diff --git a/packages/handoff/src/index.ts b/packages/handoff/src/index.ts index 1d0261f873bce9bb028e72f2c3393acf5ccfa2e9..98af756843864e3e9c33b1da45b090a864fce7d9 100644 --- a/packages/handoff/src/index.ts +++ b/packages/handoff/src/index.ts @@ -249,7 +249,7 @@ type HandoffExtraction = { */ async function extractHandoffContext( model: NonNullable, - apiKey: string, + apiKey: string | undefined, headers: Record | undefined, conversationText: string, goal: string, @@ -438,21 +438,19 @@ export default function (pi: ExtensionAPI) { if (pending) clearPending(ctx); }); - pi.on("session_switch", (_event, ctx) => { + pi.on("session_start", (event, ctx) => { if (pending) clearPending(ctx); - // A proper session switch (e.g. /new) fully resets agent state, + // A proper session switch (e.g. /new) or fork fully resets agent state, // so clear the context filter to avoid hiding new messages. - handoffTimestamp = null; + if (event.reason === "new" || event.reason === "resume" || event.reason === "fork") { + handoffTimestamp = null; + } }); pi.on("session_before_fork", (_event, ctx) => { if (pending) clearPending(ctx); }); - pi.on("session_fork", (_event, ctx) => { - if (pending) clearPending(ctx); - }); - pi.on("session_before_tree", (_event, ctx) => { if (pending) clearPending(ctx); }); @@ -476,7 +474,7 @@ export default function (pi: ExtensionAPI) { // send the handoff prompt in the next macrotask // 2. context: filter pre-handoff messages since the low-level session // switch doesn't clear agent.state.messages - // 3. session_switch (above): clears the filter on proper switches + // 3. session_start (above): clears the filter on proper switches pi.on("agent_end", (_event, ctx) => { if (!pendingHandoff) return; @@ -484,8 +482,11 @@ export default function (pi: ExtensionAPI) { const { prompt, parentSession, newModel } = pendingHandoff; pendingHandoff = null; + // Create new session first — this triggers session_start which clears + // any stale handoffTimestamp. Then set the timestamp for the current + // handoff so the context filter is active for the new session's first turn. + (ctx.sessionManager as SessionManager).newSession({ parentSession }); handoffTimestamp = Date.now(); - (ctx.sessionManager as any).newSession({ parentSession }); setTimeout(async () => { if (newModel) { @@ -501,10 +502,11 @@ export default function (pi: ExtensionAPI) { }, 0); }); - pi.on("context", (event) => { + pi.on("context", (event, _ctx) => { if (handoffTimestamp === null) return; - const newMessages = (event as any).messages.filter((m: any) => m.timestamp >= handoffTimestamp); + const cutoff = handoffTimestamp; + const newMessages = event.messages.filter((m) => m.timestamp >= cutoff); if (newMessages.length > 0) { return { messages: newMessages }; } @@ -512,7 +514,7 @@ export default function (pi: ExtensionAPI) { pi.registerTool({ name: "session_query", - label: (params) => `Session Query: ${params.question}`, + label: "Session Query", description: "Query a prior pi session file. Use when handoff prompt references a parent session and you need details.", parameters: Type.Object({ @@ -529,15 +531,15 @@ export default function (pi: ExtensionAPI) { const error = (text: string) => ({ content: [{ type: "text" as const, text }], - details: { error: true }, + details: { error: true } as const, }); const cancelled = () => ({ content: [{ type: "text" as const, text: "Session query cancelled." }], - details: { cancelled: true }, + details: { cancelled: true } as const, }); - if (signal.aborted) { + if (!signal || signal.aborted) { return cancelled(); } @@ -612,7 +614,7 @@ export default function (pi: ExtensionAPI) { const response = await complete( ctx.model, { systemPrompt: QUERY_SYSTEM_PROMPT, messages: [userMessage] }, - { apiKey: auth.apiKey, headers: auth.headers, signal }, + { apiKey: auth.apiKey, headers: auth.headers, signal: signal as AbortSignal }, ); if (response.stopReason === "aborted") { @@ -634,7 +636,7 @@ export default function (pi: ExtensionAPI) { }, }; } catch (err) { - if (signal.aborted) { + if (signal?.aborted) { return cancelled(); } if (err instanceof Error && err.name === "AbortError") { @@ -664,6 +666,7 @@ export default function (pi: ExtensionAPI) { if (!ctx.model) { return { content: [{ type: "text" as const, text: "No model selected." }], + details: undefined, }; } @@ -680,6 +683,7 @@ export default function (pi: ExtensionAPI) { text: "No conversation to hand off.", }, ], + details: undefined, }; } @@ -698,6 +702,7 @@ export default function (pi: ExtensionAPI) { text: `Failed to get API key: ${auth.error}`, }, ], + details: undefined, }; } @@ -714,7 +719,7 @@ export default function (pi: ExtensionAPI) { signal, ); } catch (err) { - if (signal.aborted) { + if (signal?.aborted) { return { content: [ { @@ -722,6 +727,7 @@ export default function (pi: ExtensionAPI) { text: "Handoff cancelled.", }, ], + details: undefined, }; } return { @@ -731,6 +737,7 @@ export default function (pi: ExtensionAPI) { text: `Handoff extraction failed: ${String(err)}`, }, ], + details: undefined, }; } @@ -742,6 +749,7 @@ export default function (pi: ExtensionAPI) { text: "Handoff extraction failed or was cancelled.", }, ], + details: undefined, }; } @@ -761,6 +769,7 @@ export default function (pi: ExtensionAPI) { text: "Handoff prepared. The session will switch after this turn completes.", }, ], + details: undefined, }; }, }); @@ -861,7 +870,7 @@ export default function (pi: ExtensionAPI) { } const next = await ctx.newSession({ - parentSession: currentSessionFile, + parentSession: currentSessionFile ?? undefined, }); if (next.cancelled) { @@ -887,9 +896,7 @@ export default function (pi: ExtensionAPI) { } const newSessionFile = ctx.sessionManager.getSessionFile(); - if (newSessionFile) { - ctx.ui.notify(`Switched to new session: ${newSessionFile}`, "info"); - } + ctx.ui.notify(`Switched to new session: ${newSessionFile ?? "(unnamed)"}`, "info"); ctx.ui.setEditorText(editedPrompt); startCountdown(ctx);