Load Pi extensions for model discovery

packages/server/src/server/agent/providers/pi-direct-agent.test.ts

@@ -1,3 +1,6 @@
+import { mkdir, mkdtemp, rm, writeFile } from "node:fs/promises";
+import { join } from "node:path";
+import { tmpdir } from "node:os";
 import { describe, expect, test, vi } from "vitest";
 import type { Api, Model } from "@mariozechner/pi-ai";
 import pino from "pino";
@@ -93,4 +96,126 @@ describe("PiDirectAgentClient", () => {
     expect(registry.getAll).not.toHaveBeenCalled();
     expect(models.map((model) => model.id)).toEqual(["anthropic/claude-opus-4-5"]);
   });
+
+  test("loads project extensions before listing available models", async () => {
+    const testRoot = await mkdtemp(join(tmpdir(), "paseo-pi-extension-"));
+    const previousAgentDir = process.env.PI_CODING_AGENT_DIR;
+
+    try {
+      const agentDir = join(testRoot, "agent");
+      const cwd = join(testRoot, "project");
+      const extensionDir = join(cwd, ".pi", "extensions");
+      process.env.PI_CODING_AGENT_DIR = agentDir;
+
+      await mkdir(extensionDir, { recursive: true });
+      await writeFile(
+        join(extensionDir, "dummy-provider.ts"),
+        `
+export default function(pi) {
+  pi.registerProvider("paseo-dummy", {
+    baseUrl: "https://example.invalid/v1",
+    apiKey: "paseo-test-key",
+    api: "openai-responses",
+    models: [
+      {
+        id: "extension-model",
+        name: "Extension Model",
+        reasoning: true,
+        input: ["text"],
+        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+        contextWindow: 128000,
+        maxTokens: 4096
+      }
+    ]
+  });
+}
+`,
+        "utf-8",
+      );
+
+      const client = new PiDirectAgentClient({
+        logger: pino({ level: "silent" }),
+      });
+
+      const models = await client.listModels({ cwd, force: false });
+
+      expect(models.map((model) => model.id)).toContain("paseo-dummy/extension-model");
+    } finally {
+      if (previousAgentDir === undefined) {
+        delete process.env.PI_CODING_AGENT_DIR;
+      } else {
+        process.env.PI_CODING_AGENT_DIR = previousAgentDir;
+      }
+      await rm(testRoot, { recursive: true, force: true });
+    }
+  });
+
+  test("creates sessions with project extension models and exposes extension commands", async () => {
+    const testRoot = await mkdtemp(join(tmpdir(), "paseo-pi-extension-session-"));
+    const previousAgentDir = process.env.PI_CODING_AGENT_DIR;
+
+    try {
+      const agentDir = join(testRoot, "agent");
+      const cwd = join(testRoot, "project");
+      const extensionDir = join(cwd, ".pi", "extensions");
+      process.env.PI_CODING_AGENT_DIR = agentDir;
+
+      await mkdir(extensionDir, { recursive: true });
+      await writeFile(
+        join(extensionDir, "dummy-command.ts"),
+        `
+export default function(pi) {
+  pi.registerProvider("paseo-dummy", {
+    baseUrl: "https://example.invalid/v1",
+    apiKey: "paseo-test-key",
+    api: "openai-responses",
+    models: [
+      {
+        id: "extension-model",
+        name: "Extension Model",
+        reasoning: true,
+        input: ["text"],
+        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+        contextWindow: 128000,
+        maxTokens: 4096
+      }
+    ]
+  });
+
+  pi.registerCommand("dummy-command", {
+    description: "Dummy extension command",
+    handler: async () => {}
+  });
+}
+`,
+        "utf-8",
+      );
+
+      const client = new PiDirectAgentClient({
+        logger: pino({ level: "silent" }),
+      });
+      const session = await client.createSession({
+        provider: "pi",
+        cwd,
+        model: "paseo-dummy/extension-model",
+      });
+
+      try {
+        await expect(session.listCommands()).resolves.toContainEqual({
+          name: "dummy-command",
+          description: "Dummy extension command",
+          argumentHint: "",
+        });
+      } finally {
+        await session.close();
+      }
+    } finally {
+      if (previousAgentDir === undefined) {
+        delete process.env.PI_CODING_AGENT_DIR;
+      } else {
+        process.env.PI_CODING_AGENT_DIR = previousAgentDir;
+      }
+      await rm(testRoot, { recursive: true, force: true });
+    }
+  });
 });

packages/server/src/server/agent/providers/pi-direct-agent.ts

