From ea63bb2473d3521fcc3d4acc66f7f75dda20cd8f Mon Sep 17 00:00:00 2001 From: Altay Date: Tue, 3 Mar 2026 10:48:48 +0300 Subject: [PATCH] fix(agents): scope exec wake dispatch to agent sessions --- src/agents/bash-tools.exec-runtime.test.ts | 64 ++++++++++++++++++++++ src/agents/bash-tools.exec-runtime.ts | 13 ++++- src/agents/bash-tools.test.ts | 53 +++++++++++++++++- 3 files changed, 127 insertions(+), 3 deletions(-) create mode 100644 src/agents/bash-tools.exec-runtime.test.ts diff --git a/src/agents/bash-tools.exec-runtime.test.ts b/src/agents/bash-tools.exec-runtime.test.ts new file mode 100644 index 00000000000..35a38b5483d --- /dev/null +++ b/src/agents/bash-tools.exec-runtime.test.ts @@ -0,0 +1,64 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; + +vi.mock("../infra/heartbeat-wake.js", () => ({ + requestHeartbeatNow: vi.fn(), +})); + +vi.mock("../infra/system-events.js", () => ({ + enqueueSystemEvent: vi.fn(), +})); + +import { requestHeartbeatNow } from "../infra/heartbeat-wake.js"; +import { enqueueSystemEvent } from "../infra/system-events.js"; +import { emitExecSystemEvent } from "./bash-tools.exec-runtime.js"; + +const requestHeartbeatNowMock = vi.mocked(requestHeartbeatNow); +const enqueueSystemEventMock = vi.mocked(enqueueSystemEvent); + +describe("emitExecSystemEvent", () => { + beforeEach(() => { + requestHeartbeatNowMock.mockClear(); + enqueueSystemEventMock.mockClear(); + }); + + it("scopes heartbeat wake to the event session key", () => { + emitExecSystemEvent("Exec finished", { + sessionKey: "agent:ops:main", + contextKey: "exec:run-1", + }); + + expect(enqueueSystemEventMock).toHaveBeenCalledWith("Exec finished", { + sessionKey: "agent:ops:main", + contextKey: "exec:run-1", + }); + expect(requestHeartbeatNowMock).toHaveBeenCalledWith({ + reason: "exec-event", + sessionKey: "agent:ops:main", + }); + }); + + it("keeps wake unscoped for non-agent session keys", () => { + emitExecSystemEvent("Exec finished", { + sessionKey: "global", + contextKey: "exec:run-global", + }); + + expect(enqueueSystemEventMock).toHaveBeenCalledWith("Exec finished", { + sessionKey: "global", + contextKey: "exec:run-global", + }); + expect(requestHeartbeatNowMock).toHaveBeenCalledWith({ + reason: "exec-event", + }); + }); + + it("ignores events without a session key", () => { + emitExecSystemEvent("Exec finished", { + sessionKey: " ", + contextKey: "exec:run-2", + }); + + expect(enqueueSystemEventMock).not.toHaveBeenCalled(); + expect(requestHeartbeatNowMock).not.toHaveBeenCalled(); + }); +}); diff --git a/src/agents/bash-tools.exec-runtime.ts b/src/agents/bash-tools.exec-runtime.ts index 22d2f14aa57..7c8e9d911cd 100644 --- a/src/agents/bash-tools.exec-runtime.ts +++ b/src/agents/bash-tools.exec-runtime.ts @@ -6,6 +6,7 @@ import { requestHeartbeatNow } from "../infra/heartbeat-wake.js"; import { isDangerousHostEnvVarName } from "../infra/host-env-security.js"; import { findPathKey, mergePathPrepend } from "../infra/path-prepend.js"; import { enqueueSystemEvent } from "../infra/system-events.js"; +import { parseAgentSessionKey } from "../routing/session-key.js"; import type { ProcessSession } from "./bash-process-registry.js"; import type { ExecToolDetails } from "./bash-tools.exec-types.js"; import type { BashSandboxConfig } from "./bash-tools.shared.js"; @@ -239,7 +240,11 @@ function maybeNotifyOnExit(session: ProcessSession, status: "completed" | "faile ? `Exec ${status} (${session.id.slice(0, 8)}, ${exitLabel}) :: ${output}` : `Exec ${status} (${session.id.slice(0, 8)}, ${exitLabel})`; enqueueSystemEvent(summary, { sessionKey }); - requestHeartbeatNow({ reason: `exec:${session.id}:exit` }); + requestHeartbeatNow( + parseAgentSessionKey(sessionKey) + ? { reason: `exec:${session.id}:exit`, sessionKey } + : { reason: `exec:${session.id}:exit` }, + ); } export function createApprovalSlug(id: string) { @@ -265,7 +270,11 @@ export function emitExecSystemEvent( return; } enqueueSystemEvent(text, { sessionKey, contextKey: opts.contextKey }); - requestHeartbeatNow({ reason: "exec-event" }); + requestHeartbeatNow( + parseAgentSessionKey(sessionKey) + ? { reason: "exec-event", sessionKey } + : { reason: "exec-event" }, + ); } export async function runExecProcess(opts: { diff --git a/src/agents/bash-tools.test.ts b/src/agents/bash-tools.test.ts index 151d705f726..368bddda9c8 100644 --- a/src/agents/bash-tools.test.ts +++ b/src/agents/bash-tools.test.ts @@ -1,5 +1,9 @@ import path from "node:path"; -import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { + resetHeartbeatWakeStateForTests, + setHeartbeatWakeHandler, +} from "../infra/heartbeat-wake.js"; import { applyPathPrepend, findPathKey } from "../infra/path-prepend.js"; import { peekSystemEvents, resetSystemEventsForTest } from "../infra/system-events.js"; import { captureEnv } from "../test-utils/env.js"; @@ -510,6 +514,14 @@ describe("exec exit codes", () => { }); describe("exec notifyOnExit", () => { + beforeEach(() => { + resetHeartbeatWakeStateForTests(); + }); + + afterEach(() => { + resetHeartbeatWakeStateForTests(); + }); + it("enqueues a system event when a backgrounded exec exits", async () => { const tool = createNotifyOnExitExecTool(); @@ -521,6 +533,45 @@ describe("exec notifyOnExit", () => { expect(hasEvent).toBe(true); }); + it("scopes notifyOnExit heartbeat wake to the exec session key", async () => { + const tool = createNotifyOnExitExecTool(); + const wakeHandler = vi.fn().mockResolvedValue({ status: "skipped", reason: "disabled" }); + const dispose = setHeartbeatWakeHandler( + wakeHandler as unknown as Parameters[0], + ); + try { + const sessionId = await startBackgroundCommand(tool, echoAfterDelay("notify")); + + await expect + .poll(() => wakeHandler.mock.calls[0]?.[0], NOTIFY_POLL_OPTIONS) + .toMatchObject({ + reason: `exec:${sessionId}:exit`, + sessionKey: DEFAULT_NOTIFY_SESSION_KEY, + }); + } finally { + dispose(); + } + }); + + it("keeps notifyOnExit heartbeat wake unscoped for non-agent session keys", async () => { + const tool = createNotifyOnExitExecTool({ sessionKey: "global" }); + const wakeHandler = vi.fn().mockResolvedValue({ status: "skipped", reason: "disabled" }); + const dispose = setHeartbeatWakeHandler( + wakeHandler as unknown as Parameters[0], + ); + try { + const sessionId = await startBackgroundCommand(tool, echoAfterDelay("notify")); + + await expect + .poll(() => wakeHandler.mock.calls[0]?.[0], NOTIFY_POLL_OPTIONS) + .toEqual({ + reason: `exec:${sessionId}:exit`, + }); + } finally { + dispose(); + } + }); + it.each(NOOP_NOTIFY_CASES)("$label", runNotifyNoopCase); });