diff --git a/apps/cli/src/lib/assistant-turn-finalize.test.ts b/apps/cli/src/lib/assistant-turn-finalize.test.ts new file mode 100644 index 000000000..fd199cace --- /dev/null +++ b/apps/cli/src/lib/assistant-turn-finalize.test.ts @@ -0,0 +1,87 @@ +import { describe, expect, it } from 'vitest'; +import type { SessionHistoryInput } from '@lody/shared'; + +import { markAssistantTurnFinished } from './assistant-turn-finalize'; + +const OPENED_AT = Date.parse('2026-01-01T00:00:00.000Z'); +const TURN_ENDED_AT = OPENED_AT + 12_000; +const APP_CLOSED_AT = OPENED_AT + 3_600_000; + +const assistantEntry = ( + overrides: Partial & { id: string } +): SessionHistoryInput => ({ + role: 'assistant', + timestamp: new Date(OPENED_AT).toISOString(), + items: [], + fileDiff: [], + ...overrides, +}); + +describe('markAssistantTurnFinished', () => { + it('stamps the open assistant entry when no turn id is given', () => { + const history = [assistantEntry({ id: 'assistant:u1' })]; + + markAssistantTurnFinished(history, { endedAt: TURN_ENDED_AT }); + + expect(history[0]).toMatchObject({ finished: true, endedAt: TURN_ENDED_AT }); + }); + + it('leaves an already finished turn alone when a later teardown finalizes', () => { + // Regression for #260: closing the app runs the no-turnId finalize for every + // live session, which used to re-stamp `endedAt = now` on a turn that ended + // an hour earlier and inflate its rendered "Worked for …". + const history = [assistantEntry({ id: 'assistant:u1' })]; + + markAssistantTurnFinished(history, { endedAt: TURN_ENDED_AT }); + markAssistantTurnFinished(history, { endedAt: APP_CLOSED_AT }); + + expect(history[0]).toMatchObject({ finished: true, endedAt: TURN_ENDED_AT }); + }); + + it('records no duration for a finished entry that never carried one', () => { + // Image-group and file entries publish `finished: true` with no `endedAt`. + const history = [assistantEntry({ id: 'assistant-image-1', finished: true })]; + + markAssistantTurnFinished(history, { endedAt: APP_CLOSED_AT }); + + expect(history[0]?.endedAt).toBeUndefined(); + }); + + it('stamps the addressed turn even when a later assistant entry follows it', () => { + const history = [ + assistantEntry({ id: 'assistant:u1' }), + assistantEntry({ id: 'assistant-image-1', finished: true }), + ]; + + markAssistantTurnFinished(history, { turnId: 'assistant:u1', endedAt: TURN_ENDED_AT }); + + expect(history[0]).toMatchObject({ finished: true, endedAt: TURN_ENDED_AT }); + expect(history[1]?.endedAt).toBeUndefined(); + }); + + it('records the permission wait when the finalizing turn measured one', () => { + const history = [assistantEntry({ id: 'assistant:u1' })]; + + markAssistantTurnFinished(history, { endedAt: TURN_ENDED_AT, permissionWaitMs: 4_000 }); + + expect(history[0]?.permissionWaitMs).toBe(4_000); + }); + + it('never stamps a user or system entry standing after the turn', () => { + const history: SessionHistoryInput[] = [ + assistantEntry({ id: 'assistant:u1' }), + { + id: 'system:1', + role: 'system', + timestamp: new Date(OPENED_AT).toISOString(), + items: [], + fileDiff: [], + }, + ]; + + markAssistantTurnFinished(history, { endedAt: TURN_ENDED_AT }); + + expect(history[0]).toMatchObject({ finished: true, endedAt: TURN_ENDED_AT }); + expect(history[1]?.finished).toBeUndefined(); + }); +}); diff --git a/apps/cli/src/lib/assistant-turn-finalize.ts b/apps/cli/src/lib/assistant-turn-finalize.ts new file mode 100644 index 000000000..ac92f6b72 --- /dev/null +++ b/apps/cli/src/lib/assistant-turn-finalize.ts @@ -0,0 +1,51 @@ +import type { SessionHistoryInput } from '@lody/shared'; + +/** + * Stamp the terminal footprint (`finished`/`endedAt`/`permissionWaitMs`) on the + * assistant entry a finalize call owns. Extracted from `finalizeACPState` so the + * one rule that matters here is testable: a terminal stamp is written once. + * + * `finalizeACPState` has a no-turnId overload used by teardown/cancel paths + * (session `exit`/`terminated`, error, cleanup). Those callers only check that + * transient state exists, so on app close they run for sessions whose turn ended + * long ago — and the loop matched the last assistant entry regardless of state, + * re-stamping `endedAt = now`. The renderer derives "Worked for …" from + * `endedAt - timestamp`, so every close inflated a finished turn's duration by + * the wall-clock time the app stayed open. + * + * Skipping an already-finished entry (rather than filling in a missing + * `endedAt`) is deliberate: `createAssistantImageGroupEntry` and + * `createAssistantFileEntry` publish assistant entries with `finished: true` and + * no `endedAt`, so an `endedAt`-only guard would still stamp `now` on an entry + * that finished whenever it finished. No duration is the honest answer there. + * + * A turn genuinely still running is never finished: the teardown stamp on an + * interrupted turn still lands, and resume clears the footprint through + * `writeAssistantEntryForTurn`'s reopen branch before streaming into the entry + * again. See `apps/cli/src/session/AGENTS.md`. + */ +export const markAssistantTurnFinished = ( + history: SessionHistoryInput[], + options: { + /** Finalize the entry with this id; absent means "whichever turn is open". */ + turnId?: string | undefined; + endedAt: number; + permissionWaitMs?: number | undefined; + } +): SessionHistoryInput[] => { + const { turnId, endedAt, permissionWaitMs } = options; + for (let i = history.length - 1; i >= 0; i--) { + const entry = history[i]; + if (entry && entry.role === 'assistant' && (!turnId || entry.id === turnId)) { + // Already finalized: its terminal timing is the truth, not this call's clock. + if (entry.finished === true) break; + entry.finished = true; + entry.endedAt = endedAt; + if (permissionWaitMs !== undefined) { + entry.permissionWaitMs = permissionWaitMs; + } + break; + } + } + return history; +}; diff --git a/apps/cli/src/lib/message-handler.ts b/apps/cli/src/lib/message-handler.ts index 50d1170ca..d7a65c8fa 100644 --- a/apps/cli/src/lib/message-handler.ts +++ b/apps/cli/src/lib/message-handler.ts @@ -262,6 +262,7 @@ import { resolveImageGenerationStatusWrite, shouldRestoreRunningAfterPermission, } from './session-activity-status'; +import { markAssistantTurnFinished } from './assistant-turn-finalize'; import type { RepoWatchHandle } from 'loro-repo'; import { resolveGitBranchName } from './git/resolve-git-branch-name'; import { @@ -5807,20 +5808,9 @@ export class MessageHandler { // Mark the owning assistant entry as finished and record timing. const sessionDoc = await this.workspaceDocument.getOrCreateSessionDoc(sessionId); - await sessionDoc.updateHistory((history) => { - for (let i = history.length - 1; i >= 0; i--) { - const entry = history[i]; - if (entry && entry.role === 'assistant' && (!turnId || entry.id === turnId)) { - entry.finished = true; - entry.endedAt = endedAt; - if (permissionWaitMs !== undefined) { - entry.permissionWaitMs = permissionWaitMs; - } - break; - } - } - return history; - }); + await sessionDoc.updateHistory((history) => + markAssistantTurnFinished(history, { turnId, endedAt, permissionWaitMs }) + ); await sessionDoc.waitUntilSynced(); } catch (error) { this.logger.error(`[${sessionId}] Failed to flush ACP updates during finalization:`, error); diff --git a/apps/cli/src/session/AGENTS.md b/apps/cli/src/session/AGENTS.md index a2b7f1adc..370195ba4 100644 --- a/apps/cli/src/session/AGENTS.md +++ b/apps/cli/src/session/AGENTS.md @@ -157,7 +157,11 @@ delegation proofs or a shared-machine gate without a new product and security de history mutation, leaving B pending and permanently unwatched. Because teardown/cancel finalize (`message-handler.ts` `finalizeACPState`, no-turnId overload) stamps `finished=true`/`endedAt` on the - in-progress entry, resume must **reopen** it: `writeAssistantEntryForTurn`'s + in-progress entry — and only that entry: `assistant-turn-finalize.ts` + `markAssistantTurnFinished` is a no-op on an already-finished one, because those + callers fire per live session at app close and a second stamp would rewrite a + long-finished turn's `endedAt` to now — resume must **reopen** it: + `writeAssistantEntryForTurn`'s existing-entry branch clears `finished`/`endedAt`/`permissionWaitMs` when re-adopting the entry for a live turn. Without that reset a machine-death-then-resume turn streams new output into a `finished=true` entry — the web renderer folds the still-streaming diff --git a/apps/cli/tests/message-handler-turn-duration.test.ts b/apps/cli/tests/message-handler-turn-duration.test.ts new file mode 100644 index 000000000..4df1d15b9 --- /dev/null +++ b/apps/cli/tests/message-handler-turn-duration.test.ts @@ -0,0 +1,148 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { LoroRepo } from 'loro-repo'; + +import type { SessionHistoryInput, SessionId, WorkspaceId } from '@lody/shared'; + +import { MessageHandler } from '../src/lib/message-handler'; +import { SessionDocument } from '../src/lib/loro/doc'; +import type { LoroDocumentManager } from '../src/lib/loro/doc'; +import type { SessionManager } from '../src/session/session-manager'; +import type { Logger } from '../src/utils/logger'; +import { loadEnv } from '../src/utils/const'; +import { createTestCloudPort } from './test-cloud-port'; + +const createSilentLogger = (): Logger => ({ + info: () => {}, + warn: () => {}, + error: () => {}, + success: () => {}, + debug: () => {}, + setLevel: () => {}, + child: () => createSilentLogger(), + close: async () => {}, +}); + +const originalLodyServerUrl = process.env.LODY_SERVER_URL; + +type MessageHandlerHost = { + finalizeACPState(sessionId: SessionId, turnId?: string): Promise; +}; + +const createHandlerHarness = async (sessionId: SessionId) => { + const logger = createSilentLogger(); + const repo = await LoroRepo.create({}); + const doc = new SessionDocument(repo, sessionId); + await doc.initOffline(); + + const workspaceDocument = { + isTransportConnected: vi.fn(() => true), + markMachineFlockDocDirty: vi.fn(), + registerMachine: vi.fn(), + repo: { + watch: vi.fn(() => ({ unsubscribe: vi.fn() })), + getDocMeta: vi.fn(async () => ({ + meta: { needToArchiveSessions: {}, needToDeleteSessions: {} }, + })), + }, + getOrCreateSessionDoc: vi.fn(async () => doc), + }; + const sessionManager = { + on: vi.fn(), + setRequestPermissionHandler: vi.fn(), + getSession: vi.fn(() => null), + }; + + const handler = new MessageHandler( + sessionManager as unknown as SessionManager, + workspaceDocument as unknown as LoroDocumentManager, + logger, + { + token: 't', + workspaceId: 'ws-1' as WorkspaceId, + userId: 'u-1', + machineId: 'm-1', + machineName: 'machine', + cliVersion: '0.0.0', + cloudPort: createTestCloudPort(), + } + ); + + return { repo, doc, handler: handler as unknown as MessageHandlerHost }; +}; + +const TURN_STARTED_AT = Date.parse('2026-01-01T00:00:00.000Z'); +const TURN_ENDED_AT = TURN_STARTED_AT + 12_000; +const APP_CLOSED_AT = TURN_STARTED_AT + 3_600_000; + +const assistantEntry = (): SessionHistoryInput => ({ + id: 'assistant:user-turn-1', + role: 'assistant', + timestamp: new Date(TURN_STARTED_AT).toISOString(), + fileDiff: [], + items: [{ type: 'text', text: 'done' }] as unknown as SessionHistoryInput['items'], +}); + +describe('MessageHandler turn duration (finalize stamps)', () => { + beforeEach(() => { + process.env.LODY_SERVER_URL = 'https://server.example.test'; + loadEnv(); + }); + + afterEach(() => { + vi.restoreAllMocks(); + vi.clearAllMocks(); + vi.unstubAllGlobals(); + if (originalLodyServerUrl === undefined) { + delete process.env.LODY_SERVER_URL; + } else { + process.env.LODY_SERVER_URL = originalLodyServerUrl; + } + loadEnv(); + }); + + it('keeps a finished turn timing when teardown finalizes again at app close', async () => { + // #260: closing the app raises exit/terminated for every live session, and + // both handlers run the no-turnId finalize whether or not the turn is over. + // The turn's `endedAt` used to be rewritten to the close-time clock, so the + // rendered "Worked for …" grew with however long the app had been open. + const sessionId = 's-duration-1' as SessionId; + const { repo, doc, handler } = await createHandlerHarness(sessionId); + + try { + await doc.updateHistory((history) => [...history, assistantEntry()]); + + const now = vi.spyOn(Date, 'now').mockReturnValue(TURN_ENDED_AT); + await handler.finalizeACPState(sessionId); + expect((await doc.getHistory())[0]).toMatchObject({ + finished: true, + endedAt: TURN_ENDED_AT, + }); + + now.mockReturnValue(APP_CLOSED_AT); + await handler.finalizeACPState(sessionId); + + expect((await doc.getHistory())[0]?.endedAt).toBe(TURN_ENDED_AT); + } finally { + await repo.destroy(); + } + }); + + it('still stamps a turn that was interrupted before it finished', async () => { + const sessionId = 's-duration-2' as SessionId; + const { repo, doc, handler } = await createHandlerHarness(sessionId); + + try { + await doc.updateHistory((history) => [...history, assistantEntry()]); + + vi.spyOn(Date, 'now').mockReturnValue(APP_CLOSED_AT); + await handler.finalizeACPState(sessionId); + + expect((await doc.getHistory())[0]).toMatchObject({ + finished: true, + endedAt: APP_CLOSED_AT, + }); + } finally { + await repo.destroy(); + } + }); +});