fix(chat): preserve mailbox cursor across waitpoints

Author: gtremper

apps/webapp/app/routes/api.v1.runs.$runFriendlyId.session-streams.wait.ts

@@ -1,6 +1,8 @@
 import { json } from "@remix-run/server-runtime";
 import {
   CreateSessionStreamWaitpointRequestBody,
+  SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE,
+  serializeSessionStreamWaitpointRecord,
   type CreateSessionStreamWaitpointResponseBody,
 } from "@trigger.dev/core/v3";
 import { WaitpointId } from "@trigger.dev/core/v3/isomorphic";
@@ -125,7 +127,8 @@ const { action, loader } = createActionApiRoute(
         addressingKey,
         body.io,
         result.waitpoint.id,
-        ttlMs && ttlMs > 0 ? ttlMs : undefined
+        ttlMs && ttlMs > 0 ? ttlMs : undefined,
+        body.responseFormat
       );
 
       // Race-check. If a record landed on the channel before this
@@ -155,8 +158,14 @@ const { action, loader } = createActionApiRoute(
               await engine.completeWaitpoint({
                 id: result.waitpoint.id,
                 output: {
-                  value: record.data,
-                  type: "application/json",
+                  value:
+                    body.responseFormat === "record-v1"
+                      ? serializeSessionStreamWaitpointRecord(record.data, record.seqNum)
+                      : record.data,
+                  type:
+                    body.responseFormat === "record-v1"
+                      ? SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE
+                      : "application/json",
                   isError: false,
                 },
               });

apps/webapp/app/routes/realtime.v1.sessions.$session.$io.append.ts

@@ -15,6 +15,7 @@ import {
   claimSessionStreamPart,
   drainSessionStreamWaitpoints,
   releaseSessionStreamPart,
+  sessionStreamWaitpointOutput,
 } from "~/services/sessionStreamWaitpointCache.server";
 import { anyResource, createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server";
 import { engine } from "~/v3/runEngine.server";
@@ -201,7 +202,7 @@ const { action, loader } = createActionApiRoute(
     // keyed on the canonical addressing key the agent registered with via
     // `sessions.open(...).in.wait()`, so writers and readers converge
     // regardless of which URL form they used.
-    const [drainError, waitpointIds] = await tryCatch(
+    const [drainError, waitpoints] = await tryCatch(
       drainSessionStreamWaitpoints(authentication.environment.id, addressingKey, params.io)
     );
     if (drainError) {
@@ -210,24 +211,20 @@ const { action, loader } = createActionApiRoute(
         io: params.io,
         error: drainError,
       });
-    } else if (waitpointIds && waitpointIds.length > 0) {
+    } else if (waitpoints && waitpoints.length > 0) {
       await Promise.all(
-        waitpointIds.map(async (waitpointId) => {
+        waitpoints.map(async (waitpoint) => {
           const [completeError] = await tryCatch(
             engine.completeWaitpoint({
-              id: waitpointId,
-              output: {
-                value: part,
-                type: "application/json",
-                isError: false,
-              },
+              id: waitpoint.id,
+              output: sessionStreamWaitpointOutput(waitpoint, part, appendSeq),
             })
           );
           if (completeError) {
             logger.error("Failed to complete session stream waitpoint", {
               addressingKey,
               io: params.io,
-              waitpointId,
+              waitpointId: waitpoint.id,
               error: completeError,
             });
           }

apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.playground.realtime.v1.sessions.$session.$io.append.ts

@@ -13,7 +13,10 @@ import {
   resolveSessionByIdOrExternalId,
 } from "~/services/realtime/sessions.server";
 import { getRealtimeStreamInstance } from "~/services/realtime/v1StreamsGlobal.server";
-import { drainSessionStreamWaitpoints } from "~/services/sessionStreamWaitpointCache.server";
+import {
+  drainSessionStreamWaitpoints,
+  sessionStreamWaitpointOutput,
+} from "~/services/sessionStreamWaitpointCache.server";
 import { requireUserId } from "~/services/session.server";
 import { EnvironmentParamSchema } from "~/utils/pathBuilder";
 import { engine } from "~/v3/runEngine.server";
@@ -114,7 +117,7 @@ export async function action({ request, params }: ActionFunctionArgs) {
 
   // Drain any waitpoints registered for this channel — same as the
   // public append. Best-effort; failure doesn't fail the append.
-  const [drainError, waitpointIds] = await tryCatch(
+  const [drainError, waitpoints] = await tryCatch(
     drainSessionStreamWaitpoints(environment.id, addressingKey, io)
   );
   if (drainError) {
@@ -123,24 +126,20 @@ export async function action({ request, params }: ActionFunctionArgs) {
       io,
       error: drainError,
     });
-  } else if (waitpointIds && waitpointIds.length > 0) {
+  } else if (waitpoints && waitpoints.length > 0) {
     await Promise.all(
-      waitpointIds.map(async (waitpointId) => {
+      waitpoints.map(async (waitpoint) => {
         const [completeError] = await tryCatch(
           engine.completeWaitpoint({
-            id: waitpointId,
-            output: {
-              value: part,
-              type: "application/json",
-              isError: false,
-            },
+            id: waitpoint.id,
+            output: sessionStreamWaitpointOutput(waitpoint, part, appendSeq ?? undefined),
           })
         );
         if (completeError) {
           logger.error("Failed to complete session stream waitpoint (playground)", {
             addressingKey,
             io,
-            waitpointId,
+            waitpointId: waitpoint.id,
             error: completeError,
           });
         }

apps/webapp/app/services/sessionStreamWaitpointCache.server.ts

@@ -1,5 +1,9 @@
 import { Redis } from "ioredis";
 import { defaultReconnectOnError } from "@internal/redis";
+import {
+  SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE,
+  serializeSessionStreamWaitpointRecord,
+} from "@trigger.dev/core/v3";
 import { env } from "~/env.server";
 import { singleton } from "~/utils/singleton";
 import { logger } from "./logger.server";
@@ -13,12 +17,35 @@ import { logger } from "./logger.server";
 // is shared — without it, two environments using the same externalId
 // would drain each other's waitpoints.
 const KEY_PREFIX = "ssw:";
+const FORMAT_KEY_PREFIX = "sswf:";
 const DEFAULT_TTL_MS = 7 * 24 * 60 * 60 * 1000; // 7 days
 
+export type SessionStreamWaitpoint = {
+  id: string;
+  responseFormat?: "record-v1";
+};
+
+export function sessionStreamWaitpointOutput(
+  waitpoint: SessionStreamWaitpoint,
+  data: string,
+  seqNum: number | undefined
+): { value: string; type: string; isError: false } {
+  const hasRecordEnvelope = waitpoint.responseFormat === "record-v1" && seqNum !== undefined;
+  return {
+    value: hasRecordEnvelope ? serializeSessionStreamWaitpointRecord(data, seqNum) : data,
+    type: hasRecordEnvelope ? SESSION_STREAM_WAITPOINT_RECORD_CONTENT_TYPE : "application/json",
+    isError: false,
+  };
+}
+
 function buildKey(environmentId: string, addressingKey: string, io: "out" | "in"): string {
   return `${KEY_PREFIX}${environmentId}:${addressingKey}:${io}`;
 }
 
+function buildFormatKey(waitpointId: string): string {
+  return `${FORMAT_KEY_PREFIX}${waitpointId}`;
+}
+
 // Pre-env-scoping key format, drained for one release so waitpoints from the
 // previous deploy still wake. Removable once this has been live > turn timeout.
 function buildLegacyKey(addressingKey: string, io: "out" | "in"): string {
@@ -81,13 +108,25 @@ export async function addSessionStreamWaitpoint(
   addressingKey: string,
   io: "out" | "in",
   waitpointId: string,
-  ttlMs?: number
+  ttlMs?: number,
+  responseFormat?: "record-v1"
 ): Promise<void> {
   if (!redis) return;
 
   try {
     const key = buildKey(environmentId, addressingKey, io);
-    await redis.eval(ADD_WAITPOINT_SCRIPT, 1, key, waitpointId, String(ttlMs ?? DEFAULT_TTL_MS));
+    const effectiveTtlMs = ttlMs ?? DEFAULT_TTL_MS;
+
+    // Keep the set member as the plain waitpoint id so an older append
+    // instance can still drain it during a rolling deploy. New instances read
+    // the optional response format from this separate, TTL-bound key.
+    if (responseFormat) {
+      await redis.set(buildFormatKey(waitpointId), responseFormat, "PX", effectiveTtlMs);
+    } else {
+      await redis.del(buildFormatKey(waitpointId));
+    }
+
+    await redis.eval(ADD_WAITPOINT_SCRIPT, 1, key, waitpointId, String(effectiveTtlMs));
   } catch (error) {
     logger.error("Failed to set session stream waitpoint cache", {
       environmentId,
@@ -107,7 +146,7 @@ export async function drainSessionStreamWaitpoints(
   environmentId: string,
   addressingKey: string,
   io: "out" | "in"
-): Promise<string[]> {
+): Promise<SessionStreamWaitpoint[]> {
   if (!redis) return [];
 
   try {
@@ -129,7 +168,34 @@ export async function drainSessionStreamWaitpoints(
       if (err || !Array.isArray(members)) continue;
       for (const m of members as string[]) ids.add(m);
     }
-    return [...ids];
+    const waitpointIds = [...ids];
+    if (waitpointIds.length === 0) return [];
+
+    let formatResults: Awaited<ReturnType<typeof pipeline.exec>> | null = null;
+    try {
+      const formatPipeline = redis.multi();
+      for (const waitpointId of waitpointIds) {
+        formatPipeline.get(buildFormatKey(waitpointId));
+        formatPipeline.del(buildFormatKey(waitpointId));
+      }
+      formatResults = await formatPipeline.exec();
+    } catch (error) {
+      // The waitpoint ids were already drained. Complete them with raw data
+      // rather than losing the wake-up because optional metadata was unavailable.
+      logger.error("Failed to read session stream waitpoint response formats", {
+        environmentId,
+        addressingKey,
+        io,
+        error,
+      });
+    }
+
+    return waitpointIds.map((id, index) => {
+      const formatEntry = formatResults?.[index * 2];
+      const responseFormat =
+        formatEntry && !formatEntry[0] && formatEntry[1] === "record-v1" ? "record-v1" : undefined;
+      return { id, responseFormat };
+    });
   } catch (error) {
     logger.error("Failed to drain session stream waitpoint cache", {
       environmentId,
@@ -240,7 +306,10 @@ export async function removeSessionStreamWaitpoint(
 
   try {
     const key = buildKey(environmentId, addressingKey, io);
-    await redis.srem(key, waitpointId);
+    const pipeline = redis.multi();
+    pipeline.srem(key, waitpointId);
+    pipeline.del(buildFormatKey(waitpointId));
+    await pipeline.exec();
   } catch (error) {
     logger.error("Failed to remove session stream waitpoint cache entry", {
       environmentId,

apps/webapp/app/v3/webhookEngine.server.ts

@@ -17,6 +17,7 @@ import {
   claimSessionStreamPart,
   drainSessionStreamWaitpoints,
   releaseSessionStreamPart,
+  sessionStreamWaitpointOutput,
 } from "~/services/sessionStreamWaitpointCache.server";
 import { getSecretStore } from "~/services/secrets/secretStore.server";
 import { singleton } from "~/utils/singleton";
@@ -229,10 +230,12 @@ function createWebhookEngine() {
           "in",
           deliveryId
         );
+        let appendSeq: number | undefined;
         if (wonClaim) {
-          const [appendError] = await tryCatch(
+          const [appendError, seqNum] = await tryCatch(
             realtimeStream.appendPartToSessionStream(part, deliveryId, addressingKey, "in")
           );
+          appendSeq = seqNum ?? undefined;
           if (appendError) {
             // Nothing landed — release the claim so a retry re-appends the same id.
             await releaseSessionStreamPart(environment.id, addressingKey, "in", deliveryId);
@@ -245,21 +248,21 @@ function createWebhookEngine() {
         }
 
         // Wake any `.in` waitpoints the run registered (best-effort; the record is durable in S2).
-        const [drainError, waitpointIds] = await tryCatch(
+        const [drainError, waitpoints] = await tryCatch(
           drainSessionStreamWaitpoints(environment.id, addressingKey, "in")
         );
         if (drainError) {
           logger.error("deliverToSession: failed to drain session waitpoints", {
             externalId,
             error: drainError,
           });
-        } else if (waitpointIds && waitpointIds.length > 0) {
+        } else if (waitpoints && waitpoints.length > 0) {
           await Promise.all(
-            waitpointIds.map((waitpointId) =>
+            waitpoints.map((waitpoint) =>
               tryCatch(
                 runEngine.completeWaitpoint({
-                  id: waitpointId,
-                  output: { value: part, type: "application/json", isError: false },
+                  id: waitpoint.id,
+                  output: sessionStreamWaitpointOutput(waitpoint, part, appendSeq),
                 })
               )
             )

packages/core/src/v3/schemas/api.ts

@@ -1669,6 +1669,8 @@ export const CreateSessionStreamWaitpointRequestBody = z.object({
    * Used to catch data that arrived before `.wait()` was called.
    */
   lastSeqNum: z.number().optional(),
+  /** Internal capability flag: return the exact record sequence on resume. */
+  responseFormat: z.literal("record-v1").optional(),
 });
 export type CreateSessionStreamWaitpointRequestBody = z.infer<
   typeof CreateSessionStreamWaitpointRequestBody

packages/core/src/v3/sessionStreams/index.ts

@@ -94,6 +94,14 @@ export class SessionStreamsAPI implements SessionStreamManager {
     this.#getManager().setLastSeqNum(sessionId, io, seqNum);
   }
 
+  public consumeRecord(sessionId: string, io: SessionChannelIO, seqNum: number): void {
+    const manager = this.#getManager();
+    if (!manager.consumeRecord) {
+      throw new Error("The configured Session stream manager does not support exact consumption");
+    }
+    manager.consumeRecord(sessionId, io, seqNum);
+  }
+
   public lastDispatchedSeqNum(sessionId: string, io: SessionChannelIO): number | undefined {
     return this.#getManager().lastDispatchedSeqNum(sessionId, io);
   }

packages/core/src/v3/sessionStreams/manager.test.ts

@@ -386,7 +386,7 @@ describe("StandardSessionStreamManager — record metadata", () => {
     manager.disconnect();
   });
 
-  it("retains cursor barriers when disconnect clears the buffer", async () => {
+  it("preserves buffered records across disconnect and consumes only the exact sequence", async () => {
     const manager = new StandardSessionStreamManager(
       singleShotApiClient([
         {
@@ -401,6 +401,12 @@ describe("StandardSessionStreamManager — record metadata", () => {
           chunk: { kind: "stop" },
           timestamp: 2000,
         },
+        {
+          id: "52",
+          recordId: "message-2",
+          chunk: { kind: "message", payload: { id: "u2" } },
+          timestamp: 3000,
+        },
       ]),
       "http://localhost"
     );
@@ -418,11 +424,16 @@ describe("StandardSessionStreamManager — record metadata", () => {
 
     expect(manager.lastDispatchedSeqNum(sessionId, io)).toBe(49);
     manager.disconnectStream(sessionId, io);
-    expect(manager.peekRecord(sessionId, io)).toBeUndefined();
+    expect(manager.peekRecord(sessionId, io)?.seqNum).toBe(50);
     expect(manager.lastDispatchedSeqNum(sessionId, io)).toBe(49);
 
-    manager.setLastDispatchedSeqNum(sessionId, io, 51);
-    expect(manager.lastDispatchedSeqNum(sessionId, io)).toBe(49);
+    manager.consumeRecord(sessionId, io, 50);
+    expect(manager.peekRecord(sessionId, io)?.seqNum).toBe(52);
+    expect(manager.lastDispatchedSeqNum(sessionId, io)).toBe(51);
+
+    manager.consumeRecord(sessionId, io, 52);
+    expect(manager.peekRecord(sessionId, io)).toBeUndefined();
+    expect(manager.lastDispatchedSeqNum(sessionId, io)).toBe(52);
 
     manager.reset();
     expect(manager.lastDispatchedSeqNum(sessionId, io)).toBeUndefined();