feat(streams): a reply generated on one machine now streams to every machine

apps/web/src/app/api/ai/chat/stream-join/[messageId]/__tests__/route.test.ts

@@ -1,6 +1,7 @@
 import { describe, it, expect, vi, beforeEach } from 'vitest';
 import { NextResponse } from 'next/server';
 import { StreamChannelRegistry } from '@/lib/ai/core/stream-channel-registry';
+import { openStreamChannel } from '@/lib/ai/core/stream-channel';
 import type { SessionAuthResult, AuthError } from '@/lib/auth';
 
 // Fresh registry per test — module-level let, updated in beforeEach
@@ -38,6 +39,19 @@ vi.mock('@pagespace/lib/audit/audit-log', () => ({
  * the registry answers, or nothing does. The cross-instance cases opt in.
  */
 const mockSessionRow = vi.fn<() => Promise<unknown[]>>();
+
+/**
+ * The remote follower, mocked at its module boundary.
+ *
+ * The route's job is to acquire one, hand its channel to the SAME subscribe/SSE code a local
+ * channel goes through, and release it on every teardown path. The follower's own polling and
+ * terminal classification live in `remote-frame-follower.test.ts`.
+ */
+const mockRemoteRelease = vi.fn();
+const mockAcquireRemoteChannel = vi.fn();
+vi.mock('@/lib/ai/core/remote-frame-follower', () => ({
+  acquireRemoteChannel: (messageId: string) => mockAcquireRemoteChannel(messageId),
+}));
 vi.mock('@pagespace/db/db', () => ({
   db: {
     select: () => ({
@@ -127,6 +141,13 @@ describe('GET /api/ai/chat/stream-join/[messageId]', () => {
     mockCanSubscribeToStream.mockResolvedValue(true);
     vi.mocked(canUserViewPage).mockResolvedValue(true);
     mockSessionRow.mockResolvedValue([]);
+    // Default: a follower that yields an already-finished, empty channel, so a case that does
+    // not care about the remote path still terminates promptly.
+    mockAcquireRemoteChannel.mockImplementation((id: string) => {
+      const channel = openStreamChannel({ messageId: id });
+      channel.finish(false);
+      return { channel, release: mockRemoteRelease };
+    });
   });
 
   describe('authentication', () => {
@@ -260,6 +281,138 @@ describe('GET /api/ai/chat/stream-join/[messageId]', () => {
     });
   });
 
+  /**
+   * SERVING a remote stream.
+   *
+   * The route body is deliberately NOT forked here: a follower tails the durable frame log and
+   * presents it as an ordinary `StreamChannel`, so the SSE framing, the ping, the recheck, the
+   * teardown and the overflow semantics below are one code path for both sources. These cases
+   * pin that the follower is wired in and released, not the follower's own behaviour
+   * (`remote-frame-follower.test.ts` covers that).
+   */
+  describe('serving a remote stream through the follower', () => {
+    const remoteRow = () => [{
+      channelId: mockPageId,
+      userId: mockUserId,
+      displayName: mockDisplayName,
+      conversationId: mockConversationId,
+      browserSessionId: mockBrowserSessionId,
+      status: 'streaming',
+    }];
+
+    it('given a remote stream, serves its frames instead of 404ing', async () => {
+      mockSessionRow.mockResolvedValue(remoteRow());
+      mockAcquireRemoteChannel.mockImplementation((id: string) => {
+        const channel = openStreamChannel({ messageId: id });
+        channel.append(textChunk('from the durable log') as never);
+        channel.finish(false);
+        return { channel, release: mockRemoteRelease };
+      });
+
+      const response = await GET(makeRequest(), makeContext(mockMessageId));
+      const body = await readSSEBody(response);
+
+      expect(response.status).toBe(200);
+      expect(body).toContain(`data: ${JSON.stringify({ seq: 0, chunk: textChunk('from the durable log') })}\n\n`);
+    });
+
+    it('labels the response with where the frames came from', async () => {
+      mockSessionRow.mockResolvedValue(remoteRow());
+
+      const response = await GET(makeRequest(), makeContext(mockMessageId));
+
+      // An N=2 smoke test has no other way to PROVE it exercised the follower rather than
+      // getting lucky with the load balancer; in production the remote share should sit near
+      // (N-1)/N.
+      expect(response.headers.get('X-Stream-Join-Source')).toBe('remote');
+    });
+
+    it('labels a locally-owned join as local', async () => {
+      testRegistry.open(mockMessageId, mockMeta);
+
+      const response = await GET(makeRequest(), makeContext(mockMessageId));
+
+      expect(response.headers.get('X-Stream-Join-Source')).toBe('local');
+    });
+
+    it('labels a terminal join as terminal, and follows it rather than 404ing', async () => {
+      // Frames are deleted on the terminal write, so a terminal row is exactly the case where
+      // the follower's honest-answer logic is needed — not a case to short-circuit.
+      mockSessionRow.mockResolvedValue([{ ...remoteRow()[0], status: 'complete' }]);
+
+      const response = await GET(makeRequest(), makeContext(mockMessageId));
+
+      expect(response.status).toBe(200);
+      expect(response.headers.get('X-Stream-Join-Source')).toBe('terminal');
+    });
+
+    it('given a truncated end, tells the client to RELOAD rather than sending a bare done', async () => {
+      mockSessionRow.mockResolvedValue(remoteRow());
+      mockAcquireRemoteChannel.mockImplementation((id: string) => {
+        const channel = openStreamChannel({ messageId: id });
+        channel.append(textChunk('a prefix') as never);
+        channel.finish(false, { truncated: true });
+        return { channel, release: mockRemoteRelease };
+      });
+
+      const body = await readSSEBody(await GET(makeRequest(), makeContext(mockMessageId)));
+
+      // A bare `done` would leave a short reply on screen looking whole. `reload` is a different
+      // answer from `resumeFromSeq`: there is no seq to resume from, only a durable message to
+      // re-read.
+      expect(body).toContain('data: {"done":true,"aborted":false,"reload":true}\n\n');
+    });
+
+    it('releases the follower reference when the stream ends', async () => {
+      mockSessionRow.mockResolvedValue(remoteRow());
+      mockAcquireRemoteChannel.mockImplementation((id: string) => {
+        const channel = openStreamChannel({ messageId: id });
+        channel.finish(false);
+        return { channel, release: mockRemoteRelease };
+      });
+
+      await readSSEBody(await GET(makeRequest(), makeContext(mockMessageId)));
+
+      // A leaked reference keeps a poller hitting Postgres for a reader that has already gone.
+      expect(mockRemoteRelease).toHaveBeenCalled();
+    });
+
+    it('releases the follower reference when the client disconnects', async () => {
+      mockSessionRow.mockResolvedValue(remoteRow());
+      const controller = new AbortController();
+      mockAcquireRemoteChannel.mockImplementation((id: string) => ({
+        channel: openStreamChannel({ messageId: id }),
+        release: mockRemoteRelease,
+      }));
+
+      const response = await GET(makeRequest(controller.signal), makeContext(mockMessageId));
+      const reader = response.body!.getReader();
+      controller.abort();
+      await reader.read().catch(() => undefined);
+
+      expect(mockRemoteRelease).toHaveBeenCalled();
+    });
+
+    it('never acquires a follower for a locally-owned stream', async () => {
+      testRegistry.open(mockMessageId, mockMeta);
+
+      await GET(makeRequest(), makeContext(mockMessageId));
+
+      expect(mockAcquireRemoteChannel).not.toHaveBeenCalled();
+    });
+
+    it('never acquires a follower for a caller who may not subscribe', async () => {
+      mockSessionRow.mockResolvedValue(remoteRow());
+      mockCanSubscribeToStream.mockResolvedValue(false);
+
+      await GET(makeRequest(), makeContext(mockMessageId));
+
+      // Acquiring first would start a poller — and a DB read loop over another member's private
+      // conversation — for a request that is about to 404.
+      expect(mockAcquireRemoteChannel).not.toHaveBeenCalled();
+    });
+  });
+
   describe('authorization', () => {
     it('given a user without view access, should return 403', async () => {
       testRegistry.open(mockMessageId, mockMeta);

apps/web/src/app/api/ai/chat/stream-join/[messageId]/route.ts

@@ -5,6 +5,7 @@ import { auditRequest } from '@pagespace/lib/audit/audit-log';
 import { parseGlobalChannelId } from '@pagespace/lib/ai/global-channel-id';
 import { canSubscribeToStream } from '@/lib/ai/core/stream-subscription-authz';
 import { resolveStreamJoinContext } from '@/lib/ai/core/stream-join-context';
+import { acquireRemoteChannel } from '@/lib/ai/core/remote-frame-follower';
 
 export const dynamic = 'force-dynamic';
 
@@ -131,22 +132,42 @@ export async function GET(
   // frames to discard — under-skipping duplicated visible text, over-skipping left a silent
   // permanent gap, and neither was detectable from either side.
   //
-  // A `remote` or `terminal` context has no channel to read from YET — serving one is the next
-  // leaf (the durable log's follower). Until then this is the same 404 the route already gave
-  // for a registry miss, reached AFTER authorization rather than instead of it. The client's
-  // poll fallback (`stream-join-poll-fallback.ts`) handles a 404 exactly as it does today, so
-  // cross-instance delivery is unchanged by this leaf — what changes is that the answer is now
-  // classified, and that a caller who may not subscribe is refused for the right reason.
-  if (joinContext.kind !== 'local') {
-    return NextResponse.json({ error: 'Stream not found' }, { status: 404 });
-  }
-  const channel = joinContext.channel;
+  // A REMOTE or TERMINAL context is served by a follower that tails the durable frame log and
+  // presents it as an ordinary `StreamChannel` — so everything below this line (the SSE framing,
+  // the ping, the recheck, the teardown, the overflow/resumeFromSeq semantics) is ONE code path
+  // for both sources. Forking the route body here is what this deliberately does not do: it
+  // would mean two copies of five behaviours, one of which nothing local exercises.
+  //
+  // `terminal` is followed too, not short-circuited. The frames are deleted on the terminal
+  // write, so a terminal row is exactly the case where the follower's honest-answer logic is
+  // needed: serve whatever the log still holds, then end — with `truncated` when the log was
+  // already released, which tells the client to reload rather than trust a short reply.
+  const remote = joinContext.kind === 'local' ? null : acquireRemoteChannel(messageId);
+  const channel = joinContext.kind === 'local' ? joinContext.channel : remote!.channel;
+  const joinSource = joinContext.kind === 'local' ? 'local' : joinContext.kind;
+  // Dropped in EVERY exit below, not only the happy one — a follower reference leaked by an
+  // early return keeps a poller alive for a reader that never arrived.
+  const releaseRemote = () => remote?.release();
 
   const requestedFromSeq = Number(new URL(request.url).searchParams.get('fromSeq') ?? '0');
   const fromSeq = Number.isFinite(requestedFromSeq) && requestedFromSeq >= 0
     ? Math.floor(requestedFromSeq)
     : 0;
 
+  // ONE teardown for both halves of this joiner's hold: the channel subscription, and — when
+  // the channel is a follower's — this reader's reference to it. Splitting them was the obvious
+  // shape and the wrong one: every path that unsubscribed would have had to remember the
+  // release too, and the one that forgot would keep a poller running against Postgres for a
+  // reader that had already gone.
+  let unsubscribeChannel: (() => void) | null = null;
+  let detached = false;
+  const detach = (): void => {
+    if (detached) return;
+    detached = true;
+    unsubscribeChannel?.();
+    releaseRemote();
+  };
+
   const unsubscribe = channel.subscribe({
     fromSeq,
     onFrame: ({ seq, chunk }) => {
@@ -163,22 +184,34 @@ export async function GET(
       // An `overflow` end means the cursor names a frame the ring no longer holds. Say so with
       // the seq that IS available rather than quietly serving a later prefix — the client can
       // then reseed deliberately instead of rendering a gap it cannot see.
+      //
+      // A `truncated` end is the OTHER thing, and they must not be conflated: it says there is
+      // no resume point at all — the durable log was released or holds a hole — so the only
+      // correct answer is `reload`, which the client turns into a read of the durably-persisted
+      // message. Sending a bare `done` there would leave a short reply on screen looking whole.
       const done = end.reason === 'overflow'
         ? encoder.encode(`data: ${JSON.stringify({ done: true, resumeFromSeq: end.resumeFromSeq })}\n\n`)
-        : encodeDoneFrame(end.aborted);
+        : encoder.encode(`data: ${JSON.stringify(
+          end.truncated === true
+            ? { done: true, aborted: end.aborted, reload: true }
+            : { done: true, aborted: end.aborted },
+        )}\n\n`);
       if (streamController) {
         streamController.enqueue(done);
         streamController.close();
       } else {
         preBuffer.push(done);
       }
       streamClosed = true;
+      detach();
     },
   });
+  unsubscribeChannel = unsubscribe;
 
   // finish() deletes entries before notifying subscribers, so subscribe() returns
   // null for both unknown and already-finished streams.
   if (unsubscribe === null) {
+    detach();
     return NextResponse.json({ error: 'Stream not found' }, { status: 404 });
   }
 
@@ -202,7 +235,7 @@ export async function GET(
       }
       if (request.signal.aborted) {
         streamClosed = true;
-        unsubscribe();
+        detach();
         controller.close();
         return;
       }
@@ -211,14 +244,14 @@ export async function GET(
         streamClosed = true;
         clearRecheckTimeout();
         clearPingInterval();
-        unsubscribe();
+        detach();
         controller.close();
       }, { once: true });
 
       const closeStreamAsDenied = (reason: string) => {
         streamClosed = true;
         clearPingInterval();
-        unsubscribe();
+        detach();
         auditRequest(request, {
           eventType: 'authz.access.denied',
           resourceType: 'ai_stream',
@@ -283,6 +316,11 @@ export async function GET(
       'Content-Type': 'text/event-stream',
       'Cache-Control': 'no-cache',
       'X-Accel-Buffering': 'no',
+      // WHERE THE FRAMES CAME FROM. An N=2 smoke test has no other way to PROVE it exercised
+      // the follower rather than getting lucky with the load balancer, and in production the
+      // remote share should sit near (N-1)/N — a free check that traffic is actually spread and
+      // that cross-instance joins are being served rather than silently falling back to polling.
+      'X-Stream-Join-Source': joinSource,
     },
   });
 }