diff --git a/apps/agent-orchestrator/src/engine/temporal-engine.test.ts b/apps/agent-orchestrator/src/engine/temporal-engine.test.ts index 179c188..602b934 100644 --- a/apps/agent-orchestrator/src/engine/temporal-engine.test.ts +++ b/apps/agent-orchestrator/src/engine/temporal-engine.test.ts @@ -149,6 +149,48 @@ describe("TemporalEngine", () => { expect(Object.values(updates[0]!)[0]).toMatchObject({ result: "done" }); }); + // Regression: stream() used to `await this.invoke(input)` before returning + // the iterable, so its own Promise didn't settle until the whole turn had + // already finished. That left server.ts's withHeartbeat wrapper unable to + // do its job -- it can't race a source it doesn't have yet -- so a + // long-running turn's SSE connection had no keep-alive bytes at all and an + // idle-connection timeout upstream would cancel it even though the turn + // kept running server-side. stream() must resolve immediately; only + // iterating it should trigger the poll loop. + it("resolves the streamed iterable before the turn finishes polling", async () => { + let resolvePending!: () => void; + const pending = new Promise((resolve) => { + resolvePending = resolve; + }); + let released = false; + const impl = vi.fn(async (url: string | URL | Request) => { + const href = String(url); + if (href.endsWith("/invoke")) { + return new Response(JSON.stringify({ id: "x", status: "pending" }), { status: 202 }); + } + await pending; // the turn "hangs" here until the test releases it + released = true; + return new Response(JSON.stringify({ id: "x", status: "succeeded", result: "done" }), { status: 200 }); + }) as unknown as typeof fetch; + const engine = new TemporalEngine({ baseUrl: BASE, fetchImpl: impl }); + + const streamPromise = engine.stream(input(), { streamMode: "updates" }); + const source = await Promise.race([ + streamPromise, + new Promise((_, reject) => setTimeout(() => reject(new Error("stream() did not resolve promptly")), 50)), + ]); + expect(released).toBe(false); // the turn is still in-flight + + const updates: Record[] = []; + const drained = (async () => { + for await (const update of source as AsyncIterable>) updates.push(update); + })(); + resolvePending(); + await drained; + expect(updates).toHaveLength(1); + expect(Object.values(updates[0]!)[0]).toMatchObject({ result: "done" }); + }); + // Without this, a streaming chat caller on this engine saw nothing at all // until the whole turn completed -- poll() ran silently, even though the // engine's own gateway already narrates in-flight turns diff --git a/apps/agent-orchestrator/src/engine/temporal-engine.ts b/apps/agent-orchestrator/src/engine/temporal-engine.ts index c7a2211..23ff425 100644 --- a/apps/agent-orchestrator/src/engine/temporal-engine.ts +++ b/apps/agent-orchestrator/src/engine/temporal-engine.ts @@ -131,17 +131,40 @@ export class TemporalEngine implements AgentGraphLike { * streaming protocol for a status line. A streaming client on this engine * therefore sees the reply rather than the running commentary — a real * difference, recorded rather than papered over. + * + * Deliberately NOT `await this.invoke(input)` before returning: an earlier + * version awaited invoke() here, so this method's own returned Promise + * didn't settle until the whole (up to `timeoutMs`, default 30 min) turn + * had already finished. server.ts's streaming handler does + * `const source = await this.graph.stream(...)` before starting its + * `withHeartbeat(source, HEARTBEAT_MS)` loop -- so that wrapper, whose only + * job is emitting an SSE keep-alive comment whenever the source stalls for + * longer than HEARTBEAT_MS, never got a chance to run: it can't race + * against a source it doesn't have yet. A long-running turn's SSE + * connection then had no guaranteed byte cadence at all (progress + * narration is opportunistic, not periodic), so an idle-connection timeout + * upstream of this process (ingress, load balancer, browser) would cancel + * it (RST_STREAM) even though the turn itself kept running and completed + * fine server-side -- the underlying AgentRun Job is unaffected either way + * (docs/adr/0033), only the chat client watching it lost its stream. + * + * An async generator's body does not start running until its first + * `.next()` call (JS semantics), so returning the iterable synchronously + * and moving the `await this.invoke(input)` inside it defers that whole + * wait until `withHeartbeat` actually asks for the next item -- which is + * exactly when its race against HEARTBEAT_MS needs to start. */ - async stream( + stream( input: AgentGraphInput, _options: { streamMode: "updates" }, ): Promise>>> { - const state = await this.invoke(input); - return { + const engine = this; + return Promise.resolve({ async *[Symbol.asyncIterator]() { + const state = await engine.invoke(input); yield { temporalEngine: state }; }, - }; + }); } private async start(input: AgentGraphInput): Promise {