Fix chat SSE timeout on long-running Temporal-engine turns

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<void>((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<string, unknown>[] = [];
+    const drained = (async () => {
+      for await (const update of source as AsyncIterable<Record<string, unknown>>) 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

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<AsyncIterable<Record<string, Partial<AgentState>>>> {
-    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<string> {