diff --git a/packages/pi-coding-agent/src/core/agent-session.ts b/packages/pi-coding-agent/src/core/agent-session.ts index 835700e18..990438c7e 100644 --- a/packages/pi-coding-agent/src/core/agent-session.ts +++ b/packages/pi-coding-agent/src/core/agent-session.ts @@ -2102,14 +2102,17 @@ export class AgentSession { runner.bindCore( { - sendMessage: (message, options) => { - this.sendCustomMessage(message, options).catch((err) => { + sendMessage: async (message, options) => { + try { + await this.sendCustomMessage(message, options); + } catch (err) { runner.emitError({ extensionPath: "", event: "send_message", error: getErrorMessage(err), }); - }); + throw err; + } }, sendUserMessage: (content, options) => { this.sendUserMessage(content, options).catch((err) => { diff --git a/packages/pi-coding-agent/src/core/extensions/loader.ts b/packages/pi-coding-agent/src/core/extensions/loader.ts index 6091889b3..19c4e2b2b 100644 --- a/packages/pi-coding-agent/src/core/extensions/loader.ts +++ b/packages/pi-coding-agent/src/core/extensions/loader.ts @@ -530,8 +530,8 @@ function createExtensionAPI( }, // Action methods - delegate to shared runtime - sendMessage(message, options): void { - runtime.sendMessage(message, options); + sendMessage(message, options): Promise { + return runtime.sendMessage(message, options); }, sendUserMessage(content, options): void { diff --git a/packages/pi-coding-agent/src/core/extensions/types.ts b/packages/pi-coding-agent/src/core/extensions/types.ts index e9078d4c9..956b47a94 100644 --- a/packages/pi-coding-agent/src/core/extensions/types.ts +++ b/packages/pi-coding-agent/src/core/extensions/types.ts @@ -1211,7 +1211,7 @@ export interface ExtensionAPI { sendMessage( message: Pick, "customType" | "content" | "display" | "details">, options?: { triggerTurn?: boolean; deliverAs?: "steer" | "followUp" | "nextTurn" }, - ): void; + ): Promise; /** * Send a user message to the agent. Always triggers a turn. @@ -1487,7 +1487,7 @@ export interface ExtensionActions { sendMessage: ( message: Pick, "customType" | "content" | "display" | "details">, options?: { triggerTurn?: boolean; deliverAs?: "steer" | "followUp" | "nextTurn" }, - ) => void; + ) => Promise; sendUserMessage: ( content: string | (TextContent | ImageContent)[], options?: { deliverAs?: "steer" | "followUp" }, diff --git a/src/resources/extensions/sf/auto-direct-dispatch.ts b/src/resources/extensions/sf/auto-direct-dispatch.ts index 11f8de4c0..f050b4341 100644 --- a/src/resources/extensions/sf/auto-direct-dispatch.ts +++ b/src/resources/extensions/sf/auto-direct-dispatch.ts @@ -350,7 +350,7 @@ export async function dispatchDirectPhase( } } try { - pi.sendMessage( + await pi.sendMessage( { customType: "sf-dispatch", content: prompt, display: false }, { triggerTurn: true }, ); diff --git a/src/resources/extensions/sf/auto/run-unit.ts b/src/resources/extensions/sf/auto/run-unit.ts index 95ca2fafe..9e57dbe73 100644 --- a/src/resources/extensions/sf/auto/run-unit.ts +++ b/src/resources/extensions/sf/auto/run-unit.ts @@ -249,7 +249,7 @@ export async function runUnit( } } try { - pi.sendMessage( + await pi.sendMessage( { customType: "sf-auto", content: prompt, display: s.verbose }, { triggerTurn: true }, ); diff --git a/src/resources/extensions/sf/guided-flow.ts b/src/resources/extensions/sf/guided-flow.ts index d0bbf1b60..c55c61ebf 100644 --- a/src/resources/extensions/sf/guided-flow.ts +++ b/src/resources/extensions/sf/guided-flow.ts @@ -548,20 +548,22 @@ async function dispatchWorkflow( join(process.env.HOME ?? "~", ".sf", "agent", "SF-WORKFLOW.md"); const workflow = readFileSync(workflowPath, "utf-8"); - pi.sendMessage( - { - customType, - content: `Read the following SF workflow protocol and execute exactly.\n\n${workflow}\n\n## Your Task\n\n${note}`, - display: false, - }, - { triggerTurn: true }, - ); - - // Restore full tool set after the message is queued. The LLM turn has - // already captured the scoped set — restoring prevents the narrowed - // tools from leaking into subsequent dispatches (#3628). - if (savedTools) { - pi.setActiveTools(savedTools); + try { + await pi.sendMessage( + { + customType, + content: `Read the following SF workflow protocol and execute exactly.\n\n${workflow}\n\n## Your Task\n\n${note}`, + display: false, + }, + { triggerTurn: true }, + ); + } finally { + // Restore full tool set after the scoped turn has been handed to the agent. + // The LLM turn has captured the scoped set by then; restoring prevents the + // narrowed tools from leaking into subsequent dispatches (#3628). + if (savedTools) { + pi.setActiveTools(savedTools); + } } } diff --git a/src/resources/extensions/sf/tests/restore-tools-after-discuss.test.ts b/src/resources/extensions/sf/tests/restore-tools-after-discuss.test.ts index 8c0daaa28..3613c527b 100644 --- a/src/resources/extensions/sf/tests/restore-tools-after-discuss.test.ts +++ b/src/resources/extensions/sf/tests/restore-tools-after-discuss.test.ts @@ -12,7 +12,7 @@ import assert from "node:assert/strict"; import { readFileSync } from "node:fs"; import { resolve } from "node:path"; -import { describe, it } from 'vitest'; +import { describe, it } from "vitest"; const src = readFileSync( resolve( @@ -27,37 +27,45 @@ const src = readFileSync( ); describe("restore tools after discuss flow scoping (#3628)", () => { - it("savedTools is declared before the discuss scoping block", () => { - // savedTools must be declared before the discuss-* check + it("savedTools is declared before shared unit tool scoping", () => { + // savedTools must be declared before per-unit scoping const savedToolsDecl = src.indexOf("let savedTools"); - const discussCheck = src.indexOf('if (unitType?.startsWith("discuss-"))'); + const scopingCall = src.indexOf( + "const scopedTools = scopeActiveToolsForUnitType", + ); assert.ok(savedToolsDecl !== -1, "savedTools variable must be declared"); - assert.ok(discussCheck !== -1, "discuss-* type check must exist"); + assert.ok(scopingCall !== -1, "shared scoping helper must be used"); assert.ok( - savedToolsDecl < discussCheck, - "savedTools must be declared before the discuss scoping block", + savedToolsDecl < scopingCall, + "savedTools must be declared before the scoping block", ); }); - it("savedTools captures current tools inside the discuss block", () => { - const discussCheck = src.indexOf('if (unitType?.startsWith("discuss-"))'); - assert.ok(discussCheck !== -1); + it("savedTools captures current tools inside the shared scoping block", () => { + const scopingCall = src.indexOf( + "const scopedTools = scopeActiveToolsForUnitType", + ); + assert.ok(scopingCall !== -1); - // Look for savedTools assignment within the discuss block - const blockAfter = src.slice(discussCheck, discussCheck + 500); + // Look for savedTools assignment within the shared scoping block + const blockAfter = src.slice(scopingCall, scopingCall + 500); assert.ok( blockAfter.includes("savedTools = currentTools"), - "savedTools must be assigned from currentTools inside the discuss block", + "savedTools must be assigned from currentTools inside the scoping block", ); }); - it("savedTools is restored after sendMessage", () => { + it("savedTools is restored after awaited sendMessage", () => { // Find the sendMessage call - const sendMsg = src.indexOf("triggerTurn: true"); + const sendMsg = src.indexOf("await pi.sendMessage"); assert.ok(sendMsg !== -1, "sendMessage with triggerTurn must exist"); // After sendMessage, savedTools should be restored via setActiveTools - const afterSend = src.slice(sendMsg, sendMsg + 500); + const afterSend = src.slice(sendMsg, sendMsg + 700); + assert.ok( + afterSend.includes("triggerTurn: true"), + "sendMessage must trigger a turn", + ); assert.ok( afterSend.includes("if (savedTools)"), "savedTools restoration guard must exist after sendMessage",