diff --git a/src/gateway/worker-environments/live-events.test.ts b/src/gateway/worker-environments/live-events.test.ts index 96f64cbbe3d5..4fa9ad3a6e82 100644 --- a/src/gateway/worker-environments/live-events.test.ts +++ b/src/gateway/worker-environments/live-events.test.ts @@ -476,6 +476,30 @@ describe("worker live events", () => { expect(deltas()).toEqual(["worker"]); }); + it("adopts a visible dispatch-owned run context so worker live events stay visible", () => { + const lifecycleGeneration = getAgentEventLifecycleGeneration(); + // A worker-routed turn keeps its dispatch-owned Control UI visibility. The + // gateway claims the run context (isControlUiVisible: true for a visible + // turn) before handing the turn to the remote worker; adopting live events + // must inherit that visibility instead of forcing the run hidden. + claimAgentRunContext(RUN, { + ...LOCAL, + isControlUiVisible: true, + lifecycleGeneration, + }); + + ack(msg(1, "worker")); + + expect(getAgentRunContext(RUN)).toMatchObject({ + ...LOCAL, + isControlUiVisible: true, + lifecycleGeneration, + projectSessionActive: true, + }); + expect(deltas()).toEqual(["worker"]); + expect(events[0]?.controlUiVisible).toBe(true); + }); + it("rejects pre-registered gateway run contexts with mismatched identity", () => { const lifecycleGeneration = getAgentEventLifecycleGeneration(); const mismatches: Array<{ @@ -486,7 +510,6 @@ describe("worker live events", () => { { name: "session-key", context: { ...LOCAL, sessionKey: `${KEY}-other` } }, { name: "agent-id", context: { ...LOCAL, agentId: "other" } }, { name: "lifecycle", context: { ...LOCAL, lifecycleGeneration: "other-lifecycle" } }, - { name: "visibility", context: { ...LOCAL, isControlUiVisible: true } }, ]; for (const mismatch of mismatches) { @@ -502,14 +525,10 @@ describe("worker live events", () => { expect(events).toEqual([]); }); - it("keeps run ids exclusive", () => { - const local = "run-local-first"; - claimAgentRunContext(local, LOCAL); - fail(msg(1, "blocked", 0, local), "invalid-event"); - clearAgentRunContext(local); - + it("keeps a claimed run id exclusive against a later untracked local claim", () => { const worker = "run-worker-first"; ack(msg(1, "worker", 0, worker)); + // A same-identity untracked claim cannot hijack a run live events already own. claimAgentRunContext(worker, LOCAL); clearAgentRunContext(worker); emitAgentEvent({ diff --git a/src/gateway/worker-environments/live-events.ts b/src/gateway/worker-environments/live-events.ts index e0e6a1472b6b..be9a632bda5a 100644 --- a/src/gateway/worker-environments/live-events.ts +++ b/src/gateway/worker-environments/live-events.ts @@ -569,16 +569,19 @@ export function createWorkerLiveEventReceiver(options: WorkerLiveEventReceiverOp } const lifecycleGeneration = getAgentEventLifecycleGeneration(); const existingContext = getAgentRunContext(runId); - // Turn placement owns wider visibility; otherwise scope to this session. - const controlUiVisible = false; + // A dispatch-owned turn context (e.g. a worker-routed turn) owns the run's + // Control UI visibility; adopt it so worker live events keep reaching the + // visible clients that started the turn. Identity still has to match, so a + // foreign run is rejected; only the visibility preference is inherited. With + // no pre-existing turn context we scope live events to this session. + const controlUiVisible = existingContext?.isControlUiVisible ?? false; const adoptExistingUnowned = existingContext !== undefined; if ( existingContext && (existingContext.sessionId !== window.sessionId || existingContext.sessionKey !== window.target.sessionKey || existingContext.agentId !== window.target.agentId || - existingContext.lifecycleGeneration !== lifecycleGeneration || - existingContext.isControlUiVisible !== controlUiVisible) + existingContext.lifecycleGeneration !== lifecycleGeneration) ) { return invalidEvent(); } diff --git a/src/worker/worker.fault-injection.test.ts b/src/worker/worker.fault-injection.test.ts index d039ea4729bf..3765521d67af 100644 --- a/src/worker/worker.fault-injection.test.ts +++ b/src/worker/worker.fault-injection.test.ts @@ -41,7 +41,13 @@ import { } from "../gateway/worker-environments/store.js"; import { createWorkerTranscriptCommitStore } from "../gateway/worker-environments/transcript-commit-store.js"; import { createWorkerTranscriptCommitter } from "../gateway/worker-environments/transcript-commit.js"; -import { onAgentRuntimeEvent } from "../infra/agent-events.js"; +import { + claimAgentRunContext, + clearAgentRunContext, + getAgentEventLifecycleGeneration, + getAgentRunContext, + onAgentRuntimeEvent, +} from "../infra/agent-events.js"; import { rawDataToString } from "../infra/ws.js"; import type { WorkerProvider, WorkerSshEndpoint } from "../plugins/types.js"; import { @@ -920,6 +926,43 @@ describe("cloud worker milestone 2 fault injection", () => { expect(SessionManager.open(harness.sessionFile).getEntries()).toHaveLength(1); }); + it("advances a worker live stream whose run context is dispatch-owned and visible", async () => { + const current = harness.createClients(); + clients.push(current); + // A visible turn's run context is claimed by the gateway dispatch before the + // turn hands off to the worker. The worker's first live event must adopt that + // dispatch-owned context (seq advances from 1) and keep the run visible. + const lifecycleGeneration = getAgentEventLifecycleGeneration(); + claimAgentRunContext(RUN_ID, { + agentId: "main", + sessionId: SESSION_ID, + sessionKey: SESSION_KEY, + isControlUiVisible: true, + lifecycleGeneration, + }); + try { + await current.connection.start(); + + await expect( + Promise.all( + ["one", "two"].map((delta) => + current.live.emit(RUN_ID, { kind: "assistant", payload: { text: delta, delta } }), + ), + ), + ).resolves.toHaveLength(2); + + expect(harness.liveDeltas).toEqual(["one", "two"]); + expect( + harness + .requestParams("worker.live-event") + .map((request) => (request as WorkerLiveEventParams).seq), + ).toEqual([1, 2]); + expect(getAgentRunContext(RUN_ID)?.isControlUiVisible).toBe(true); + } finally { + clearAgentRunContext(RUN_ID); + } + }); + it("settles stop during an in-flight commit without retrying or spinning", async () => { const current = harness.createClients(); clients.push(current);