From 96ffeee2500c7edcda09319ad91cd8a84b7312ef Mon Sep 17 00:00:00 2001 From: Jinjing <6427696+AmethystLiang@users.noreply.github.com> Date: Sun, 5 Jul 2026 15:45:19 -0700 Subject: [PATCH] feat(native-chat): route live-session hook through owner transport (U3) Co-authored-by: Orca --- .../use-native-chat-live-session.test.ts | 186 ++++++++++++++++-- .../use-native-chat-live-session.ts | 37 ++-- 2 files changed, 194 insertions(+), 29 deletions(-) diff --git a/src/renderer/src/components/native-chat/use-native-chat-live-session.test.ts b/src/renderer/src/components/native-chat/use-native-chat-live-session.test.ts index f92937ceea8..539882c26c9 100644 --- a/src/renderer/src/components/native-chat/use-native-chat-live-session.test.ts +++ b/src/renderer/src/components/native-chat/use-native-chat-live-session.test.ts @@ -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 + subscribe: ReturnType + unsubscribe: ReturnType + } + const transports = new Map() + 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 { + await act(async () => { + await Promise.resolve() + await Promise.resolve() + }) + } + + async function render(props: UseNativeChatLiveSessionArgs): Promise { + 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 { + 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() + }) +}) diff --git a/src/renderer/src/components/native-chat/use-native-chat-live-session.ts b/src/renderer/src/components/native-chat/use-native-chat-live-session.ts index 4ba2ce73190..64a2cb66cf8 100644 --- a/src/renderer/src/components/native-chat/use-native-chat-live-session.ts +++ b/src/renderer/src/components/native-chat/use-native-chat-live-session.ts @@ -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({ 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