feat(native-chat): route live-session hook through owner transport (U3)

Co-authored-by: Orca <help@stably.ai>
2026-10-17 00:03:08 +00:00 · 2026-07-05 15:45:19 -07:00
parent 94da6eef94
commit 96ffeee250
2 changed files with 194 additions and 29 deletions
@@ -1,20 +1,54 @@
-import { describe, expect, it } from 'vitest'
+// @vitest-environment happy-dom
+
+import { act, createElement } from 'react'
+import { createRoot, type Root } from 'react-dom/client'
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 import type { NativeChatMessage } from '../../../../shared/native-chat-types'
+import { useAppStore } from '@/store'
 import { mergeNativeChatLiveSession } from './native-chat-live-status'

-// These exercise the hook's merge core (mergeNativeChatLiveSession), which is
-// the pure decision the hook delegates to after doing IPC + store reads. The
-// React wiring (subscribe/unsubscribe lifecycle) is covered by the main-side
-// transcript-watch + IPC tests; here we pin the live-status precedence rules.
+// Mock the session transport so the hook's IO is observable and controllable
+// per owner id. Each distinct owner gets its own read/subscribe/unsubscribe mocks.
+const { transportFactory, getMockTransport, resetMockTransports } = vi.hoisted(() => {
+  type MockTransport = {
+    readSession: ReturnType<typeof vi.fn>
+    subscribe: ReturnType<typeof vi.fn>
+    unsubscribe: ReturnType<typeof vi.fn>
+  }
+  const transports = new Map<string | null, MockTransport>()
+  const getMockTransport = (ownerId: string | null): MockTransport => {
+    let transport = transports.get(ownerId)
+    if (!transport) {
+      const unsubscribe = vi.fn()
+      transport = {
+        unsubscribe,
+        readSession: vi.fn().mockResolvedValue({ messages: [] }),
+        subscribe: vi.fn(() => unsubscribe)
+      }
+      transports.set(ownerId, transport)
+    }
+    return transport
+  }
+  return {
+    getMockTransport,
+    resetMockTransports: () => transports.clear(),
+    transportFactory: vi.fn((ownerId: string | null) => getMockTransport(ownerId))
+  }
+})
+
+vi.mock('./native-chat-session-transport', () => ({
+  getNativeChatSessionTransport: transportFactory
+}))
+
+// Imported after vi.mock is hoisted, so it binds to the mocked transport.
+import {
+  useNativeChatLiveSession,
+  type NativeChatLiveSession,
+  type UseNativeChatLiveSessionArgs
+} from './use-native-chat-live-session'

 function assistant(id: string, text: string): NativeChatMessage {
-  return {
-    id,
-    role: 'assistant',
-    blocks: [{ type: 'text', text }],
-    timestamp: 2,
-    source: 'transcript'
-  }
+  return { id, role: 'assistant', blocks: [{ type: 'text', text }], timestamp: 2, source: 'transcript' }
 }

 function user(id: string, text: string): NativeChatMessage {
@@ -23,8 +57,6 @@ function user(id: string, text: string): NativeChatMessage {

 describe('mergeNativeChatLiveSession', () => {
  it("surfaces live 'working' before the assistant turn lands in the transcript", () => {
-    // The user prompt is in the transcript but the assistant has not replied
-    // yet; the hook reports the agent is working.
    const session = mergeNativeChatLiveSession({
      sources: { transcript: [user('u-1', 'do a thing')] },
      sessionId: 'sess',
@@ -36,8 +68,6 @@ describe('mergeNativeChatLiveSession', () => {
  })

  it("clears 'working' once the assistant message flushes to the transcript", () => {
-    // Same pane, but now the transcript's last message is the assistant reply —
-    // the turn is visible, so 'working' is superseded by the derived 'ready'.
    const session = mergeNativeChatLiveSession({
      sources: { transcript: [user('u-1', 'do a thing'), assistant('a-1', 'done')] },
      sessionId: 'sess',
@@ -89,3 +119,127 @@ describe('mergeNativeChatLiveSession', () => {
    expect(session.status).toBe('empty')
  })
 })
+
+describe('useNativeChatLiveSession — transport routing', () => {
+  const AGENT = 'claude' as const
+  const SESSION = 'sess-1'
+  const PANE = 'pane-1'
+  const roots: Root[] = []
+  let latest: NativeChatLiveSession | null = null
+
+  function Probe(props: UseNativeChatLiveSessionArgs): null {
+    latest = useNativeChatLiveSession(props)
+    return null
+  }
+
+  async function flush(): Promise<void> {
+    await act(async () => {
+      await Promise.resolve()
+      await Promise.resolve()
+    })
+  }
+
+  async function render(props: UseNativeChatLiveSessionArgs): Promise<Root> {
+    const container = document.createElement('div')
+    const root = createRoot(container)
+    roots.push(root)
+    await act(async () => {
+      root.render(createElement(Probe, props))
+    })
+    await flush()
+    return root
+  }
+
+  async function rerender(root: Root, props: UseNativeChatLiveSessionArgs): Promise<void> {
+    await act(async () => {
+      root.render(createElement(Probe, props))
+    })
+    await flush()
+  }
+
+  beforeEach(() => {
+    useAppStore.setState({ agentStatusByPaneKey: {} })
+  })
+
+  afterEach(() => {
+    for (const root of roots.splice(0)) {
+      act(() => root.unmount())
+    }
+    latest = null
+    vi.clearAllMocks()
+    resetMockTransports()
+  })
+
+  it('routes the initial read and subscribe through the runtime transport', async () => {
+    await render({ paneKey: PANE, agent: AGENT, sessionId: SESSION, runtimeEnvironmentId: 'env-1' })
+
+    expect(transportFactory).toHaveBeenCalledWith('env-1')
+    const transport = getMockTransport('env-1')
+    expect(transport.readSession).toHaveBeenCalledWith(AGENT, SESSION, expect.any(Number), undefined)
+    expect(transport.subscribe).toHaveBeenCalledOnce()
+  })
+
+  it('re-subscribes against the new owner on an owner flip (R5)', async () => {
+    const root = await render({
+      paneKey: PANE,
+      agent: AGENT,
+      sessionId: SESSION,
+      runtimeEnvironmentId: 'env-1'
+    })
+    const first = getMockTransport('env-1')
+
+    await rerender(root, {
+      paneKey: PANE,
+      agent: AGENT,
+      sessionId: SESSION,
+      runtimeEnvironmentId: 'env-2'
+    })
+    const second = getMockTransport('env-2')
+
+    expect(first.unsubscribe).toHaveBeenCalledOnce()
+    expect(second.subscribe).toHaveBeenCalledOnce()
+  })
+
+  it('tears down the subscription on unmount (no watcher leak)', async () => {
+    const root = await render({
+      paneKey: PANE,
+      agent: AGENT,
+      sessionId: SESSION,
+      runtimeEnvironmentId: 'env-1'
+    })
+    const transport = getMockTransport('env-1')
+
+    await act(async () => {
+      root.unmount()
+    })
+
+    expect(transport.unsubscribe).toHaveBeenCalledOnce()
+  })
+
+  it('surfaces a runtime read error in the error phase (R4 end-to-end)', async () => {
+    getMockTransport('env-1').readSession.mockResolvedValueOnce({ error: 'runtime too old' })
+
+    await render({ paneKey: PANE, agent: AGENT, sessionId: SESSION, runtimeEnvironmentId: 'env-1' })
+    await flush()
+
+    expect(latest?.status).toBe('error')
+    expect(latest?.error).toBe('runtime too old')
+  })
+
+  it('never calls the transport when there is no session id', async () => {
+    await render({ paneKey: PANE, agent: AGENT, sessionId: null, runtimeEnvironmentId: 'env-1' })
+
+    const transport = getMockTransport('env-1')
+    expect(transport.readSession).not.toHaveBeenCalled()
+    expect(transport.subscribe).not.toHaveBeenCalled()
+  })
+
+  it('uses the local transport when the owner is null (unchanged behavior, R6)', async () => {
+    await render({ paneKey: PANE, agent: AGENT, sessionId: SESSION })
+
+    expect(transportFactory).toHaveBeenCalledWith(null)
+    const transport = getMockTransport(null)
+    expect(transport.readSession).toHaveBeenCalledOnce()
+    expect(transport.subscribe).toHaveBeenCalledOnce()
+  })
+})
@@ -93,17 +93,26 @@ type ReadState =
 * filled the window. Read results replace the base list (they are an ordered
 * tail), while live appends accumulate separately so a re-read never drops them.
 *
- * Remote/SSH: `nativeChat.readSession`/`subscribe` are main-process IPC, so the
- * transcript is read against the runtime's home dir (local or server-side) on
- * main — the renderer is transport-agnostic and needs no remote branch here.
+ * Transport: IO goes through a per-owner session transport selected by
+ * getNativeChatSessionTransport. A runtime-owned pane (Model B) reads/tails the
+ * REMOTE runtime host via the runtime RPCs; local- and ssh-owned panes keep the
+ * local IPC path. The transport preserves the NativeChatApi read/subscribe shape,
+ * so everything below (merge, assembler, pagination) is unchanged.
 *
- * Teardown: the subscription is closed on unmount and whenever agent/sessionId
- * change, so a toggle back to terminal or a session swap never leaks a watcher.
+ * Teardown: the subscription is closed on unmount and whenever the owner, agent,
+ * or sessionId change, so a toggle back to terminal, a session swap, or an
+ * owner-flip never leaks a watcher (remote or local).
 */
 export function useNativeChatLiveSession(
  args: UseNativeChatLiveSessionArgs
 ): NativeChatLiveSession {
-  const { paneKey, agent, sessionId, transcriptPath } = args
+  const { paneKey, agent, sessionId, transcriptPath, runtimeEnvironmentId } = args
+  // Stable per owner id, so a re-render without an owner flip keeps the same
+  // transport identity and doesn't re-subscribe.
+  const transport = useMemo(
+    () => getNativeChatSessionTransport(runtimeEnvironmentId ?? null),
+    [runtimeEnvironmentId]
+  )
  const [read, setRead] = useState<ReadState>({ phase: 'loading' })
  const [hasMore, setHasMore] = useState(false)
  const [loadingEarlier, setLoadingEarlier] = useState(false)
@@ -154,8 +163,8 @@ export function useNativeChatLiveSession(
    setAppended([])
    setHasMore(false)

-    void window.api?.nativeChat
-      ?.readSession(agent, sessionId, limitRef.current, transcriptPath ?? undefined)
+    void transport
+      .readSession(agent, sessionId, limitRef.current, transcriptPath ?? undefined)
      .then((result) => {
        if (cancelled) {
          return
@@ -175,7 +184,7 @@ export function useNativeChatLiveSession(
      })

    const subscriptionId = nextSubscriptionId()
-    const unsubscribe = window.api?.nativeChat?.subscribe?.(
+    const unsubscribe = transport.subscribe(
      { subscriptionId, agent, sessionId, transcriptPath: transcriptPath ?? undefined },
      (messages) => {
        if (!cancelled) {
@@ -210,7 +219,9 @@ export function useNativeChatLiveSession(
        })
      }
    }
-  }, [agent, sessionId, transcriptPath])
+    // `transport` identity changes on an owner flip, re-running this effect to
+    // tear down the old host's subscription and open one against the new host.
+  }, [agent, sessionId, transcriptPath, transport])

  const loadEarlier = useCallback(() => {
    if (!sessionId || loadingEarlier || !hasMore || read.phase !== 'ready') {
@@ -218,8 +229,8 @@ export function useNativeChatLiveSession(
    }
    const nextLimit = nextNativeChatLimit(limitRef.current)
    setLoadingEarlier(true)
-    void window.api?.nativeChat
-      ?.readSession(agent, sessionId, nextLimit, transcriptPath ?? undefined)
+    void transport
+      .readSession(agent, sessionId, nextLimit, transcriptPath ?? undefined)
      .then((result) => {
        // Ignore a stale resolve from a session that swapped underneath us.
        if (latestSessionId.current !== sessionId) {
@@ -240,7 +251,7 @@ export function useNativeChatLiveSession(
        // APPLYING the result above is gated on the session-id match.
        setLoadingEarlier(false)
      })
-  }, [agent, sessionId, transcriptPath, hasMore, loadingEarlier, read.phase])
+  }, [agent, sessionId, transcriptPath, transport, hasMore, loadingEarlier, read.phase])

  // Assembled messages reuse the incremental assembler across appends. Computed
  // outside the status memo: hookState changes only the status override, not the