From 091350544dbe2e051e63aac7074ecfbe8d113f98 Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Sat, 16 May 2026 10:12:55 +0700 Subject: [PATCH] Fix finish notification steering dispatch --- .../src/server/agent/agent-prompt.test.ts | 39 +++++++++++-------- .../server/src/server/agent/agent-prompt.ts | 26 +++---------- 2 files changed, 28 insertions(+), 37 deletions(-) diff --git a/packages/server/src/server/agent/agent-prompt.test.ts b/packages/server/src/server/agent/agent-prompt.test.ts index 597fce9e11..053569f923 100644 --- a/packages/server/src/server/agent/agent-prompt.test.ts +++ b/packages/server/src/server/agent/agent-prompt.test.ts @@ -10,6 +10,9 @@ import { } from "./agent-prompt.js"; import type { AgentManagerEvent, ManagedAgent } from "./agent-manager.js"; +const CHILD_AGENT_ID = "11111111-1111-4111-8111-111111111111"; +const CALLER_AGENT_ID = "22222222-2222-4222-8222-222222222222"; + test("isSystemInjectedEnvelope matches the envelope formatSystemNotificationPrompt produces", () => { expect(isSystemInjectedEnvelope(formatSystemNotificationPrompt("child finished"))).toBe(true); expect(isSystemInjectedEnvelope("hello world")).toBe(false); @@ -19,12 +22,12 @@ it("does not notify archived callers", async () => { let subscriber: ((event: AgentManagerEvent) => void) | null = null; const childAgent: ManagedAgent = Object.create(null); - Reflect.set(childAgent, "id", "child-agent"); + Reflect.set(childAgent, "id", CHILD_AGENT_ID); Reflect.set(childAgent, "lifecycle", "idle"); Reflect.set(childAgent, "config", { title: "Child Agent" }); const callerAgent: ManagedAgent = Object.create(null); - Reflect.set(callerAgent, "id", "caller-agent"); + Reflect.set(callerAgent, "id", CALLER_AGENT_ID); Reflect.set(callerAgent, "lifecycle", "idle"); Reflect.set(callerAgent, "config", { title: "Caller Agent" }); @@ -33,10 +36,10 @@ it("does not notify archived callers", async () => { agentManager, "getAgent", vi.fn((agentId: string) => { - if (agentId === "child-agent") { + if (agentId === CHILD_AGENT_ID) { return childAgent; } - if (agentId === "caller-agent") { + if (agentId === CALLER_AGENT_ID) { return callerAgent; } return null; @@ -59,7 +62,7 @@ it("does not notify archived callers", async () => { Reflect.set(agentManager, "startAgentRun", startAgentRunSpy); const agentStorageGetSpy = vi.fn(async (agentId: string) => - agentId === "caller-agent" ? { archivedAt: "2024-01-01" } : null, + agentId === CALLER_AGENT_ID ? { archivedAt: "2024-01-01" } : null, ); const agentStorage: AgentStorage = Object.create(AgentStorage.prototype); Reflect.set(agentStorage, "get", agentStorageGetSpy); @@ -67,8 +70,8 @@ it("does not notify archived callers", async () => { setupFinishNotification({ agentManager, agentStorage, - childAgentId: "child-agent", - callerAgentId: "caller-agent", + childAgentId: CHILD_AGENT_ID, + callerAgentId: CALLER_AGENT_ID, logger: createTestLogger(), }); @@ -87,7 +90,7 @@ it("does not notify archived callers", async () => { }); await vi.waitFor(() => { - expect(agentStorageGetSpy).toHaveBeenCalledWith("caller-agent"); + expect(agentStorageGetSpy).toHaveBeenCalledWith(CALLER_AGENT_ID); }); expect(startAgentRunSpy).not.toHaveBeenCalled(); @@ -97,12 +100,12 @@ it("uses AgentManager startAgentRun for finish notifications", async () => { let subscriber: ((event: AgentManagerEvent) => void) | null = null; const childAgent: ManagedAgent = Object.create(null); - Reflect.set(childAgent, "id", "child-agent"); + Reflect.set(childAgent, "id", CHILD_AGENT_ID); Reflect.set(childAgent, "lifecycle", "idle"); Reflect.set(childAgent, "config", { title: "Child Agent" }); const callerAgent: ManagedAgent = Object.create(null); - Reflect.set(callerAgent, "id", "caller-agent"); + Reflect.set(callerAgent, "id", CALLER_AGENT_ID); Reflect.set(callerAgent, "lifecycle", "idle"); Reflect.set(callerAgent, "config", { title: "Caller Agent" }); @@ -116,10 +119,10 @@ it("uses AgentManager startAgentRun for finish notifications", async () => { agentManager, "getAgent", vi.fn((agentId: string) => { - if (agentId === "child-agent") { + if (agentId === CHILD_AGENT_ID) { return childAgent; } - if (agentId === "caller-agent") { + if (agentId === CALLER_AGENT_ID) { return callerAgent; } return null; @@ -141,14 +144,16 @@ it("uses AgentManager startAgentRun for finish notifications", async () => { Reflect.set( agentStorage, "get", - vi.fn(async () => null), + vi.fn(async (agentId: string) => + agentId === CHILD_AGENT_ID ? { title: "Child Agent" } : null, + ), ); setupFinishNotification({ agentManager, agentStorage, - childAgentId: "child-agent", - callerAgentId: "caller-agent", + childAgentId: CHILD_AGENT_ID, + callerAgentId: CALLER_AGENT_ID, logger: createTestLogger(), }); @@ -168,8 +173,8 @@ it("uses AgentManager startAgentRun for finish notifications", async () => { await vi.waitFor(() => { expect(startAgentRunSpy).toHaveBeenCalledWith( - "caller-agent", - "\nAgent child-agent (Child Agent) finished.\n", + CALLER_AGENT_ID, + `\nAgent ${CHILD_AGENT_ID} (Child Agent) finished.\n`, { replaceRunning: true }, ); }); diff --git a/packages/server/src/server/agent/agent-prompt.ts b/packages/server/src/server/agent/agent-prompt.ts index 492e4ffd8e..762f41083d 100644 --- a/packages/server/src/server/agent/agent-prompt.ts +++ b/packages/server/src/server/agent/agent-prompt.ts @@ -1,19 +1,11 @@ import type { Logger } from "pino"; import type { AgentPromptInput, AgentRunOptions } from "./agent-sdk-types.js"; -import type { AgentManager, ManagedAgent } from "./agent-manager.js"; +import type { AgentManager, ManagedAgent, StartAgentRunOptions } from "./agent-manager.js"; import type { AgentStorage } from "./agent-storage.js"; import { ensureAgentLoaded } from "./agent-loading.js"; -export type AgentRunController = Pick< - AgentManager, - "getAgent" | "tryRunOutOfBand" | "hasInFlightRun" | "replaceAgentRun" | "streamAgent" ->; - -export interface StartAgentRunOptions { - replaceRunning?: boolean; - runOptions?: AgentRunOptions; -} +export type AgentRunController = Pick; export function startAgentRun( agentManager: AgentRunController, @@ -35,23 +27,17 @@ export function startAgentRun( }, "agent.session.start_stream.request", ); - // Out-of-band commands (e.g. /goal pause) must run WITHOUT canceling an - // in-flight turn — replaceAgentRun would interrupt the running turn. The - // intercept lives at this layer so it covers every prompt entrypoint. - if (agentManager.tryRunOutOfBand(agentId, prompt)) { + + const result = agentManager.startAgentRun(agentId, prompt, options); + if (result.outOfBand) { return { outOfBand: true }; } - const shouldReplace = Boolean(options?.replaceRunning && agentManager.hasInFlightRun(agentId)); - const runOptions = options?.runOptions; - const iterator = shouldReplace - ? agentManager.replaceAgentRun(agentId, prompt, runOptions) - : agentManager.streamAgent(agentId, prompt, runOptions); + const iterator = result.events; logger.trace( { agentId, provider: snapshot?.provider, providerSessionId: snapshot?.persistence?.sessionId ?? undefined, - shouldReplace, }, "agent.session.start_stream.iterator_returned", );