feat: Pi extension uses Pi's internal complete() for LLM analysis

- PiLLMClient now imports @mariozechner/pi-ai's complete() function
- Respects user's /model selection and /login auth
- No external fetch() — Pi handles transport internally
- Added src/types/pi-ai.d.ts for TypeScript declarations
- Extension reverted to real LLM calls (initProject, patchFile, reinitPath)
- Tests mock @mariozechner/pi-ai for verification

pi-extension.ts

@@ -1,39 +1,20 @@
 import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
 import { Type } from "typebox";
-import { readFileSync, readdirSync, statSync, writeFileSync } from "fs";
+import { readFileSync, readdirSync, statSync } from "fs";
-import { join, relative, dirname } from "path";
+import { join, relative } from "path";
-import { validateMaps } from "./src/index.js";
+import { initProject, patchFile, validateMaps, reinitPath } from "./src/index.js";
-import { discoverProject } from "./src/discover.js";
+import { createLLMClient } from "./src/llm-client.js";
-import { extractFileAST } from "./src/ast-extract.js";
+import { LLMError } from "./src/llm-error.js";
-import { extractFileHeuristic, extractPackageHeuristic } from "./src/llm-extract.js";
-import { mergeFileData } from "./src/merge.js";
-import { renderPackageMap, type FileEntry } from "./src/format.js";
 
 /**
- * Generate maps using heuristics + AST only (no LLM).
+ * Get the LLM client for Pi runtime.
- * Inside Pi, we let Pi's agent handle semantic analysis.
+ *
- * The tool does structural work; the agent can refine later.
+ * When running inside Pi, we ALWAYS use Pi's native LLM via
+ * @mariozechner/pi-ai's complete() function. This respects the
+ * user's /model selection and /login auth. No external fetch().
  */
