fix: close Pi sessions through runtime owner (#863)

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<void>): PiDirectSessionAdapter {
   };
 }
 
+function createPiRuntime(
+  session: PiDirectSessionAdapter,
+  dispose: () => Promise<void> = vi.fn(async () => undefined),
+): PiDirectSessionRuntimeAdapter {
+  return {
+    session,
+    dispose,
+  };
+}
+
 function createPiModel(provider: string, id: string): Model<Api> {
   return {
     provider,
@@ -61,7 +72,7 @@ function createPiModel(provider: string, id: string): Model<Api> {
 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) {

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<void>;
+}
+
 type PiDirectModelRegistry = Pick<ModelRegistry, "find" | "getAll">;
 
 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<void> {
-    this.session.dispose();
+    await this.runtime.dispose();
   }
 
   async listCommands(): Promise<AgentSlashCommand[]> {
@@ -1350,9 +1363,10 @@ export class PiDirectAgentClient implements AgentClient {
     this.runtimeSettings = options.runtimeSettings;
   }
 
-  private async getSessionServices(cwd: string): Promise<AgentSessionServices> {
+  private async getSessionServices(cwd: string, agentDir?: string): Promise<AgentSessionServices> {
     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<PiDirectSessionResult> {
+  private async createSdkRuntime(
-    const thinkingLevel =
+    config: AgentSessionConfig,
-      normalizePiThinkingOption(config.thinkingOptionId) ?? DEFAULT_PI_THINKING_LEVEL;
+    sessionManager: SessionManager,
-    const services = await this.getSessionServices(config.cwd);
+    options: { defaultThinkingLevel?: ThinkingLevel } = {},
-    const model = this.resolveConfiguredModel(services.modelRegistry, config.model);
+  ): Promise<PiDirectSessionResult> {
+    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({
+      return {
-      services,
+        ...(await createAgentSessionFromServices({
-      sessionManager: SessionManager.create(config.cwd),
+          services,
-      thinkingLevel,
+          sessionManager: runtimeSessionManager,
-      ...(model ? { model } : {}),
+          sessionStartEvent,
+          ...(thinkingLevel ? { thinkingLevel } : {}),
+          ...(model ? { model } : {}),
+        })),
+        services,
+        diagnostics: services.diagnostics,
+      };
+    };
+
+    const runtime = await createAgentSessionRuntime(createRuntime, {
+      cwd: sessionManager.getCwd(),
+      agentDir: getAgentDir(),
+      sessionManager,
     });
-    await session.bindExtensions({});
+    await runtime.session.bindExtensions({});
-    applySystemPrompt(session, config.systemPrompt);
+    applySystemPrompt(runtime.session, config.systemPrompt);
-    return { session, modelRegistry: services.modelRegistry };
+    return { runtime, modelRegistry: runtime.services.modelRegistry };
   }
 
   async createSession(
     config: AgentSessionConfig,
     _launchContext?: AgentLaunchContext,
   ): Promise<AgentSession> {
-    const { session, modelRegistry } = await this.createSdkSession(config);
+    const { runtime, modelRegistry } = await this.createSdkRuntime(
-    return new PiDirectAgentSession(session, modelRegistry, config);
+      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 { runtime, modelRegistry } = await this.createSdkRuntime(mergedConfig, resumedManager);
-    const model = this.resolveConfiguredModel(services.modelRegistry, mergedConfig.model);
+    return new PiDirectAgentSession(runtime, modelRegistry, mergedConfig);
-    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);
   }
 
   async listModels(options: ListModelsOptions): Promise<AgentModelDefinition[]> {