diff --git a/packages/server/src/server/agent/providers/pi-direct-agent.test.ts b/packages/server/src/server/agent/providers/pi-direct-agent.test.ts index 07ccecfb0..8445c9e7f 100644 --- a/packages/server/src/server/agent/providers/pi-direct-agent.test.ts +++ b/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 }); + } + }); }); diff --git a/packages/server/src/server/agent/providers/pi-direct-agent.ts b/packages/server/src/server/agent/providers/pi-direct-agent.ts index 8c5c422c6..cfaed4432 100644 --- a/packages/server/src/server/agent/providers/pi-direct-agent.ts +++ b/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 { - if (!this.modelRegistry) { - this.modelRegistry = ModelRegistry.create(AuthStorage.create()); - } - return this.modelRegistry; + private async getSessionServices(cwd: string): Promise { + return createAgentSessionServices({ + cwd, + ...(this.modelRegistry ? { modelRegistry: this.modelRegistry } : {}), + }); } - private resolveConfiguredModel(modelId: string | null | undefined): Model | undefined { + private resolveConfiguredModel( + registry: PiDirectModelRegistry, + modelId: string | null | undefined, + ): Model | undefined { const parsedReference = parseModelReference(modelId ?? null); if (!parsedReference) { return undefined; } - const registry = this.getModelRegistry(); return findModelInRegistry(registry, parsedReference); } - private async createSdkSession(config: AgentSessionConfig): Promise { + private async createSdkSession(config: AgentSessionConfig): Promise { const thinkingLevel = normalizePiThinkingOption(config.thinkingOptionId) ?? DEFAULT_PI_THINKING_LEVEL; - const modelRegistry = this.getModelRegistry(); - const model = this.resolveConfiguredModel(config.model); + const services = await this.getSessionServices(config.cwd); + const model = this.resolveConfiguredModel(services.modelRegistry, config.model); - const sessionOptions = { - cwd: config.cwd, - modelRegistry, + const { session } = await createAgentSessionFromServices({ + services, 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 { - const session = await this.createSdkSession(config); - return new PiDirectAgentSession(session, this.getModelRegistry(), config); + const { session, modelRegistry } = await this.createSdkSession(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 = { - cwd: mergedConfig.cwd, - modelRegistry: this.getModelRegistry(), + const { session } = await createAgentSessionFromServices({ + services, 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 { - // Pi Direct uses an in-process global registry; cwd/force are intentionally irrelevant. - const models = this.getModelRegistry() - .getAvailable() - .map((model) => ({ - provider: PI_PROVIDER, - id: `${model.provider}/${model.id}`, - label: `${model.provider}/${model.name}`, - description: `${model.provider}/${model.id}`, - metadata: { - provider: model.provider, - modelId: model.id, - } satisfies AgentMetadata, - thinkingOptions: model.reasoning ? PI_THINKING_OPTIONS.map(mapThinkingOption) : undefined, - defaultThinkingOptionId: model.reasoning ? DEFAULT_PI_THINKING_LEVEL : undefined, - })); + async listModels(options: ListModelsOptions): Promise { + const services = await this.getSessionServices(options.cwd); + const models = services.modelRegistry.getAvailable().map((model) => ({ + provider: PI_PROVIDER, + id: `${model.provider}/${model.id}`, + label: `${model.provider}/${model.name}`, + description: `${model.provider}/${model.id}`, + metadata: { + provider: model.provider, + modelId: model.id, + } satisfies AgentMetadata, + thinkingOptions: model.reasoning ? PI_THINKING_OPTIONS.map(mapThinkingOption) : undefined, + defaultThinkingOptionId: model.reasoning ? DEFAULT_PI_THINKING_LEVEL : undefined, + })); return transformPiModels(models); }