fix(chat): harden append sequence recovery

Author: gtremper

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

@@ -100,10 +100,12 @@ type S2ReadRecord = {
   timestamp: number;
   headers?: Array<[string, string]>;
 };
+type S2TailResponse = {
+  tail: { seq_num: number; timestamp: number };
+};
 
 const PART_RECOVERY_CLOCK_SKEW_MS = 10_000;
 const PART_RECOVERY_PAGE_SIZE = 1_000;
-const PART_RECOVERY_MAX_PAGES = 32;
 
 export class S2RealtimeStreams implements StreamResponder, StreamIngestor {
   private readonly basin: string;
@@ -276,8 +278,8 @@ export class S2RealtimeStreams implements StreamResponder, StreamIngestor {
 
   /**
    * Find an append whose Redis claim was left pending after S2 accepted it.
-   * Reads forward from the claim timestamp and stops at the first response's
-   * tail snapshot, so an active stream cannot make recovery chase new data.
+   * Snapshots the stream tail, then reads forward from the claim timestamp.
+   * The fixed tail keeps an active stream from extending the recovery search.
    */
   async findSessionStreamPartSequence(
     friendlyId: string,
@@ -286,10 +288,33 @@ export class S2RealtimeStreams implements StreamResponder, StreamIngestor {
     claimedAt?: number
   ): Promise<number | undefined> {
     const s2Stream = this.toSessionStreamName(friendlyId, io);
+    const tailResponse = await fetch(
+      `${this.baseUrl}/streams/${encodeURIComponent(s2Stream)}/records/tail`,
+      {
+        method: "GET",
+        headers: {
+          Authorization: `Bearer ${this.token}`,
+          Accept: "application/json",
+          "S2-Basin": this.basin,
+        },
+      }
+    );
+
+    if (!tailResponse.ok) {
+      if (tailResponse.status === 404) return undefined;
+      const text = await tailResponse.text().catch(() => "");
+      throw new Error(
+        `S2 findSessionStreamPartSequence tail check failed: ${tailResponse.status} ${tailResponse.statusText} ${text}`
+      );
+    }
+
+    const { tail } = (await tailResponse.json()) as S2TailResponse;
+    const snapshotTail = tail.seq_num;
+    if (snapshotTail === 0) return undefined;
+
     let nextSeq: number | undefined;
-    let snapshotTail: number | undefined;
 
-    for (let page = 0; page < PART_RECOVERY_MAX_PAGES; page++) {
+    while (true) {
       const qs = new URLSearchParams();
       if (nextSeq !== undefined) {
         qs.set("seq_num", String(nextSeq));
@@ -326,13 +351,11 @@ export class S2RealtimeStreams implements StreamResponder, StreamIngestor {
 
       const json = (await res.json()) as {
         records?: S2ReadRecord[];
-        tail?: { seq_num: number };
       };
       const records = json.records ?? [];
-      snapshotTail ??= json.tail?.seq_num;
 
       for (const record of records) {
-        if (snapshotTail !== undefined && record.seq_num >= snapshotTail) break;
+        if (record.seq_num >= snapshotTail) break;
         if (record.headers?.[0]?.[0] === "") continue;
         try {
           const envelope = JSON.parse(record.body) as { id?: unknown };
@@ -345,12 +368,8 @@ export class S2RealtimeStreams implements StreamResponder, StreamIngestor {
       const lastRecord = records.at(-1);
       if (!lastRecord) return undefined;
       nextSeq = lastRecord.seq_num + 1;
-      if (snapshotTail !== undefined && nextSeq >= snapshotTail) return undefined;
+      if (nextSeq >= snapshotTail) return undefined;
     }
-
-    // An incomplete search must never be treated as proof that the append is
-    // absent; the caller will keep the claim pending and retry recovery.
-    throw new Error(`S2 part recovery exceeded ${PART_RECOVERY_MAX_PAGES} pages`);
   }
 
   async #readRecordsByName(s2Stream: string, afterSeqNum?: number): Promise<StreamRecord[]> {

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

@@ -212,36 +212,50 @@ export async function claimSessionStreamPart(
   if (!redis) return { status: "claimed", claimValue: undefined };
 
   try {
-    const claimedAt = Date.now();
-    const pendingValue = `${APPEND_DEDUPE_PENDING_PREFIX}${claimedAt}:${randomUUID()}`;
-    // SET NX is the atomic claim: "OK" when set (we won), null when the key
-    // already exists (someone else owns this id).
     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", claimValue: pendingValue };
-    }
+    for (let attempt = 0; attempt < 3; attempt++) {
+      const claimedAt = Date.now();
+      const pendingValue = `${APPEND_DEDUPE_PENDING_PREFIX}${claimedAt}:${randomUUID()}`;
+      // 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(key, pendingValue, "EX", APPEND_DEDUPE_TTL_SECONDS, "NX");
+      if (result === "OK") {
+        return { status: "claimed", claimValue: pendingValue };
+      }
 
-    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 };
+      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 };
+        }
       }
-    }
 
-    if (existing) {
-      const pendingClaimedAt = parsePendingClaimedAt(existing);
-      return {
-        status: "pending",
-        claimValue: existing,
-        claimedAt: pendingClaimedAt,
-      };
+      if (existing) {
+        let pendingClaimedAt = parsePendingClaimedAt(existing);
+        if (pendingClaimedAt === undefined) {
+          // Claims from the previous release stored only "1". Estimate their
+          // creation time from the remaining TTL so recovery starts near the
+          // original append during a rolling deploy.
+          const ttlMs = await redis.pttl(key);
+          if (ttlMs >= 0) {
+            pendingClaimedAt = Date.now() - Math.max(0, APPEND_DEDUPE_TTL_SECONDS * 1000 - ttlMs);
+          }
+        }
+        return {
+          status: "pending",
+          claimValue: existing,
+          claimedAt: pendingClaimedAt,
+        };
+      }
+
+      // The key expired or was evicted between SET NX and GET. Retry the
+      // reservation a small number of times without growing the call stack.
     }
 
-    // The key expired between SET NX and GET. Let this request retry the
-    // claim instead of treating a vanished owner as pending.
-    return claimSessionStreamPart(environmentId, addressingKey, io, partId);
+    // Redis kept losing the reservation. Fail open like other cache errors so
+    // this append degrades to at-least-once instead of failing the request.
+    return { status: "claimed", claimValue: undefined };
   } catch (error) {
     logger.error("Failed to claim session stream append part", {
       environmentId,

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

@@ -413,7 +413,72 @@ describe("session stream e2e", () => {
     }
   });
 
-  it("E14 subscribe with an invalid token is rejected", async () => {
+  it("E14 in/append recovers a sequence from a legacy claim", async () => {
+    const { addressingKey, environmentId, token, inProducer, baseUrl } = await setupSession();
+    const redis = new Redis({ ...server.redis, keyPrefix: "tr:" });
+
+    try {
+      const payload = JSON.stringify({ kind: "message", text: "rolling deploy retry" });
+      const partId = `legacy-${randomBytes(6).toString("hex")}`;
+      const claimKey = `ssa:${encodeURIComponent(environmentId)}:${encodeURIComponent(
+        addressingKey
+      )}:in:${encodeURIComponent(partId)}`;
+
+      // The previous release stored literal "1" for both pending and accepted
+      // appends. Its remaining TTL supplies the recovery window after deploy.
+      await redis.set(claimKey, "1", "EX", 5 * 60);
+      const originalSeq = await inProducer.appendData(payload, partId);
+
+      const retry = await appendInput({ baseUrl, addressingKey, token, partId, body: payload });
+      expect(retry.status).toBe(200);
+      expect(retry.json).toEqual({ ok: true, seq: originalSeq });
+      expect(await redis.get(claimKey)).toBe(`seq:${originalSeq}`);
+
+      const { parts } = await collectSessionOut({
+        baseUrl,
+        addressingKey,
+        token,
+        io: "in",
+        timeoutInSeconds: 1,
+        maxMs: 5_000,
+      });
+      expect(parts.filter((part) => part.chunk != null)).toHaveLength(1);
+    } finally {
+      await redis.quit();
+    }
+  });
+
+  it("E15 in/append recovery searches beyond 32 full read pages", async () => {
+    const { addressingKey, environmentId, token, inProducer, baseUrl } = await setupSession();
+    const redis = new Redis({ ...server.redis, keyPrefix: "tr:" });
+
+    try {
+      const payload = JSON.stringify({ kind: "message", text: "after large records" });
+      const partId = `deep-${randomBytes(6).toString("hex")}`;
+      const claimKey = `ssa:${encodeURIComponent(environmentId)}:${encodeURIComponent(
+        addressingKey
+      )}:in:${encodeURIComponent(partId)}`;
+
+      await redis.set(claimKey, `pending:${Date.now()}:crashed-owner`, "EX", 5 * 60);
+
+      // Each record is over half the 1 MiB read budget, forcing one record per
+      // page. The target therefore sits on page 33 of the recovery scan.
+      const largeDecoy = "x".repeat(530_000);
+      for (let index = 0; index < 32; index++) {
+        await inProducer.appendData(largeDecoy, `decoy-${index}`);
+      }
+      const originalSeq = await inProducer.appendData(payload, partId);
+
+      const retry = await appendInput({ baseUrl, addressingKey, token, partId, body: payload });
+      expect(retry.status).toBe(200);
+      expect(retry.json).toEqual({ ok: true, seq: originalSeq });
+      expect(await redis.get(claimKey)).toBe(`seq:${originalSeq}`);
+    } finally {
+      await redis.quit();
+    }
+  });
+
+  it("E16 subscribe with an invalid token is rejected", async () => {
     const { addressingKey, baseUrl } = await setupSession();
 
     const { status } = await openChannelRaw({