feat(sdk,core): close resumed chat streams promptly when caught up

.changeset/chat-agent-preload-oom-recovery.md

@@ -0,0 +1,5 @@
+---
+"@trigger.dev/sdk": patch
+---
+
+Fix a preloaded `chat.agent` run dropping an in-flight message when it retries after an out-of-memory error. The message being processed when the run hit the OOM is now recovered and re-run on the retry, instead of being skipped while the run waited for a new message.

.changeset/chat-session-caught-up-resume.md

@@ -0,0 +1,6 @@
+---
+"@trigger.dev/core": patch
+"@trigger.dev/sdk": patch
+---
+
+Chat sessions now close a resumed stream as soon as it has caught up to the latest output, instead of holding the connection open for the full long-poll window. Reloading or reconnecting to an idle chat settles faster.

.github/workflows/e2e-webapp.yml

@@ -80,6 +80,8 @@ jobs:
           docker pull postgres:14
           docker pull redis:7.2
           docker pull testcontainers/ryuk:0.11.0
+          docker pull ghcr.io/s2-streamstore/s2:0.40.0@sha256:b26249e2ede0949755f5af8028185dc2bcfc3aa2db21eb9610543d144eb6ee9d
+          docker pull minio/minio:latest
           echo "Image pre-pull complete"
 
       - name: 📥 Download deps
@@ -91,6 +93,9 @@ jobs:
       - name: 🏗️ Build Webapp
         run: pnpm run build --filter webapp
 
+      - name: 🎭 Install Playwright Chromium
+        run: cd apps/webapp && pnpm exec playwright install chromium
+
       - name: 🧪 Run Webapp E2E Tests
         run: cd apps/webapp && pnpm exec vitest run --config vitest.e2e.config.ts --reporter=default
         env:

apps/webapp/app/services/realtime/s2realtimeStreams.server.ts

