diff --git a/apps/mobile/src/app/(app)/agent-chat/[session-id].tsx b/apps/mobile/src/app/(app)/agent-chat/[session-id].tsx index e2017c75cc..ea6364d00a 100644 --- a/apps/mobile/src/app/(app)/agent-chat/[session-id].tsx +++ b/apps/mobile/src/app/(app)/agent-chat/[session-id].tsx @@ -14,7 +14,6 @@ import { Button } from '@/components/ui/button'; import { Text } from '@/components/ui/text'; import { shouldRetryNotFoundOnSpawnedRoute } from '@/lib/spawned-not-found-retry'; import { useTRPC } from '@/lib/trpc'; -import { useAckSessionAttentionOnOpen } from '@/lib/session-attention'; export default function SessionDetailScreen() { const { @@ -40,7 +39,6 @@ export default function SessionDetailScreen() { }>(); const trpc = useTRPC(); const router = useRouter(); - useAckSessionAttentionOnOpen(sessionId); const sessionQuery = useQuery({ ...trpc.cliSessionsV2.get.queryOptions( { session_id: sessionId }, diff --git a/apps/mobile/src/components/agents/session-detail-content.tsx b/apps/mobile/src/components/agents/session-detail-content.tsx index 68f7c77ce8..40a4b8a01f 100644 --- a/apps/mobile/src/components/agents/session-detail-content.tsx +++ b/apps/mobile/src/components/agents/session-detail-content.tsx @@ -136,6 +136,7 @@ export function SessionDetailContent({ const remoteModelState = useAtomValue(manager.atoms.remoteModelState); const observedModel = useAtomValue(manager.atoms.observedModel); const remoteModelOverride = useAtomValue(manager.atoms.remoteModelOverride); + const cloudAgentModelOverride = useAtomValue(manager.atoms.cloudAgentModelOverride); const availableCommands = useAtomValue(manager.atoms.availableCommands); const remoteCommandState = useAtomValue(manager.atoms.remoteCommandState); const contextUsage = useAtomValue(manager.atoms.contextUsage); @@ -162,6 +163,7 @@ export function SessionDetailContent({ handleRespondToPermission, } = useInteractionHandlers({ manager, + kiloSessionId: sessionId, activeQuestion, activePermission, surface: analyticsSurface, @@ -271,6 +273,7 @@ export function SessionDetailContent({ modelOptions, selectedModel: sessionModels.selectedValue, selectedVariant: sessionModels.selectedVariant, + cloudAgentModelOverride, }); const viewTrackedRef = useRef(null); @@ -413,6 +416,10 @@ export function SessionDetailContent({ return; } + manager.setCloudAgentModelOverride({ + model: value, + ...(variant ? { variant } : {}), + }); setCurrentModel(value); setCurrentVariant(variant); savePersistedModel(organizationId, { model: value, variant }); diff --git a/apps/mobile/src/components/agents/use-interaction-handlers.test.ts b/apps/mobile/src/components/agents/use-interaction-handlers.test.ts new file mode 100644 index 0000000000..575b1cd3a5 --- /dev/null +++ b/apps/mobile/src/components/agents/use-interaction-handlers.test.ts @@ -0,0 +1,268 @@ +import * as React from 'react'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +import { + __peekSessionAttentionForTests, + __resetSessionAttentionForTests, + isAttentionAcked, + shouldShowNeedsInput, +} from '@/lib/session-attention'; + +import { useInteractionHandlers } from './use-interaction-handlers'; + +const toastError = vi.hoisted(() => vi.fn()); +const captureEvent = vi.hoisted(() => vi.fn()); + +vi.mock('sonner-native', () => ({ + toast: { error: (...args: unknown[]) => toastError(...args) }, +})); + +vi.mock('@/lib/analytics/posthog', () => ({ + captureEvent: (...args: unknown[]) => captureEvent(...args), + PERMISSION_RESPONDED_EVENT: 'permission_responded', + QUESTION_ANSWERED_EVENT: 'question_answered', +})); + +type ReactInternals = { + __CLIENT_INTERNALS_DO_NOT_USE_OR_WARN_USERS_THEY_CANNOT_UPGRADE: { + H: unknown; + }; +}; + +type HookDispatcher = { + useCallback: (callback: T) => T; + useState: (initialValue: T) => [T, (value: T | ((previous: T) => T)) => void]; +}; + +type InteractionHandlersArgs = Parameters[0]; +type InteractionHandlersResult = ReturnType; + +function renderInteractionHandlers(args: InteractionHandlersArgs) { + const reactInternals = React as typeof React & ReactInternals; + const hookState: unknown[] = []; + let hookIndex = 0; + + const dispatcher: HookDispatcher = { + useCallback: hookCallback => { + hookIndex += 1; + return hookCallback; + }, + useState: initialValue => { + const stateIndex = hookIndex; + hookIndex += 1; + if (hookState[stateIndex] === undefined) { + hookState[stateIndex] = initialValue; + } + const setState = ( + value: typeof initialValue | ((previous: typeof initialValue) => typeof initialValue) + ) => { + hookState[stateIndex] = + typeof value === 'function' + ? (value as (previous: typeof initialValue) => typeof initialValue)( + hookState[stateIndex] as typeof initialValue + ) + : value; + }; + return [hookState[stateIndex] as typeof initialValue, setState]; + }, + }; + + function render(): InteractionHandlersResult { + const previousDispatcher = + reactInternals.__CLIENT_INTERNALS_DO_NOT_USE_OR_WARN_USERS_THEY_CANNOT_UPGRADE.H; + hookIndex = 0; + reactInternals.__CLIENT_INTERNALS_DO_NOT_USE_OR_WARN_USERS_THEY_CANNOT_UPGRADE.H = dispatcher; + try { + // Alias avoids rules-of-hooks lexical false positives under the fake dispatcher. + const mountHandlers = useInteractionHandlers; + return mountHandlers(args); + } finally { + reactInternals.__CLIENT_INTERNALS_DO_NOT_USE_OR_WARN_USERS_THEY_CANNOT_UPGRADE.H = + previousDispatcher; + } + } + + return { render }; +} + +describe('useInteractionHandlers attention ack', () => { + beforeEach(() => { + __resetSessionAttentionForTests(); + toastError.mockReset(); + captureEvent.mockReset(); + }); + + it('acks the kilo session id after a successful answer', async () => { + const manager = { + answerQuestion: vi.fn().mockResolvedValue(undefined), + rejectQuestion: vi.fn(), + respondToPermission: vi.fn(), + }; + const { render } = renderInteractionHandlers({ + manager: manager as never, + kiloSessionId: 'kilo-session-1', + activeQuestion: { requestId: 'q1' }, + activePermission: null, + surface: 'remote-session', + }); + + await render().handleAnswerQuestion([['yes']]); + + expect(manager.answerQuestion).toHaveBeenCalledWith('q1', [['yes']]); + expect(isAttentionAcked('kilo-session-1', 'R1')).toBe(true); + expect( + shouldShowNeedsInput({ + status: 'question', + raiseId: 'R1', + isAcked: isAttentionAcked('kilo-session-1', 'R1'), + }) + ).toBe(false); + expect(__peekSessionAttentionForTests('kilo-session-1')).toEqual({ raiseId: null }); + // Must key by kilo id, never the cloud-agent id. + expect(isAttentionAcked('cloud-agent-id', 'R1')).toBe(false); + expect(captureEvent).toHaveBeenCalledWith('question_answered', { + surface: 'remote-session', + skipped: false, + }); + expect(toastError).not.toHaveBeenCalled(); + }); + + it('acks after a successful skip/reject', async () => { + const manager = { + answerQuestion: vi.fn(), + rejectQuestion: vi.fn().mockResolvedValue(undefined), + respondToPermission: vi.fn(), + }; + const { render } = renderInteractionHandlers({ + manager: manager as never, + kiloSessionId: 'kilo-session-1', + activeQuestion: { requestId: 'q1' }, + activePermission: null, + surface: 'cloud-agent', + }); + + await render().handleRejectQuestion(); + + expect(manager.rejectQuestion).toHaveBeenCalledWith('q1'); + expect(isAttentionAcked('kilo-session-1', 'R1')).toBe(true); + expect(captureEvent).toHaveBeenCalledWith('question_answered', { + surface: 'cloud-agent', + skipped: true, + }); + }); + + it('acks after a successful permission response', async () => { + const manager = { + answerQuestion: vi.fn(), + rejectQuestion: vi.fn(), + respondToPermission: vi.fn().mockResolvedValue(undefined), + }; + const { render } = renderInteractionHandlers({ + manager: manager as never, + kiloSessionId: 'kilo-session-1', + activeQuestion: null, + activePermission: { requestId: 'p1' }, + surface: 'remote-session', + }); + + await render().handleRespondToPermission('once'); + + expect(manager.respondToPermission).toHaveBeenCalledWith('p1', 'once'); + expect(isAttentionAcked('kilo-session-1', 'R1')).toBe(true); + expect(captureEvent).toHaveBeenCalledWith('permission_responded', { + surface: 'remote-session', + response: 'once', + }); + }); + + it('does not ack when answer submit fails', async () => { + const manager = { + answerQuestion: vi.fn().mockRejectedValue(new Error('network')), + rejectQuestion: vi.fn(), + respondToPermission: vi.fn(), + }; + const { render } = renderInteractionHandlers({ + manager: manager as never, + kiloSessionId: 'kilo-session-1', + activeQuestion: { requestId: 'q1' }, + activePermission: null, + surface: 'remote-session', + }); + + await render().handleAnswerQuestion([['yes']]); + + expect(isAttentionAcked('kilo-session-1', 'R1')).toBe(false); + expect(__peekSessionAttentionForTests('kilo-session-1')).toBeUndefined(); + expect( + shouldShowNeedsInput({ + status: 'question', + raiseId: 'R1', + isAcked: isAttentionAcked('kilo-session-1', 'R1'), + }) + ).toBe(true); + expect(toastError).toHaveBeenCalledWith('Failed to submit answer'); + expect(captureEvent).not.toHaveBeenCalled(); + }); + + it('does not ack when skip submit fails', async () => { + const manager = { + answerQuestion: vi.fn(), + rejectQuestion: vi.fn().mockRejectedValue(new Error('network')), + respondToPermission: vi.fn(), + }; + const { render } = renderInteractionHandlers({ + manager: manager as never, + kiloSessionId: 'kilo-session-1', + activeQuestion: { requestId: 'q1' }, + activePermission: null, + surface: 'remote-session', + }); + + await render().handleRejectQuestion(); + + expect(isAttentionAcked('kilo-session-1', 'R1')).toBe(false); + expect(toastError).toHaveBeenCalledWith('Failed to skip question'); + expect(captureEvent).not.toHaveBeenCalled(); + }); + + it('does not ack when permission submit fails', async () => { + const manager = { + answerQuestion: vi.fn(), + rejectQuestion: vi.fn(), + respondToPermission: vi.fn().mockRejectedValue(new Error('network')), + }; + const { render } = renderInteractionHandlers({ + manager: manager as never, + kiloSessionId: 'kilo-session-1', + activeQuestion: null, + activePermission: { requestId: 'p1' }, + surface: 'remote-session', + }); + + await render().handleRespondToPermission('reject'); + + expect(isAttentionAcked('kilo-session-1', 'R1')).toBe(false); + expect(toastError).toHaveBeenCalledWith('Failed to respond to permission request'); + expect(captureEvent).not.toHaveBeenCalled(); + }); + + it('does not ack a different session id than the one supplied', async () => { + const manager = { + answerQuestion: vi.fn().mockResolvedValue(undefined), + rejectQuestion: vi.fn(), + respondToPermission: vi.fn(), + }; + const { render } = renderInteractionHandlers({ + manager: manager as never, + kiloSessionId: 'session-a', + activeQuestion: { requestId: 'q1' }, + activePermission: null, + surface: 'remote-session', + }); + + await render().handleAnswerQuestion([['ok']]); + + expect(isAttentionAcked('session-a', 'R1')).toBe(true); + expect(isAttentionAcked('session-b', 'R1')).toBe(false); + }); +}); diff --git a/apps/mobile/src/components/agents/use-interaction-handlers.ts b/apps/mobile/src/components/agents/use-interaction-handlers.ts index d633fc71d0..b7e1171abe 100644 --- a/apps/mobile/src/components/agents/use-interaction-handlers.ts +++ b/apps/mobile/src/components/agents/use-interaction-handlers.ts @@ -7,11 +7,17 @@ import { PERMISSION_RESPONDED_EVENT, QUESTION_ANSWERED_EVENT, } from '@/lib/analytics/posthog'; +import { ackSessionAttention } from '@/lib/session-attention'; import { type useSessionManager } from './session-provider'; type InteractionHandlersArgs = { manager: ReturnType; + /** + * Kilo session id (route `session-id` / list `session.id`). Keys the + * attention ack store. Must not be the cloud-agent session id. + */ + kiloSessionId: string; activeQuestion: { requestId: string; questions?: unknown[] } | null; activePermission: { requestId: string } | null; surface: AnalyticsSurface; @@ -19,6 +25,7 @@ type InteractionHandlersArgs = { export function useInteractionHandlers({ manager, + kiloSessionId, activeQuestion, activePermission, surface, @@ -34,6 +41,7 @@ export function useInteractionHandlers({ setIsAnswering(true); try { await manager.answerQuestion(activeQuestion.requestId, answers); + ackSessionAttention(kiloSessionId); captureEvent(QUESTION_ANSWERED_EVENT, { surface, skipped: false }); } catch { toast.error('Failed to submit answer'); @@ -41,7 +49,7 @@ export function useInteractionHandlers({ setIsAnswering(false); } }, - [manager, activeQuestion, surface] + [manager, kiloSessionId, activeQuestion, surface] ); const handleRejectQuestion = useCallback(async () => { @@ -51,13 +59,14 @@ export function useInteractionHandlers({ setIsAnswering(true); try { await manager.rejectQuestion(activeQuestion.requestId); + ackSessionAttention(kiloSessionId); captureEvent(QUESTION_ANSWERED_EVENT, { surface, skipped: true }); } catch { toast.error('Failed to skip question'); } finally { setIsAnswering(false); } - }, [manager, activeQuestion, surface]); + }, [manager, kiloSessionId, activeQuestion, surface]); const handleRespondToPermission = useCallback( async (response: 'once' | 'always' | 'reject') => { @@ -67,6 +76,7 @@ export function useInteractionHandlers({ setIsRespondingToPermission(true); try { await manager.respondToPermission(activePermission.requestId, response); + ackSessionAttention(kiloSessionId); captureEvent(PERMISSION_RESPONDED_EVENT, { surface, response }); } catch { toast.error('Failed to respond to permission request'); @@ -74,7 +84,7 @@ export function useInteractionHandlers({ setIsRespondingToPermission(false); } }, - [manager, activePermission, surface] + [manager, kiloSessionId, activePermission, surface] ); return { diff --git a/apps/mobile/src/components/agents/use-session-config-sync.test.ts b/apps/mobile/src/components/agents/use-session-config-sync.test.ts index 0dd5ce3ad9..69f1d73454 100644 --- a/apps/mobile/src/components/agents/use-session-config-sync.test.ts +++ b/apps/mobile/src/components/agents/use-session-config-sync.test.ts @@ -43,4 +43,46 @@ describe('resolveSessionConfigSelection', () => { }) ).toEqual({ model: 'gateway/first', variant: 'high' }); }); + + it('prefers the cloud-agent model override over stored session config', () => { + expect( + resolveSessionConfigSelection({ + activeSessionType: 'cloud-agent', + fetchedData: { model: 'stored/from-fetch', variant: 'low' }, + sessionConfig: { model: 'stored/from-session', variant: 'medium' }, + modelOptions: gatewayModels, + selectedModel: '', + selectedVariant: '', + cloudAgentModelOverride: { model: 'user/picked', variant: 'high' }, + }) + ).toEqual({ model: 'user/picked', variant: 'high' }); + }); + + it('falls back to stored session model when there is no cloud-agent override', () => { + expect( + resolveSessionConfigSelection({ + activeSessionType: 'cloud-agent', + fetchedData: { model: 'stored/from-fetch', variant: 'low' }, + sessionConfig: { model: 'stored/from-session', variant: 'medium' }, + modelOptions: gatewayModels, + selectedModel: '', + selectedVariant: '', + cloudAgentModelOverride: null, + }) + ).toEqual({ model: 'stored/from-session', variant: 'medium' }); + }); + + it('ignores cloud-agent override on remote sessions', () => { + expect( + resolveSessionConfigSelection({ + activeSessionType: 'remote', + fetchedData: {}, + sessionConfig: null, + modelOptions: gatewayModels, + selectedModel: 'remote/selected', + selectedVariant: 'max', + cloudAgentModelOverride: { model: 'should/not/win', variant: 'high' }, + }) + ).toEqual({ model: 'remote/selected', variant: 'max' }); + }); }); diff --git a/apps/mobile/src/components/agents/use-session-config-sync.ts b/apps/mobile/src/components/agents/use-session-config-sync.ts index 7172e16e97..f8b2070bd9 100644 --- a/apps/mobile/src/components/agents/use-session-config-sync.ts +++ b/apps/mobile/src/components/agents/use-session-config-sync.ts @@ -11,6 +11,11 @@ type SessionConfigSnapshot = { variant?: string | null; }; +type CloudAgentModelOverrideSnapshot = { + model: string; + variant?: string; +} | null; + type ResolveSessionConfigSelectionOptions = { activeSessionType: ResolvedSession['type'] | null; fetchedData: SessionConfigSnapshot | null; @@ -18,6 +23,8 @@ type ResolveSessionConfigSelectionOptions = { modelOptions: SessionModelOption[]; selectedModel: string; selectedVariant: string; + /** Manager-held cloud-agent pick; wins over session/fetched config on non-remote. */ + cloudAgentModelOverride?: CloudAgentModelOverrideSnapshot; }; type UseSessionConfigSyncOptions = ResolveSessionConfigSelectionOptions; @@ -38,11 +45,21 @@ export function resolveSessionConfigSelection({ modelOptions, selectedModel, selectedVariant, + cloudAgentModelOverride = null, }: ResolveSessionConfigSelectionOptions): { model: string; variant: string } { if (activeSessionType === 'remote') { return { model: selectedModel, variant: selectedVariant }; } + // Cloud-agent in-session override must beat stored session config so the + // sync effect cannot revert a user pick before send. + if (cloudAgentModelOverride?.model) { + return { + model: cloudAgentModelOverride.model, + variant: cloudAgentModelOverride.variant ?? '', + }; + } + const configuredModel = sessionConfig?.model ?? fetchedData?.model ?? ''; if (configuredModel) { return { @@ -68,6 +85,7 @@ export function useSessionConfigSync({ modelOptions, selectedModel, selectedVariant, + cloudAgentModelOverride = null, }: UseSessionConfigSyncOptions): UseSessionConfigSyncResult { const initialSelection = resolveSessionConfigSelection({ activeSessionType, @@ -76,6 +94,7 @@ export function useSessionConfigSync({ modelOptions, selectedModel, selectedVariant, + cloudAgentModelOverride, }); const [currentMode, setCurrentMode] = useState(() => normalizeAgentMode(fetchedData?.mode) @@ -98,12 +117,14 @@ export function useSessionConfigSync({ modelOptions, selectedModel, selectedVariant, + cloudAgentModelOverride, }); const isAutoSelectingFirstModel = activeSessionType === 'cloud-agent' && fetchedData !== null && !sessionConfig?.model && !fetchedData.model && + !cloudAgentModelOverride?.model && selection.model === modelOptions[0]?.id; if (isAutoSelectingFirstModel && currentModel) { return; @@ -117,6 +138,7 @@ export function useSessionConfigSync({ modelOptions, selectedModel, selectedVariant, + cloudAgentModelOverride, currentModel, ]); diff --git a/apps/mobile/src/lib/active-sessions-live-sync.attention.test.ts b/apps/mobile/src/lib/active-sessions-live-sync.attention.test.ts new file mode 100644 index 0000000000..859f6b8001 --- /dev/null +++ b/apps/mobile/src/lib/active-sessions-live-sync.attention.test.ts @@ -0,0 +1,144 @@ +import { describe, expect, it, type Mock } from 'vitest'; + +import { + ActiveSessionsLiveSync, + makeCached, + makeConnection, + makeFakeQueryClient, + makeQueryFn, + QUERY_KEY, + setupTimers, +} from '@/lib/active-sessions-live-sync.test-helpers'; + +setupTimers(); + +describe('ActiveSessionsLiveSync — session.status.updated', () => { + it('applies the lightweight sessionId payload and clears attention', async () => { + const conn = makeConnection(); + const qc = makeFakeQueryClient(); + qc.__setCached({ + sessions: [makeCached({ id: 'ses-1', status: 'question', connectionId: 'c1' })], + }); + const sync = new ActiveSessionsLiveSync({ + connection: conn, + queryClient: qc, + queryKey: QUERY_KEY, + queryFn: makeQueryFn(), + }); + sync.attach(); + conn.__fireSystem({ + event: 'session.status.updated', + data: { + source: 'v2', + sessionId: 'ses-1', + previousStatus: 'question', + status: 'idle', + statusUpdatedAt: 'now', + changedAt: 'now', + }, + }); + await sync.getWriteQueue(); + expect(qc.__getCached()?.sessions[0]?.status).toBe('idle'); + }); + + it('applies the full session-row payload into attention', async () => { + const conn = makeConnection(); + const qc = makeFakeQueryClient(); + qc.__setCached({ + sessions: [makeCached({ id: 'ses-2', status: 'busy', connectionId: 'c1' })], + }); + const sync = new ActiveSessionsLiveSync({ + connection: conn, + queryClient: qc, + queryKey: QUERY_KEY, + queryFn: makeQueryFn(), + }); + sync.attach(); + conn.__fireSystem({ + event: 'session.status.updated', + data: { + source: 'v2', + session: { + source: 'v2', + sessionId: 'ses-2', + createdAt: 'now', + updatedAt: 'now', + title: 't', + createdOnPlatform: null, + organizationId: null, + gitUrl: null, + gitBranch: null, + parentSessionId: null, + status: 'permission', + statusUpdatedAt: 'now', + }, + previousStatus: 'busy', + status: 'permission', + statusUpdatedAt: 'now', + changedAt: 'now', + }, + }); + await sync.getWriteQueue(); + expect(qc.__getCached()?.sessions[0]?.status).toBe('permission'); + }); + + it('ignores a malformed status payload', async () => { + const conn = makeConnection(); + const qc = makeFakeQueryClient(); + qc.__setCached({ + sessions: [makeCached({ id: 'ses-1', status: 'question' })], + }); + const setQueryDataCalls = qc.setQueryData as Mock; + setQueryDataCalls.mockClear(); + const sync = new ActiveSessionsLiveSync({ + connection: conn, + queryClient: qc, + queryKey: QUERY_KEY, + queryFn: makeQueryFn(), + }); + sync.attach(); + conn.__fireSystem({ + event: 'session.status.updated', + data: { sessionId: 'ses-1' }, + }); + await Promise.resolve(); + expect(setQueryDataCalls).not.toHaveBeenCalled(); + expect(qc.__getCached()?.sessions[0]?.status).toBe('question'); + }); + + it('keeps attention sticky across a busy heartbeat after status.updated sets it', async () => { + const conn = makeConnection(); + const qc = makeFakeQueryClient(); + qc.__setCached({ + sessions: [makeCached({ id: 'ses-1', status: 'busy', connectionId: 'c1', title: 'A' })], + }); + const sync = new ActiveSessionsLiveSync({ + connection: conn, + queryClient: qc, + queryKey: QUERY_KEY, + queryFn: makeQueryFn(), + }); + sync.attach(); + conn.__fireSystem({ + event: 'session.status.updated', + data: { + source: 'v2', + sessionId: 'ses-1', + previousStatus: 'busy', + status: 'question', + statusUpdatedAt: 'now', + changedAt: 'now', + }, + }); + await sync.getWriteQueue(); + conn.__fireSystem({ + event: 'sessions.heartbeat', + data: { + connectionId: 'c1', + sessions: [{ id: 'ses-1', status: 'busy', title: 'A' }], + }, + }); + await sync.getWriteQueue(); + expect(qc.__getCached()?.sessions[0]?.status).toBe('question'); + }); +}); diff --git a/apps/mobile/src/lib/active-sessions-live-sync.ts b/apps/mobile/src/lib/active-sessions-live-sync.ts index d235d40ca6..85861d3e9a 100644 --- a/apps/mobile/src/lib/active-sessions-live-sync.ts +++ b/apps/mobile/src/lib/active-sessions-live-sync.ts @@ -1,23 +1,9 @@ /** - * App-level owner for the active-sessions live-sync. The owner: - * - * - retains the shared `UserWebConnection` while mounted; - * - subscribes to `onSystemEvent` and applies the WS payloads to the - * shared `trpc.activeSessions.list` cache through ONE serialized - * pipeline (`cancelQueries` + `setQueryData`); the pipeline never - * awaits a network fetch, so a stalled tRPC refetch cannot block - * later heartbeats; - * - requests explicit refreshes (`cli.connected`, `cli.disconnected`, - * reconnect, enrichment) through a coalescing `scheduleRefresh` with - * durable per-reason pending state. Refreshes call - * `queryClient.fetchQuery({ queryKey, queryFn, staleTime: 0 })` so - * the network call is forced even after a preceding `setQueryData`. - * - observes the connection-state rising edge and triggers exactly - * one reconnect refresh per disconnect → connect transition. - * - * This module is framework-agnostic: it does not import React or the - * `UserWebConnectionProvider`. The thin React glue that wires it into the - * provider lives in `active-sessions-live-sync-mount.tsx`. + * App-level owner for active-sessions live-sync: retains UserWebConnection, + * applies onSystemEvent payloads to trpc.activeSessions.list via a serialized + * cancelQueries+setQueryData pipeline, and coalesces refreshes (cli.connected / + * disconnected, reconnect, enrichment) with fetchQuery staleTime:0. + * Framework-agnostic; React glue is active-sessions-live-sync-mount.tsx. */ import { type QueryClient, type QueryFunction, type QueryKey } from '@tanstack/react-query'; @@ -26,13 +12,7 @@ import { type CachedActiveSession, type CachedActiveSessionsData, hasUnenrichedLiveId, - mergeHeartbeatForActiveSessions, - mergeSnapshotForActiveSessions, - parseCliConnectionPayload, - parseHeartbeatPayload, - parseSessionsListPayload, - removeActiveSessionsForConnection, - selectRootWsSessions, + planLiveSystemEventActions, } from './active-sessions-live'; import { type UserWebConnection, type UserWebSystemEvent } from 'cloud-agent-sdk'; @@ -45,11 +25,7 @@ type SystemEvent = UserWebSystemEvent; type WriteUpdater = (current: CachedActiveSession[]) => CachedActiveSession[]; -/** - * Minimal contract this owner needs from the SDK. Mirrors the public - * surface of `UserWebConnection`; the test double in - * `mobile-session-manager.test.ts` already conforms (see S2). - */ +/** Minimal UserWebConnection surface used by this owner. */ export type LiveSyncConnection = Pick< UserWebConnection, 'retain' | 'isConnected' | 'onConnectionChange' | 'onSystemEvent' @@ -68,11 +44,7 @@ type CreateLiveSyncOptions = { now?: () => number; }; -/** - * The owner as a plain class. Exposed so the test suite can exercise - * the serialized pipeline, pending-reason state, and reconnect - * detection without a React renderer. - */ +/** Testable owner: serialized pipeline, pending reasons, reconnect edge. */ export class ActiveSessionsLiveSync { private readonly connection: LiveSyncConnection; private readonly queryClient: LiveSyncQueryClient; @@ -111,12 +83,7 @@ export class ActiveSessionsLiveSync { this.lastConnectedState = this.connection.isConnected(); } - /** - * Subscribes to WS events, retains the connection, and tracks the - * initial connection state for the reconnect-rising-edge detector. - * Returns a detach function that releases all listeners and the - * retain. - */ + /** Subscribe, retain; detach releases listeners + retain. */ attach(): () => void { if (this.releaseRetain) { throw new Error('ActiveSessionsLiveSync already attached'); @@ -188,48 +155,17 @@ export class ActiveSessionsLiveSync { } private handleSystemEvent(event: SystemEvent): void { - if (event.event === 'sessions.list') { - const sessions = parseSessionsListPayload(event.data); - if (sessions) { - const roots = selectRootWsSessions(sessions); - this.enqueueWrite(current => mergeSnapshotForActiveSessions(current, roots)); - } - return; - } - if (event.event === 'sessions.heartbeat') { - const payload = parseHeartbeatPayload(event.data); - if (payload) { - const roots = selectRootWsSessions(payload.sessions); - this.enqueueWrite(current => - mergeHeartbeatForActiveSessions(current, { - connectionId: payload.connectionId, - sessions: roots, - }) - ); + for (const action of planLiveSystemEventActions(event)) { + if (action.type === 'write') { + this.enqueueWrite(current => action.updater(current)); + } else { + this.scheduleRefresh(action.reason); } - return; - } - if (event.event === 'cli.disconnected') { - const payload = parseCliConnectionPayload(event.data); - if (payload) { - this.enqueueWrite(current => - removeActiveSessionsForConnection(current, payload.connectionId) - ); - this.scheduleRefresh('cli-disconnected'); - } - return; - } - if (event.event === 'cli.connected' && parseCliConnectionPayload(event.data)) { - this.scheduleRefresh('cli-connected'); } } private handleConnectionChange(connected: boolean): void { - // Rising-edge detector: a false → true transition triggers exactly - // one reconnect refresh. A true → false transition does not - // schedule a refresh (the conditional `refetchInterval` covers - // the WS-down window; re-scheduling here would either duplicate - // the work or be absorbed by the next refresh regardless). + // Rising edge only; disconnect relies on refetchInterval. if (!this.lastConnectedState && connected) { this.scheduleRefresh('reconnect'); } @@ -241,18 +177,13 @@ export class ActiveSessionsLiveSync { return; } const attachmentEpoch = this.attachmentEpoch; - // Serialize ALL cache writes (cancel + setQueryData) on one queue. - // The pipeline never awaits a network fetch — that is the - // cancel-based fencing model. + // Serialized cancel+setQueryData; never awaits network. this.writeQueue = (async () => { await this.writeQueue; if (attachmentEpoch !== this.attachmentEpoch) { return; } - // A write always cancels the in-flight fetch so the new cache - // state can never be overwritten by a stale result. Record that - // this cancellation was intentional, so the fetch queue can retry - // immediately rather than waiting for the next external trigger. + // Cancel in-flight fetch so stale results cannot overwrite. if (this.isFetchInFlight) { this.inFlightFetchCanceled = true; } @@ -282,15 +213,13 @@ export class ActiveSessionsLiveSync { this.pendingReasons.delete('enrichment'); return; } - if (this.inFlightReasons?.has('enrichment')) { return; } - - if ( + const due = this.lastEnrichmentAttemptAt === null || - this.now() - this.lastEnrichmentAttemptAt >= ENRICHMENT_RETRY_MIN_INTERVAL_MS - ) { + this.now() - this.lastEnrichmentAttemptAt >= ENRICHMENT_RETRY_MIN_INTERVAL_MS; + if (due) { this.pendingReasons.add('enrichment'); } else { this.pendingReasons.delete('enrichment'); @@ -319,12 +248,8 @@ export class ActiveSessionsLiveSync { return; } const attachmentEpoch = this.attachmentEpoch; - // Cancel any in-flight fetch so a newly scheduled refresh can start - // immediately instead of being queued behind a stale one. + // Cancel in-flight fetch so a new refresh starts immediately. if (this.isFetchInFlight) { - // Record that this cancellation was intentional, so the fetch - // queue can retry immediately rather than waiting for the next - // external trigger. this.inFlightFetchCanceled = true; void this.queryClient.cancelQueries({ queryKey: this.queryKey }); } @@ -342,8 +267,6 @@ export class ActiveSessionsLiveSync { return; } this.isFetchInFlight = true; - // Reset the cancellation flag for THIS fetch instance. Any cancel - // that targets this fetch will set it back to true before the catch. this.inFlightFetchCanceled = false; const inFlightReasons = new Set(this.pendingReasons); this.inFlightReasons = inFlightReasons; @@ -351,24 +274,18 @@ export class ActiveSessionsLiveSync { try { await this.queryClient.cancelQueries({ queryKey: this.queryKey }); if (attachmentEpoch === this.attachmentEpoch) { - // fetchQuery with staleTime: 0 forces a network call regardless - // of any preceding setQueryData. + // staleTime:0 forces a network call after setQueryData. const fetchPromise = this.queryClient.fetchQuery({ queryKey: this.queryKey, queryFn: this.queryFn, staleTime: 0, }); - // The fetch is now in flight: let tests waiting on getFetchQueue() - // observe the pending state and/or issue cancellations. this.notifyFetchStart(); await fetchPromise; success = true; } } catch { - // Either the query was canceled by a later WS write, or the - // network call itself failed. In either case, the reasons stay - // pending so a future scheduleRefresh (or the trailing - // re-kick) re-attempts. + // Canceled or network failure: keep reasons pending for retry. } finally { this.isFetchInFlight = false; } @@ -387,16 +304,11 @@ export class ActiveSessionsLiveSync { this.updateEnrichmentReason(); } const hasNewReasons = [...this.pendingReasons].some(reason => !inFlightReasons.has(reason)); - // Read via helper so control-flow analysis does not treat the field as - // stuck at the `false` written above — other methods flip it during await. + // Helper so CFA does not treat the field as stuck at false after await. const wasCanceled = this.readInFlightFetchCanceled(); this.inFlightReasons = null; this.notifyFetchCompletion(); - // Re-kick immediately only when a replacement fetch is genuinely - // warranted: either the in-flight fetch was intentionally canceled by - // newer work, or new reasons were raised while it was in flight. On a - // genuine failure without either condition, stay quiet and let the - // next scheduled trigger (WS event / fallback poll) retry. + // Re-kick only after intentional cancel or new reasons mid-flight. if (this.pendingReasons.size > 0 && (wasCanceled || hasNewReasons)) { this.kickFetch(); } diff --git a/apps/mobile/src/lib/active-sessions-live.attention.test.ts b/apps/mobile/src/lib/active-sessions-live.attention.test.ts new file mode 100644 index 0000000000..39e9b876ab --- /dev/null +++ b/apps/mobile/src/lib/active-sessions-live.attention.test.ts @@ -0,0 +1,106 @@ +import { describe, expect, it } from 'vitest'; + +import { + applySessionStatusUpdated, + type CachedActiveSession, + effectiveStatus, + isAttentionStatus, + mergeHeartbeatForActiveSessions, + mergeSnapshotForActiveSessions, +} from '@/lib/active-sessions-live'; + +function makeCached(over: Partial = {}): CachedActiveSession { + return { + id: 'a1', + status: 'running', + title: 'test', + connectionId: 'c1', + ...over, + }; +} + +describe('isAttentionStatus / effectiveStatus', () => { + it('treats question and permission as attention', () => { + expect(isAttentionStatus('question')).toBe(true); + expect(isAttentionStatus('permission')).toBe(true); + expect(isAttentionStatus('busy')).toBe(false); + expect(isAttentionStatus('idle')).toBe(false); + expect(isAttentionStatus(null)).toBe(false); + }); + + it('prefers stored attention over live idle/busy', () => { + expect(effectiveStatus('busy', 'question')).toBe('question'); + expect(effectiveStatus('idle', 'permission')).toBe('permission'); + }); + + it('yields to live when stored is not attention', () => { + expect(effectiveStatus('busy', 'idle')).toBe('busy'); + expect(effectiveStatus('idle', null)).toBe('idle'); + expect(effectiveStatus('busy', undefined)).toBe('busy'); + }); +}); + +describe('sticky attention on snapshot / heartbeat merge', () => { + it('does not clear held attention when the snapshot reports busy', () => { + const current = [makeCached({ id: 'a', status: 'question', connectionId: 'c1' })]; + const snapshot = [{ id: 'a', status: 'busy', title: 'A', connectionId: 'c1' }]; + expect(mergeSnapshotForActiveSessions(current, snapshot)[0]?.status).toBe('question'); + }); + + it('does not clear held permission when the snapshot reports idle', () => { + const current = [makeCached({ id: 'a', status: 'permission', connectionId: 'c1' })]; + const snapshot = [{ id: 'a', status: 'idle', title: 'A', connectionId: 'c1' }]; + expect(mergeSnapshotForActiveSessions(current, snapshot)[0]?.status).toBe('permission'); + }); + + it('takes the snapshot status when the cache is not in attention', () => { + const current = [makeCached({ id: 'a', status: 'idle', connectionId: 'c1' })]; + const snapshot = [{ id: 'a', status: 'busy', title: 'A', connectionId: 'c1' }]; + expect(mergeSnapshotForActiveSessions(current, snapshot)[0]?.status).toBe('busy'); + }); + + it('keeps held attention across repeated non-attention heartbeats (no flicker)', () => { + const current = [makeCached({ id: 'a', status: 'question', connectionId: 'c1', title: 'A' })]; + let next = current; + for (const heartbeatStatus of ['busy', 'idle', 'busy'] as const) { + next = mergeHeartbeatForActiveSessions(next, { + connectionId: 'c1', + sessions: [{ id: 'a', status: heartbeatStatus, title: 'A' }], + }); + expect(next[0]?.status).toBe('question'); + } + }); + + it('does not overwrite non-attention rows with sticky logic', () => { + const current = [makeCached({ id: 'a', status: 'idle', connectionId: 'c1' })]; + const result = mergeHeartbeatForActiveSessions(current, { + connectionId: 'c1', + sessions: [{ id: 'a', status: 'busy', title: 'A' }], + }); + expect(result[0]?.status).toBe('busy'); + }); +}); + +describe('applySessionStatusUpdated', () => { + it('applies a transition out of attention', () => { + const current = [makeCached({ id: 'a', status: 'question' })]; + expect(applySessionStatusUpdated(current, 'a', 'idle')[0]?.status).toBe('idle'); + }); + + it('applies a transition into attention', () => { + const current = [makeCached({ id: 'a', status: 'busy' })]; + expect(applySessionStatusUpdated(current, 'a', 'permission')[0]?.status).toBe('permission'); + }); + + it('ignores unknown session ids', () => { + const current = [makeCached({ id: 'a', status: 'question' })]; + const result = applySessionStatusUpdated(current, 'other', 'idle'); + expect(result).toEqual(current); + expect(result[0]?.status).toBe('question'); + }); + + it('accepts an empty status string to clear attention', () => { + const current = [makeCached({ id: 'a', status: 'question' })]; + expect(applySessionStatusUpdated(current, 'a', '')[0]?.status).toBe(''); + }); +}); diff --git a/apps/mobile/src/lib/active-sessions-live.test.ts b/apps/mobile/src/lib/active-sessions-live.test.ts index 00f847c9f4..e6dcdc81fd 100644 --- a/apps/mobile/src/lib/active-sessions-live.test.ts +++ b/apps/mobile/src/lib/active-sessions-live.test.ts @@ -4,6 +4,7 @@ import { parseCliConnectionPayload, parseHeartbeatPayload, parseSessionsListPayload, + parseSessionStatusUpdatedPayload, selectRootWsSessions, } from '@/lib/active-sessions-live'; @@ -91,3 +92,49 @@ describe('parseCliConnectionPayload', () => { expect(parseCliConnectionPayload({})).toBeNull(); }); }); + +describe('parseSessionStatusUpdatedPayload', () => { + it('parses the lightweight sessionId shape', () => { + expect( + parseSessionStatusUpdatedPayload({ + source: 'v2', + sessionId: 'ses-1', + previousStatus: 'question', + status: 'idle', + statusUpdatedAt: 'now', + changedAt: 'now', + }) + ).toEqual({ sessionId: 'ses-1', status: 'idle' }); + }); + + it('parses the full session-row shape', () => { + expect( + parseSessionStatusUpdatedPayload({ + source: 'v2', + session: { + source: 'v2', + sessionId: 'ses-2', + createdAt: 'now', + updatedAt: 'now', + title: 't', + createdOnPlatform: null, + organizationId: null, + gitUrl: null, + gitBranch: null, + parentSessionId: null, + status: 'permission', + statusUpdatedAt: 'now', + }, + previousStatus: 'busy', + status: 'permission', + statusUpdatedAt: 'now', + changedAt: 'now', + }) + ).toEqual({ sessionId: 'ses-2', status: 'permission' }); + }); + + it('rejects a malformed payload', () => { + expect(parseSessionStatusUpdatedPayload({ sessionId: 'ses-1' })).toBeNull(); + expect(parseSessionStatusUpdatedPayload(null)).toBeNull(); + }); +}); diff --git a/apps/mobile/src/lib/active-sessions-live.ts b/apps/mobile/src/lib/active-sessions-live.ts index 043eda556e..fa3fd2e209 100644 --- a/apps/mobile/src/lib/active-sessions-live.ts +++ b/apps/mobile/src/lib/active-sessions-live.ts @@ -8,6 +8,11 @@ * between CLI connections. The functions here never touch React, the * network, or a QueryClient — they are pure and exhaustively unit-tested * alongside this file. + * + * Status resolution for live rows: CLI heartbeats/snapshots often report + * only idle/busy while `cli_sessions_v2` holds question/permission. A + * held attention status is sticky across WS snapshots and heartbeats; + * only an explicit `session.status.updated` (or disconnect) clears it. */ import { @@ -17,15 +22,13 @@ import { heartbeatDataSchema, type SessionsListData, sessionsListDataSchema, + type SessionStatusUpdatedPayload, + sessionStatusUpdatedPayloadSchema, } from 'cloud-agent-sdk/schemas'; import { type ActiveSession } from '@/lib/hooks/use-agent-sessions'; -/** - * Incoming WS session row (per the SDK schemas). Carries - * `parentSessionId` so the root filter can drop subagent sessions; the - * cached `ActiveSession` (from the tRPC router) does not. - */ +/** Incoming WS row; carries `parentSessionId` for the root filter. */ type IncomingWsSession = { id: string; status: string; @@ -36,26 +39,35 @@ type IncomingWsSession = { connectionId?: string; }; -/** - * Cached active session: `ActiveSession` (tRPC output) plus the - * enrichment fields `createdOnPlatform` / `createdAt` / `updatedAt` that - * the live-sync owner preserves across WS updates. - */ +/** Cached active session (tRPC output); enrichment fields preserved across WS. */ export type CachedActiveSession = ActiveSession; export type CachedActiveSessionsData = { sessions: CachedActiveSession[]; }; -/** - * The three enrichment fields that the live-sync owner preserves across - * WS updates. Every other field comes from the latest WS payload (this - * includes `connectionId`, so ownership transfer between CLIs lands - * correctly on the next heartbeat/snapshot). - */ const ENRICHMENT_FIELDS = ['createdOnPlatform', 'createdAt', 'updatedAt'] as const; type EnrichmentField = (typeof ENRICHMENT_FIELDS)[number]; +/** Structured question/permission — the Active Now "NEEDS INPUT" badge. */ +export function isAttentionStatus(status: string | null | undefined): boolean { + return status === 'question' || status === 'permission'; +} + +/** + * Prefer stored attention over live idle/busy so released-CLI heartbeats + * do not clear NEEDS INPUT. Non-attention stored values yield to live. + */ +export function effectiveStatus( + live: string | null | undefined, + stored: string | null | undefined +): string { + if (isAttentionStatus(stored) && stored != null) { + return stored; + } + return live ?? ''; +} + function isRootWsSession(session: IncomingWsSession): boolean { return !session.parentSessionId; } @@ -98,6 +110,30 @@ export function parseCliConnectionPayload(value: unknown): CliConnectionData | n return parsed.data; } +/** + * Dual-shaped `session.status.updated` (full session row vs lightweight + * sessionId). Null payload status becomes `''` for the cache string field. + */ +export function parseSessionStatusUpdatedPayload( + value: unknown +): { sessionId: string; status: string } | null { + const parsed = sessionStatusUpdatedPayloadSchema.safeParse(value); + if (!parsed.success) { + return null; + } + const data: SessionStatusUpdatedPayload = parsed.data; + if ('session' in data) { + return { + sessionId: data.session.sessionId, + status: data.status ?? data.session.status ?? '', + }; + } + return { + sessionId: data.sessionId, + status: data.status ?? '', + }; +} + // ── Enrichment-preserving merge helpers ────────────────────────────── function readEnrichment( @@ -121,7 +157,9 @@ function withEnrichmentAndConnectionId( const enrichment = readEnrichment(current); return { id: row.id, - status: row.status, + // Sticky attention: a non-attention WS status must not clear a held + // question/permission. "stored" for WS paths is the cached row status. + status: effectiveStatus(row.status, current?.status), title: row.title, gitUrl: row.gitUrl, gitBranch: row.gitBranch, @@ -132,9 +170,10 @@ function withEnrichmentAndConnectionId( /** * Replace the entire cache with the snapshot. Rows whose id is in both - * the snapshot and the cache keep ONLY the three enrichment fields from - * the cache; every other field (including `connectionId`) comes from the - * snapshot. Rows absent from the snapshot are dropped. + * the snapshot and the cache keep the three enrichment fields and any + * held attention status from the cache; every other field (including + * `connectionId`) comes from the snapshot. Rows absent from the snapshot + * are dropped. */ export function mergeSnapshotForActiveSessions( current: readonly CachedActiveSession[], @@ -163,6 +202,9 @@ export function mergeSnapshotForActiveSessions( * the payload under a DIFFERENT connectionId — so ownership transfer * between CLIs (same session id, new owner) reflects the new owner on * the next heartbeat without leaving a stale copy under the old one. + * + * A non-attention heartbeat status does not overwrite a currently-held + * attention status (sticky overlay for released CLIs). */ export function mergeHeartbeatForActiveSessions( current: readonly CachedActiveSession[], @@ -194,6 +236,20 @@ export function mergeHeartbeatForActiveSessions( return result; } +/** + * Apply an explicit status transition (including leaving attention). + * Unknown session ids are ignored — live cache only holds active rows. + */ +export function applySessionStatusUpdated( + current: readonly CachedActiveSession[], + sessionId: string, + status: string +): CachedActiveSession[] { + return current.map(row => + row.id === sessionId && row.status !== status ? { ...row, status } : row + ); +} + export function removeActiveSessionsForConnection( current: readonly CachedActiveSession[], connectionId: string @@ -218,3 +274,72 @@ export function isEnriched(row: CachedActiveSession): boolean { export function hasUnenrichedLiveId(rows: readonly CachedActiveSession[]): boolean { return rows.some(row => !isEnriched(row)); } + +/** Actions produced by routing a live-sync system event. */ +type LiveSystemEventAction = + | { type: 'write'; updater: (current: readonly CachedActiveSession[]) => CachedActiveSession[] } + | { type: 'refresh'; reason: 'cli-connected' | 'cli-disconnected' }; + +/** + * Pure routing for ActiveSessionsLiveSync system events. session.status.updated + * is included here so the owner can handle it via onSystemEvent only. + */ +export function planLiveSystemEventActions(event: { + event: string; + data: unknown; +}): LiveSystemEventAction[] { + if (event.event === 'sessions.list') { + const sessions = parseSessionsListPayload(event.data); + if (!sessions) { + return []; + } + const roots = selectRootWsSessions(sessions); + return [{ type: 'write', updater: current => mergeSnapshotForActiveSessions(current, roots) }]; + } + if (event.event === 'sessions.heartbeat') { + const payload = parseHeartbeatPayload(event.data); + if (!payload) { + return []; + } + const roots = selectRootWsSessions(payload.sessions); + return [ + { + type: 'write', + updater: current => + mergeHeartbeatForActiveSessions(current, { + connectionId: payload.connectionId, + sessions: roots, + }), + }, + ]; + } + if (event.event === 'session.status.updated') { + const payload = parseSessionStatusUpdatedPayload(event.data); + if (!payload) { + return []; + } + return [ + { + type: 'write', + updater: current => applySessionStatusUpdated(current, payload.sessionId, payload.status), + }, + ]; + } + if (event.event === 'cli.disconnected') { + const payload = parseCliConnectionPayload(event.data); + if (!payload) { + return []; + } + return [ + { + type: 'write', + updater: current => removeActiveSessionsForConnection(current, payload.connectionId), + }, + { type: 'refresh', reason: 'cli-disconnected' }, + ]; + } + if (event.event === 'cli.connected' && parseCliConnectionPayload(event.data)) { + return [{ type: 'refresh', reason: 'cli-connected' }]; + } + return []; +} diff --git a/apps/mobile/src/lib/pr-review/mutation-error-display.test.ts b/apps/mobile/src/lib/pr-review/mutation-error-display.test.ts index 524ce4e61c..5c32b9b022 100644 --- a/apps/mobile/src/lib/pr-review/mutation-error-display.test.ts +++ b/apps/mobile/src/lib/pr-review/mutation-error-display.test.ts @@ -39,7 +39,7 @@ describe('mutationErrorDisplay', () => { expect(mutationErrorDisplay('submit', classification)).toEqual({ kind: 'bad-request', message: - "This review can't be submitted as is. The PR may have changed, or you can't approve your own pull request.", + "This review can't be submitted as is. The PR may have changed, or you can't review your own pull request.", }); }); diff --git a/apps/mobile/src/lib/pr-review/mutation-error-display.ts b/apps/mobile/src/lib/pr-review/mutation-error-display.ts index f03d10f85b..e887721bde 100644 --- a/apps/mobile/src/lib/pr-review/mutation-error-display.ts +++ b/apps/mobile/src/lib/pr-review/mutation-error-display.ts @@ -19,7 +19,7 @@ type MutationErrorDisplay = { const COMPOSER_BAD_REQUEST = "This comment can't be posted. The selected line may have changed, or the PR may have been updated."; const SUBMIT_BAD_REQUEST = - "This review can't be submitted as is. The PR may have changed, or you can't approve your own pull request."; + "This review can't be submitted as is. The PR may have changed, or you can't review your own pull request."; const COMPOSER_RETRYABLE_FALLBACK = 'Could not post comment.'; const SUBMIT_RETRYABLE = 'Could not submit review. Check your connection and try again.'; const RECONNECT_MESSAGE = 'GitHub connection expired.'; diff --git a/apps/mobile/src/lib/session-attention.test.ts b/apps/mobile/src/lib/session-attention.test.ts index 1147da8acb..f5c170fefb 100644 --- a/apps/mobile/src/lib/session-attention.test.ts +++ b/apps/mobile/src/lib/session-attention.test.ts @@ -95,16 +95,41 @@ describe('ack store state machine', () => { expect(isAttentionAcked('s1', 'R2')).toBe(false); }); - it('re-opening a session with a stale resolved ack re-pends and hides the new raise', () => { + it('a later successful answer re-pends a resolved ack and hides the new raise', () => { ackSessionAttention('s1'); reconcileSessionAttention('s1', 'question', 'R1'); // new raise R2 arrives and is not acked expect(isAttentionAcked('s1', 'R2')).toBe(false); - // user opens the session again → ack overwrites with pending + // user answers again → ack overwrites with pending ackSessionAttention('s1'); expect(isAttentionAcked('s1', 'R2')).toBe(true); }); + it('without an ack, opening alone leaves the badge visible for any raise', () => { + // DEF-4: no on-mount ack — viewing a blocked session must not hide NEEDS INPUT + expect(isAttentionAcked('s1', 'R1')).toBe(false); + expect( + shouldShowNeedsInput({ + status: 'question', + raiseId: 'R1', + isAcked: isAttentionAcked('s1', 'R1'), + }) + ).toBe(true); + }); + + it('acking one session never suppresses a raise on a different session', () => { + ackSessionAttention('s1'); + expect(isAttentionAcked('s1', 'R1')).toBe(true); + expect(isAttentionAcked('s2', 'R1')).toBe(false); + expect( + shouldShowNeedsInput({ + status: 'question', + raiseId: 'R1', + isAcked: isAttentionAcked('s2', 'R1'), + }) + ).toBe(true); + }); + it('reconcile deletes the entry on a non-attention status, so the next raise shows', () => { ackSessionAttention('s1'); reconcileSessionAttention('s1', 'question', 'R1'); @@ -133,8 +158,8 @@ describe('ack store state machine', () => { expect(isAttentionAcked('s1', 'question')).toBe(false); }); - it('frozen-return sequence: pending entry resolves via reconcile, then next raise is not absorbed', () => { - // raise R1 observed, then user opens the session (ack → pending) + it('answer-then-clear sequence: pending entry resolves via reconcile, then next raise is not absorbed', () => { + // user answers (ack → pending), then status still reports the same raise ackSessionAttention('s1'); // reconcile resolves the pending entry to R1 reconcileSessionAttention('s1', 'question', 'R1'); @@ -217,7 +242,7 @@ describe('revision snapshot and listener notification', () => { const listener = vi.fn<() => void>(); const unsubscribe = subscribe(listener); - // Repeated open of the same still-pending session is a no-op. + // Repeated successful answers while still pending is a no-op. ackSessionAttention('s1'); ackSessionAttention('s1'); diff --git a/apps/mobile/src/lib/session-attention.ts b/apps/mobile/src/lib/session-attention.ts index c686aa1555..5608b9223f 100644 --- a/apps/mobile/src/lib/session-attention.ts +++ b/apps/mobile/src/lib/session-attention.ts @@ -1,13 +1,15 @@ -import { useEffect, useSyncExternalStore } from 'react'; +import { useSyncExternalStore } from 'react'; /** * Pure session-attention derivation + in-memory ack store for the mobile * Agents session list "needs input" indicator. * - * The detail screen is the only ack writer. Acks are intentionally NOT - * persisted across app restarts. Raise identity is `statusUpdatedAt ?? status` - * (stored rows carry server `status_updated_at`; remote active-only rows - * carry none so identity degrades to the status string). + * Acks are written only when the user successfully answers, skips, or + * responds to a permission — never on merely opening the detail screen. + * Acks are intentionally NOT persisted across app restarts. Raise identity + * is `statusUpdatedAt ?? status` (stored rows carry server + * `status_updated_at`; remote active-only rows carry none so identity + * degrades to the status string). * * No backend, tRPC, or shared-package imports: this is a mobile-local * module so the web client can keep its own copy. @@ -67,10 +69,15 @@ function getServerSnapshot(): number { return 0; } +/** + * Hide the needs-input badge immediately after a successful answer / skip / + * permission response, before the server status round-trip lands. + * `reconcileSessionAttention` clears the entry once status leaves attention. + * + * If the entry is already pending, nothing changes — skip the bump so we + * don't fire a redundant global re-render. + */ export function ackSessionAttention(sessionId: string): void { - // Opening a session always leaves the entry pending. If it is already - // pending, nothing changes — skip the bump so we don't fire a redundant - // global re-render (e.g. React Strict Mode's double effect invocation). if (store.entries.get(sessionId)?.raiseId === null) { return; } @@ -135,16 +142,6 @@ export function useSessionAttentionRevision(): number { return useSyncExternalStore(subscribe, getRevisionSnapshot, getServerSnapshot); } -/** - * Ack a session's attention indicator when the detail screen opens. - * Re-runs if `sessionId` changes (e.g. switching sessions). - */ -export function useAckSessionAttentionOnOpen(sessionId: string): void { - useEffect(() => { - ackSessionAttention(sessionId); - }, [sessionId]); -} - /** * Test-only: clear all acks and reset the revision counter so each * test starts from a known state. Not for production use. diff --git a/apps/web/src/lib/cloud-agent-sdk/index.ts b/apps/web/src/lib/cloud-agent-sdk/index.ts index 58b97dd009..8209604ef0 100644 --- a/apps/web/src/lib/cloud-agent-sdk/index.ts +++ b/apps/web/src/lib/cloud-agent-sdk/index.ts @@ -3,6 +3,7 @@ export { createSessionManager } from './session-manager'; export { CLI_MODEL_ID, cliModelLabel } from './cli-model'; export type { ActiveSessionType, + CloudAgentModelOverride, SessionManager, SessionManagerConfig, SessionManagerAtoms, diff --git a/apps/web/src/lib/cloud-agent-sdk/session-manager.test.ts b/apps/web/src/lib/cloud-agent-sdk/session-manager.test.ts index d8fdcd52c7..9554ba3220 100644 --- a/apps/web/src/lib/cloud-agent-sdk/session-manager.test.ts +++ b/apps/web/src/lib/cloud-agent-sdk/session-manager.test.ts @@ -1485,6 +1485,46 @@ describe('createSessionManager', () => { }); }); + it('uses cloud-agent model override on send and clears it on switchSession', async () => { + const config = createMockConfig(); + const mgr = createSessionManager(config); + + // Mock connect auto-resolves as cloud-agent. + await mgr.switchSession(kiloId('ses-1')); + mgr.setCloudAgentModelOverride({ model: 'openai/gpt-5', variant: 'high' }); + expect(atomValue(config.store, mgr.atoms.cloudAgentModelOverride)).toEqual({ + model: 'openai/gpt-5', + variant: 'high', + }); + + mockSession.send.mockResolvedValue(undefined); + await mgr.send({ + payload: { + type: 'prompt', + prompt: 'use override', + mode: 'code', + // Stale composer payload must not win over the manager override. + model: 'stale/composer-model', + variant: 'stale', + }, + }); + + expect(mockSession.send).toHaveBeenLastCalledWith({ + messageId: expect.stringMatching(/^msg_/), + payload: { + type: 'prompt', + prompt: 'use override', + mode: 'code', + model: { providerID: 'kilo', modelID: 'openai/gpt-5' }, + variant: 'high', + }, + images: undefined, + }); + + await mgr.switchSession(kiloId('ses-2')); + expect(atomValue(config.store, mgr.atoms.cloudAgentModelOverride)).toBeNull(); + }); + it('sends only the explicit remote override and omits stale session model fields after clear', async () => { const config = createMockConfig(); const mgr = createSessionManager(config); diff --git a/apps/web/src/lib/cloud-agent-sdk/session-manager.ts b/apps/web/src/lib/cloud-agent-sdk/session-manager.ts index 6243164675..e23aa6e6c7 100644 --- a/apps/web/src/lib/cloud-agent-sdk/session-manager.ts +++ b/apps/web/src/lib/cloud-agent-sdk/session-manager.ts @@ -65,6 +65,11 @@ import { CLI_MODEL_ID, cliModelLabel } from './cli-model'; type StoredMessage = { info: MessageInfo; parts: Part[] }; type SessionManagerPromptPayload = Omit & { model?: string }; type SessionManagerSendPayload = SessionManagerPromptPayload | SendCommandPayload; +/** In-session cloud-agent model pick. Separate from remoteModelOverride — no remote clear rules. */ +type CloudAgentModelOverride = { + model: string; + variant?: string; +}; type SessionStatusIndicator = { type: 'error' | 'warning' | 'info' | 'progress'; message: string; @@ -217,6 +222,8 @@ type SessionManagerAtoms = { remoteCommandState: W; observedModel: W; remoteModelOverride: W; + /** Session-scoped cloud-agent model pick; cleared on switchSession. Not remote. */ + cloudAgentModelOverride: W; canSend: W; canInterrupt: W; statusIndicator: W; @@ -286,6 +293,7 @@ type SessionManager = { attachmentParts?: RemoteAttachmentPart[]; }): Promise; setRemoteModelOverride(override: RemoteModelOverride | null): void; + setCloudAgentModelOverride(override: CloudAgentModelOverride | null): void; retryRemoteModels(): void; retryRemoteCommands(): void; createRemoteSession(): Promise; @@ -425,6 +433,7 @@ function createSessionManager(config: SessionManagerConfig): SessionManager { const remoteCommandStateAtom = atom(EMPTY_REMOTE_COMMAND_STATE); const observedModelAtom = atom(null); const remoteModelOverrideAtom = atom(null); + const cloudAgentModelOverrideAtom = atom(null); const canSendAtom = atom(false); const canInterruptAtom = atom(false); const statusIndicatorAtom = atom(null); @@ -574,6 +583,7 @@ function createSessionManager(config: SessionManagerConfig): SessionManager { observedModelSource = null; remoteHistoryReplaying = true; store.set(remoteModelOverrideAtom, null); + store.set(cloudAgentModelOverrideAtom, null); store.set(canSendAtom, false); store.set(canInterruptAtom, false); store.set(statusIndicatorAtom, null); @@ -1296,6 +1306,7 @@ function createSessionManager(config: SessionManagerConfig): SessionManager { ? `/${input.payload.command}${input.payload.arguments ? ` ${input.payload.arguments}` : ''}` : input.payload.prompt; const remoteModelOverride = store.get(remoteModelOverrideAtom); + const cloudAgentModelOverride = store.get(cloudAgentModelOverrideAtom); let transportPayload: TransportSendPayload; if (input.payload.type === 'command') { transportPayload = input.payload; @@ -1314,14 +1325,18 @@ function createSessionManager(config: SessionManagerConfig): SessionManager { : {}), }; } else { + // Prefer the in-session cloud-agent override over the payload so a stale + // composer model cannot bypass the manager's single source of truth. + const cloudModel = cloudAgentModelOverride?.model ?? input.payload.model; + const cloudVariant = cloudAgentModelOverride + ? cloudAgentModelOverride.variant + : input.payload.variant; transportPayload = { type: 'prompt', prompt: input.payload.prompt, ...(input.payload.mode ? { mode: input.payload.mode } : {}), - ...(input.payload.model - ? { model: { providerID: 'kilo', modelID: input.payload.model } } - : {}), - ...(input.payload.model && input.payload.variant ? { variant: input.payload.variant } : {}), + ...(cloudModel ? { model: { providerID: 'kilo', modelID: cloudModel } } : {}), + ...(cloudModel && cloudVariant ? { variant: cloudVariant } : {}), }; } @@ -1440,6 +1455,10 @@ function createSessionManager(config: SessionManagerConfig): SessionManager { store.set(remoteModelOverrideAtom, override); } + function setCloudAgentModelOverride(override: CloudAgentModelOverride | null): void { + store.set(cloudAgentModelOverrideAtom, override); + } + function retryRemoteModels(): void { currentSession?.retryRemoteModels(); } @@ -1485,6 +1504,7 @@ function createSessionManager(config: SessionManagerConfig): SessionManager { loadOlderMessages, send, setRemoteModelOverride, + setCloudAgentModelOverride, retryRemoteModels, retryRemoteCommands, createRemoteSession, @@ -1511,6 +1531,7 @@ function createSessionManager(config: SessionManagerConfig): SessionManager { remoteCommandState: remoteCommandStateAtom, observedModel: observedModelAtom, remoteModelOverride: remoteModelOverrideAtom, + cloudAgentModelOverride: cloudAgentModelOverrideAtom, canSend: canSendAtom, canInterrupt: canInterruptAtom, statusIndicator: statusIndicatorAtom, @@ -1553,6 +1574,7 @@ function createSessionManager(config: SessionManagerConfig): SessionManager { export { CLI_MODEL_ID, cliModelLabel, createSessionManager, formatError }; export type { ActiveSessionType, + CloudAgentModelOverride, SessionManager, SessionManagerConfig, SessionManagerAtoms, diff --git a/apps/web/src/routers/active-sessions-router.list.test.ts b/apps/web/src/routers/active-sessions-router.list.test.ts index 54691d5581..f98b8587d9 100644 --- a/apps/web/src/routers/active-sessions-router.list.test.ts +++ b/apps/web/src/routers/active-sessions-router.list.test.ts @@ -101,6 +101,92 @@ describe('active-sessions-router.list', () => { } }); + it('overlays stored question status over a live busy heartbeat status', async () => { + const sessionId = 'ses_active_attention_question_1234'; + await db.insert(cli_sessions_v2).values({ + session_id: sessionId, + kilo_user_id: regularUser.id, + created_on_platform: 'cli', + status: 'question', + }); + + fetchSpy = mockWorkerSessions([ + { + id: sessionId, + status: 'busy', + title: 'needs input', + connectionId: 'conn-attn', + }, + ]); + + try { + const caller = await createCallerForUser(regularUser.id); + const result = await caller.activeSessions.list(); + + expect(result.sessions).toHaveLength(1); + expect(result.sessions[0]?.status).toBe('question'); + expect(result.sessions[0]?.title).toBe('needs input'); + } finally { + await db.delete(cli_sessions_v2).where(eq(cli_sessions_v2.session_id, sessionId)); + } + }); + + it('overlays stored permission status over a live idle heartbeat status', async () => { + const sessionId = 'ses_active_attention_permission_1234'; + await db.insert(cli_sessions_v2).values({ + session_id: sessionId, + kilo_user_id: regularUser.id, + created_on_platform: 'cli', + status: 'permission', + }); + + fetchSpy = mockWorkerSessions([ + { + id: sessionId, + status: 'idle', + title: 'needs permission', + connectionId: 'conn-perm', + }, + ]); + + try { + const caller = await createCallerForUser(regularUser.id); + const result = await caller.activeSessions.list(); + + expect(result.sessions[0]?.status).toBe('permission'); + } finally { + await db.delete(cli_sessions_v2).where(eq(cli_sessions_v2.session_id, sessionId)); + } + }); + + it('keeps the live status when the stored DB status is not attention', async () => { + const sessionId = 'ses_active_attention_non_attn_1234'; + await db.insert(cli_sessions_v2).values({ + session_id: sessionId, + kilo_user_id: regularUser.id, + created_on_platform: 'cli', + status: 'idle', + }); + + fetchSpy = mockWorkerSessions([ + { + id: sessionId, + status: 'busy', + title: 'working', + connectionId: 'conn-busy', + }, + ]); + + try { + const caller = await createCallerForUser(regularUser.id); + const result = await caller.activeSessions.list(); + + expect(result.sessions[0]?.status).toBe('busy'); + } finally { + await db.delete(cli_sessions_v2).where(eq(cli_sessions_v2.session_id, sessionId)); + } + }); + it('passes sessions with undefined enrichment fields when no matching row exists', async () => { const unmatchedId = 'ses_active_enrich_unmatched_1234'; diff --git a/apps/web/src/routers/active-sessions-router.schema.test.ts b/apps/web/src/routers/active-sessions-router.schema.test.ts index c14abf8cda..e27d67205e 100644 --- a/apps/web/src/routers/active-sessions-router.schema.test.ts +++ b/apps/web/src/routers/active-sessions-router.schema.test.ts @@ -1,5 +1,18 @@ import { describe, it, expect } from '@jest/globals'; -import { activeSessionSchema } from './active-sessions-router'; +import { activeSessionSchema, resolveActiveSessionStatus } from './active-sessions-router'; + +describe('resolveActiveSessionStatus', () => { + it('prefers stored question/permission over live', () => { + expect(resolveActiveSessionStatus('busy', 'question')).toBe('question'); + expect(resolveActiveSessionStatus('idle', 'permission')).toBe('permission'); + }); + + it('keeps live when stored is not attention', () => { + expect(resolveActiveSessionStatus('busy', 'idle')).toBe('busy'); + expect(resolveActiveSessionStatus('idle', null)).toBe('idle'); + expect(resolveActiveSessionStatus('busy', undefined)).toBe('busy'); + }); +}); describe('activeSessionSchema capabilities', () => { it('accepts a session row with capabilities.attachments: true', () => { diff --git a/apps/web/src/routers/active-sessions-router.ts b/apps/web/src/routers/active-sessions-router.ts index e61edffd74..40d37f301f 100644 --- a/apps/web/src/routers/active-sessions-router.ts +++ b/apps/web/src/routers/active-sessions-router.ts @@ -47,6 +47,25 @@ const connectedInstancesResponseSchema = z.object({ export type ActiveSession = z.infer; export type ConnectedInstance = z.infer; +/** + * Overlay stored attention (question/permission) onto a live heartbeat + * status. Non-attention DB values yield to live so busy/idle remain + * authoritative while the CLI is connected. + * + * Must run in the router: client fetchQuery replaces the cache wholesale, + * so sticky attention held only in client helpers is wiped on every + * enrichment / reconnect / cli.connected refresh. + */ +export function resolveActiveSessionStatus( + liveStatus: string, + storedStatus: string | null | undefined +): string { + if (storedStatus === 'question' || storedStatus === 'permission') { + return storedStatus; + } + return liveStatus; +} + export const activeSessionsRouter = createTRPCRouter({ getToken: baseProcedure.query(async ({ ctx }) => { const token = generateInternalServiceToken(ctx.user.id); @@ -99,6 +118,7 @@ export const activeSessionsRouter = createTRPCRouter({ created_on_platform: string | null; created_at: string; updated_at: string; + status: string | null; }> = []; try { rows = await db @@ -107,6 +127,7 @@ export const activeSessionsRouter = createTRPCRouter({ created_on_platform: cli_sessions_v2.created_on_platform, created_at: cli_sessions_v2.created_at, updated_at: cli_sessions_v2.updated_at, + status: cli_sessions_v2.status, }) .from(cli_sessions_v2) .where( @@ -129,8 +150,12 @@ export const activeSessionsRouter = createTRPCRouter({ // Explicit snake_case → camelCase mapping: the mobile client only // reads createdOnPlatform/createdAt/updatedAt, so we do not spread // the DB row (which carries snake_case keys it never uses). + // Status: prefer DB attention (question/permission) over the live + // heartbeat so a released CLI that only reports idle/busy still + // surfaces NEEDS INPUT after tRPC seed / enrichment / reconnect. return { ...session, + status: resolveActiveSessionStatus(session.status, row.status), createdOnPlatform: row.created_on_platform ?? undefined, createdAt: row.created_at, updatedAt: row.updated_at, diff --git a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts index 9ad8fbbdc5..c737843ba5 100644 --- a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts +++ b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts @@ -93,6 +93,8 @@ import { renderExecutionTurnContent } from '../execution/types.js'; import type { Env as WorkerEnv, SandboxId } from '../types.js'; import { deriveSharedSandboxId, generateSandboxId } from '../sandbox-id.js'; import { recordSharedSandboxFailover } from '../shared-sandbox-route.js'; +import { nextMetadataAfterAdmittedAgentModel } from './persist-admitted-agent-model.js'; +import { dispatchedKilocodeModelId } from './model-utils.js'; import { resolveSecret, validateStreamTicket } from '../auth.js'; import { resolveTerminalWrapperClient, type TerminalWrapperClient } from '../terminal/access.js'; @@ -3231,7 +3233,52 @@ export class CloudAgentSession extends DurableObject { request: SubmittedSessionMessageRequest ): Promise { const deletionPending = await this.deletionPendingAdmissionFailure(); - return deletionPending ?? this.getSessionMessageQueue().admitSubmittedMessage(request); + if (deletionPending) return deletionPending; + const result = await this.getSessionMessageQueue().admitSubmittedMessage(request); + if (result.success) { + await this.persistAdmittedAgentModelIfChanged(request); + } + return result; + } + + /** + * After a successful admit, update stored agent.model/variant when the run's + * resolved selection differs. Single owning write site for post-registration + * model persistence — do not duplicate from the message queue. + */ + private async persistAdmittedAgentModelIfChanged( + request: SubmittedSessionMessageRequest + ): Promise { + const metadata = await this.getMetadata(); + if (!metadata?.agent) return; + + // Mirror the queue's resolve at admit time (read-only there): requested + // override wins, else stored default. Normalize like the queue so we store + // the same dispatched model id the run actually uses. + const resolvedModel = dispatchedKilocodeModelId(request.agent?.model ?? metadata.agent.model); + if (!resolvedModel) return; + const resolvedVariant = request.agent?.variant ?? metadata.agent.variant; + + const next = nextMetadataAfterAdmittedAgentModel(metadata, { + model: resolvedModel, + variant: resolvedVariant, + }); + if (!next) return; + + try { + await this.updateMetadata(next); + } catch (err) { + // Admission already succeeded; do not fail the client for a metadata + // bookkeeping write. Cold relaunch may show the previous model until the + // next successful persist. + logger + .withFields({ + error: err instanceof Error ? err.message : String(err), + model: resolvedModel, + variant: resolvedVariant, + }) + .warn('Failed to persist admitted agent model on session metadata'); + } } async replayPreparedInitialMessage( diff --git a/services/cloud-agent-next/src/persistence/persist-admitted-agent-model.test.ts b/services/cloud-agent-next/src/persistence/persist-admitted-agent-model.test.ts new file mode 100644 index 0000000000..d0146e385b --- /dev/null +++ b/services/cloud-agent-next/src/persistence/persist-admitted-agent-model.test.ts @@ -0,0 +1,94 @@ +import { describe, expect, it } from 'vitest'; + +import { nextMetadataAfterAdmittedAgentModel } from './persist-admitted-agent-model.js'; +import type { SessionMetadata } from './session-metadata.js'; + +function baseMetadata(agent: SessionMetadata['agent']): SessionMetadata { + return { + metadataSchemaVersion: 2, + identity: { + sessionId: 'agent_test', + userId: 'user_test', + orgId: 'org_test', + }, + auth: { + kiloSessionId: 'ses_test_aaaaaaaaaaaaaaaaaaaaaaaa', + kilocodeToken: 'token', + }, + agent, + lifecycle: { + version: 1, + timestamp: 1, + }, + }; +} + +describe('nextMetadataAfterAdmittedAgentModel', () => { + it('returns null when metadata has no agent block', () => { + expect( + nextMetadataAfterAdmittedAgentModel(baseMetadata(undefined), { + model: 'anthropic/claude-sonnet-4', + variant: 'high', + }) + ).toBeNull(); + }); + + it('returns null when model and variant already match', () => { + const metadata = baseMetadata({ + mode: 'code', + model: 'anthropic/claude-sonnet-4', + variant: 'high', + }); + expect( + nextMetadataAfterAdmittedAgentModel(metadata, { + model: 'anthropic/claude-sonnet-4', + variant: 'high', + }) + ).toBeNull(); + }); + + it('returns updated metadata when model differs', () => { + const metadata = baseMetadata({ + mode: 'code', + model: 'anthropic/claude-sonnet-4', + variant: 'high', + }); + const next = nextMetadataAfterAdmittedAgentModel(metadata, { + model: 'openai/gpt-5', + variant: 'high', + }); + expect(next).not.toBeNull(); + expect(next?.agent).toEqual({ + mode: 'code', + model: 'openai/gpt-5', + variant: 'high', + }); + expect(next?.lifecycle.version).toBeGreaterThan(metadata.lifecycle.version); + }); + + it('updates variant together with model when only variant differs', () => { + const metadata = baseMetadata({ + mode: 'code', + model: 'anthropic/claude-sonnet-4', + variant: 'low', + }); + const next = nextMetadataAfterAdmittedAgentModel(metadata, { + model: 'anthropic/claude-sonnet-4', + variant: 'high', + }); + expect(next?.agent?.model).toBe('anthropic/claude-sonnet-4'); + expect(next?.agent?.variant).toBe('high'); + }); + + it('clears variant when admitted run has no variant and stored had one', () => { + const metadata = baseMetadata({ + mode: 'code', + model: 'anthropic/claude-sonnet-4', + variant: 'high', + }); + const next = nextMetadataAfterAdmittedAgentModel(metadata, { + model: 'anthropic/claude-sonnet-4', + }); + expect(next?.agent?.variant).toBeUndefined(); + }); +}); diff --git a/services/cloud-agent-next/src/persistence/persist-admitted-agent-model.ts b/services/cloud-agent-next/src/persistence/persist-admitted-agent-model.ts new file mode 100644 index 0000000000..5e56507ffb --- /dev/null +++ b/services/cloud-agent-next/src/persistence/persist-admitted-agent-model.ts @@ -0,0 +1,34 @@ +import type { SessionMetadata } from './session-metadata.js'; + +/** + * Build next session metadata after a successful message admission when the + * run's model/variant differs from the stored agent defaults. + * + * Returns null when nothing should be written (no agent block, or same values). + * Call only after admission succeeded — never at resolve-before-admit time. + */ +export function nextMetadataAfterAdmittedAgentModel( + metadata: SessionMetadata, + admitted: { model: string; variant?: string } +): SessionMetadata | null { + if (!metadata.agent) return null; + + const nextModel = admitted.model; + const nextVariant = admitted.variant; + if (metadata.agent.model === nextModel && metadata.agent.variant === nextVariant) { + return null; + } + + return { + ...metadata, + agent: { + ...metadata.agent, + model: nextModel, + variant: nextVariant, + }, + lifecycle: { + ...metadata.lifecycle, + version: Date.now(), + }, + }; +} diff --git a/services/session-ingest/src/dos/SessionIngestDO.ts b/services/session-ingest/src/dos/SessionIngestDO.ts index d8834d8441..b7f1f515fc 100644 --- a/services/session-ingest/src/dos/SessionIngestDO.ts +++ b/services/session-ingest/src/dos/SessionIngestDO.ts @@ -38,6 +38,7 @@ import { readKiloSdkSessionSnapshot, type KiloSdkSessionSnapshotRead, } from './kilo-sdk-materialization'; +import { resetAttentionStatusOnCliDisconnect } from '../ingest/metadata'; type IngestMetaKey = | ExtractableMetaKey @@ -448,6 +449,15 @@ export class SessionIngestDO extends DurableObject { ); } + /** + * Clear a stored attention status when the owning CLI disconnects. + * Metadata owner entry point — UserConnectionDO has no Postgres write path. + * Attention check and conditional write live in `resetAttentionStatusOnCliDisconnect`. + */ + async resetAttentionStatusOnCliDisconnect(kiloUserId: string, sessionId: string): Promise { + await resetAttentionStatusOnCliDisconnect(this.env, kiloUserId, sessionId, this.ctx); + } + /** Builds a text excerpt for a completed assistant message from its already-ingested text parts. */ private buildAssistantExcerptForMessage(messageId: string): string { const range = getPartItemIdentityRange(messageId); diff --git a/services/session-ingest/src/dos/UserConnectionDO.test.ts b/services/session-ingest/src/dos/UserConnectionDO.test.ts index f38ad526c7..6b07ff7b94 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.test.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.test.ts @@ -12,6 +12,16 @@ vi.mock('cloudflare:workers', () => ({ }, })); +const sessionIngestMocks = vi.hoisted(() => ({ + resetAttentionStatusOnCliDisconnect: vi.fn(async () => undefined), + claimSessionReadyPush: vi.fn(async () => undefined), + getSessionIngestDO: vi.fn(), +})); + +vi.mock('./SessionIngestDO', () => ({ + getSessionIngestDO: sessionIngestMocks.getSessionIngestDO, +})); + import { MAX_CATALOG_RESULT_BYTES, UserConnectionDO } from './UserConnectionDO'; // --------------------------------------------------------------------------- @@ -195,15 +205,18 @@ function addCliSocket( title: string; platform?: string; }> = [], - instance?: { name: string; projectName: string; version?: string } + instance?: { name: string; projectName: string; version?: string }, + kiloUserId?: string ): MockWS { const attachment: { role: 'cli'; connectionId: string; sessions: typeof sessions; instance?: typeof instance; + kiloUserId?: string; } = { role: 'cli', connectionId, sessions }; if (instance) attachment.instance = instance; + if (kiloUserId) attachment.kiloUserId = kiloUserId; const ws = createMockWs(['cli'], attachment); mockCtx.addSocket(ws); return ws; @@ -308,14 +321,14 @@ function createUtf8OversizedResult(): { padding: string } { return result; } -/** Trigger CLI disconnect */ -function disconnectCli(doInstance: UserConnectionDO, cliWs: MockWS) { - doInstance.webSocketClose(cliWs as never, 0, '', false); +/** Trigger CLI disconnect (awaits attention reset before broadcast). */ +async function disconnectCli(doInstance: UserConnectionDO, cliWs: MockWS) { + await doInstance.webSocketClose(cliWs as never, 0, '', false); } /** Trigger web disconnect */ function disconnectWeb(doInstance: UserConnectionDO, webWs: MockWS) { - doInstance.webSocketClose(webWs as never, 0, '', false); + void doInstance.webSocketClose(webWs as never, 0, '', false); } // =========================================================================== @@ -325,6 +338,14 @@ function disconnectWeb(doInstance: UserConnectionDO, webWs: MockWS) { describe('UserConnectionDO', () => { beforeEach(() => { vi.restoreAllMocks(); + sessionIngestMocks.resetAttentionStatusOnCliDisconnect.mockReset(); + sessionIngestMocks.resetAttentionStatusOnCliDisconnect.mockResolvedValue(undefined); + sessionIngestMocks.claimSessionReadyPush.mockReset(); + sessionIngestMocks.getSessionIngestDO.mockReset(); + sessionIngestMocks.getSessionIngestDO.mockReturnValue({ + resetAttentionStatusOnCliDisconnect: sessionIngestMocks.resetAttentionStatusOnCliDisconnect, + claimSessionReadyPush: sessionIngestMocks.claimSessionReadyPush, + }); }); afterEach(() => { @@ -377,7 +398,7 @@ describe('UserConnectionDO', () => { }); describe('hasActiveCliSession', () => { - it('tracks whether a connected CLI heartbeat currently owns the session', () => { + it('tracks whether a connected CLI heartbeat currently owns the session', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); @@ -388,7 +409,7 @@ describe('UserConnectionDO', () => { expect(doInstance.hasActiveCliSession('ses_1')).toBe(true); mockCtx.removeSocket(cliWs); - disconnectCli(doInstance, cliWs); + await disconnectCli(doInstance, cliWs); expect(doInstance.hasActiveCliSession('ses_1')).toBe(false); }); @@ -1168,7 +1189,7 @@ describe('UserConnectionDO', () => { // ------------------------------------------------------------------------- describe('CLI disconnect', () => { - it('cleans up session ownership and broadcasts cli.disconnected', () => { + it('cleans up session ownership and broadcasts cli.disconnected', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -1178,7 +1199,7 @@ describe('UserConnectionDO', () => { // Remove from sockets before disconnect (simulates runtime closing) mockCtx.removeSocket(cliWs); - disconnectCli(doInstance, cliWs); + await disconnectCli(doInstance, cliWs); // Web receives cli.disconnected expect(webWs.send).toHaveBeenCalled(); @@ -1198,7 +1219,7 @@ describe('UserConnectionDO', () => { expect(parseSent(web2)).toMatchObject({ type: 'response', error: 'Session owner not found' }); }); - it('sends error responses for pending commands on disconnect', () => { + it('sends error responses for pending commands on disconnect', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -1211,7 +1232,7 @@ describe('UserConnectionDO', () => { // CLI disconnects mockCtx.removeSocket(cliWs); - disconnectCli(doInstance, cliWs); + await disconnectCli(doInstance, cliWs); // Web receives error response with original id const msgs = allSent(webWs); @@ -1221,7 +1242,7 @@ describe('UserConnectionDO', () => { expect(errorResp).toMatchObject({ type: 'response', id: 'cmd-1', error: 'CLI disconnected' }); }); - it('reports owner change when an owner-fenced command target disconnects', () => { + it('reports owner change when an owner-fenced command target disconnects', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -1236,7 +1257,7 @@ describe('UserConnectionDO', () => { webWs.send.mockClear(); mockCtx.removeSocket(cliWs); - disconnectCli(doInstance, cliWs); + await disconnectCli(doInstance, cliWs); expect(parseSent(webWs)).toEqual({ type: 'response', @@ -1278,7 +1299,7 @@ describe('UserConnectionDO', () => { }); }); - it('sends error for connection-routed pending commands on CLI disconnect', () => { + it('sends error for connection-routed pending commands on CLI disconnect', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -1295,7 +1316,7 @@ describe('UserConnectionDO', () => { // CLI disconnects before responding mockCtx.removeSocket(cliWs); - disconnectCli(doInstance, cliWs); + await disconnectCli(doInstance, cliWs); const msgs = allSent(webWs); const errorResp = msgs.find( @@ -1308,7 +1329,7 @@ describe('UserConnectionDO', () => { }); }); - it('sends error for fallback-routed pending commands on CLI disconnect', () => { + it('sends error for fallback-routed pending commands on CLI disconnect', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -1320,7 +1341,7 @@ describe('UserConnectionDO', () => { webWs.send.mockClear(); mockCtx.removeSocket(cliWs); - disconnectCli(doInstance, cliWs); + await disconnectCli(doInstance, cliWs); const msgs = allSent(webWs); const errorResp = msgs.find( @@ -1333,7 +1354,7 @@ describe('UserConnectionDO', () => { }); }); - it('reconnecting CLI — old socket close does not destroy state', () => { + it('reconnecting CLI — old socket close does not destroy state', async () => { const { doInstance, mockCtx } = setup(); const cli1 = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -1348,7 +1369,7 @@ describe('UserConnectionDO', () => { // DON'T remove cli2 from sockets — cli2 is the replacement // Just remove cli1 to simulate it being closed mockCtx.removeSocket(cli1); - disconnectCli(doInstance, cli1); + await disconnectCli(doInstance, cli1); // State should NOT be cleaned up — cli2 is live // Verify by routing a command to s1 — should reach cli2 @@ -1358,7 +1379,7 @@ describe('UserConnectionDO', () => { expect(parseSent(cli2)).toEqual({ type: 'subscribe', sessionId: 's1' }); }); - it('reconnecting CLI — commands sent to replacement socket are not spuriously failed', () => { + it('reconnecting CLI — commands sent to replacement socket are not spuriously failed', async () => { const { doInstance, mockCtx } = setup(); const cli1 = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -1383,7 +1404,7 @@ describe('UserConnectionDO', () => { webWs.send.mockClear(); // Now cli1's close event fires (stale socket teardown) - disconnectCli(doInstance, cli1); + await disconnectCli(doInstance, cli1); // Web should NOT have received an error for cmd-new — it was sent to cli2, not cli1 const errorMsgs = allSent(webWs).filter( @@ -1403,7 +1424,7 @@ describe('UserConnectionDO', () => { }); }); - it('reconnecting CLI — pending commands from old socket get error responses', () => { + it('reconnecting CLI — pending commands from old socket get error responses', async () => { const { doInstance, mockCtx } = setup(); const cli1 = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -1420,7 +1441,7 @@ describe('UserConnectionDO', () => { // cli1's close event fires — cmd-1 was sent on cli1's wire, cli2 never saw it mockCtx.removeSocket(cli1); - disconnectCli(doInstance, cli1); + await disconnectCli(doInstance, cli1); // Web should receive an error for the stranded command const msgs = allSent(webWs); @@ -1433,6 +1454,171 @@ describe('UserConnectionDO', () => { error: 'CLI disconnected', }); }); + + it('resets attention status for owned sessions before broadcasting cli.disconnected', async () => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1', [], undefined, 'usr_1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [ + makeSession('s-question', 'question'), + makeSession('s-busy', 'busy'), + ]); + webWs.send.mockClear(); + + const callOrder: string[] = []; + sessionIngestMocks.resetAttentionStatusOnCliDisconnect.mockImplementation(async () => { + callOrder.push('reset'); + // Disconnect must not have been broadcast yet (ordering guarantee). + expect( + allSent(webWs).some(m => m.type === 'system' && m.event === 'cli.disconnected') + ).toBe(false); + }); + + // Leave the socket in getWebSockets() — matches workerd during webSocketClose. + await disconnectCli(doInstance, cliWs); + + callOrder.push('disconnect'); + + expect(sessionIngestMocks.getSessionIngestDO).toHaveBeenCalledWith(expect.anything(), { + kiloUserId: 'usr_1', + sessionId: 's-question', + }); + expect(sessionIngestMocks.getSessionIngestDO).toHaveBeenCalledWith(expect.anything(), { + kiloUserId: 'usr_1', + sessionId: 's-busy', + }); + // Both owned sessions are delegated; attention-only filtering is on the metadata side. + expect(sessionIngestMocks.resetAttentionStatusOnCliDisconnect).toHaveBeenCalledTimes(2); + expect(sessionIngestMocks.resetAttentionStatusOnCliDisconnect).toHaveBeenCalledWith( + 'usr_1', + 's-question' + ); + expect(sessionIngestMocks.resetAttentionStatusOnCliDisconnect).toHaveBeenCalledWith( + 'usr_1', + 's-busy' + ); + expect(callOrder.filter(step => step === 'reset')).toHaveLength(2); + expect(callOrder.at(-1)).toBe('disconnect'); + expect(allSent(webWs).some(m => m.type === 'system' && m.event === 'cli.disconnected')).toBe( + true + ); + }); + + it('resets attention when the closing socket is still listed in getWebSockets (workerd)', async () => { + // Production wrangler/workerd keeps the closing WebSocket in getWebSockets() + // during webSocketClose. Matching connectionId without excluding self would + // treat every disconnect as a stale reconnect and skip the attention reset. + // Prior unit tests always called removeSocket first, so they never caught this. + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1', [], undefined, 'usr_1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s-question', 'question')]); + webWs.send.mockClear(); + sessionIngestMocks.getSessionIngestDO.mockClear(); + sessionIngestMocks.resetAttentionStatusOnCliDisconnect.mockClear(); + + // Do NOT removeSocket — mirrors workerd during webSocketClose. + expect(mockCtx.sockets).toContain(cliWs); + await disconnectCli(doInstance, cliWs); + + expect(sessionIngestMocks.resetAttentionStatusOnCliDisconnect).toHaveBeenCalledTimes(1); + expect(sessionIngestMocks.resetAttentionStatusOnCliDisconnect).toHaveBeenCalledWith( + 'usr_1', + 's-question' + ); + expect(allSent(webWs).some(m => m.type === 'system' && m.event === 'cli.disconnected')).toBe( + true + ); + }); + + it('does not reset attention when kiloUserId is missing on the attachment', async () => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1', 'question')]); + webWs.send.mockClear(); + + mockCtx.removeSocket(cliWs); + await disconnectCli(doInstance, cliWs); + + expect(sessionIngestMocks.getSessionIngestDO).not.toHaveBeenCalled(); + expect(sessionIngestMocks.resetAttentionStatusOnCliDisconnect).not.toHaveBeenCalled(); + expect(warn).toHaveBeenCalledWith( + 'Skipping attention status reset on CLI disconnect: missing kiloUserId on attachment', + { ownedSessionCount: 1 } + ); + expect(allSent(webWs).some(m => m.type === 'system' && m.event === 'cli.disconnected')).toBe( + true + ); + }); + + it('does not reset attention for sessions owned by another live connection', async () => { + const { doInstance, mockCtx } = setup(); + const cli1 = addCliSocket(mockCtx, 'cli-1', [], undefined, 'usr_1'); + const cli2 = addCliSocket(mockCtx, 'cli-2', [], undefined, 'usr_1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cli1, [makeSession('s1', 'question')]); + // cli2 takes ownership of s1 + sendHeartbeat(doInstance, cli2, [makeSession('s1', 'question')]); + webWs.send.mockClear(); + sessionIngestMocks.getSessionIngestDO.mockClear(); + sessionIngestMocks.resetAttentionStatusOnCliDisconnect.mockClear(); + + mockCtx.removeSocket(cli1); + await disconnectCli(doInstance, cli1); + + // cli1 no longer owns s1, so no reset for that session + expect(sessionIngestMocks.resetAttentionStatusOnCliDisconnect).not.toHaveBeenCalled(); + expect(allSent(webWs).some(m => m.type === 'system' && m.event === 'cli.disconnected')).toBe( + true + ); + }); + + it('stale reconnect close does not reset attention status', async () => { + const { doInstance, mockCtx } = setup(); + const cli1 = addCliSocket(mockCtx, 'cli-1', [], undefined, 'usr_1'); + const cli2 = addCliSocket(mockCtx, 'cli-1', [], undefined, 'usr_1'); + + sendHeartbeat(doInstance, cli1, [makeSession('s1', 'question')]); + sendHeartbeat(doInstance, cli2, [makeSession('s1', 'question')]); + sessionIngestMocks.getSessionIngestDO.mockClear(); + sessionIngestMocks.resetAttentionStatusOnCliDisconnect.mockClear(); + + mockCtx.removeSocket(cli1); + await disconnectCli(doInstance, cli1); + + expect(sessionIngestMocks.getSessionIngestDO).not.toHaveBeenCalled(); + expect(sessionIngestMocks.resetAttentionStatusOnCliDisconnect).not.toHaveBeenCalled(); + }); + + it('still broadcasts cli.disconnected when attention reset fails', async () => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1', [], undefined, 'usr_1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + const error = vi.spyOn(console, 'error').mockImplementation(() => undefined); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1', 'question')]); + webWs.send.mockClear(); + sessionIngestMocks.resetAttentionStatusOnCliDisconnect.mockRejectedValueOnce( + new Error('db down') + ); + + // Leave socket listed (workerd close semantics). + await disconnectCli(doInstance, cliWs); + + expect(allSent(webWs).some(m => m.type === 'system' && m.event === 'cli.disconnected')).toBe( + true + ); + expect(error).toHaveBeenCalledWith( + 'Failed to reset attention status on CLI disconnect', + expect.objectContaining({ error: 'db down' }) + ); + }); }); // ------------------------------------------------------------------------- @@ -3539,7 +3725,7 @@ describe('UserConnectionDO', () => { expect(JSON.stringify(warn.mock.calls)).not.toContain(secret); }); - it('webSocketError triggers webSocketClose', () => { + it('webSocketError triggers webSocketClose', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -3549,7 +3735,7 @@ describe('UserConnectionDO', () => { // Remove CLI so disconnect can clean up mockCtx.removeSocket(cliWs); - doInstance.webSocketError(cliWs as never); + await doInstance.webSocketError(cliWs as never); // Should broadcast cli.disconnected const msgs = allSent(webWs); @@ -3575,13 +3761,11 @@ describe('UserConnectionDO', () => { const mockCtx = createMockCtx(); const ctx = mockCtx.build(); const claimSessionReadyPush = vi.fn(async () => {}); - const env = { - SESSION_INGEST_DO: { - idFromName: vi.fn((name: string) => name), - get: vi.fn(() => ({ claimSessionReadyPush })), - }, - }; - const doInstance = new UserConnectionDO(ctx as never, env as never); + sessionIngestMocks.getSessionIngestDO.mockReturnValue({ + claimSessionReadyPush, + resetAttentionStatusOnCliDisconnect: sessionIngestMocks.resetAttentionStatusOnCliDisconnect, + }); + const doInstance = new UserConnectionDO(ctx as never, {} as never); return { doInstance, mockCtx, claimSessionReadyPush }; } diff --git a/services/session-ingest/src/dos/UserConnectionDO.ts b/services/session-ingest/src/dos/UserConnectionDO.ts index ae8df340e0..c18d39cd66 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.ts @@ -331,26 +331,32 @@ export class UserConnectionDO extends DurableObject { } } - webSocketClose(ws: WebSocket, _code: number, _reason: string, _wasClean: boolean): void { + webSocketClose( + ws: WebSocket, + _code: number, + _reason: string, + _wasClean: boolean + ): void | Promise { this.ensureState(); const attachment = ws.deserializeAttachment() as WSAttachment | null; if (!attachment) return; if (attachment.role === 'cli') { - this.handleCliDisconnect(ws, attachment); - } else { - this.handleWebDisconnect(ws); + // Await attention resets so cli.disconnected is not broadcast until the + // stored status write has committed (mobile history refetch races otherwise). + return this.handleCliDisconnect(ws, attachment); } + this.handleWebDisconnect(ws); } - webSocketError(ws: WebSocket): void { + webSocketError(ws: WebSocket): void | Promise { const attachment = ws.deserializeAttachment() as WSAttachment | null; console.error('WebSocket error', { role: attachment?.role ?? 'unknown', connectionId: attachment?.connectionId ?? 'unknown', }); - this.webSocketClose(ws, 0, '', false); + return this.webSocketClose(ws, 0, '', false); } async alarm(): Promise { @@ -832,15 +838,19 @@ export class UserConnectionDO extends DurableObject { // Disconnect handling // --------------------------------------------------------------------------- - private handleCliDisconnect( + private async handleCliDisconnect( disconnectedWs: WebSocket, attachment: WSAttachment & { role: 'cli' } - ): void { + ): Promise { const { connectionId } = attachment; // If another CLI socket already has this connectionId, this is a stale // close from a reconnect — the replacement socket is already active. + // Exclude the closing socket: under wrangler/workerd, getWebSockets() still + // includes it during webSocketClose, so matching self would always look "replaced" + // and skip ownership cleanup + attention reset (DEF-5 E2E failure). const replaced = this.ctx.getWebSockets('cli').some(ws => { + if (ws === disconnectedWs) return false; const att = ws.deserializeAttachment() as WSAttachment | null; return att?.role === 'cli' && att.connectionId === connectionId; }); @@ -875,6 +885,12 @@ export class UserConnectionDO extends DurableObject { // Leave webSubscriptions intact — a reconnecting CLI can resume + // Reset stored attention before broadcasting disconnect so the mobile + // departure refetch observes `retry` rather than a stuck `question`. + // kiloUserId comes from the CLI attachment (authenticated /user/cli route); + // without it we cannot safely target rows and must no-op. + await this.resetOwnedSessionAttentionOnDisconnect(attachment.kiloUserId, ownedSessions); + this.broadcastToWeb({ type: 'system', event: 'cli.disconnected', @@ -882,6 +898,40 @@ export class UserConnectionDO extends DurableObject { }); } + /** + * Commit attention clears for owned sessions before `cli.disconnected`. + * Identity: attachment `kiloUserId` only — never guess from DO name. + */ + private async resetOwnedSessionAttentionOnDisconnect( + kiloUserId: string | undefined, + ownedSessions: ReadonlySet + ): Promise { + if (ownedSessions.size === 0) return; + + if (!kiloUserId) { + console.warn( + 'Skipping attention status reset on CLI disconnect: missing kiloUserId on attachment', + { ownedSessionCount: ownedSessions.size } + ); + return; + } + + const results = await Promise.allSettled( + [...ownedSessions].map(async sessionId => { + const stub = getSessionIngestDO(this.env, { kiloUserId, sessionId }); + await stub.resetAttentionStatusOnCliDisconnect(kiloUserId, sessionId); + }) + ); + + for (const result of results) { + if (result.status === 'rejected') { + console.error('Failed to reset attention status on CLI disconnect', { + error: result.reason instanceof Error ? result.reason.message : String(result.reason), + }); + } + } + } + private handleWebDisconnect(ws: WebSocket): void { const attachment = ws.deserializeAttachment() as WSAttachment | null; const connectionId = attachment?.role === 'web' ? attachment.connectionId : 'unknown'; diff --git a/services/session-ingest/src/ingest/metadata.test.ts b/services/session-ingest/src/ingest/metadata.test.ts new file mode 100644 index 0000000000..e8cfc88602 --- /dev/null +++ b/services/session-ingest/src/ingest/metadata.test.ts @@ -0,0 +1,210 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +vi.mock('cloudflare:workers', () => ({ + DurableObject: class { + ctx: unknown; + env: unknown; + constructor(ctx: unknown, env: unknown) { + this.ctx = ctx; + this.env = env; + } + }, +})); + +vi.mock('@kilocode/db/client', () => ({ + getWorkerDb: vi.fn(), +})); + +vi.mock('../dos/SessionAccessCacheDO', () => ({ + getSessionAccessCacheDO: vi.fn(), +})); + +vi.mock('../session-events', () => ({ + mapSessionEventRow: vi.fn((row: { session_id: string; status: string | null }) => ({ + source: 'v2' as const, + sessionId: row.session_id, + status: row.status, + statusUpdatedAt: '2026-07-25T00:00:00.000Z', + updatedAt: '2026-07-25T00:00:00.000Z', + })), + notifyUserSessionEvent: vi.fn(), +})); + +import { getWorkerDb } from '@kilocode/db/client'; +import { notifyUserSessionEvent } from '../session-events'; +import { + CLI_DISCONNECT_ATTENTION_RESET_STATUS, + resetAttentionStatusOnCliDisconnect, +} from './metadata'; + +type StatusRow = { status: string | null }; + +function createTransactionDb(options: { + initialStatus: string | null; + /** After the conditional update, status read-back (simulates concurrent overwrite). */ + persistedStatus?: string | null; + rowMissing?: boolean; +}) { + const updateWhere = vi.fn(async () => undefined); + const updateSet = vi.fn(() => ({ where: updateWhere })); + // Named without the substring "update" so oxlint drizzle rules do not flag test spies. + const applyUpdate = vi.fn(() => ({ set: updateSet })); + + let selectCall = 0; + + function rowsForSelect(): unknown[] { + selectCall += 1; + if (options.rowMissing) return []; + if (selectCall === 1) { + return [{ status: options.initialStatus } satisfies StatusRow]; + } + const status = + options.persistedStatus !== undefined + ? options.persistedStatus + : options.initialStatus !== null && + (options.initialStatus === 'question' || options.initialStatus === 'permission') + ? CLI_DISCONNECT_ATTENTION_RESET_STATUS + : options.initialStatus; + return [ + { + session_id: 'ses_1', + created_at: '2026-01-01T00:00:00.000Z', + updated_at: '2026-01-01T00:00:01.000Z', + title: 'T', + created_on_platform: 'cli', + organization_id: null, + git_url: null, + git_branch: null, + parent_session_id: null, + status, + status_updated_at: '2026-07-25T00:00:00.000Z', + }, + ]; + } + + /** Thenable that also supports `.for('update')` (first select locks; second does not). */ + function limitResult() { + const promise = Promise.resolve(rowsForSelect()); + return Object.assign(promise, { + for: vi.fn(() => promise), + }); + } + + const select = vi.fn(() => ({ + from: vi.fn(() => ({ + where: vi.fn(() => ({ + limit: vi.fn(() => limitResult()), + })), + })), + })); + + const transaction = vi.fn(async (fn: (tx: unknown) => Promise) => + fn({ select, update: applyUpdate }) + ); + + return { transaction, select, applyUpdate, updateSet, updateWhere }; +} + +describe('resetAttentionStatusOnCliDisconnect', () => { + beforeEach(() => { + vi.mocked(getWorkerDb).mockReset(); + vi.mocked(notifyUserSessionEvent).mockReset(); + }); + + it('writes retry and notifies when stored status is question', async () => { + const db = createTransactionDb({ initialStatus: 'question' }); + vi.mocked(getWorkerDb).mockReturnValue(db as never); + + const env = { HYPERDRIVE: { connectionString: 'postgres://unused' } } as never; + await resetAttentionStatusOnCliDisconnect(env, 'usr_1', 'ses_1'); + + expect(db.applyUpdate).toHaveBeenCalled(); + expect(db.updateSet).toHaveBeenCalledWith({ + status: CLI_DISCONNECT_ATTENTION_RESET_STATUS, + status_updated_at: expect.any(String), + }); + expect(notifyUserSessionEvent).toHaveBeenCalledWith( + env, + 'usr_1', + expect.objectContaining({ + type: 'session.status.updated', + data: expect.objectContaining({ + previousStatus: 'question', + status: CLI_DISCONNECT_ATTENTION_RESET_STATUS, + }), + }), + undefined + ); + }); + + it('writes retry and notifies when stored status is permission', async () => { + const db = createTransactionDb({ initialStatus: 'permission' }); + vi.mocked(getWorkerDb).mockReturnValue(db as never); + + const env = { HYPERDRIVE: { connectionString: 'postgres://unused' } } as never; + await resetAttentionStatusOnCliDisconnect(env, 'usr_1', 'ses_1'); + + expect(db.applyUpdate).toHaveBeenCalled(); + expect(notifyUserSessionEvent).toHaveBeenCalledWith( + env, + 'usr_1', + expect.objectContaining({ + type: 'session.status.updated', + data: expect.objectContaining({ + previousStatus: 'permission', + status: CLI_DISCONNECT_ATTENTION_RESET_STATUS, + }), + }), + undefined + ); + }); + + it.each(['busy', 'idle', 'retry', null] as const)( + 'no-ops without write or notify when stored status is %s', + async status => { + const db = createTransactionDb({ initialStatus: status }); + vi.mocked(getWorkerDb).mockReturnValue(db as never); + + await resetAttentionStatusOnCliDisconnect( + { HYPERDRIVE: { connectionString: 'postgres://unused' } } as never, + 'usr_1', + 'ses_1' + ); + + expect(db.applyUpdate).not.toHaveBeenCalled(); + expect(notifyUserSessionEvent).not.toHaveBeenCalled(); + } + ); + + it('no-ops when the session row is missing', async () => { + const db = createTransactionDb({ initialStatus: 'question', rowMissing: true }); + vi.mocked(getWorkerDb).mockReturnValue(db as never); + + await resetAttentionStatusOnCliDisconnect( + { HYPERDRIVE: { connectionString: 'postgres://unused' } } as never, + 'usr_1', + 'ses_missing' + ); + + expect(db.applyUpdate).not.toHaveBeenCalled(); + expect(notifyUserSessionEvent).not.toHaveBeenCalled(); + }); + + it('does not notify when a concurrent write wins the conditional update', async () => { + const db = createTransactionDb({ + initialStatus: 'question', + // Conditional WHERE matched nothing; row still shows busy after the race. + persistedStatus: 'busy', + }); + vi.mocked(getWorkerDb).mockReturnValue(db as never); + + await resetAttentionStatusOnCliDisconnect( + { HYPERDRIVE: { connectionString: 'postgres://unused' } } as never, + 'usr_1', + 'ses_1' + ); + + expect(db.applyUpdate).toHaveBeenCalled(); + expect(notifyUserSessionEvent).not.toHaveBeenCalled(); + }); +}); diff --git a/services/session-ingest/src/ingest/metadata.ts b/services/session-ingest/src/ingest/metadata.ts index d222defbbb..b33a95c86a 100644 --- a/services/session-ingest/src/ingest/metadata.ts +++ b/services/session-ingest/src/ingest/metadata.ts @@ -1,13 +1,17 @@ -import { and, eq, sql } from 'drizzle-orm'; +import { and, eq, inArray, sql } from 'drizzle-orm'; import { getWorkerDb } from '@kilocode/db/client'; import { cli_sessions_v2 } from '@kilocode/db/schema'; import { normalizeGitUrl, withDORetry } from '@kilocode/worker-utils'; import type { Env } from '../env'; import { getSessionAccessCacheDO } from '../dos/SessionAccessCacheDO'; +import { isNeedsInputStatus } from '../dos/session-ingest-attention'; import { mapSessionEventRow, notifyUserSessionEvent } from '../session-events'; import { SessionStatusSchema } from '../types/user-connection-protocol'; +/** Stored status written when a CLI disconnects while the session is waiting on input. */ +export const CLI_DISCONNECT_ATTENTION_RESET_STATUS = 'retry' as const; + type SessionMetadataUpdates = Partial< Pick< typeof cli_sessions_v2.$inferInsert, @@ -245,3 +249,99 @@ export async function flushPartialMetadataChanges( }); } } + +/** + * Clear a stored attention status when the owning CLI disconnects. + * + * Only rows currently in `question`/`permission` are updated (to `retry`). Uses a + * conditional write so concurrent non-attention updates are not overwritten. Emits + * `session.status.updated` via the metadata path only — never enters the ingest + * completion pipeline, so no "Task completed" push can fire. + */ +export async function resetAttentionStatusOnCliDisconnect( + env: Env, + kiloUserId: string, + sessionId: string, + ctx?: { waitUntil(promise: Promise): void } +): Promise { + const db = getWorkerDb(env.HYPERDRIVE.connectionString); + const statusUpdatedAt = new Date().toISOString(); + + const notification = await db.transaction(async tx => { + const [statusRow] = await tx + .select({ status: cli_sessions_v2.status }) + .from(cli_sessions_v2) + .where( + and(eq(cli_sessions_v2.session_id, sessionId), eq(cli_sessions_v2.kilo_user_id, kiloUserId)) + ) + .limit(1) + .for('update'); + + if (!statusRow) return null; + + const previousStatus = SessionStatusSchema.nullable().parse(statusRow.status); + if (!isNeedsInputStatus(previousStatus)) return null; + + await tx + .update(cli_sessions_v2) + .set({ + status: CLI_DISCONNECT_ATTENTION_RESET_STATUS, + status_updated_at: statusUpdatedAt, + }) + .where( + and( + eq(cli_sessions_v2.session_id, sessionId), + eq(cli_sessions_v2.kilo_user_id, kiloUserId), + // Re-check in WHERE so a concurrent non-attention write wins. + inArray(cli_sessions_v2.status, ['question', 'permission']) + ) + ); + + const [persistedRow] = await tx + .select({ + session_id: cli_sessions_v2.session_id, + created_at: cli_sessions_v2.created_at, + updated_at: cli_sessions_v2.updated_at, + title: cli_sessions_v2.title, + created_on_platform: cli_sessions_v2.created_on_platform, + organization_id: cli_sessions_v2.organization_id, + git_url: cli_sessions_v2.git_url, + git_branch: cli_sessions_v2.git_branch, + parent_session_id: cli_sessions_v2.parent_session_id, + status: cli_sessions_v2.status, + status_updated_at: cli_sessions_v2.status_updated_at, + }) + .from(cli_sessions_v2) + .where( + and(eq(cli_sessions_v2.session_id, sessionId), eq(cli_sessions_v2.kilo_user_id, kiloUserId)) + ) + .limit(1); + + if (!persistedRow) return null; + if (persistedRow.status !== CLI_DISCONNECT_ATTENTION_RESET_STATUS) return null; + + return { + previousStatus, + session: mapSessionEventRow(persistedRow), + }; + }); + + if (!notification) return; + + notifyUserSessionEvent( + env, + kiloUserId, + { + type: 'session.status.updated', + data: { + source: 'v2', + session: notification.session, + previousStatus: notification.previousStatus, + status: notification.session.status, + statusUpdatedAt: notification.session.statusUpdatedAt, + changedAt: notification.session.updatedAt, + }, + }, + ctx + ); +}