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 f773e1000..20de3cbbb 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,4 +1,4 @@ -import { mkdir, mkdtemp, rm, writeFile } from "node:fs/promises"; +import { mkdir, mkdtemp, readFile, rm, writeFile } from "node:fs/promises"; import { join } from "node:path"; import { tmpdir } from "node:os"; import { describe, expect, test, vi } from "vitest"; @@ -9,6 +9,7 @@ import type { AgentStreamEvent } from "../agent-sdk-types.js"; import { PiDirectAgentClient, PiDirectAgentSession, + type PiDirectSessionRuntimeAdapter, type PiDirectSessionAdapter, } from "./pi-direct-agent.js"; @@ -43,6 +44,16 @@ function createPiSession(prompt: () => Promise): PiDirectSessionAdapter { }; } +function createPiRuntime( + session: PiDirectSessionAdapter, + dispose: () => Promise = vi.fn(async () => undefined), +): PiDirectSessionRuntimeAdapter { + return { + session, + dispose, + }; +} + function createPiModel(provider: string, id: string): Model { return { provider, @@ -61,7 +72,7 @@ function createPiModel(provider: string, id: string): Model { describe("PiDirectAgentSession", () => { test("treats SDK request abort rejections as turn cancellations", async () => { const session = new PiDirectAgentSession( - createPiSession(() => Promise.reject(new Error("Request was aborted."))), + createPiRuntime(createPiSession(() => Promise.reject(new Error("Request was aborted.")))), { find: vi.fn(), getAll: vi.fn(() => []) }, { provider: "pi", @@ -87,7 +98,7 @@ describe("PiDirectAgentSession", () => { test("setModel creates a minimal model for new ids under a known provider", async () => { const sdkSession = createPiSession(async () => undefined); const session = new PiDirectAgentSession( - sdkSession, + createPiRuntime(sdkSession), { find: vi.fn(() => undefined), getAll: vi.fn(() => [createPiModel("openrouter", "known-model")]), @@ -196,12 +207,15 @@ export default function(pi) { const agentDir = join(testRoot, "agent"); const cwd = join(testRoot, "project"); const extensionDir = join(cwd, ".pi", "extensions"); + const shutdownMarker = join(testRoot, "shutdown.txt"); process.env.PI_CODING_AGENT_DIR = agentDir; await mkdir(extensionDir, { recursive: true }); await writeFile( join(extensionDir, "dummy-command.ts"), ` +import { writeFileSync } from "node:fs"; + export default function(pi) { pi.registerProvider("paseo-dummy", { baseUrl: "https://example.invalid/v1", @@ -224,6 +238,10 @@ export default function(pi) { description: "Dummy extension command", handler: async () => {} }); + + pi.on("session_shutdown", async () => { + writeFileSync(${JSON.stringify(shutdownMarker)}, "closed"); + }); } `, "utf-8", @@ -237,6 +255,7 @@ export default function(pi) { cwd, model: "paseo-dummy/extension-model", }); + let closed = false; try { await expect(session.listCommands()).resolves.toContainEqual({ @@ -244,8 +263,13 @@ export default function(pi) { description: "Dummy extension command", argumentHint: "", }); - } finally { await session.close(); + closed = true; + await expect(readFile(shutdownMarker, "utf-8")).resolves.toBe("closed"); + } finally { + if (!closed) { + await session.close(); + } } } finally { if (previousAgentDir === undefined) { 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 a07f3ada6..3501631b4 100644 --- a/packages/server/src/server/agent/providers/pi-direct-agent.ts +++ b/packages/server/src/server/agent/providers/pi-direct-agent.ts @@ -8,11 +8,14 @@ import { ModelRegistry, SessionManager, createAgentSessionFromServices, + createAgentSessionRuntime, createAgentSessionServices, + getAgentDir, type AgentSession as PiAgentSession, type AgentSessionEvent, type AgentSessionServices, type BashToolInput, + type CreateAgentSessionRuntimeFactory, type EditToolInput, type FindToolInput, type GrepToolInput, @@ -104,7 +107,7 @@ interface StartTurnResult { } interface PiDirectSessionResult { - session: PiAgentSession; + runtime: PiDirectSessionRuntimeAdapter; modelRegistry: ModelRegistry; } @@ -128,6 +131,11 @@ export type PiDirectSessionAdapter = Pick< | "thinkingLevel" >; +export interface PiDirectSessionRuntimeAdapter { + readonly session: PiDirectSessionAdapter; + dispose(): Promise; +} + type PiDirectModelRegistry = Pick; interface PiToolResultObject { @@ -835,18 +843,23 @@ export class PiDirectAgentSession implements AgentSession { private latestUsage: AgentUsage | undefined; constructor( - private readonly session: PiDirectSessionAdapter, + private readonly runtime: PiDirectSessionRuntimeAdapter, private readonly modelRegistry: PiDirectModelRegistry, private readonly config: AgentSessionConfig, ) { + const session = this.session; this.lastKnownThinkingOptionId = normalizePiThinkingOption(config.thinkingOptionId) ?? session.thinkingLevel ?? null; - this.session.subscribe((event) => { + session.subscribe((event) => { this.handleSessionEvent(event); }); } + private get session(): PiDirectSessionAdapter { + return this.runtime.session; + } + get id(): string | null { return this.session.sessionId; } @@ -1307,7 +1320,7 @@ export class PiDirectAgentSession implements AgentSession { } async close(): Promise { - this.session.dispose(); + await this.runtime.dispose(); } async listCommands(): Promise { @@ -1350,9 +1363,10 @@ export class PiDirectAgentClient implements AgentClient { this.runtimeSettings = options.runtimeSettings; } - private async getSessionServices(cwd: string): Promise { + private async getSessionServices(cwd: string, agentDir?: string): Promise { return createAgentSessionServices({ cwd, + ...(agentDir ? { agentDir } : {}), ...(this.modelRegistry ? { modelRegistry: this.modelRegistry } : {}), }); } @@ -1369,29 +1383,57 @@ export class PiDirectAgentClient implements AgentClient { return findModelInRegistry(registry, parsedReference); } - private async createSdkSession(config: AgentSessionConfig): Promise { - const thinkingLevel = - normalizePiThinkingOption(config.thinkingOptionId) ?? DEFAULT_PI_THINKING_LEVEL; - const services = await this.getSessionServices(config.cwd); - const model = this.resolveConfiguredModel(services.modelRegistry, config.model); + private async createSdkRuntime( + config: AgentSessionConfig, + sessionManager: SessionManager, + options: { defaultThinkingLevel?: ThinkingLevel } = {}, + ): Promise { + const createRuntime: CreateAgentSessionRuntimeFactory = async ({ + cwd, + agentDir, + sessionManager: runtimeSessionManager, + sessionStartEvent, + }) => { + const thinkingLevel = + normalizePiThinkingOption(config.thinkingOptionId) ?? options.defaultThinkingLevel; + const services = await this.getSessionServices(cwd, agentDir); + const model = this.resolveConfiguredModel(services.modelRegistry, config.model); - const { session } = await createAgentSessionFromServices({ - services, - sessionManager: SessionManager.create(config.cwd), - thinkingLevel, - ...(model ? { model } : {}), + return { + ...(await createAgentSessionFromServices({ + services, + sessionManager: runtimeSessionManager, + sessionStartEvent, + ...(thinkingLevel ? { thinkingLevel } : {}), + ...(model ? { model } : {}), + })), + services, + diagnostics: services.diagnostics, + }; + }; + + const runtime = await createAgentSessionRuntime(createRuntime, { + cwd: sessionManager.getCwd(), + agentDir: getAgentDir(), + sessionManager, }); - await session.bindExtensions({}); - applySystemPrompt(session, config.systemPrompt); - return { session, modelRegistry: services.modelRegistry }; + await runtime.session.bindExtensions({}); + applySystemPrompt(runtime.session, config.systemPrompt); + return { runtime, modelRegistry: runtime.services.modelRegistry }; } async createSession( config: AgentSessionConfig, _launchContext?: AgentLaunchContext, ): Promise { - const { session, modelRegistry } = await this.createSdkSession(config); - return new PiDirectAgentSession(session, modelRegistry, config); + const { runtime, modelRegistry } = await this.createSdkRuntime( + config, + SessionManager.create(config.cwd), + { + defaultThinkingLevel: DEFAULT_PI_THINKING_LEVEL, + }, + ); + return new PiDirectAgentSession(runtime, modelRegistry, config); } async resumeSession( @@ -1404,9 +1446,10 @@ export class PiDirectAgentClient implements AgentClient { throw new Error("Pi resume requires a native session file handle"); } - const resumedManager = SessionManager.open(sessionFile); + const initialManager = SessionManager.open(sessionFile); const persistenceMetadata = parsePersistenceMetadata(handle.metadata); - const cwd = overrides?.cwd ?? persistenceMetadata.cwd ?? resumedManager.getCwd(); + const cwd = overrides?.cwd ?? persistenceMetadata.cwd ?? initialManager.getCwd(); + const resumedManager = SessionManager.open(sessionFile, undefined, cwd); const mergedConfig: AgentSessionConfig = { provider: PI_PROVIDER, cwd, @@ -1425,18 +1468,8 @@ export class PiDirectAgentClient implements AgentClient { modeId: overrides?.modeId, }; - const services = await this.getSessionServices(mergedConfig.cwd); - const model = this.resolveConfiguredModel(services.modelRegistry, mergedConfig.model); - const thinkingLevel = normalizePiThinkingOption(mergedConfig.thinkingOptionId); - const { session } = await createAgentSessionFromServices({ - services, - sessionManager: resumedManager, - ...(model ? { model } : {}), - ...(thinkingLevel ? { thinkingLevel } : {}), - }); - await session.bindExtensions({}); - applySystemPrompt(session, mergedConfig.systemPrompt); - return new PiDirectAgentSession(session, services.modelRegistry, mergedConfig); + const { runtime, modelRegistry } = await this.createSdkRuntime(mergedConfig, resumedManager); + return new PiDirectAgentSession(runtime, modelRegistry, mergedConfig); } async listModels(options: ListModelsOptions): Promise {