fix: bind extensions for web sessions

Fixes #19
This commit is contained in:
Federico Jaramillo Martinez
2026-06-10 20:32:25 +02:00
parent 313dbea34e
commit 824b7a0a2f
3 changed files with 90 additions and 3 deletions
+5
View File
@@ -0,0 +1,5 @@
---
"@jmfederico/pi-web": patch
---
Initialize Pi extensions for web-managed sessions so `session_start` handlers, extension resources, and startup-dependent tools run correctly.
+60 -1
View File
@@ -45,8 +45,9 @@ function sessionRecord(id: string, cwd = "/workspace") {
function fakeRuntime(sessionId = "session-1", patch: Partial<TestSession> = {}) {
const promptCalls: { text: string; options: unknown }[] = [];
const bindExtensionCalls: unknown[] = [];
const listeners: ((event: unknown) => void)[] = [];
const calls = { abort: 0, clearQueue: 0, dispose: 0, prompt: promptCalls };
const calls = { abort: 0, bindExtensions: bindExtensionCalls, clearQueue: 0, dispose: 0, prompt: promptCalls };
const session: TestSession = {
sessionId,
sessionFile: `/tmp/${sessionId}.jsonl`,
@@ -71,6 +72,10 @@ function fakeRuntime(sessionId = "session-1", patch: Partial<TestSession> = {})
if (index !== -1) listeners.splice(index, 1);
};
},
bindExtensions: (bindings: unknown) => {
calls.bindExtensions.push(bindings);
return Promise.resolve();
},
getSessionStats: () => ({ sessionId, totalMessages: 0, userMessages: 0, assistantMessages: 0, toolCalls: 0, tokens: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, cost: 0 }),
getContextUsage: () => undefined,
prompt: (text: string, options: unknown) => {
@@ -146,6 +151,7 @@ describe("PiSessionService", () => {
const session = await service.start("/workspace");
expect(createCalls).toBe(1);
expect(fake.calls.bindExtensions).toHaveLength(1);
expect(session).toMatchObject({ id: "session-1", cwd: "/workspace", messageCount: 0 });
expect(service.activeCount()).toBe(1);
expect(hub.globalEvents.some((event) => event.type === "status.update" && event.status.sessionId === "session-1")).toBe(true);
@@ -155,6 +161,59 @@ describe("PiSessionService", () => {
expect(fake.calls.dispose).toBe(1);
});
it("binds extensions again when the SDK runtime replaces the active session", async () => {
const hub = new CapturingSessionEventHub();
const fake = fakeRuntime("session-1");
const replacement = fakeRuntime("session-2");
let rebindSession: ((session: PiAgentSession) => Promise<void>) | undefined;
fake.runtime.setRebindSession = (callback) => { rebindSession = callback; };
const service = new PiSessionService(hub, {
createAgentRuntime: runtimeCreator(fake.runtime),
sessionManager: sessionGateway([]),
heartbeatIntervalMs: 60_000,
});
await service.start("/workspace");
Object.defineProperty(fake.runtime, "session", { configurable: true, value: replacement.session });
await rebindSession?.(replacement.session);
expect(fake.calls.bindExtensions).toHaveLength(1);
expect(replacement.calls.bindExtensions).toHaveLength(1);
expect(service.activeCount()).toBe(1);
expect(await service.status("session-2")).toMatchObject({ sessionId: "session-2" });
await service.dispose();
});
it("publishes extension errors reported while binding session extensions", async () => {
const hub = new CapturingSessionEventHub();
const fake = fakeRuntime("extension-session", {
bindExtensions: (bindings) => {
bindings.onError?.({ extensionPath: "pi-mcp-adapter", event: "session_start", error: "MCP failed" });
return Promise.resolve();
},
});
const service = new PiSessionService(hub, {
createAgentRuntime: runtimeCreator(fake.runtime),
sessionManager: sessionGateway([]),
heartbeatIntervalMs: 60_000,
});
await service.start("/workspace");
expect(hub.sessionEvents).toContainEqual({
sessionId: "extension-session",
event: { type: "session.error", message: "pi-mcp-adapter: MCP failed" },
});
const extensionErrorActivity = hub.globalEvents.find((event) => event.type === "activity.update" && event.activity.sessionId === "extension-session");
expect(extensionErrorActivity).toMatchObject({
type: "activity.update",
activity: { sessionId: "extension-session", phase: "error", label: "extension error", detail: "pi-mcp-adapter: MCP failed" },
});
await service.dispose();
});
it("clears stale active activity once a previously active session becomes idle", async () => {
vi.useFakeTimers();
let service: PiSessionService | undefined;
+25 -2
View File
@@ -89,6 +89,17 @@ export interface PiSessionManagerGateway {
open(path: string): PiSessionManager;
}
interface PiExtensionError {
extensionPath: string;
event: string;
error: string;
stack?: string;
}
interface PiExtensionBindings {
onError?: (error: PiExtensionError) => void;
}
export interface PiAgentSession {
modelRegistry: ModelRegistryInstance;
sessionManager: PiSessionManager;
@@ -107,6 +118,7 @@ export interface PiAgentSession {
promptTemplates: readonly { name: string; description?: string }[];
resourceLoader: { getSkills(): { skills: readonly { name: string; description?: string }[] } };
subscribe(listener: (event: unknown) => void): () => void;
bindExtensions(bindings: PiExtensionBindings): Promise<void>;
compact(instructions?: string): Promise<{ summary: string; tokensBefore: number }>;
getUserMessagesForForking(): readonly { entryId: string; text: string }[];
getSessionStats(): { sessionId: string; totalMessages: number; userMessages: number; assistantMessages: number; toolCalls: number; tokens: ClientSessionStatus["tokens"]; cost: number };
@@ -642,17 +654,28 @@ export class PiSessionService {
private async create(sessionManager: PiSessionManager, cwd: string): Promise<ActiveSession<PiSessionRuntime>> {
const runtime = await this.createAgentRuntime(this.createRuntime, { cwd, agentDir: this.agentDir, sessionManager });
await this.bindSessionExtensions(runtime.session);
const active: ActiveSession<PiSessionRuntime> = { runtime, unsubscribe: noop };
this.bindRuntime(active);
runtime.setRebindSession(() => {
runtime.setRebindSession(async (session) => {
await this.bindSessionExtensions(session);
this.bindRuntime(active);
return Promise.resolve();
});
this.active.set(runtime.session.sessionId, active);
this.publishStatus(runtime.session);
return active;
}
private async bindSessionExtensions(session: PiAgentSession): Promise<void> {
await session.bindExtensions({
onError: (error) => {
const message = `${error.extensionPath}: ${error.error}`;
this.publishActivity(session, "extension error", "error", message);
this.events.publish(session.sessionId, { type: "session.error", message });
},
});
}
private bindRuntime(active: ActiveSession<PiSessionRuntime>): void {
active.unsubscribe();
const { session } = active.runtime;