@@ -4,12 +4,13 @@ import { join } from "node:path";
 import { homedir } from "node:os";
 import type { Logger } from "pino";
 import {
-  AuthStorage,
   ModelRegistry,
   SessionManager,
-  createAgentSession,
+  createAgentSessionFromServices,
+  createAgentSessionServices,
   type AgentSession as PiAgentSession,
   type AgentSessionEvent,
+  type AgentSessionServices,
   type BashToolInput,
   type EditToolInput,
   type FindToolInput,
@@ -101,6 +102,11 @@ interface StartTurnResult {
   turnId: string;
 }
 
+interface PiDirectSessionResult {
+  session: PiAgentSession;
+  modelRegistry: ModelRegistry;
+}
+
 export type PiDirectSessionAdapter = Pick<
   PiAgentSession,
   | "abort"
@@ -1359,48 +1365,48 @@ export class PiDirectAgentClient implements AgentClient {
     this.runtimeSettings = options.runtimeSettings;
   }
 
-  private getModelRegistry(): ModelRegistry {
+  private async getSessionServices(cwd: string): Promise<AgentSessionServices> {
-    if (!this.modelRegistry) {
+    return createAgentSessionServices({
-      this.modelRegistry = ModelRegistry.create(AuthStorage.create());
+      cwd,
-    }
+      ...(this.modelRegistry ? { modelRegistry: this.modelRegistry } : {}),
-    return this.modelRegistry;
+    });
   }
 
-  private resolveConfiguredModel(modelId: string | null | undefined): Model<Api> | undefined {
+  private resolveConfiguredModel(
+    registry: PiDirectModelRegistry,
+    modelId: string | null | undefined,
+  ): Model<Api> | undefined {
     const parsedReference = parseModelReference(modelId ?? null);
     if (!parsedReference) {
       return undefined;
     }
 
-    const registry = this.getModelRegistry();
     return findModelInRegistry(registry, parsedReference);
   }
 
-  private async createSdkSession(config: AgentSessionConfig): Promise<PiAgentSession> {
+  private async createSdkSession(config: AgentSessionConfig): Promise<PiDirectSessionResult> {
     const thinkingLevel =
       normalizePiThinkingOption(config.thinkingOptionId) ?? DEFAULT_PI_THINKING_LEVEL;
-    const modelRegistry = this.getModelRegistry();
+    const services = await this.getSessionServices(config.cwd);
-    const model = this.resolveConfiguredModel(config.model);
+    const model = this.resolveConfiguredModel(services.modelRegistry, config.model);
 
-    const sessionOptions = {
+    const { session } = await createAgentSessionFromServices({
-      cwd: config.cwd,
+      services,
-      modelRegistry,
       sessionManager: SessionManager.create(config.cwd),
       thinkingLevel,
       ...(model ? { model } : {}),
-    };
+    });
-    const { session } = await createAgentSession(sessionOptions);
     await session.bindExtensions({});
     applySystemPrompt(session, config.systemPrompt);
-    return session;
+    return { session, modelRegistry: services.modelRegistry };
   }
 
   async createSession(
     config: AgentSessionConfig,
     _launchContext?: AgentLaunchContext,
   ): Promise<AgentSession> {
-    const session = await this.createSdkSession(config);
+    const { session, modelRegistry } = await this.createSdkSession(config);
-    return new PiDirectAgentSession(session, this.getModelRegistry(), config);
+    return new PiDirectAgentSession(session, modelRegistry, config);
   }
 
   async resumeSession(
@@ -1434,37 +1440,34 @@ export class PiDirectAgentClient implements AgentClient {
       modeId: overrides?.modeId,
     };
 
-    const model = this.resolveConfiguredModel(mergedConfig.model);
+    const services = await this.getSessionServices(mergedConfig.cwd);
+    const model = this.resolveConfiguredModel(services.modelRegistry, mergedConfig.model);
     const thinkingLevel = normalizePiThinkingOption(mergedConfig.thinkingOptionId);
-    const sessionOptions = {
+    const { session } = await createAgentSessionFromServices({
-      cwd: mergedConfig.cwd,
+      services,
-      modelRegistry: this.getModelRegistry(),
       sessionManager: resumedManager,
       ...(model ? { model } : {}),
       ...(thinkingLevel ? { thinkingLevel } : {}),
-    };
+    });
-    const { session } = await createAgentSession(sessionOptions);
     await session.bindExtensions({});
     applySystemPrompt(session, mergedConfig.systemPrompt);
-    return new PiDirectAgentSession(session, this.getModelRegistry(), mergedConfig);
+    return new PiDirectAgentSession(session, services.modelRegistry, mergedConfig);
   }
 
-  async listModels(_options: ListModelsOptions): Promise<AgentModelDefinition[]> {
+  async listModels(options: ListModelsOptions): Promise<AgentModelDefinition[]> {
-    // Pi Direct uses an in-process global registry; cwd/force are intentionally irrelevant.
+    const services = await this.getSessionServices(options.cwd);
-    const models = this.getModelRegistry()
+    const models = services.modelRegistry.getAvailable().map((model) => ({
-      .getAvailable()
+      provider: PI_PROVIDER,
-      .map((model) => ({
+      id: `${model.provider}/${model.id}`,
-        provider: PI_PROVIDER,
+      label: `${model.provider}/${model.name}`,
-        id: `${model.provider}/${model.id}`,
+      description: `${model.provider}/${model.id}`,
-        label: `${model.provider}/${model.name}`,
+      metadata: {
-        description: `${model.provider}/${model.id}`,
+        provider: model.provider,
-        metadata: {
+        modelId: model.id,
-          provider: model.provider,
+      } satisfies AgentMetadata,
-          modelId: model.id,
+      thinkingOptions: model.reasoning ? PI_THINKING_OPTIONS.map(mapThinkingOption) : undefined,
-        } satisfies AgentMetadata,
+      defaultThinkingOptionId: model.reasoning ? DEFAULT_PI_THINKING_LEVEL : undefined,
-        thinkingOptions: model.reasoning ? PI_THINKING_OPTIONS.map(mapThinkingOption) : undefined,
+    }));
-        defaultThinkingOptionId: model.reasoning ? DEFAULT_PI_THINKING_LEVEL : undefined,
-      }));
 
     return transformPiModels(models);
   }