fix(desktop): retain registry session ownership

This commit is contained in:
arya
2026-08-26 03:07:14 +10:00
committed by Teknium
parent 1ec32e7385
commit 26932ea0bf
5 changed files with 144 additions and 4 deletions
+7 -2
View File
@@ -276,7 +276,8 @@ import {
findRemoteOwnerProfileForSession, findRemoteOwnerProfileForSession,
mergeProfileSessionWindow, mergeProfileSessionWindow,
type RegistrySessionSource, type RegistrySessionSource,
spliceRegistrySessionRows spliceRegistrySessionRows,
tagRegistrySessionResponse
} from './profile-session-routing' } from './profile-session-routing'
import { createQuickEntryShortcut, quickEntryWindowBounds, sanitizeQuickEntrySettings } from './quick-entry' import { createQuickEntryShortcut, quickEntryWindowBounds, sanitizeQuickEntrySettings } from './quick-entry'
import { type ActiveWork, mergeActiveWork, normalizeActiveWork, quitPromptFor } from './quit-guard' import { type ActiveWork, mergeActiveWork, normalizeActiveWork, quitPromptFor } from './quit-guard'
@@ -14347,12 +14348,16 @@ async function dispatchRegistryApiRequest(
const requestPath = pathForRegistryBackendRequest(request.path, requestProfile, connection) const requestPath = pathForRegistryBackendRequest(request.path, requestProfile, connection)
return fetchJsonForBackend(connection, requestPath, { const response = await fetchJsonForBackend(connection, requestPath, {
method: request?.method, method: request?.method,
body: request?.body, body: request?.body,
upload: request?.upload, upload: request?.upload,
timeoutMs: resolveTimeoutMs(request?.timeoutMs, DEFAULT_FETCH_TIMEOUT_MS) timeoutMs: resolveTimeoutMs(request?.timeoutMs, DEFAULT_FETCH_TIMEOUT_MS)
}) })
return (request?.method || 'GET').toUpperCase() === 'GET'
? tagRegistrySessionResponse(requestPath, response, registryConnectionId)
: response
} }
function registryConnectionKind(connectionId) { function registryConnectionKind(connectionId) {
@@ -9,7 +9,8 @@ import {
fetchRemoteProfileSessions, fetchRemoteProfileSessions,
findRemoteOwnerProfileForSession, findRemoteOwnerProfileForSession,
mergeProfileSessionWindow, mergeProfileSessionWindow,
spliceRegistrySessionRows spliceRegistrySessionRows,
tagRegistrySessionResponse
} from './profile-session-routing' } from './profile-session-routing'
test('remote sidebar slices all follow the selected profile', () => { 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 () => { test('registry sources: an older shared host without the aggregator falls back to its flat list', async () => {
const calls: string[] = [] const calls: string[] = []
@@ -20,6 +20,52 @@ function rowsOf(data: unknown): unknown[] {
return Array.isArray(data.sessions) ? data.sessions : [] 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<string, unknown>
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<string, unknown>
for (const key of ['recents', 'cron', 'messaging']) {
tagRowsWithConnection(rowsOf(response[key]), connectionId)
}
return data
}
if (/^\/api\/sessions\/[^/]+$/.test(pathname)) {
const session = data as Record<string, unknown>
session.connection_id = connectionId
}
return data
}
function sessionId(row: unknown): string | null { function sessionId(row: unknown): string | null {
if (!row || typeof row !== 'object' || !('id' in row)) { if (!row || typeof row !== 'object' || !('id' in row)) {
return null return null
@@ -2093,6 +2093,49 @@ describe('resumeSession warm-cache mapping integrity', () => {
expect(ambientRequest).not.toHaveBeenCalled() 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<unknown>) | null = null
render(<ResumeHarness onReady={ready => (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 () => { 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 // A recycled runtime id ('rt-recycled') is mapped to 'stored-A', but its
// cached state actually belongs to a DIFFERENT session ('stored-B') — the // cached state actually belongs to a DIFFERENT session ('stored-B') — the
@@ -847,7 +847,12 @@ export function useSessionActions({
connectionId: ownerRoute.connectionId, connectionId: ownerRoute.connectionId,
profile: ownerRoute.targetProfile || ownerRoute.profile 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 // Re-check after the profile-resolve / gateway-swap awaits above: the
// cache may have changed, and takeWarmCache re-validates belongs-to and // cache may have changed, and takeWarmCache re-validates belongs-to and