@@ -107,8 +107,8 @@ export class S2RealtimeStreams implements StreamResponder, StreamIngestor {
 
   constructor(opts: S2RealtimeStreamsOptions) {
     this.basin = opts.basin;
-    this.baseUrl = opts.endpoint ?? `https://${this.basin}.b.aws.s2.dev/v1`;
-    this.accountUrl = opts.endpoint ?? `https://aws.s2.dev/v1`;
+    this.baseUrl = opts.endpoint ?? `https://${this.basin}.b.s2.dev/v1`;
+    this.accountUrl = opts.endpoint ?? `https://a.s2.dev/v1`;
     this.endpoint = opts.endpoint;
     this.token = opts.accessToken;
     this.streamPrefix = opts.streamPrefix ?? "";

apps/webapp/app/services/realtime/streamBasinProvisioner.server.ts

@@ -185,7 +185,7 @@ type CreateBasinOptions = {
 };
 
 async function s2CreateBasin(name: string, opts: CreateBasinOptions): Promise<void> {
-  const url = `https://aws.s2.dev/v1/basins`;
+  const url = `https://a.s2.dev/v1/basins`;
   const body = {
     basin: name,
     config: {
@@ -222,7 +222,7 @@ type ReconfigureBasinOptions = {
 };
 
 async function s2ReconfigureBasin(name: string, opts: ReconfigureBasinOptions): Promise<void> {
-  const url = `https://aws.s2.dev/v1/basins/${encodeURIComponent(name)}`;
+  const url = `https://a.s2.dev/v1/basins/${encodeURIComponent(name)}`;
   const body = {
     default_stream_config: {
       retention_policy: { age: parseDuration(opts.retentionPolicy) },

apps/webapp/package.json

@@ -103,7 +103,7 @@
     "@remix-run/react": "2.17.5",
     "@remix-run/router": "^1.23.3",
     "@remix-run/server-runtime": "2.17.5",
-    "@s2-dev/streamstore": "^0.22.10",
+    "@s2-dev/streamstore": "^0.25.0",
     "@sentry/remix": "9.46.0",
     "@slack/web-api": "7.16.0",
     "@socket.io/redis-adapter": "^8.3.0",
@@ -219,6 +219,7 @@
     "@internal/clickhouse": "workspace:*",
     "@internal/replication": "workspace:*",
     "@internal/testcontainers": "workspace:*",
+    "@playwright/test": "^1.36.2",
     "@remix-run/dev": "2.17.5",
     "@remix-run/testing": "^2.17.5",
     "@sentry/cli": "2.50.2",

apps/webapp/test/helpers/agentHarness.ts

@@ -0,0 +1,195 @@
+import { apiClientManager, resourceCatalog } from "@trigger.dev/core/v3";
+import type { LocalsKey } from "@trigger.dev/core/v3";
+import type { LanguageModel } from "ai";
+import {
+  installSessionWaitpointBackend,
+  runInMockTaskContext,
+  StandardSessionStreamManager,
+} from "@trigger.dev/core/v3/test";
+
+export type RunRealChatAgentOptions = {
+  agentId: string;
+  baseUrl: string;
+  addressingKey: string;
+  /**
+   * The environment secret key. The agent writes `.out` and reads `.in` as the
+   * backend (PRIVATE auth) — the `.out` channel rejects client session tokens.
+   */
+  secretKey: string;
+  model: LanguageModel;
+  modelLocal: LocalsKey<LanguageModel>;
+  runId?: string;
+  /**
+   * Boot as a continuation of a previous run for the same session. Gates the
+   * snapshot + `.out`/`.in` replay boot path, so the agent restores prior
+   * history instead of treating the chat as brand new.
+   */
+  continuation?: boolean;
+  previousRunId?: string;
+  /**
+   * Idle window (seconds) before the turn loop falls through from the SSE
+   * once() to the suspending `session.in.wait()`. Set this low to force the
+   * suspend/resume path in a test.
+   */
+  idleTimeoutInSeconds?: number;
+  /**
+   * `ctx.attempt.number`. A value greater than 1 makes the boot treat the run
+   * as a retry (`couldHavePriorState`), restoring from the snapshot + `.in`
+   * replay. Used to model an OOM retry re-dispatch.
+   */
+  attemptNumber?: number;
+};
+
+export type RunningAgent = {
+  done: Promise<void>;
+  close: () => Promise<void>;
+};
+
+/**
+ * Run the real `chat.agent` turn loop in-process, wired to a running webapp:
+ * `apiClientManager` + a real `StandardSessionStreamManager` point the agent's
+ * `.in`/`.out` at the webapp's Session streams (real S2 + SSE), the model is
+ * injected via locals (so it survives without serialization), and turns are
+ * driven by appending to `.in` over HTTP. Callers keep each message inside the
+ * idle window and `close()` promptly so the run-engine suspend path is never
+ * reached.
+ */
+export function runRealChatAgent(opts: RunRealChatAgentOptions): RunningAgent {
+  apiClientManager.setGlobalAPIClientConfiguration({
+    baseURL: opts.baseUrl,
+    accessToken: opts.secretKey,
+  });
+  const apiClient = apiClientManager.clientOrThrow();
+  const manager = new StandardSessionStreamManager(apiClient, opts.baseUrl);
+  const { runtimeManager, restore } = installSessionWaitpointBackend(apiClient);
+
+  const taskEntry = resourceCatalog.getTask(opts.agentId);
+  if (!taskEntry) {
+    restore();
+    throw new Error(`runRealChatAgent: agent "${opts.agentId}" is not registered`);
+  }
+  const runFn = taskEntry.fns.run as (
+    payload: unknown,
+    params: { ctx: unknown; signal: AbortSignal }
+  ) => Promise<unknown>;
+
+  const runSignal = new AbortController();
+  const runId = opts.runId ?? `run_${opts.addressingKey}`;
+
+  const idle =
+    opts.idleTimeoutInSeconds !== undefined
+      ? { idleTimeoutInSeconds: opts.idleTimeoutInSeconds }
+      : {};
+
+  const done = (
+    runInMockTaskContext(
+      async (drivers) => {
+        drivers.locals.set(opts.modelLocal, opts.model);
+        const payload = opts.continuation
+          ? {
+              chatId: opts.addressingKey,
+              continuation: true,
+              metadata: {},
+              ...idle,
+              ...(opts.previousRunId ? { previousRunId: opts.previousRunId } : {}),
+            }
+          : { chatId: opts.addressingKey, trigger: "preload", metadata: {}, ...idle };
+        await runFn(payload, { ctx: drivers.ctx, signal: runSignal.signal });
+      },
+      {
+        ctx: {
+          run: { id: runId },
+          ...(opts.attemptNumber !== undefined ? { attempt: { number: opts.attemptNumber } } : {}),
+        },
+        sessionStreamManager: manager,
+        runtimeManager,
+      }
+    ) as Promise<void>
+  ).finally(restore);
+
+  return {
+    done,
+    close: async () => {
+      try {
+        await fetch(
+          `${opts.baseUrl}/realtime/v1/sessions/${encodeURIComponent(opts.addressingKey)}/in/append`,
+          {
+            method: "POST",
+            headers: {
+              Authorization: `Bearer ${opts.secretKey}`,
+              "Content-Type": "application/json",
+              "X-Part-Id": "close",
+            },
+            body: JSON.stringify({
+              kind: "message",
+              payload: { chatId: opts.addressingKey, trigger: "close" },
+            }),
+          }
+        );
+      } catch {}
+      runSignal.abort();
+      await Promise.race([
+        done.catch(() => {}),
+        new Promise((resolve) => setTimeout(resolve, 10_000)),
+      ]);
+    },
+  };
+}
+
+export type ChatAgentSessionOptions = Omit<
+  RunRealChatAgentOptions,
+  "runId" | "continuation" | "previousRunId"
+>;
+
+export type ChatAgentSession = {
+  /** How many runs the session has spawned so far (1 fresh + N continuations). */
+  runCount: () => number;
+  /** Close the currently-active run and stop spawning continuations. */
+  close: () => Promise<void>;
+};
+
+/**
+ * A session-scoped orchestrator that stands in for the run-engine's run
+ * lifecycle: it starts a run, and whenever that run exits on its own
+ * (`chat.endRun()` / `chat.requestUpgrade()`), spawns the next run as a
+ * continuation (new run id, `continuation: true`, `previousRunId` threaded)
+ * for the same session. That mirrors the server triggering a fresh run on the
+ * next append after the previous run went terminal, and lets each continuation
+ * restore prior history from the persisted snapshot. Runs never overlap: the
+ * next spawn is chained on the previous run's `done` (after its manager
+ * teardown), so the process-global managers are never installed twice at once.
+ */
+export function runChatAgentSession(opts: ChatAgentSessionOptions): ChatAgentSession {
+  let closed = false;
+  let index = 0;
+  let current: RunningAgent | undefined;
+  let previousRunId: string | undefined;
+
+  const spawn = () => {
+    index += 1;
+    const runId = `run_${opts.addressingKey}_${index}`;
+    current = runRealChatAgent({
+      ...opts,
+      runId,
+      continuation: index > 1,
+      previousRunId,
+    });
+    previousRunId = runId;
+    const settle = () => {
+      if (!closed) {
+        spawn();
+      }
+    };
+    current.done.then(settle, settle);
+  };
+
+  spawn();
+
+  return {
+    runCount: () => index,
+    close: async () => {
+      closed = true;
+      await current?.close();
+    },
+  };
+}