From 064f65680fda9b7b940e1eee0658afe2af7cf5f5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 28 Jul 2026 15:37:17 +0200 Subject: [PATCH 1/3] fix(session-ingest): resolve forwarded exit_cli when ownership ends A successful remote /exit was reported to the app as "Session owner changed": the CLI ACKs exit_cli only after the detach heartbeat that removes the session, so the relay always failed the pending command first. For a forwarded exit_cli, ownership ending is the requested outcome, so resolve it with result: {} (the CLI's own ACK shape) when the owner-change path sees no next owner or when the owning socket is really gone. Genuine takeovers, reconnect replacements, and every other command keep failing with SESSION_OWNER_CHANGED. --- .../src/dos/UserConnectionDO.test.ts | 165 ++++++++++++++++++ .../src/dos/UserConnectionDO.ts | 45 +++-- 2 files changed, 194 insertions(+), 16 deletions(-) diff --git a/services/session-ingest/src/dos/UserConnectionDO.test.ts b/services/session-ingest/src/dos/UserConnectionDO.test.ts index e2f3863a3eb..5f6ac2017c1 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.test.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.test.ts @@ -3045,6 +3045,171 @@ describe('UserConnectionDO', () => { expect(parseSent(webWs)).toEqual({ type: 'response', id: 'cmd-1', result }); }); + + it('resolves exit_cli successfully when heartbeat drops the session', () => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + cliWs.send.mockClear(); + webWs.send.mockClear(); + sendCommand(doInstance, webWs, { + id: 'cmd-1', + command: 'exit_cli', + sessionId: 's1', + connectionId: 'cli-1', + data: { protocolVersion: 1 }, + }); + const correlationId = getCorrelationId(cliWs); + webWs.send.mockClear(); + + // Exit's own effect: CLI heartbeat no longer lists the session. + sendHeartbeat(doInstance, cliWs, []); + + expect(allSent(webWs).find(m => m.type === 'response' && m.id === 'cmd-1')).toEqual({ + type: 'response', + id: 'cmd-1', + result: {}, + }); + expect( + allSent(webWs).some( + m => + m.type === 'response' && + m.id === 'cmd-1' && + isRecord(m.error) && + m.error.code === 'SESSION_OWNER_CHANGED' + ) + ).toBe(false); + + // CLI's late ACK must not produce a second cmd-1 response. + const responsesBeforeLateAck = allSent(webWs).filter( + m => m.type === 'response' && m.id === 'cmd-1' + ).length; + sendCliResponse(doInstance, cliWs, { id: correlationId, result: {} }); + expect(allSent(webWs).filter(m => m.type === 'response' && m.id === 'cmd-1')).toHaveLength( + responsesBeforeLateAck + ); + }); + + it('resolves exit_cli successfully when the owning socket closes', async () => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendCommand(doInstance, webWs, { + id: 'cmd-1', + command: 'exit_cli', + sessionId: 's1', + connectionId: 'cli-1', + data: { protocolVersion: 1 }, + }); + webWs.send.mockClear(); + + mockCtx.removeSocket(cliWs); + await disconnectCli(doInstance, cliWs); + + expect(parseSent(webWs)).toEqual({ + type: 'response', + id: 'cmd-1', + result: {}, + }); + }); + + it.each(['list_models', 'send_message'] as const)( + 'still fails %s with SESSION_OWNER_CHANGED when heartbeat drops the session', + command => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + cliWs.send.mockClear(); + webWs.send.mockClear(); + sendCommand(doInstance, webWs, { + id: 'cmd-1', + command, + sessionId: 's1', + connectionId: 'cli-1', + }); + webWs.send.mockClear(); + + sendHeartbeat(doInstance, cliWs, []); + + expect(allSent(webWs).find(m => m.type === 'response' && m.id === 'cmd-1')).toEqual({ + type: 'response', + id: 'cmd-1', + error: { + source: 'relay', + code: 'SESSION_OWNER_CHANGED', + message: 'Session owner changed', + }, + }); + } + ); + + it('still fails exit_cli with SESSION_OWNER_CHANGED on genuine takeover', () => { + const { doInstance, mockCtx } = setup(); + const firstOwner = addCliSocket(mockCtx, 'cli-1'); + const nextOwner = addCliSocket(mockCtx, 'cli-2'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, firstOwner, [makeSession('s1')]); + sendHeartbeat(doInstance, nextOwner, []); + firstOwner.send.mockClear(); + webWs.send.mockClear(); + sendCommand(doInstance, webWs, { + id: 'cmd-1', + command: 'exit_cli', + sessionId: 's1', + connectionId: 'cli-1', + data: { protocolVersion: 1 }, + }); + webWs.send.mockClear(); + + sendHeartbeat(doInstance, nextOwner, [makeSession('s1')]); + + expect(allSent(webWs).find(m => m.type === 'response' && m.id === 'cmd-1')).toEqual({ + type: 'response', + id: 'cmd-1', + error: { + source: 'relay', + code: 'SESSION_OWNER_CHANGED', + message: 'Session owner changed', + }, + }); + }); + + it('still fails exit_cli with SESSION_OWNER_CHANGED when socket is replaced by reconnect', () => { + const { doInstance, mockCtx } = setup(); + const firstCli = connectCliSocket(doInstance, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, firstCli, [makeSession('s1')]); + firstCli.send.mockClear(); + sendCommand(doInstance, webWs, { + id: 'cmd-1', + command: 'exit_cli', + sessionId: 's1', + connectionId: 'cli-1', + data: { protocolVersion: 1 }, + }); + webWs.send.mockClear(); + + connectCliSocket(doInstance, 'cli-1'); + + expect(firstCli.close).toHaveBeenCalledWith(1000, 'replaced by reconnect'); + expect(parseSent(webWs)).toEqual({ + type: 'response', + id: 'cmd-1', + error: { + source: 'relay', + code: 'SESSION_OWNER_CHANGED', + message: 'Session owner changed', + }, + }); + }); }); // ------------------------------------------------------------------------- diff --git a/services/session-ingest/src/dos/UserConnectionDO.ts b/services/session-ingest/src/dos/UserConnectionDO.ts index ab1c52306f5..0197e787358 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.ts @@ -860,7 +860,7 @@ export class UserConnectionDO extends DurableObject { }); // Fail pending commands that targeted this specific socket - this.failPendingCommandsForSocket(disconnectedWs); + this.failPendingCommandsForSocket(disconnectedWs, !replaced); if (replaced) { console.log('Stale CLI socket closed (already replaced)', { connectionId }); @@ -1083,7 +1083,7 @@ export class UserConnectionDO extends DurableObject { const att = ws.deserializeAttachment() as WSAttachment | null; if (att?.role === 'cli' && att.connectionId === connectionId) { console.log('Closing stale CLI socket for reconnect', { connectionId }); - this.failPendingCommandsForSocket(ws); + this.failPendingCommandsForSocket(ws, false); // Preserve session ownership — the reconnecting CLI still owns these sessions ws.close(1000, 'replaced by reconnect'); return true; @@ -1132,16 +1132,25 @@ export class UserConnectionDO extends DurableObject { return undefined; } - private failPendingCommandsForSocket(targetWs: WebSocket): void { + private failPendingCommandsForSocket(targetWs: WebSocket, cliGone: boolean): void { for (const [id, entry] of this.pendingCommands) { - if (entry.targetCliWs === targetWs) { - this.sendToWeb(entry.ws, { - type: 'response', - id: entry.originalId, - error: entry.expectedOwnerConnectionId ? SESSION_OWNER_CHANGED_ERROR : 'CLI disconnected', - }); - this.pendingCommands.delete(id); - } + if (entry.targetCliWs !== targetWs) continue; + // The owning CLI is really gone, so a forwarded `exit_cli` got what it + // asked for: this session is no longer owned by anyone. + const exited = cliGone && entry.command === 'exit_cli'; + this.sendToWeb( + entry.ws, + exited + ? { type: 'response', id: entry.originalId, result: {} } + : { + type: 'response', + id: entry.originalId, + error: entry.expectedOwnerConnectionId + ? SESSION_OWNER_CHANGED_ERROR + : 'CLI disconnected', + } + ); + this.pendingCommands.delete(id); } } @@ -1154,11 +1163,15 @@ export class UserConnectionDO extends DurableObject { continue; } this.pendingCommands.delete(id); - this.sendToWeb(entry.ws, { - type: 'response', - id: entry.originalId, - error: SESSION_OWNER_CHANGED_ERROR, - }); + // `exit_cli` asked for exactly this: the session is no longer owned. + // Ownership moving to another CLI is a genuine owner change and still fails. + const exited = entry.command === 'exit_cli' && nextOwnerConnectionId === undefined; + this.sendToWeb( + entry.ws, + exited + ? { type: 'response', id: entry.originalId, result: {} } + : { type: 'response', id: entry.originalId, error: SESSION_OWNER_CHANGED_ERROR } + ); } } From 08dc4250b30fc0619009f2d189ffed2f3d224a8d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 28 Jul 2026 15:37:32 +0200 Subject: [PATCH 2/3] fix(mobile): mention updating the CLI in the picker empty state A CLI too old to advertise itself is connected but never offered as a Run on target, and the empty state only said nothing was connected. Name both recoveries: run kilo remote, or update Kilo CLI if one is already running. --- apps/mobile/src/app/(app)/agent-chat/instance-picker.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/mobile/src/app/(app)/agent-chat/instance-picker.tsx b/apps/mobile/src/app/(app)/agent-chat/instance-picker.tsx index 8cda6869ddc..36e33b777e2 100644 --- a/apps/mobile/src/app/(app)/agent-chat/instance-picker.tsx +++ b/apps/mobile/src/app/(app)/agent-chat/instance-picker.tsx @@ -241,7 +241,7 @@ export default function InstancePickerScreen() { icon={Server} placement="top" title="No CLI instances connected" - description="Run `kilo remote` in a project on your computer to connect one." + description="Run `kilo remote` in a project on your computer, or update Kilo CLI if one is already running." action={