jb/paseo
1
0
Fork 0
Code Issues 1 Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Mohamed Boudra
2026-05-10 12:28:29 +08:00
committed by GitHub
parent ab4525db5b
commit 7d823fd3ec
2 changed files with 95 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/server/src/server/agent/providers/pi-direct-agent.test.ts
+28 -4
View File
@@ -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
+67 -34
View File
@@ -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> {
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<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({
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<AgentSession> {
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<AgentModelDefinition[]> {