diff --git a/apps/desktop/electron/main.ts b/apps/desktop/electron/main.ts index 0846a0b6ce..00ca47f5dd 100644 --- a/apps/desktop/electron/main.ts +++ b/apps/desktop/electron/main.ts @@ -276,7 +276,8 @@ import { findRemoteOwnerProfileForSession, mergeProfileSessionWindow, type RegistrySessionSource, - spliceRegistrySessionRows + spliceRegistrySessionRows, + tagRegistrySessionResponse } from './profile-session-routing' import { createQuickEntryShortcut, quickEntryWindowBounds, sanitizeQuickEntrySettings } from './quick-entry' import { type ActiveWork, mergeActiveWork, normalizeActiveWork, quitPromptFor } from './quit-guard' @@ -14347,12 +14348,16 @@ async function dispatchRegistryApiRequest( const requestPath = pathForRegistryBackendRequest(request.path, requestProfile, connection) - return fetchJsonForBackend(connection, requestPath, { + const response = await fetchJsonForBackend(connection, requestPath, { method: request?.method, body: request?.body, upload: request?.upload, timeoutMs: resolveTimeoutMs(request?.timeoutMs, DEFAULT_FETCH_TIMEOUT_MS) }) + + return (request?.method || 'GET').toUpperCase() === 'GET' + ? tagRegistrySessionResponse(requestPath, response, registryConnectionId) + : response } function registryConnectionKind(connectionId) { diff --git a/apps/desktop/electron/profile-session-routing.test.ts b/apps/desktop/electron/profile-session-routing.test.ts index 8cc530ff99..7cf49c76bd 100644 --- a/apps/desktop/electron/profile-session-routing.test.ts +++ b/apps/desktop/electron/profile-session-routing.test.ts @@ -9,7 +9,8 @@ import { fetchRemoteProfileSessions, findRemoteOwnerProfileForSession, mergeProfileSessionWindow, - spliceRegistrySessionRows + spliceRegistrySessionRows, + tagRegistrySessionResponse } from './profile-session-routing' test('remote sidebar slices all follow the selected profile', () => { @@ -301,6 +302,46 @@ test('registry sources: shared remote hosts read the cross-profile aggregate onc ) }) +test('registry-pinned session responses retain their owning connection', () => { + const sidebar = tagRegistrySessionResponse( + '/api/profiles/sessions/sidebar?recents_profile=default', + { + recents: { sessions: [{ id: 'remote-chat', profile: 'default' }] }, + cron: { sessions: [{ id: 'remote-cron', profile: 'default' }] }, + messaging: { sessions: [] } + }, + 'test-amnezia' + ) as any + + assert.equal(sidebar.recents.sessions[0].connection_id, 'test-amnezia') + assert.equal(sidebar.cron.sessions[0].connection_id, 'test-amnezia') + + const aggregate = tagRegistrySessionResponse( + '/api/profiles/sessions?profile=all', + { sessions: [{ id: 'remote-profile-chat', profile: 'research' }] }, + 'test-amnezia' + ) as any + + assert.equal(aggregate.sessions[0].connection_id, 'test-amnezia') + + const single = tagRegistrySessionResponse( + '/api/sessions/remote-chat?profile=default', + { id: 'remote-chat', profile: 'default' }, + 'test-amnezia' + ) as any + + assert.equal(single.connection_id, 'test-amnezia') +}) + +test('registry response ownership tagging ignores non-session payloads and transcript messages', () => { + const status = { ok: true } + const messages = { messages: [{ id: 'message-1' }], session_id: 'remote-chat' } + + assert.equal(tagRegistrySessionResponse('/api/status', status, 'test-amnezia'), status) + assert.equal(tagRegistrySessionResponse('/api/sessions/remote-chat/messages', messages, 'test-amnezia'), messages) + assert.equal((messages.messages[0] as any).connection_id, undefined) +}) + test('registry sources: an older shared host without the aggregator falls back to its flat list', async () => { const calls: string[] = [] diff --git a/apps/desktop/electron/profile-session-routing.ts b/apps/desktop/electron/profile-session-routing.ts index c669401021..52677ba60c 100644 --- a/apps/desktop/electron/profile-session-routing.ts +++ b/apps/desktop/electron/profile-session-routing.ts @@ -20,6 +20,52 @@ function rowsOf(data: unknown): unknown[] { return Array.isArray(data.sessions) ? data.sessions : [] } +function tagRowsWithConnection(rows: unknown[], connectionId: string): void { + for (const row of rows) { + if (row && typeof row === 'object') { + const session = row as Record + session.connection_id = connectionId + } + } +} + +/** Preserve the registry source that served a session REST response. + * + * A registry-pinned request is dispatched directly to that remote host, so its + * own session rows naturally omit Desktop's synthetic `connection_id`. Without + * restoring that provenance, a `profile: "default"` row later resumes through + * the legacy local primary instead of the active registry gateway. */ +export function tagRegistrySessionResponse(path: string, data: unknown, connectionId: string): unknown { + if (!data || typeof data !== 'object') { + return data + } + + const pathname = path.split('?', 1)[0].replace(/\/+$/, '') + + if (pathname === '/api/sessions' || pathname === '/api/profiles/sessions') { + tagRowsWithConnection(rowsOf(data), connectionId) + + return data + } + + if (pathname === '/api/profiles/sessions/sidebar') { + const response = data as Record + + for (const key of ['recents', 'cron', 'messaging']) { + tagRowsWithConnection(rowsOf(response[key]), connectionId) + } + + return data + } + + if (/^\/api\/sessions\/[^/]+$/.test(pathname)) { + const session = data as Record + session.connection_id = connectionId + } + + return data +} + function sessionId(row: unknown): string | null { if (!row || typeof row !== 'object' || !('id' in row)) { return null diff --git a/apps/desktop/src/app/session/hooks/use-session-actions.test.tsx b/apps/desktop/src/app/session/hooks/use-session-actions.test.tsx index c7e1adac50..0ca358f86f 100644 --- a/apps/desktop/src/app/session/hooks/use-session-actions.test.tsx +++ b/apps/desktop/src/app/session/hooks/use-session-actions.test.tsx @@ -2093,6 +2093,49 @@ describe('resumeSession warm-cache mapping integrity', () => { expect(ambientRequest).not.toHaveBeenCalled() }) + it('keeps a registry-tagged cached session on its owning connection without an explicit route', async () => { + setSessions([ + storedSession({ + connection_id: 'test-amnezia', + id: 'stored-registry', + profile: 'default' + }) + ]) + vi.mocked(getSession).mockImplementation(async id => + storedSession({ connection_id: 'test-amnezia', id, profile: 'default' }) + ) + vi.mocked(getLatestSessionMessages).mockResolvedValue({ messages: [], session_id: 'stored-registry' } as never) + vi.mocked(requestGatewayForAgent).mockImplementation(async (_connectionId, _profile, method, params) => { + if (method === 'session.resume') { + return { + info: {}, + messages: [], + resumed: params?.session_id, + session_id: 'runtime-registry' + } as never + } + + return {} as never + }) + + const ambientRequest = vi.fn(async () => ({}) as never) + let resume: ((storedSessionId: string, replaceRoute?: boolean) => Promise) | null = null + + render( (resume = ready)} requestGateway={ambientRequest} />) + await waitFor(() => expect(resume).not.toBeNull()) + await resume!('stored-registry', true) + + const restScope = { connectionId: 'test-amnezia', profile: 'default' } + expect(getLatestSessionMessages).toHaveBeenCalledWith('stored-registry', restScope) + expect(requestGatewayForAgent).toHaveBeenCalledWith( + 'test-amnezia', + 'default', + 'session.resume', + expect.objectContaining({ session_id: 'stored-registry' }) + ) + expect(ambientRequest).not.toHaveBeenCalled() + }) + it('rejects a cross-wired runtime mapping and falls through to a full resume', async () => { // A recycled runtime id ('rt-recycled') is mapped to 'stored-A', but its // cached state actually belongs to a DIFFERENT session ('stored-B') — the diff --git a/apps/desktop/src/app/session/hooks/use-session-actions/index.ts b/apps/desktop/src/app/session/hooks/use-session-actions/index.ts index d5e21f57fa..431be7cc11 100644 --- a/apps/desktop/src/app/session/hooks/use-session-actions/index.ts +++ b/apps/desktop/src/app/session/hooks/use-session-actions/index.ts @@ -847,7 +847,12 @@ export function useSessionActions({ connectionId: ownerRoute.connectionId, profile: ownerRoute.targetProfile || ownerRoute.profile } - : sessionProfile + : storedForProfile?.connection_id + ? { + connectionId: storedForProfile.connection_id, + profile: sessionProfile || 'default' + } + : sessionProfile // Re-check after the profile-resolve / gateway-swap awaits above: the // cache may have changed, and takeWarmCache re-validates belongs-to and