-async function generateMapHeuristic(rootPath: string): Promise<void> {
+function getPiLLMClient(ctx: any) {
-	const entries = discoverProject(rootPath);
+	return createLLMClient("pi", { extensionContext: ctx });
-	for (const entry of entries) {
-		const fileData: FileEntry[] = [];
-		for (const file of entry.files) {
-			const filePath = join(entry.dirPath, file);
-			const astData = await extractFileAST(filePath);
-			const heuristicData = await extractFileHeuristic(filePath);
-			fileData.push(mergeFileData(file, heuristicData, astData));
-		}
-		const packageData = await extractPackageHeuristic(entry.relativePath, fileData);
-		const mapData = {
-			path: entry.relativePath,
-			role: packageData.role,
-			files: fileData,
-			arch: packageData.arch,
-			dirty: "-",
-		};
-		writeFileSync(join(entry.dirPath, ".pi-map.md"), renderPackageMap(mapData));
-	}
 }
 
 function findPiMapFiles(cwd: string): string[] {
@@ -91,18 +72,19 @@ export default function (pi: ExtensionAPI) {
 		async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
 			try {
 				const targetPath = params.path || ctx.cwd;
-				await generateMapHeuristic(targetPath);
+				const client = getPiLLMClient(ctx);
+				await initProject(targetPath, { verbose: false, llmClient: client });
 				return {
 					content: [
 						{
 							type: "text",
-							text: `Generated heuristic .pi-map.md files for ${targetPath}.`,
+							text: `Generated .pi-map.md files for ${targetPath}`,
 						},
 					],
 					details: { success: true, cwd: ctx.cwd },
 				};
 			} catch (err: any) {
-				const msg = String(err);
+				const msg = err instanceof LLMError ? err.message : String(err);
 				return {
 					content: [{ type: "text", text: `Error: ${msg}` }],
 					details: { success: false, error: msg },
@@ -128,8 +110,8 @@ export default function (pi: ExtensionAPI) {
 		}),
 		async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
 			try {
-				const dirPath = dirname(params.file_path);
+				const client = getPiLLMClient(ctx);
-				await generateMapHeuristic(dirPath);
+				await patchFile(params.file_path, client);
 				return {
 					content: [
 						{
@@ -140,7 +122,7 @@ export default function (pi: ExtensionAPI) {
 					details: { success: true },
 				};
 			} catch (err: any) {
-				const msg = String(err);
+				const msg = err instanceof LLMError ? err.message : String(err);
 				return {
 					content: [{ type: "text", text: `Error: ${msg}` }],
 					details: { success: false, error: msg },
@@ -183,7 +165,7 @@ export default function (pi: ExtensionAPI) {
 					details: { success: true, clean: result.clean },
 				};
 			} catch (err: any) {
-				const msg = String(err);
+				const msg = err instanceof LLMError ? err.message : String(err);
 				return {
 					content: [{ type: "text", text: `Error: ${msg}` }],
 					details: { success: false, error: msg },
@@ -211,18 +193,19 @@ export default function (pi: ExtensionAPI) {
 		async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
 			try {
 				const targetPath = params.path || ctx.cwd;
-				await generateMapHeuristic(targetPath);
+				const client = getPiLLMClient(ctx);
+				await reinitPath(targetPath, { verbose: false, llmClient: client });
 				return {
 					content: [
 						{
 							type: "text",
-							text: `Regenerated heuristic maps for ${targetPath}`,
+							text: `Regenerated maps for ${targetPath}`,
 						},
 					],
 					details: { success: true },
 				};
 			} catch (err: any) {
-				const msg = String(err);
+				const msg = err instanceof LLMError ? err.message : String(err);
 				return {
 					content: [{ type: "text", text: `Error: ${msg}` }],
 					details: { success: false, error: msg },

src/pi-llm-client.ts

@@ -2,10 +2,11 @@ import { LLMError } from "./llm-error.js";
 import type { LLMClient } from "./llm-client.js";
 
 /**
- * Pi LLM Client — uses Pi's built-in model when running inside the Pi runtime.
+ * Pi LLM Client — uses Pi's built-in `complete()` from `@mariozechner/pi-ai`.
  *
- * This client ALWAYS attempts to use Pi's native LLM and throws a hard error
+ * This runs inside the Pi extension runtime and calls Pi's internal AI layer,
- * if Pi's model is not accessible. It never falls back to external APIs.
+ * which respects the user's configured model (/model) and auth (/login).
+ * No external fetch() — Pi handles transport, retries, and token accounting.
  */
 export class PiLLMClient implements LLMClient {
 	constructor(private extensionContext?: unknown) {}
@@ -16,57 +17,66 @@ export class PiLLMClient implements LLMClient {
 		if (!ctx) {
 			throw new LLMError(
 				"Pi LLM not accessible: no ExtensionContext provided. " +
-					"This extension must run inside Pi.",
+					"This tool must run inside Pi.",
 			);
 		}
 
-		// Try to access Pi's model or modelRegistry
+		// Get the active model from Pi's runtime
 		const model = ctx.model ?? ctx.modelRegistry?.get?.();
 		if (!model) {
 			throw new LLMError(
-				"Pi LLM not accessible: no model configured in Pi runtime. " +
+				"Pi LLM not accessible: no model configured. " +
-					"Configure a model (e.g., via /model) before using project-map tools.",
+					"Set a model via /model before using project-map tools.",
 			);
 		}
 
-		// Attempt to call the model via Pi's provider interface.
-		// Pi exposes an OpenAI-compatible completions endpoint through the
-		// active provider. We construct a minimal chat-completion payload
-		// and send it through the model's fetch wrapper.
 		try {
-			const provider = ctx.modelRegistry?.getProvider?.(model.provider);
+			// Dynamically import Pi's AI module (available in the Pi runtime)
-			if (!provider || typeof provider.complete !== "function") {
+			const { complete } = await import("@mariozechner/pi-ai");
-				// Fallback: try using Pi's internal agent session to get a completion
+
-				return await this.callViaAgentSession(ctx, prompt);
+			const response = await complete(
+				model,
+				{
+					systemPrompt:
+						"You are a code analysis assistant. Analyze the provided file and respond with concise, structured information.",
+					messages: [
+						{
+							role: "user",
+							content: prompt,
+							timestamp: Date.now(),
+						},
+					],
+				},
+				{
+					temperature: 0.1,
+					maxTokens: 256,
+				},
+			);
+
+			if (response.errorMessage) {
+				throw new LLMError(`Pi LLM error: ${response.errorMessage}`);
 			}
-			return await provider.complete(prompt);
+
+			// Extract text content from AssistantMessage
+			const text = response.content
+				.filter((c: any) => c.type === "text")
+				.map((c: any) => c.text)
+				.join("")
+				.trim();
+
+			return text;
 		} catch (err: any) {
 			if (err instanceof LLMError) throw err;
+			if (err.code === "MODULE_NOT_FOUND") {
+				throw new LLMError(
+					"Pi LLM not accessible: @mariozechner/pi-ai is not available. " +
+						"This extension must run inside the Pi runtime.",
+				);
+			}
 			throw new LLMError(
-				`Pi LLM call failed: ${err.message || String(err)}. ` +
+				`Pi LLM call failed: ${err.message || String(err)}`,
-					"Ensure Pi has a working model configured.",
 				err,
 			);
 		}
 	}
-
-	private async callViaAgentSession(ctx: any, prompt: string): Promise<string> {
-		// Pi's ExtensionContext may expose a session or agent that can
-		// make model calls. This is the fallback when direct provider
-		// access is not available.
-		if (typeof ctx.sendMessage === "function") {
-			// Not ideal — sendMessage injects into conversation.
-			// Reserved for future Pi SDK improvement.
-			throw new LLMError(
-				"Pi LLM direct call not yet supported by this Pi version. " +
-					"The extension requires a Pi runtime with exposed model provider API.",
-			);
-		}
-
-		throw new LLMError(
-			"Pi LLM not accessible: the Pi runtime does not expose a callable model API. " +
-				"This extension requires Pi to provide either ctx.modelRegistry.getProvider() " +
-				"or a direct model completion interface.",
-		);
-	}
 }

src/types/pi-ai.d.ts

@@ -0,0 +1,24 @@
+// Pi AI module — available only inside the Pi runtime
+// Declared here so TypeScript doesn't error during `npm run build`
+
+declare module "@mariozechner/pi-ai" {
+	export function complete(
+		model: any,
+		context: {
+			systemPrompt?: string;
+			messages: Array<{
+				role: "user" | "assistant" | "system" | "toolResult";
+				content: string;
+				timestamp?: number;
+			}>;
+		},
+		options?: {
+			temperature?: number;
+			maxTokens?: number;
+		},
+	): Promise<{
+		role: "assistant";
+		content: Array<{ type: string; text?: string }>;
+		errorMessage?: string;
+	}>;
+}

tests/pi-extension.test.ts

@@ -15,14 +15,17 @@ vi.mock("typebox", () => ({
 	},
 }));
 
+// Mock Pi's AI module so PiLLMClient works in tests
+vi.mock("@mariozechner/pi-ai", () => ({
+	complete: vi.fn(async (_model: any, _context: any) => ({
+		role: "assistant" as const,
+		content: [{ type: "text", text: "PURPOSE: Test file\nDEPS: none\nCONCEPTS: testing" }],
+	})),
+}));
+
 import extension from "../pi-extension.js";
 
-const CACHE_FILE = join(
+const CACHE_FILE = join(homedir(), ".cache", "pi-project-map", "llm-cache.json");
-	homedir(),
-	".cache",
-	"pi-project-map",
-	"llm-cache.json",
-);
 function clearCache() {
 	if (existsSync(CACHE_FILE)) unlinkSync(CACHE_FILE);
 }
@@ -40,7 +43,10 @@ describe("pi-extension", () => {
 		mockNotify = vi.fn();
 		mockCtx = {
 			cwd: "/home/project",
-			modelRegistry: {},
+			model: { provider: "openai", id: "gpt-4o-mini", api: "openai-completions" },
+			modelRegistry: {
+				getApiKeyAndHeaders: vi.fn(async () => ({ apiKey: "test-key", headers: {} })),
+			},
 			ui: { notify: mockNotify },
 		};
 
@@ -72,7 +78,7 @@ describe("pi-extension", () => {
 	});
 
 	describe("project_map_init tool", () => {
-		it("generates heuristic maps successfully", async () => {
+		it("generates maps with mocked Pi LLM", async () => {
 			const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
 			writeFileSync(join(dir, "test.ts"), `export const x = ${Date.now()};`);
 			mockCtx.cwd = dir;
@@ -80,12 +86,24 @@ describe("pi-extension", () => {
 			const tool = registeredTools.project_map_init;
 			const result = await tool.execute("tool-1", {}, null, null, mockCtx);
 			expect(result.details.success).toBe(true);
-			expect(result.content[0].text).toContain("Generated heuristic");
+			expect(result.content[0].text).toContain("Generated");
+		});
+
+		it("returns error when no Pi model configured", async () => {
+			const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
+			writeFileSync(join(dir, "test.ts"), `export const x = ${Date.now()};`);
+			mockCtx.cwd = dir;
+			mockCtx.model = null;
+
+			const tool = registeredTools.project_map_init;
+			const result = await tool.execute("tool-1", {}, null, null, mockCtx);
+			expect(result.details.success).toBe(false);
+			expect(result.content[0].text).toContain("no model configured");
 		});
 	});
 
 	describe("project_map_patch tool", () => {
-		it("patches map successfully using heuristics", async () => {
+		it("patches map with mocked Pi LLM", async () => {
 			const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
 			const file = join(dir, "test.ts");
 			writeFileSync(file, `export const x = ${Date.now()};`);
@@ -132,7 +150,7 @@ describe("pi-extension", () => {
 	});
 
 	describe("project_map_reinit tool", () => {
-		it("regenerates heuristic maps successfully", async () => {
+		it("regenerates maps with mocked Pi LLM", async () => {
 			const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
 			writeFileSync(join(dir, "test.ts"), `export const x = ${Date.now()};`);
 			mockCtx.cwd = dir;
@@ -140,7 +158,7 @@ describe("pi-extension", () => {
 			const tool = registeredTools.project_map_reinit;
 			const result = await tool.execute("tool-1", {}, null, null, mockCtx);
 			expect(result.details.success).toBe(true);
-			expect(result.content[0].text).toContain("Regenerated heuristic");
+			expect(result.content[0].text).toContain("Regenerated");
 		});
 	});