address comments

Author: icecrasher321

apps/sim/app/api/mothership/chats/[chatId]/route.test.ts

@@ -7,13 +7,13 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'
 
 const {
   mockGetAccessibleCopilotChat,
-  mockGetActiveChatStreamIds,
+  mockReconcileChatStreamMarkers,
   mockReadEvents,
   mockReadFilePreviewSessions,
   mockGetLatestRunForStream,
 } = vi.hoisted(() => ({
   mockGetAccessibleCopilotChat: vi.fn(),
-  mockGetActiveChatStreamIds: vi.fn(),
+  mockReconcileChatStreamMarkers: vi.fn(),
   mockReadEvents: vi.fn(),
   mockReadFilePreviewSessions: vi.fn(),
   mockGetLatestRunForStream: vi.fn(),
@@ -50,8 +50,8 @@ vi.mock('@/lib/copilot/chat/lifecycle', () => ({
   getAccessibleCopilotChat: mockGetAccessibleCopilotChat,
 }))
 
-vi.mock('@/lib/copilot/request/session/abort', () => ({
-  getActiveChatStreamIds: mockGetActiveChatStreamIds,
+vi.mock('@/lib/copilot/chat/stream-liveness', () => ({
+  reconcileChatStreamMarkers: mockReconcileChatStreamMarkers,
 }))
 
 vi.mock('@/lib/copilot/request/session/buffer', () => ({
@@ -105,7 +105,19 @@ describe('GET /api/mothership/chats/[chatId]', () => {
       userId: 'user-1',
       isAuthenticated: true,
     })
-    mockGetActiveChatStreamIds.mockResolvedValue(new Set<string>())
+    mockReconcileChatStreamMarkers.mockImplementation(
+      async (candidates: Array<{ chatId: string; streamId: string | null }>) =>
+        new Map(
+          candidates.map((candidate) => [
+            candidate.chatId,
+            {
+              chatId: candidate.chatId,
+              streamId: candidate.streamId,
+              status: candidate.streamId ? 'active' : 'inactive',
+            },
+          ])
+        )
+    )
     mockReadEvents.mockResolvedValue([])
     mockReadFilePreviewSessions.mockResolvedValue([])
     mockGetLatestRunForStream.mockResolvedValue(null)
@@ -122,13 +134,18 @@ describe('GET /api/mothership/chats/[chatId]', () => {
       createdAt: new Date('2026-05-11T12:00:00Z'),
       updatedAt: new Date('2026-05-11T12:00:00Z'),
     })
-    mockGetActiveChatStreamIds.mockResolvedValueOnce(new Set<string>())
+    mockReconcileChatStreamMarkers.mockResolvedValueOnce(
+      new Map([['chat-stuck', { chatId: 'chat-stuck', streamId: null, status: 'inactive' }]])
+    )
 
     const response = await GET(createRequest('chat-stuck'), makeContext('chat-stuck'))
     expect(response.status).toBe(200)
     const body = await response.json()
 
-    expect(mockGetActiveChatStreamIds).toHaveBeenCalledWith(['chat-stuck'])
+    expect(mockReconcileChatStreamMarkers).toHaveBeenCalledWith(
+      [{ chatId: 'chat-stuck', streamId: 'stream-orphaned' }],
+      { repairVerifiedStaleMarkers: true }
+    )
     expect(body.success).toBe(true)
     expect(body.chat.conversationId).toBeNull()
     expect(body.chat.streamSnapshot).toBeUndefined()
@@ -146,7 +163,6 @@ describe('GET /api/mothership/chats/[chatId]', () => {
       createdAt: new Date('2026-05-11T12:00:00Z'),
       updatedAt: new Date('2026-05-11T12:00:00Z'),
     })
-    mockGetActiveChatStreamIds.mockResolvedValueOnce(new Set(['chat-live']))
     mockGetLatestRunForStream.mockResolvedValueOnce({ status: 'active' })
 
     const response = await GET(createRequest('chat-live'), makeContext('chat-live'))
@@ -159,7 +175,32 @@ describe('GET /api/mothership/chats/[chatId]', () => {
     expect(body.chat.streamSnapshot.status).toBe('active')
   })
 
-  it('skips the reconciliation lookup when conversationId is already null', async () => {
+  it('uses the Redis lock owner when it differs from a stale conversationId', async () => {
+    mockGetAccessibleCopilotChat.mockResolvedValueOnce({
+      id: 'chat-mismatch',
+      type: 'mothership',
+      title: 'Mismatch',
+      messages: [],
+      resources: [],
+      conversationId: 'stream-stale',
+      createdAt: new Date('2026-05-11T12:00:00Z'),
+      updatedAt: new Date('2026-05-11T12:00:00Z'),
+    })
+    mockReconcileChatStreamMarkers.mockResolvedValueOnce(
+      new Map([
+        ['chat-mismatch', { chatId: 'chat-mismatch', streamId: 'stream-live', status: 'active' }],
+      ])
+    )
+
+    const response = await GET(createRequest('chat-mismatch'), makeContext('chat-mismatch'))
+    expect(response.status).toBe(200)
+    const body = await response.json()
+
+    expect(body.chat.conversationId).toBe('stream-live')
+    expect(mockReadEvents).toHaveBeenCalledWith('stream-live', '0')
+  })
+
+  it('returns null when conversationId is already null', async () => {
     mockGetAccessibleCopilotChat.mockResolvedValueOnce({
       id: 'chat-idle',
       type: 'mothership',
@@ -174,7 +215,10 @@ describe('GET /api/mothership/chats/[chatId]', () => {
     const response = await GET(createRequest('chat-idle'), makeContext('chat-idle'))
     expect(response.status).toBe(200)
 
-    expect(mockGetActiveChatStreamIds).not.toHaveBeenCalled()
+    expect(mockReconcileChatStreamMarkers).toHaveBeenCalledWith(
+      [{ chatId: 'chat-idle', streamId: null }],
+      { repairVerifiedStaleMarkers: true }
+    )
     const body = await response.json()
     expect(body.chat.conversationId).toBeNull()
   })
@@ -184,7 +228,7 @@ describe('GET /api/mothership/chats/[chatId]', () => {
 
     const response = await GET(createRequest('chat-missing'), makeContext('chat-missing'))
     expect(response.status).toBe(404)
-    expect(mockGetActiveChatStreamIds).not.toHaveBeenCalled()
+    expect(mockReconcileChatStreamMarkers).not.toHaveBeenCalled()
   })
 
   it('returns 401 when unauthenticated', async () => {
@@ -196,6 +240,6 @@ describe('GET /api/mothership/chats/[chatId]', () => {
     const response = await GET(createRequest('chat-x'), makeContext('chat-x'))
     expect(response.status).toBe(401)
     expect(mockGetAccessibleCopilotChat).not.toHaveBeenCalled()
-    expect(mockGetActiveChatStreamIds).not.toHaveBeenCalled()
+    expect(mockReconcileChatStreamMarkers).not.toHaveBeenCalled()
   })
 })

apps/sim/app/api/mothership/chats/[chatId]/route.ts

@@ -14,13 +14,13 @@ import { getLatestRunForStream } from '@/lib/copilot/async-runs/repository'
 import { buildEffectiveChatTranscript } from '@/lib/copilot/chat/effective-transcript'
 import { getAccessibleCopilotChat } from '@/lib/copilot/chat/lifecycle'
 import { normalizeMessage } from '@/lib/copilot/chat/persisted-message'
+import { reconcileChatStreamMarkers } from '@/lib/copilot/chat/stream-liveness'
 import {
   authenticateCopilotRequestSessionOnly,
   createInternalServerErrorResponse,
   createUnauthorizedResponse,
 } from '@/lib/copilot/request/http'
 import type { FilePreviewSession } from '@/lib/copilot/request/session'
-import { getActiveChatStreamIds } from '@/lib/copilot/request/session/abort'
 import { readEvents } from '@/lib/copilot/request/session/buffer'
 import { readFilePreviewSessions } from '@/lib/copilot/request/session/file-preview-session'
 import { type StreamBatchEvent, toStreamBatchEvent } from '@/lib/copilot/request/session/types'
@@ -53,16 +53,11 @@ export const GET = withRouteHandler(
         status: string
       } | null = null
 
-      // Reconcile the persisted stream marker against the canonical Redis
-      // lock. If `conversation_id` is set but no lock is held, the stream
-      // is no longer running (process died before finalize) — treat the
-      // marker as null so the client doesn't try to reconnect to a dead
-      // stream. Mirrors the same reconciliation in the task list route.
-      const activeIds = chat.conversationId
-        ? await getActiveChatStreamIds([chat.id])
-        : new Set<string>()
-      const liveConversationId =
-        chat.conversationId && activeIds.has(chat.id) ? chat.conversationId : null
+      const reconciledMarkers = await reconcileChatStreamMarkers(
+        [{ chatId: chat.id, streamId: chat.conversationId }],
+        { repairVerifiedStaleMarkers: true }
+      )
+      const liveConversationId = reconciledMarkers.get(chat.id)?.streamId ?? null
 
       if (liveConversationId) {
         try {

apps/sim/app/api/mothership/chats/route.test.ts

@@ -5,13 +5,13 @@ import { copilotHttpMock, copilotHttpMockFns, permissionsMock } from '@sim/testi
 import { NextRequest } from 'next/server'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
-const { mockSelect, mockFrom, mockWhere, mockOrderBy, mockGetActiveChatStreamIds } = vi.hoisted(
+const { mockSelect, mockFrom, mockWhere, mockOrderBy, mockReconcileChatStreamMarkers } = vi.hoisted(
   () => ({
     mockSelect: vi.fn(),
     mockFrom: vi.fn(),
     mockWhere: vi.fn(),
     mockOrderBy: vi.fn(),
-    mockGetActiveChatStreamIds: vi.fn(),
+    mockReconcileChatStreamMarkers: vi.fn(),
   })
 )
 
@@ -43,8 +43,8 @@ vi.mock('drizzle-orm', () => ({
 vi.mock('@/lib/copilot/request/http', () => copilotHttpMock)
 vi.mock('@/lib/workspaces/permissions/utils', () => permissionsMock)
 
-vi.mock('@/lib/copilot/request/session/abort', () => ({
-  getActiveChatStreamIds: mockGetActiveChatStreamIds,
+vi.mock('@/lib/copilot/chat/stream-liveness', () => ({
+  reconcileChatStreamMarkers: mockReconcileChatStreamMarkers,
 }))
 
 vi.mock('@/lib/copilot/tasks', () => ({
@@ -77,7 +77,19 @@ describe('GET /api/mothership/chats', () => {
     mockFrom.mockReturnValue({ where: mockWhere })
     mockSelect.mockReturnValue({ from: mockFrom })
 
-    mockGetActiveChatStreamIds.mockResolvedValue(new Set<string>())
+    mockReconcileChatStreamMarkers.mockImplementation(
+      async (candidates: Array<{ chatId: string; streamId: string | null }>) =>
+        new Map(
+          candidates.map((candidate) => [
+            candidate.chatId,
+            {
+              chatId: candidate.chatId,
+              streamId: candidate.streamId,
+              status: candidate.streamId ? 'active' : 'inactive',
+            },
+          ])
+        )
+    )
   })
 
   it('clears activeStreamId on chats whose redis lock has expired (stuck-yellow bug)', async () => {
@@ -105,13 +117,26 @@ describe('GET /api/mothership/chats', () => {
         lastSeenAt: null,
       },
     ])
-    mockGetActiveChatStreamIds.mockResolvedValueOnce(new Set(['chat-live']))
+    mockReconcileChatStreamMarkers.mockResolvedValueOnce(
+      new Map([
+        ['chat-stuck', { chatId: 'chat-stuck', streamId: null, status: 'inactive' }],
+        ['chat-live', { chatId: 'chat-live', streamId: 'stream-live', status: 'active' }],
+        ['chat-idle', { chatId: 'chat-idle', streamId: null, status: 'inactive' }],
+      ])
+    )
 
     const response = await GET(createRequest('ws-1'))
     expect(response.status).toBe(200)
     const body = await response.json()
 
-    expect(mockGetActiveChatStreamIds).toHaveBeenCalledWith(['chat-stuck', 'chat-live'])
+    expect(mockReconcileChatStreamMarkers).toHaveBeenCalledWith(
+      [
+        { chatId: 'chat-stuck', streamId: 'stream-orphaned' },
+        { chatId: 'chat-live', streamId: 'stream-live' },
+        { chatId: 'chat-idle', streamId: null },
+      ],
+      { repairVerifiedStaleMarkers: true }
+    )
     expect(body.success).toBe(true)
     expect(body.data).toEqual([
       expect.objectContaining({ id: 'chat-stuck', activeStreamId: null }),
@@ -120,7 +145,7 @@ describe('GET /api/mothership/chats', () => {
     ])
   })
 
-  it('issues no Redis MGET when no chat has a stream marker set (empty candidateIds)', async () => {
+  it('preserves chats when no chat has a stream marker set', async () => {
     const now = new Date('2026-05-11T12:00:00Z')
     mockOrderBy.mockResolvedValueOnce([
       { id: 'chat-1', title: null, updatedAt: now, activeStreamId: null, lastSeenAt: null },
@@ -130,11 +155,18 @@ describe('GET /api/mothership/chats', () => {
     const response = await GET(createRequest('ws-1'))
     expect(response.status).toBe(200)
 
-    expect(mockGetActiveChatStreamIds).toHaveBeenCalledWith([])
+    expect(mockReconcileChatStreamMarkers).toHaveBeenCalledWith(
+      [
+        { chatId: 'chat-1', streamId: null },
+        { chatId: 'chat-2', streamId: null },
+      ],
+      { repairVerifiedStaleMarkers: true }
+    )
     const body = await response.json()
-    expect(
-      body.data.every((c: { activeStreamId: string | null }) => c.activeStreamId === null)
-    ).toBe(true)
+    expect(body.data).toEqual([
+      expect.objectContaining({ id: 'chat-1', activeStreamId: null }),
+      expect.objectContaining({ id: 'chat-2', activeStreamId: null }),
+    ])
   })
 
   it('leaves activeStreamId untouched when redis confirms every lock is live', async () => {
@@ -143,7 +175,6 @@ describe('GET /api/mothership/chats', () => {
       { id: 'chat-a', title: null, updatedAt: now, activeStreamId: 'stream-a', lastSeenAt: null },
       { id: 'chat-b', title: null, updatedAt: now, activeStreamId: 'stream-b', lastSeenAt: null },
     ])
-    mockGetActiveChatStreamIds.mockResolvedValueOnce(new Set(['chat-a', 'chat-b']))
 
     const response = await GET(createRequest('ws-1'))
     const body = await response.json()
@@ -154,6 +185,32 @@ describe('GET /api/mothership/chats', () => {
     ])
   })
 
+  it('uses Redis lock owner when it differs from a stale activeStreamId', async () => {
+    const now = new Date('2026-05-11T12:00:00Z')
+    mockOrderBy.mockResolvedValueOnce([
+      {
+        id: 'chat-mismatch',
+        title: null,
+        updatedAt: now,
+        activeStreamId: 'stream-stale',
+        lastSeenAt: null,
+      },
+    ])
+    mockReconcileChatStreamMarkers.mockResolvedValueOnce(
+      new Map([
+        ['chat-mismatch', { chatId: 'chat-mismatch', streamId: 'stream-live', status: 'active' }],
+      ])
+    )
+
+    const response = await GET(createRequest('ws-1'))
+    expect(response.status).toBe(200)
+    const body = await response.json()
+
+    expect(body.data).toEqual([
+      expect.objectContaining({ id: 'chat-mismatch', activeStreamId: 'stream-live' }),
+    ])
+  })
+
   it('returns 401 when unauthenticated', async () => {
     copilotHttpMockFns.mockAuthenticateCopilotRequestSessionOnly.mockResolvedValueOnce({
       userId: null,
@@ -163,6 +220,6 @@ describe('GET /api/mothership/chats', () => {
     const response = await GET(createRequest('ws-1'))
     expect(response.status).toBe(401)
     expect(mockSelect).not.toHaveBeenCalled()
-    expect(mockGetActiveChatStreamIds).not.toHaveBeenCalled()
+    expect(mockReconcileChatStreamMarkers).not.toHaveBeenCalled()
   })
 })

apps/sim/app/api/mothership/chats/route.ts

@@ -8,12 +8,12 @@ import {
   listMothershipChatsContract,
 } from '@/lib/api/contracts/mothership-tasks'
 import { parseRequest } from '@/lib/api/server'
+import { reconcileChatStreamMarkers } from '@/lib/copilot/chat/stream-liveness'
 import {
   authenticateCopilotRequestSessionOnly,
   createInternalServerErrorResponse,
   createUnauthorizedResponse,
 } from '@/lib/copilot/request/http'
-import { getActiveChatStreamIds } from '@/lib/copilot/request/session/abort'
 import { taskPubSub } from '@/lib/copilot/tasks'
 import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
 import { captureServerEvent } from '@/lib/posthog/server'
@@ -56,17 +56,14 @@ export const GET = withRouteHandler(async (request: NextRequest) => {
       )
       .orderBy(desc(copilotChats.updatedAt))
 
-    // Reconcile the persisted stream marker against the canonical Redis
-    // lock. `conversation_id` is set when a stream starts and cleared on
-    // the finalize/stop paths — but if those never run (pod crash, OOM,
-    // throw before callback), the column is orphaned and the task renders
-    // yellow forever. The Redis lock self-heals via its 60s TTL, so a
-    // missing lock means the stream is no longer running.
-    const candidateIds = chats.filter((c) => c.activeStreamId !== null).map((c) => c.id)
-    const activeIds = await getActiveChatStreamIds(candidateIds)
-    const reconciled = chats.map((c) =>
-      c.activeStreamId !== null && !activeIds.has(c.id) ? { ...c, activeStreamId: null } : c
+    const streamMarkers = await reconcileChatStreamMarkers(
+      chats.map((c) => ({ chatId: c.id, streamId: c.activeStreamId })),
+      { repairVerifiedStaleMarkers: true }
     )
+    const reconciled = chats.map((c) => {
+      const activeStreamId = streamMarkers.get(c.id)?.streamId ?? null
+      return activeStreamId === c.activeStreamId ? c : { ...c, activeStreamId }
+    })
 
     return NextResponse.json({ success: true, data: reconciled })
   } catch (error) {