fix(chat): retain input sequence for append retries

Author: gtremper

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

@@ -13,6 +13,7 @@ import { getRealtimeStreamInstance } from "~/services/realtime/v1StreamsGlobal.s
 import { stripClientWebhookActionSource } from "~/services/realtime/sanitizeSessionInput.server";
 import {
   claimSessionStreamPart,
+  commitSessionStreamPart,
   drainSessionStreamWaitpoints,
   releaseSessionStreamPart,
 } from "~/services/sessionStreamWaitpointCache.server";
@@ -147,22 +148,28 @@ const { action, loader } = createActionApiRoute(
     const partId = clientPartId ?? nanoid(7);
 
     // Idempotency on client-supplied part ids: atomically claim the id before
-    // appending. A concurrent or retried POST that loses the claim skips the
-    // append (no duplicate record) but still falls through to the drain below,
-    // so a retry whose first attempt died before waking the waitpoint can still
-    // recover it. The claim is released on append failure so a genuine retry
-    // can re-claim and proceed.
-    const wonClaim = clientPartId
+    // appending, then store the resulting sequence with the claim. A retry
+    // returns that original sequence and still falls through to the drain, so
+    // a first attempt that lost its HTTP response cannot lose turn correlation
+    // or strand a waitpoint. The claim is released on append failure.
+    const claim = clientPartId
       ? await claimSessionStreamPart(
           authentication.environment.id,
           addressingKey,
           params.io,
           clientPartId
         )
-      : true;
+      : ({ status: "claimed", claimToken: undefined } as const);
 
-    let appendSeq: number | undefined;
-    if (wonClaim) {
+    if (claim.status === "pending") {
+      return json(
+        { ok: false, error: "This append is still in progress, please retry." },
+        { status: 409, headers: { "Retry-After": "1" } }
+      );
+    }
+
+    let appendSeq = claim.status === "committed" ? claim.seq : undefined;
+    if (claim.status === "claimed") {
       const [appendError, seq] = await tryCatch(
         realtimeStream.appendPartToSessionStream(part, partId, addressingKey, params.io)
       );
@@ -174,7 +181,8 @@ const { action, loader } = createActionApiRoute(
             authentication.environment.id,
             addressingKey,
             params.io,
-            clientPartId
+            clientPartId,
+            claim.claimToken
           );
         }
         if (appendError instanceof ServiceValidationError) {
@@ -193,6 +201,17 @@ const { action, loader } = createActionApiRoute(
           { status: 500 }
         );
       }
+
+      if (clientPartId && appendSeq !== undefined) {
+        await commitSessionStreamPart(
+          authentication.environment.id,
+          addressingKey,
+          params.io,
+          clientPartId,
+          claim.claimToken,
+          appendSeq
+        );
+      }
     }
 
     // Fire any run-scoped waitpoints registered against this channel. Best

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

@@ -1,3 +1,4 @@
+import { randomUUID } from "node:crypto";
 import { Redis } from "ioredis";
 import { defaultReconnectOnError } from "@internal/redis";
 import { env } from "~/env.server";
@@ -146,15 +147,36 @@ export async function drainSessionStreamWaitpoints(
  * where `.wait()` completed the waitpoint from pre-arrived data.
  */
 // "ssa" — session-stream-append. Idempotency claim for the append route:
-// when a caller supplies an `X-Part-Id`, the first request atomically claims
-// the key (SET NX) before appending; a concurrent or retried POST with the
-// same id fails the claim and skips the append, so it never produces a
-// duplicate record (or double-fires the waitpoint drain). The claim is
-// released if the append fails, so a retry of a genuinely failed append
-// still goes through. 5-minute window — covers retry storms, not a
-// permanent idempotency store.
+// when a caller supplies an `X-Part-Id`, the first request atomically stores a
+// tokened pending claim before appending, then replaces it with the committed
+// S2 sequence. A retry can therefore return the original sequence without
+// producing a duplicate record. The token makes commit/release compare-and-
+// set operations safe if a slow owner outlives the claim TTL. 5-minute window
+// — covers retry storms, not a permanent idempotency store.
 const APPEND_DEDUPE_PREFIX = "ssa:";
 const APPEND_DEDUPE_TTL_SECONDS = 5 * 60;
+const APPEND_DEDUPE_PENDING_PREFIX = "pending:";
+const APPEND_DEDUPE_SEQUENCE_PREFIX = "seq:";
+
+const COMMIT_APPEND_DEDUPE_SCRIPT = `
+  if redis.call("GET", KEYS[1]) ~= ARGV[1] then
+    return 0
+  end
+  redis.call("SET", KEYS[1], ARGV[2], "EX", ARGV[3])
+  return 1
+`;
+
+const RELEASE_APPEND_DEDUPE_SCRIPT = `
+  if redis.call("GET", KEYS[1]) ~= ARGV[1] then
+    return 0
+  end
+  return redis.call("DEL", KEYS[1])
+`;
+
+type SessionStreamPartClaim =
+  | { status: "claimed"; claimToken: string | undefined }
+  | { status: "committed"; seq: number }
+  | { status: "pending" };
 
 function buildAppendDedupeKey(
   environmentId: string,
@@ -171,40 +193,84 @@ function buildAppendDedupeKey(
 }
 
 /**
- * Atomically claim a part id before appending. Returns true if this caller
- * won the claim (first to see this id) and should perform the append, false
- * if the id was already claimed (a concurrent or retried POST) and the append
- * should be skipped. Fails open (returns true) when Redis is unavailable —
- * appends degrade to at-least-once, never to dropped.
+ * Atomically claim a part id before appending. Returns the caller's claim
+ * token, the sequence from an already committed append, or pending when the
+ * winning request has not committed yet. Fails open as an untracked claim
+ * when Redis is unavailable — appends degrade to at-least-once, never to
+ * dropped.
  */
 export async function claimSessionStreamPart(
   environmentId: string,
   addressingKey: string,
   io: "out" | "in",
   partId: string
-): Promise<boolean> {
-  if (!redis) return true;
+): Promise<SessionStreamPartClaim> {
+  if (!redis) return { status: "claimed", claimToken: undefined };
 
   try {
+    const claimToken = randomUUID();
+    const pendingValue = `${APPEND_DEDUPE_PENDING_PREFIX}${claimToken}`;
     // SET NX is the atomic claim: "OK" when set (we won), null when the key
     // already exists (someone else owns this id).
-    const result = await redis.set(
+    const key = buildAppendDedupeKey(environmentId, addressingKey, io, partId);
+    const result = await redis.set(key, pendingValue, "EX", APPEND_DEDUPE_TTL_SECONDS, "NX");
+    if (result === "OK") {
+      return { status: "claimed", claimToken };
+    }
+
+    const existing = await redis.get(key);
+    if (existing?.startsWith(APPEND_DEDUPE_SEQUENCE_PREFIX)) {
+      const seq = Number(existing.slice(APPEND_DEDUPE_SEQUENCE_PREFIX.length));
+      if (Number.isSafeInteger(seq) && seq >= 0) {
+        return { status: "committed", seq };
+      }
+    }
+
+    // Includes claims from an older deploy (literal "1"). They suppress a
+    // duplicate append but cannot safely recover its sequence, so callers
+    // must retry instead of returning a successful response without a cursor.
+    return { status: "pending" };
+  } catch (error) {
+    logger.error("Failed to claim session stream append part", {
+      environmentId,
+      addressingKey,
+      io,
+      partId,
+      error,
+    });
+    return { status: "claimed", claimToken: undefined };
+  }
+}
+
+/** Replace this caller's pending claim with the committed S2 sequence. */
+export async function commitSessionStreamPart(
+  environmentId: string,
+  addressingKey: string,
+  io: "out" | "in",
+  partId: string,
+  claimToken: string | undefined,
+  seq: number
+): Promise<void> {
+  if (!redis || !claimToken) return;
+
+  try {
+    await redis.eval(
+      COMMIT_APPEND_DEDUPE_SCRIPT,
+      1,
       buildAppendDedupeKey(environmentId, addressingKey, io, partId),
-      "1",
-      "EX",
-      APPEND_DEDUPE_TTL_SECONDS,
-      "NX"
+      `${APPEND_DEDUPE_PENDING_PREFIX}${claimToken}`,
+      `${APPEND_DEDUPE_SEQUENCE_PREFIX}${seq}`,
+      String(APPEND_DEDUPE_TTL_SECONDS)
     );
-    return result === "OK";
   } catch (error) {
-    logger.error("Failed to claim session stream append part", {
+    logger.error("Failed to commit session stream append part", {
       environmentId,
       addressingKey,
       io,
       partId,
+      seq,
       error,
     });
-    return true;
   }
 }
 
@@ -213,12 +279,18 @@ export async function releaseSessionStreamPart(
   environmentId: string,
   addressingKey: string,
   io: "out" | "in",
-  partId: string
+  partId: string,
+  claimToken: string | undefined
 ): Promise<void> {
-  if (!redis) return;
+  if (!redis || !claimToken) return;
 
   try {
-    await redis.del(buildAppendDedupeKey(environmentId, addressingKey, io, partId));
+    await redis.eval(
+      RELEASE_APPEND_DEDUPE_SCRIPT,
+      1,
+      buildAppendDedupeKey(environmentId, addressingKey, io, partId),
+      `${APPEND_DEDUPE_PENDING_PREFIX}${claimToken}`
+    );
   } catch (error) {
     logger.error("Failed to release session stream append part", {
       environmentId,

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

@@ -15,6 +15,7 @@ import { findOrCreateSession, findSessionByExternalId } from "~/services/realtim
 import { getRealtimeStreamInstance } from "~/services/realtime/v1StreamsGlobal.server";
 import {
   claimSessionStreamPart,
+  commitSessionStreamPart,
   drainSessionStreamWaitpoints,
   releaseSessionStreamPart,
 } from "~/services/sessionStreamWaitpointCache.server";
@@ -221,27 +222,38 @@ function createWebhookEngine() {
               };
         const part = JSON.stringify({ kind: "message", payload });
 
-        // deliveryId as the part id → a deliver-job retry re-claims the same id and skips a duplicate
-        // append. The S2 record is durable, so a run that boots later still reads it.
-        const wonClaim = await claimSessionStreamPart(
-          environment.id,
-          addressingKey,
-          "in",
-          deliveryId
-        );
-        if (wonClaim) {
-          const [appendError] = await tryCatch(
+        // deliveryId as the part id → a deliver-job retry reuses the committed claim and skips a
+        // duplicate append. The S2 record is durable, so a run that boots later still reads it.
+        const claim = await claimSessionStreamPart(environment.id, addressingKey, "in", deliveryId);
+        if (claim.status === "claimed") {
+          const [appendError, appendSeq] = await tryCatch(
             realtimeStream.appendPartToSessionStream(part, deliveryId, addressingKey, "in")
           );
           if (appendError) {
             // Nothing landed — release the claim so a retry re-appends the same id.
-            await releaseSessionStreamPart(environment.id, addressingKey, "in", deliveryId);
+            await releaseSessionStreamPart(
+              environment.id,
+              addressingKey,
+              "in",
+              deliveryId,
+              claim.claimToken
+            );
             // A ServiceValidationError (e.g. record too large) is terminal; anything else is transient.
             if (appendError instanceof ServiceValidationError) {
               return { success: false, error: appendError.message };
             }
             throw appendError;
           }
+          if (typeof appendSeq === "number") {
+            await commitSessionStreamPart(
+              environment.id,
+              addressingKey,
+              "in",
+              deliveryId,
+              claim.claimToken,
+              appendSeq
+            );
+          }
         }
 
         // Wake any `.in` waitpoints the run registered (best-effort; the record is durable in S2).

apps/webapp/test/session-stream.e2e.test.ts

@@ -311,7 +311,38 @@ describe("session stream e2e", () => {
     expect(String(got?.chunk)).toContain("hello from client");
   });
 
-  it("E12 subscribe with an invalid token is rejected", async () => {
+  it("E12 in/append retry returns the original seq without duplicating the record", async () => {
+    const { addressingKey, token, baseUrl } = await setupSession();
+
+    const payload = JSON.stringify({ kind: "message", text: "lost response" });
+    const partId = `retry-${randomBytes(6).toString("hex")}`;
+
+    // The first append commits. Model its response being lost by issuing the
+    // same idempotent request again; keep the first seq only as a test oracle.
+    const first = await appendInput({ baseUrl, addressingKey, token, partId, body: payload });
+    expect(first.status).toBe(200);
+    expect(first.json).toEqual({ ok: true, seq: expect.any(Number) });
+    const firstSeq = (first.json as { seq: number }).seq;
+
+    const retry = await appendInput({ baseUrl, addressingKey, token, partId, body: payload });
+    expect(retry.status).toBe(200);
+    expect(retry.json).toEqual({ ok: true, seq: firstSeq });
+    const retrySeq = (retry.json as { seq: number }).seq;
+
+    const { parts } = await collectSessionOut({
+      baseUrl,
+      addressingKey,
+      token,
+      io: "in",
+      timeoutInSeconds: 1,
+      maxMs: 5_000,
+    });
+    const dataParts = parts.filter((part) => part.chunk != null);
+    expect(dataParts).toHaveLength(1);
+    expect(Number(dataParts[0]!.id)).toBe(retrySeq);
+  });
+
+  it("E13 subscribe with an invalid token is rejected", async () => {
     const { addressingKey, baseUrl } = await setupSession();
 
     const { status } = await openChannelRaw({