(null);
- const [menuFlip, setMenuFlip] = useState<{ vertical?: boolean; horizontal?: boolean }>({});
+ const [menuPos, setMenuPos] = useState<{
+ top: number;
+ left: number;
+ placement: 'up' | 'down';
+ } | null>(null);
const { copiedId, copy } = useCopyToClipboard();
const messageContent = item.type === 'message' ? item.content : null;
const isCopied = copiedId === `msg-${item.id}`;
- // Dynamically flip menu if it would overflow the viewport
- useLayoutEffect(() => {
- if (!rollbackMenuOpen || !rollbackMenuRef.current || !rollbackBtnRef.current) return;
- const menuRect = rollbackMenuRef.current.getBoundingClientRect();
- const flip: { vertical?: boolean; horizontal?: boolean } = {};
- // Account for title bar overlay (~36px on Windows, ~28px on macOS)
- const safeTop = 40;
- if (menuRect.top < safeTop) {
- flip.vertical = true;
+ const updateMenuPos = useCallback(() => {
+ if (!rollbackBtnRef.current) return;
+ const rect = rollbackBtnRef.current.getBoundingClientRect();
+ const GAP = 4;
+ let top = rect.top - MENU_HEIGHT_EST - GAP;
+ let placement: 'up' | 'down' = 'up';
+ if (top < 40) {
+ top = rect.bottom + GAP;
+ placement = 'down';
}
- if (menuRect.right > window.innerWidth) {
- flip.horizontal = true;
+ const left = Math.max(4, rect.right - MENU_WIDTH);
+ setMenuPos({ top, left, placement });
+ }, []);
+
+ useLayoutEffect(() => {
+ if (!rollbackMenuOpen) {
+ setMenuPos(null);
+ return;
}
- setMenuFlip(flip);
- }, [rollbackMenuOpen]);
+ updateMenuPos();
+ }, [rollbackMenuOpen, updateMenuPos]);
+
+ useEffect(() => {
+ if (!rollbackMenuOpen) return;
+ const handler = () => updateMenuPos();
+ window.addEventListener('resize', handler);
+ window.addEventListener('scroll', handler, true);
+ return () => {
+ window.removeEventListener('resize', handler);
+ window.removeEventListener('scroll', handler, true);
+ };
+ }, [rollbackMenuOpen, updateMenuPos]);
if (item.type === 'message') {
const content = item.content;
@@ -79,39 +104,64 @@ const MessageItem = memo(function MessageItem({
>
↩
- {rollbackMenuOpen && (
-
- {onRollbackHere && (
-
- )}
- {onForkFromHere && (
-
- )}
-
- )}
+ {rollbackMenuOpen &&
+ menuPos &&
+ createPortal(
+
+ {onRollbackHere && (
+
+ )}
+ {onForkFromHere && (
+
+ )}
+
,
+ document.body
+ )}
)}
diff --git a/packages/desktop/src/stores/agent.store.ts b/packages/desktop/src/stores/agent.store.ts
index 8b17bbef..512872aa 100644
--- a/packages/desktop/src/stores/agent.store.ts
+++ b/packages/desktop/src/stores/agent.store.ts
@@ -62,6 +62,7 @@ interface AgentActions {
threadId: string,
usage: { prompt: number; completion: number; total: number }
) => void;
+ clearThreadUsage: (threadId: string) => void;
loadThreads: (threads: Thread[]) => void;
updateToolCallStatus: (
threadId: string,
@@ -156,6 +157,11 @@ export const useAgentStore = create()(
s.usageByThreadId[threadId] = usage;
}),
+ clearThreadUsage: (threadId) =>
+ set((s) => {
+ delete s.usageByThreadId[threadId];
+ }),
+
loadThreads: (threads) => {
const incomingIds = new Set(threads.map((t) => t.id));
set((s) => {
diff --git a/packages/desktop/test/compact-usage-reset.test.ts b/packages/desktop/test/compact-usage-reset.test.ts
new file mode 100644
index 00000000..b394a84e
--- /dev/null
+++ b/packages/desktop/test/compact-usage-reset.test.ts
@@ -0,0 +1,154 @@
+/**
+ * @vitest-environment jsdom
+ */
+import { describe, it, expect, beforeEach } from 'vitest';
+import { useAgentStore } from '../src/stores/agent.store';
+import { useWorkspaceStore } from '../src/stores/workspace.store';
+
+// Reconstruct the side-effect block from useAgentCore.streamChunkToItem for the
+// 'reactive_compact' case. This mirrors the actual implementation so we can
+// exercise it without rendering the full hook.
+function handleReactiveCompact(threadId: string, event: { promptEstimate: number }): void {
+ const contextUsage = useAgentStore.getState().contextUsage;
+ if (contextUsage) {
+ useAgentStore.getState().setContextUsage({
+ used: event.promptEstimate,
+ contextWindow: contextUsage.contextWindow,
+ });
+ }
+ useAgentStore.getState().clearThreadUsage(threadId);
+}
+
+// Reconstruct the manual /compact handler from AgentWorkspace.ContextIndicator.
+// Mirrors the onClick body so we can drive it without rendering the component.
+async function runManualCompact(
+ threadId: string,
+ response: { didCompress: boolean; promptEstimate: number; released: number }
+): Promise {
+ // mirrors: if (res.didCompress && contextUsage) { setContextUsage(...); clearThreadUsage(threadId); }
+ const contextUsage = useAgentStore.getState().contextUsage;
+ if (response.didCompress && contextUsage) {
+ useAgentStore.getState().setContextUsage({
+ used: response.promptEstimate,
+ contextWindow: contextUsage.contextWindow,
+ });
+ useAgentStore.getState().clearThreadUsage(threadId);
+ }
+}
+
+beforeEach(() => {
+ useAgentStore.setState({
+ currentThreadId: 'thread-1',
+ threads: {
+ 'thread-1': {
+ id: 'thread-1',
+ projectId: '',
+ title: 't1',
+ cwd: '/test/cwd',
+ turns: [],
+ createdAt: 0,
+ updatedAt: 0,
+ },
+ },
+ approvalPolicy: 'ask-all',
+ model: 'model-1',
+ models: [{ id: 'model-1', provider: 'p', name: 'm1', context_window: 128000 } as any],
+ contextUsage: { used: 50000, contextWindow: 128000 },
+ todoByThreadId: {},
+ pendingInput: null,
+ usageByThreadId: {},
+ isCompressing: false,
+ automations: [],
+ });
+ useWorkspaceStore.setState({
+ rootPath: '/test/cwd',
+ name: 'test',
+ projects: [],
+ currentProjectId: '',
+ git: { branch: 'main', isDirty: false, staged: [], unstaged: [] },
+ });
+});
+
+describe('reactive_compact streaming handler', () => {
+ it('clears usageByThreadId for the affected thread only', () => {
+ useAgentStore
+ .getState()
+ .setThreadUsage('thread-1', { prompt: 3000, completion: 500, total: 3500 });
+ useAgentStore
+ .getState()
+ .setThreadUsage('thread-2', { prompt: 800, completion: 400, total: 1200 });
+
+ handleReactiveCompact('thread-1', { promptEstimate: 1200 });
+
+ expect(useAgentStore.getState().usageByThreadId['thread-1']).toBeUndefined();
+ expect(useAgentStore.getState().usageByThreadId['thread-2']).toEqual({
+ prompt: 800,
+ completion: 400,
+ total: 1200,
+ });
+ });
+
+ it('updates contextUsage.used to the new promptEstimate', () => {
+ useAgentStore.getState().setContextUsage({ used: 95000, contextWindow: 128000 });
+
+ handleReactiveCompact('thread-1', { promptEstimate: 1200 });
+
+ expect(useAgentStore.getState().contextUsage).toEqual({
+ used: 1200,
+ contextWindow: 128000,
+ });
+ });
+
+ it('does not throw when contextUsage is null (e.g., model not loaded)', () => {
+ useAgentStore.getState().setContextUsage(null);
+ useAgentStore
+ .getState()
+ .setThreadUsage('thread-1', { prompt: 3000, completion: 500, total: 3500 });
+
+ expect(() => handleReactiveCompact('thread-1', { promptEstimate: 1200 })).not.toThrow();
+ expect(useAgentStore.getState().usageByThreadId['thread-1']).toBeUndefined();
+ });
+});
+
+describe('manual /compact button handler (ContextIndicator)', () => {
+ it('clears usageByThreadId when didCompress is true', async () => {
+ useAgentStore
+ .getState()
+ .setThreadUsage('thread-1', { prompt: 3000, completion: 500, total: 3500 });
+
+ await runManualCompact('thread-1', {
+ didCompress: true,
+ promptEstimate: 1200,
+ released: 5000,
+ });
+
+ expect(useAgentStore.getState().usageByThreadId['thread-1']).toBeUndefined();
+ expect(useAgentStore.getState().contextUsage).toEqual({
+ used: 1200,
+ contextWindow: 128000,
+ });
+ });
+
+ it('does NOT update state when didCompress is false (below threshold)', async () => {
+ useAgentStore
+ .getState()
+ .setThreadUsage('thread-1', { prompt: 3000, completion: 500, total: 3500 });
+ useAgentStore.getState().setContextUsage({ used: 50000, contextWindow: 128000 });
+
+ await runManualCompact('thread-1', {
+ didCompress: false,
+ promptEstimate: 45000,
+ released: 0,
+ });
+
+ expect(useAgentStore.getState().usageByThreadId['thread-1']).toEqual({
+ prompt: 3000,
+ completion: 500,
+ total: 3500,
+ });
+ expect(useAgentStore.getState().contextUsage).toEqual({
+ used: 50000,
+ contextWindow: 128000,
+ });
+ });
+});
diff --git a/packages/desktop/test/fork-button-portal.test.tsx b/packages/desktop/test/fork-button-portal.test.tsx
new file mode 100644
index 00000000..925e946e
--- /dev/null
+++ b/packages/desktop/test/fork-button-portal.test.tsx
@@ -0,0 +1,303 @@
+/**
+ * @vitest-environment jsdom
+ */
+import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
+import { act, render, cleanup, fireEvent } from '@testing-library/react';
+import '@testing-library/jest-dom/vitest';
+import { useAgentStore } from '../src/stores/agent.store';
+import MessageStream from '../src/agent/MessageStream';
+import type { Turn } from '../shared/types';
+
+const forkThreadMock = vi.fn();
+const previewRollbackMock = vi.fn();
+
+vi.mock('@tanstack/react-virtual', () => ({
+ useVirtualizer: (options: { count: number }) => {
+ const count = options.count;
+ return {
+ getTotalSize: () => count * 60,
+ getVirtualItems: () =>
+ Array.from({ length: count }, (_, index) => ({
+ key: `row-${index}`,
+ index,
+ start: index * 60,
+ size: 60,
+ })),
+ measureElement: vi.fn(),
+ scrollToEnd: vi.fn(),
+ };
+ },
+}));
+
+vi.mock('../src/hooks/useAgent', () => ({
+ useAgentApproval: () => ({ approveTool: vi.fn(), rejectTool: vi.fn() }),
+ useAgentRollback: () => ({
+ loadCheckpointDiff: vi.fn().mockResolvedValue({ turnId: 0, files: [] }),
+ revertFile: vi.fn(),
+ revertFiles: vi.fn(),
+ previewRollback: (...args: unknown[]) => {
+ previewRollbackMock(...args);
+ return Promise.resolve({});
+ },
+ rollbackCtx: vi.fn(),
+ rollbackBoth: vi.fn(),
+ undoCodeRollback: vi.fn(),
+ forkThread: (...args: unknown[]) => forkThreadMock(...args),
+ initRollbackState: vi.fn(),
+ deleteThread: vi.fn(),
+ revertedFilesByTurnId: {},
+ }),
+}));
+
+vi.mock('../src/hooks/useCopyToClipboard', () => ({
+ useCopyToClipboard: () => ({ copiedId: null, copy: vi.fn() }),
+}));
+
+function makeTurn(id: string, items: Turn['items']): Turn {
+ return { id, items, status: 'completed' };
+}
+
+function setThread(threadId: string, turns: Turn[]) {
+ act(() => {
+ useAgentStore.setState((s) => {
+ s.threads[threadId] = {
+ id: threadId,
+ projectId: '',
+ title: threadId,
+ cwd: '/test/cwd',
+ turns,
+ createdAt: Date.now(),
+ updatedAt: Date.now(),
+ };
+ });
+ });
+}
+
+function mockGetBoundingClientRect(rect: Partial) {
+ const baseRect: DOMRect = {
+ top: 0,
+ left: 0,
+ right: 0,
+ bottom: 0,
+ width: 0,
+ height: 0,
+ x: 0,
+ y: 0,
+ toJSON() {
+ return {};
+ },
+ };
+ const merged = { ...baseRect, ...rect };
+ vi.spyOn(HTMLElement.prototype, 'getBoundingClientRect').mockImplementation(function () {
+ return merged;
+ });
+}
+
+beforeEach(() => {
+ cleanup();
+ vi.restoreAllMocks();
+ forkThreadMock.mockReset();
+ previewRollbackMock.mockReset();
+ forkThreadMock.mockResolvedValue('new-session-id-1234');
+ Object.defineProperty(window, 'innerWidth', { value: 1000, writable: true });
+ Object.defineProperty(window, 'innerHeight', { value: 800, writable: true });
+ useAgentStore.setState({
+ currentThreadId: null,
+ threads: {},
+ approvalPolicy: 'ask-all',
+ model: '',
+ models: [],
+ contextUsage: null,
+ todoByThreadId: {},
+ pendingInput: null,
+ usageByThreadId: {},
+ isCompressing: false,
+ automations: [],
+ });
+});
+
+afterEach(() => {
+ cleanup();
+ vi.restoreAllMocks();
+});
+
+describe('fork button via portal', () => {
+ it('renders the rollback menu in document.body (not inside the virtual row)', async () => {
+ setThread('t1', [
+ makeTurn('1', [
+ { id: 'u1', type: 'message', role: 'user', content: 'hi' },
+ { id: 'a1', type: 'message', role: 'assistant', content: 'hello' },
+ ]),
+ ]);
+ mockGetBoundingClientRect({ top: 200, right: 600, bottom: 220, width: 20, height: 20 });
+
+ const { container } = render();
+ const triggerBtn = container.querySelector('button[title="回退到此"]') as HTMLElement;
+ expect(triggerBtn).toBeTruthy();
+
+ await act(async () => {
+ fireEvent.click(triggerBtn);
+ });
+
+ const menu = document.body.querySelector('[data-testid="rollback-menu"]') as HTMLElement;
+ expect(menu).toBeTruthy();
+ expect(menu.style.position).toBe('fixed');
+ expect(menu.style.zIndex).toBe('100');
+ });
+
+ it('places menu at expected fixed coordinates from getBoundingClientRect', async () => {
+ setThread('t1', [
+ makeTurn('1', [
+ { id: 'u1', type: 'message', role: 'user', content: 'hi' },
+ { id: 'a1', type: 'message', role: 'assistant', content: 'hello' },
+ ]),
+ ]);
+ mockGetBoundingClientRect({ top: 200, right: 600, bottom: 220, width: 20, height: 20 });
+
+ const { container } = render();
+ const triggerBtn = container.querySelector('button[title="回退到此"]') as HTMLElement;
+ await act(async () => {
+ fireEvent.click(triggerBtn);
+ });
+
+ const menu = document.body.querySelector('[data-testid="rollback-menu"]') as HTMLElement;
+ expect(menu).toBeTruthy();
+ expect(menu.style.position).toBe('fixed');
+ expect(menu.getAttribute('data-placement')).toBe('up');
+ });
+
+ it('flips menu below the trigger when there is no room above', async () => {
+ setThread('t1', [
+ makeTurn('1', [
+ { id: 'u1', type: 'message', role: 'user', content: 'hi' },
+ { id: 'a1', type: 'message', role: 'assistant', content: 'hello' },
+ ]),
+ ]);
+ mockGetBoundingClientRect({ top: 10, right: 600, bottom: 30, width: 20, height: 20 });
+
+ const { container } = render();
+ const triggerBtn = container.querySelector('button[title="回退到此"]') as HTMLElement;
+ await act(async () => {
+ fireEvent.click(triggerBtn);
+ });
+
+ const menu = document.body.querySelector('[data-testid="rollback-menu"]') as HTMLElement;
+ expect(menu).toBeTruthy();
+ expect(menu.getAttribute('data-placement')).toBe('down');
+ });
+
+ it('updates menu position on scroll events', async () => {
+ setThread('t1', [
+ makeTurn('1', [
+ { id: 'u1', type: 'message', role: 'user', content: 'hi' },
+ { id: 'a1', type: 'message', role: 'assistant', content: 'hello' },
+ ]),
+ ]);
+ let rect: Partial = { top: 200, right: 600, bottom: 220, width: 20, height: 20 };
+ vi.spyOn(HTMLElement.prototype, 'getBoundingClientRect').mockImplementation(function () {
+ return {
+ ...{
+ top: 0,
+ left: 0,
+ right: 0,
+ bottom: 0,
+ width: 0,
+ height: 0,
+ x: 0,
+ y: 0,
+ toJSON: () => ({}),
+ },
+ ...rect,
+ } as DOMRect;
+ });
+
+ const { container } = render();
+ const triggerBtn = container.querySelector('button[title="回退到此"]') as HTMLElement;
+ await act(async () => {
+ fireEvent.click(triggerBtn);
+ });
+ const menu = document.body.querySelector('[data-testid="rollback-menu"]') as HTMLElement;
+ const initialTop = menu.style.top;
+ expect(initialTop).toBeTruthy();
+
+ rect = { top: 400, right: 600, bottom: 420, width: 20, height: 20 };
+ await act(async () => {
+ window.dispatchEvent(new Event('scroll'));
+ });
+
+ const updatedMenu = document.body.querySelector('[data-testid="rollback-menu"]') as HTMLElement;
+ expect(updatedMenu.style.top).not.toBe(initialTop);
+ });
+
+ it('clicking fork triggers forkThread with the correct threadId and numeric turnId', async () => {
+ setThread('t1', [
+ makeTurn('1', [
+ { id: 'u1', type: 'message', role: 'user', content: 'hello world' },
+ { id: 'a1', type: 'message', role: 'assistant', content: 'reply' },
+ ]),
+ ]);
+ mockGetBoundingClientRect({ top: 200, right: 600, bottom: 220, width: 20, height: 20 });
+
+ const { container } = render();
+ const triggerBtn = container.querySelector('button[title="回退到此"]') as HTMLElement;
+ await act(async () => {
+ fireEvent.click(triggerBtn);
+ });
+ const forkBtn = document.body.querySelector('[data-testid="fork-menu-item"]') as HTMLElement;
+ expect(forkBtn).toBeTruthy();
+ expect(forkBtn.textContent).toBe('Fork from here');
+
+ await act(async () => {
+ fireEvent.click(forkBtn);
+ });
+
+ expect(forkThreadMock).toHaveBeenCalledTimes(1);
+ expect(forkThreadMock).toHaveBeenCalledWith('t1', 1);
+ });
+
+ it('does not throw when forkThread rejects (try/catch surfaces error to console)', async () => {
+ const errSpy = vi.spyOn(console, 'error').mockImplementation(() => {});
+ forkThreadMock.mockRejectedValueOnce(new Error('network down'));
+
+ setThread('t1', [
+ makeTurn('1', [
+ { id: 'u1', type: 'message', role: 'user', content: 'hi' },
+ { id: 'a1', type: 'message', role: 'assistant', content: 'hello' },
+ ]),
+ ]);
+ mockGetBoundingClientRect({ top: 200, right: 600, bottom: 220, width: 20, height: 20 });
+
+ const { container } = render();
+ const triggerBtn = container.querySelector('button[title="回退到此"]') as HTMLElement;
+ await act(async () => {
+ fireEvent.click(triggerBtn);
+ });
+ const forkBtn = document.body.querySelector('[data-testid="fork-menu-item"]') as HTMLElement;
+ await act(async () => {
+ fireEvent.click(forkBtn);
+ });
+
+ expect(errSpy).toHaveBeenCalled();
+ errSpy.mockRestore();
+ });
+
+ it('menu is in document.body, not inside MessageStream DOM tree', async () => {
+ setThread('t1', [
+ makeTurn('1', [
+ { id: 'u1', type: 'message', role: 'user', content: 'hi' },
+ { id: 'a1', type: 'message', role: 'assistant', content: 'hello' },
+ ]),
+ ]);
+ mockGetBoundingClientRect({ top: 200, right: 600, bottom: 220, width: 20, height: 20 });
+
+ const { container } = render();
+ const triggerBtn = container.querySelector('button[title="回退到此"]') as HTMLElement;
+ await act(async () => {
+ fireEvent.click(triggerBtn);
+ });
+
+ const menu = document.body.querySelector('[data-testid="rollback-menu"]');
+ expect(container.contains(menu)).toBe(false);
+ expect(document.body.contains(menu)).toBe(true);
+ });
+});
diff --git a/packages/desktop/test/global-store.test.ts b/packages/desktop/test/global-store.test.ts
index a6bd9f38..b2277764 100644
--- a/packages/desktop/test/global-store.test.ts
+++ b/packages/desktop/test/global-store.test.ts
@@ -450,11 +450,9 @@ describe('global store - per-thread isStreaming derivation', () => {
useAgentStore.getState().startTurn(threadB, { id: 'turn-b', items: [], status: 'running' });
const isStreamingA = () =>
- useAgentStore.getState().threads[threadA]?.turns.some((t) => t.status === 'running') ??
- false;
+ useAgentStore.getState().threads[threadA]?.turns.some((t) => t.status === 'running') ?? false;
const isStreamingB = () =>
- useAgentStore.getState().threads[threadB]?.turns.some((t) => t.status === 'running') ??
- false;
+ useAgentStore.getState().threads[threadB]?.turns.some((t) => t.status === 'running') ?? false;
expect(isStreamingA()).toBe(true);
expect(isStreamingB()).toBe(true);
@@ -469,9 +467,8 @@ describe('global store - per-thread isStreaming derivation', () => {
it('thread with no running turns is not streaming', () => {
const threadId = 'thread-x';
const isStreaming = () =>
- useAgentStore
- .getState()
- .threads[threadId]?.turns.some((t) => t.status === 'running') ?? false;
+ useAgentStore.getState().threads[threadId]?.turns.some((t) => t.status === 'running') ??
+ false;
// Thread not yet created
expect(isStreaming()).toBe(false);
@@ -592,6 +589,30 @@ describe('global store - token usage', () => {
useAgentStore.getState().setCurrentThread('t1');
expect(useAgentStore.getState().contextUsage).toBeNull();
});
+
+ it('clearThreadUsage removes the entry for a single thread', () => {
+ useAgentStore.getState().setThreadUsage('t1', { prompt: 1000, completion: 500, total: 1500 });
+ useAgentStore.getState().setThreadUsage('t2', { prompt: 800, completion: 400, total: 1200 });
+ useAgentStore.getState().clearThreadUsage('t1');
+ expect(useAgentStore.getState().usageByThreadId['t1']).toBeUndefined();
+ expect(useAgentStore.getState().usageByThreadId['t2']).toEqual({
+ prompt: 800,
+ completion: 400,
+ total: 1200,
+ });
+ expect('t1' in useAgentStore.getState().usageByThreadId).toBe(false);
+ expect('t2' in useAgentStore.getState().usageByThreadId).toBe(true);
+ });
+
+ it('clearThreadUsage is a no-op for a threadId with no entry', () => {
+ useAgentStore.getState().setThreadUsage('t1', { prompt: 1000, completion: 500, total: 1500 });
+ useAgentStore.getState().clearThreadUsage('t-does-not-exist');
+ expect(useAgentStore.getState().usageByThreadId['t1']).toEqual({
+ prompt: 1000,
+ completion: 500,
+ total: 1500,
+ });
+ });
});
describe('global store - compressing state', () => {
@@ -614,9 +635,9 @@ describe('global store - compressing state', () => {
describe('global store - loadThreads orphan data cleanup', () => {
it('cleans up todoByThreadId for deleted threads', () => {
- useAgentStore.getState().applyTodoUpdate('deleted-thread', [
- { id: '1', text: 'todo', status: 'in_progress' },
- ]);
+ useAgentStore
+ .getState()
+ .applyTodoUpdate('deleted-thread', [{ id: '1', text: 'todo', status: 'in_progress' }]);
expect(useAgentStore.getState().todoByThreadId['deleted-thread']).toBeDefined();
useAgentStore.getState().loadThreads([]);
@@ -624,11 +645,19 @@ describe('global store - loadThreads orphan data cleanup', () => {
});
it('preserves todoByThreadId for threads still in the list', () => {
- useAgentStore.getState().applyTodoUpdate('kept-thread', [
- { id: '1', text: 'todo', status: 'in_progress' },
- ]);
+ useAgentStore
+ .getState()
+ .applyTodoUpdate('kept-thread', [{ id: '1', text: 'todo', status: 'in_progress' }]);
useAgentStore.getState().loadThreads([
- { id: 'kept-thread', projectId: '', title: 'test', cwd: '/x', turns: [], createdAt: 1, updatedAt: 2 },
+ {
+ id: 'kept-thread',
+ projectId: '',
+ title: 'test',
+ cwd: '/x',
+ turns: [],
+ createdAt: 1,
+ updatedAt: 2,
+ },
]);
expect(useAgentStore.getState().todoByThreadId['kept-thread']).toBeDefined();
});
@@ -644,7 +673,8 @@ describe('global store - loadThreads orphan data cleanup', () => {
it('cleans up checkpointDiffByTurnId for deleted threads', () => {
useRollbackStore.getState().setCheckpointDiff('deleted-thread', '1', {
- turnId: 1, files: [],
+ turnId: 1,
+ files: [],
} as any);
useAgentStore.getState().loadThreads([]);
expect(useRollbackStore.getState().checkpointDiffByTurnId['deleted-thread:1']).toBeUndefined();
@@ -668,13 +698,22 @@ describe('global store - loadThreads orphan data cleanup', () => {
code: { canUndoLast: false, lastEntry: null, revertedFiles: [], lastEntryId: '' },
} as any);
useRollbackStore.getState().setCheckpointDiff('kept-thread', '1', {
- turnId: 1, files: [],
+ turnId: 1,
+ files: [],
} as any);
useRollbackStore.getState().markFileReverted('kept-thread', '1', '/a.ts');
useRollbackStore.getState().setTurnCheckpointMapping('kept-thread', 1, 'ui-1');
useAgentStore.getState().loadThreads([
- { id: 'kept-thread', projectId: '', title: 'test', cwd: '/x', turns: [], createdAt: 1, updatedAt: 2 },
+ {
+ id: 'kept-thread',
+ projectId: '',
+ title: 'test',
+ cwd: '/x',
+ turns: [],
+ createdAt: 1,
+ updatedAt: 2,
+ },
]);
expect(useRollbackStore.getState().rollbackStateByThreadId['kept-thread']).toBeDefined();
diff --git a/packages/desktop/test/rollback-usage-reset.test.ts b/packages/desktop/test/rollback-usage-reset.test.ts
new file mode 100644
index 00000000..99e72129
--- /dev/null
+++ b/packages/desktop/test/rollback-usage-reset.test.ts
@@ -0,0 +1,287 @@
+/**
+ * @vitest-environment jsdom
+ */
+import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
+import { act, renderHook } from '@testing-library/react';
+import '@testing-library/jest-dom/vitest';
+import { useAgentStore } from '../src/stores/agent.store';
+import { useWorkspaceStore } from '../src/stores/workspace.store';
+import { useAgentRollback } from '../src/hooks/useAgent';
+
+const { rollbackContextMock, rollbackBothToTurnMock } = vi.hoisted(() => ({
+ rollbackContextMock: vi.fn(),
+ rollbackBothToTurnMock: vi.fn(),
+}));
+
+vi.mock('../src/lib/core-api', () => ({
+ rollbackContext: rollbackContextMock,
+ rollbackBothToTurn: rollbackBothToTurnMock,
+ deleteSession: vi.fn(),
+ listSessions: vi.fn(),
+ getCheckpointDiff: vi.fn(),
+ revertFile: vi.fn(),
+ revertFiles: vi.fn(),
+ previewRollbackDiff: vi.fn(),
+ rollbackCodeToTurn: vi.fn(),
+ undoLastCodeRollback: vi.fn(),
+ getRollbackState: vi.fn(),
+ forkSession: vi.fn(),
+ listModels: vi.fn(),
+ switchModel: vi.fn(),
+ listAgents: vi.fn(),
+ createSession: vi.fn(),
+ getSessionHistory: vi.fn(),
+ resumeSession: vi.fn(),
+ setSessionPermissionMode: vi.fn(),
+ sendApprovalResponse: vi.fn(),
+ getMemoryConfig: vi.fn(),
+ setMemoryEnabled: vi.fn(),
+ setMemoryTypeDisabled: vi.fn(),
+ createMemoryExtraType: vi.fn(),
+ updateMemoryExtraType: vi.fn(),
+ deleteMemoryExtraType: vi.fn(),
+ setMemoryModel: vi.fn(),
+ setAgentConfig: vi.fn(),
+ getAgentConfig: vi.fn(),
+ setCompactionModel: vi.fn(),
+ listMcpServers: vi.fn(),
+ setMcpDisabled: vi.fn(),
+ resetMcpDisabled: vi.fn(),
+ createMcpServer: vi.fn(),
+ updateMcpServer: vi.fn(),
+ deleteMcpServer: vi.fn(),
+ setAgentDisabled: vi.fn(),
+ resetAgentDisabled: vi.fn(),
+ createAgent: vi.fn(),
+ updateAgent: vi.fn(),
+ deleteAgent: vi.fn(),
+ getSubagentEnabled: vi.fn(),
+ setSubagentEnabled: vi.fn(),
+ resetSubagentEnabled: vi.fn(),
+ listSkills: vi.fn(),
+ toggleSkill: vi.fn(),
+ listHooks: vi.fn(),
+ createHook: vi.fn(),
+ updateHook: vi.fn(),
+ deleteHook: vi.fn(),
+ listAutomations: vi.fn(),
+ createAutomation: vi.fn(),
+ updateAutomation: vi.fn(),
+ runAutomationOnce: vi.fn(),
+}));
+
+function resetStores() {
+ useAgentStore.setState({
+ currentThreadId: 'thread-1',
+ threads: {
+ 'thread-1': {
+ id: 'thread-1',
+ projectId: '',
+ title: 't1',
+ cwd: '/test/cwd',
+ turns: [],
+ createdAt: 0,
+ updatedAt: 0,
+ },
+ },
+ approvalPolicy: 'ask-all',
+ model: 'model-1',
+ models: [{ id: 'model-1', provider: 'p', name: 'm1', context_window: 128000 } as any],
+ contextUsage: null,
+ todoByThreadId: {},
+ pendingInput: null,
+ usageByThreadId: {},
+ isCompressing: false,
+ automations: [],
+ });
+ useWorkspaceStore.setState({
+ rootPath: '/test/cwd',
+ name: 'test',
+ projects: [],
+ currentProjectId: '',
+ git: { branch: 'main', isDirty: false, staged: [], unstaged: [] },
+ });
+}
+
+beforeEach(() => {
+ rollbackContextMock.mockReset();
+ rollbackBothToTurnMock.mockReset();
+ resetStores();
+});
+
+afterEach(() => {
+ vi.restoreAllMocks();
+});
+
+describe('useAgentRollback().rollbackCtx - per-thread usage from server', () => {
+ it('adopts the server usage when the rolled-back state still has prior assistant usage', async () => {
+ act(() => {
+ useAgentStore.getState().setThreadUsage('thread-1', {
+ prompt: 3000,
+ completion: 500,
+ total: 3500,
+ });
+ });
+
+ rollbackContextMock.mockResolvedValue({
+ turns: [],
+ rolledBackMessage: null,
+ promptEstimate: 1200,
+ usage: { prompt: 800, completion: 400, total: 1200 },
+ });
+
+ const { result } = renderHook(() => useAgentRollback());
+ await act(async () => {
+ await result.current.rollbackCtx('thread-1', 2);
+ });
+
+ expect(useAgentStore.getState().usageByThreadId['thread-1']).toEqual({
+ prompt: 800,
+ completion: 400,
+ total: 1200,
+ });
+ expect(useAgentStore.getState().contextUsage).toEqual({
+ used: 1200,
+ contextWindow: 128000,
+ });
+ });
+
+ it('falls back to zeros when the server returns no usage (first-round rollback)', async () => {
+ act(() => {
+ useAgentStore.getState().setThreadUsage('thread-1', {
+ prompt: 3000,
+ completion: 500,
+ total: 3500,
+ });
+ });
+
+ rollbackContextMock.mockResolvedValue({
+ turns: [],
+ rolledBackMessage: 'first prompt',
+ promptEstimate: 0,
+ });
+
+ const { result } = renderHook(() => useAgentRollback());
+ await act(async () => {
+ await result.current.rollbackCtx('thread-1', 1);
+ });
+
+ expect(useAgentStore.getState().usageByThreadId['thread-1']).toEqual({
+ prompt: 0,
+ completion: 0,
+ total: 0,
+ });
+ expect(useAgentStore.getState().contextUsage).toEqual({
+ used: 0,
+ contextWindow: 128000,
+ });
+ });
+
+ it('uses promptEstimate for contextUsage.used when usage is also provided', async () => {
+ rollbackContextMock.mockResolvedValue({
+ turns: [],
+ rolledBackMessage: null,
+ promptEstimate: 1234,
+ usage: { prompt: 800, completion: 400, total: 1200 },
+ });
+
+ const { result } = renderHook(() => useAgentRollback());
+ await act(async () => {
+ await result.current.rollbackCtx('thread-1', 2);
+ });
+
+ expect(useAgentStore.getState().contextUsage?.used).toBe(1234);
+ expect(useAgentStore.getState().usageByThreadId['thread-1']).toEqual({
+ prompt: 800,
+ completion: 400,
+ total: 1200,
+ });
+ });
+
+ it('refills the rolled-back message into pendingInput', async () => {
+ rollbackContextMock.mockResolvedValue({
+ turns: [],
+ rolledBackMessage: 'first prompt',
+ promptEstimate: 0,
+ usage: undefined,
+ });
+
+ const { result } = renderHook(() => useAgentRollback());
+ await act(async () => {
+ await result.current.rollbackCtx('thread-1', 1);
+ });
+
+ expect(useAgentStore.getState().pendingInput).toBe('first prompt');
+ });
+});
+
+describe('useAgentRollback().rollbackBoth - per-thread usage from server', () => {
+ it('adopts the server usage when the rolled-back state still has prior assistant usage', async () => {
+ act(() => {
+ useAgentStore.getState().setThreadUsage('thread-1', {
+ prompt: 3000,
+ completion: 500,
+ total: 3500,
+ });
+ });
+
+ rollbackBothToTurnMock.mockResolvedValue({
+ turns: [],
+ rolledBackMessage: null,
+ codeResult: {
+ reverted: false,
+ throughTurnId: 0,
+ affectedTurns: [],
+ selectedFiles: [],
+ restoreEntry: null,
+ },
+ promptEstimate: 1200,
+ usage: { prompt: 800, completion: 400, total: 1200 },
+ });
+
+ const { result } = renderHook(() => useAgentRollback());
+ await act(async () => {
+ await result.current.rollbackBoth('thread-1', 2);
+ });
+
+ expect(useAgentStore.getState().usageByThreadId['thread-1']).toEqual({
+ prompt: 800,
+ completion: 400,
+ total: 1200,
+ });
+ expect(useAgentStore.getState().contextUsage).toEqual({
+ used: 1200,
+ contextWindow: 128000,
+ });
+ });
+
+ it('falls back to zeros when the server returns no usage', async () => {
+ rollbackBothToTurnMock.mockResolvedValue({
+ turns: [],
+ rolledBackMessage: null,
+ codeResult: {
+ reverted: false,
+ throughTurnId: 0,
+ affectedTurns: [],
+ selectedFiles: [],
+ restoreEntry: null,
+ },
+ promptEstimate: 0,
+ });
+
+ const { result } = renderHook(() => useAgentRollback());
+ await act(async () => {
+ await result.current.rollbackBoth('thread-1', 1);
+ });
+
+ expect(useAgentStore.getState().usageByThreadId['thread-1']).toEqual({
+ prompt: 0,
+ completion: 0,
+ total: 0,
+ });
+ expect(useAgentStore.getState().contextUsage).toEqual({
+ used: 0,
+ contextWindow: 128000,
+ });
+ });
+});
diff --git a/packages/desktop/test/sidebar-thread-switch.test.tsx b/packages/desktop/test/sidebar-thread-switch.test.tsx
new file mode 100644
index 00000000..3571a4e7
--- /dev/null
+++ b/packages/desktop/test/sidebar-thread-switch.test.tsx
@@ -0,0 +1,129 @@
+/**
+ * @vitest-environment jsdom
+ */
+import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
+import { act, render, fireEvent, cleanup } from '@testing-library/react';
+import '@testing-library/jest-dom/vitest';
+import { useAgentStore } from '../src/stores/agent.store';
+import { useUIStore } from '../src/stores/ui.store';
+import { useWorkspaceStore } from '../src/stores/workspace.store';
+import AgentSidebar from '../src/agent/AgentSidebar';
+
+const deleteThreadMock = vi.fn();
+
+vi.mock('../src/hooks/useAgent', () => ({
+ useAgentRollback: () => ({
+ deleteThread: deleteThreadMock,
+ }),
+ useAgentApproval: () => ({ approveTool: vi.fn(), rejectTool: vi.fn() }),
+}));
+
+function resetStores({ currentThreadId = null as string | null } = {}) {
+ useAgentStore.setState({
+ currentThreadId,
+ threads: {},
+ approvalPolicy: 'ask-all',
+ model: '',
+ models: [],
+ contextUsage: null,
+ todoByThreadId: {},
+ pendingInput: null,
+ usageByThreadId: {},
+ isCompressing: false,
+ automations: [],
+ });
+ useUIStore.setState({
+ view: 'agent',
+ sidebarCollapsed: false,
+ settingsInitialTab: null,
+ });
+ useWorkspaceStore.setState({
+ rootPath: '/test/cwd',
+ name: 'test',
+ projects: [],
+ currentProjectId: '',
+ git: { branch: 'main', isDirty: false, staged: [], unstaged: [] },
+ });
+}
+
+function seedThreads() {
+ act(() => {
+ useAgentStore.setState((s) => {
+ s.threads = {
+ 't-1': {
+ id: 't-1',
+ projectId: '',
+ title: 'first',
+ cwd: '/test/cwd',
+ turns: [],
+ createdAt: 1000,
+ updatedAt: 1000,
+ },
+ 't-2': {
+ id: 't-2',
+ projectId: '',
+ title: 'second',
+ cwd: '/test/cwd',
+ turns: [],
+ createdAt: 2000,
+ updatedAt: 2000,
+ },
+ };
+ });
+ });
+}
+
+beforeEach(() => {
+ cleanup();
+ deleteThreadMock.mockReset();
+ resetStores();
+});
+
+afterEach(() => {
+ cleanup();
+ vi.restoreAllMocks();
+});
+
+describe('AgentSidebar session switching and new session', () => {
+ it('clicking the "新对话" button sets currentThreadId to null', () => {
+ seedThreads();
+ act(() => {
+ useAgentStore.setState((s) => {
+ s.currentThreadId = 't-1';
+ });
+ });
+ const { getByText } = render();
+ const newBtn = getByText('新对话');
+ act(() => {
+ fireEvent.click(newBtn);
+ });
+ expect(useAgentStore.getState().currentThreadId).toBeNull();
+ });
+
+ it('clicking a session item in the list sets currentThreadId to that session', () => {
+ seedThreads();
+ act(() => {
+ useAgentStore.setState((s) => {
+ s.currentThreadId = 't-1';
+ });
+ });
+ const { getByText } = render();
+ const secondSessionBtn = getByText('second');
+ act(() => {
+ fireEvent.click(secondSessionBtn);
+ });
+ expect(useAgentStore.getState().currentThreadId).toBe('t-2');
+ });
+
+ it('clicking a session item does not throw and does not invoke deleteThread', () => {
+ seedThreads();
+ const { getByText } = render();
+ const sessionBtn = getByText('first');
+ expect(() => {
+ act(() => {
+ fireEvent.click(sessionBtn);
+ });
+ }).not.toThrow();
+ expect(deleteThreadMock).not.toHaveBeenCalled();
+ });
+});
diff --git a/packages/desktop/test/thread-delete.test.ts b/packages/desktop/test/thread-delete.test.ts
new file mode 100644
index 00000000..9a2b290e
--- /dev/null
+++ b/packages/desktop/test/thread-delete.test.ts
@@ -0,0 +1,212 @@
+/**
+ * @vitest-environment jsdom
+ */
+import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
+import { act, renderHook } from '@testing-library/react';
+import '@testing-library/jest-dom/vitest';
+import { useAgentStore } from '../src/stores/agent.store';
+import { useWorkspaceStore } from '../src/stores/workspace.store';
+import { useAgentRollback } from '../src/hooks/useAgent';
+
+const { deleteSessionMock, listSessionsMock } = vi.hoisted(() => ({
+ deleteSessionMock: vi.fn(),
+ listSessionsMock: vi.fn(),
+}));
+
+vi.mock('../src/lib/core-api', () => ({
+ deleteSession: deleteSessionMock,
+ listSessions: listSessionsMock,
+ // other named exports are stubs (useAgentRollback only uses these two in deleteThread)
+ getCheckpointDiff: vi.fn(),
+ revertFile: vi.fn(),
+ revertFiles: vi.fn(),
+ previewRollbackDiff: vi.fn(),
+ rollbackCodeToTurn: vi.fn(),
+ rollbackContext: vi.fn(),
+ rollbackBothToTurn: vi.fn(),
+ undoLastCodeRollback: vi.fn(),
+ getRollbackState: vi.fn(),
+ forkSession: vi.fn(),
+ listModels: vi.fn(),
+ switchModel: vi.fn(),
+ listAgents: vi.fn(),
+ createSession: vi.fn(),
+ getSessionHistory: vi.fn(),
+ resumeSession: vi.fn(),
+ setSessionPermissionMode: vi.fn(),
+ sendApprovalResponse: vi.fn(),
+ getMemoryConfig: vi.fn(),
+ setMemoryEnabled: vi.fn(),
+ setMemoryTypeDisabled: vi.fn(),
+ createMemoryExtraType: vi.fn(),
+ updateMemoryExtraType: vi.fn(),
+ deleteMemoryExtraType: vi.fn(),
+ setMemoryModel: vi.fn(),
+ setAgentConfig: vi.fn(),
+ getAgentConfig: vi.fn(),
+ setCompactionModel: vi.fn(),
+ listMcpServers: vi.fn(),
+ setMcpDisabled: vi.fn(),
+ resetMcpDisabled: vi.fn(),
+ createMcpServer: vi.fn(),
+ updateMcpServer: vi.fn(),
+ deleteMcpServer: vi.fn(),
+ setAgentDisabled: vi.fn(),
+ resetAgentDisabled: vi.fn(),
+ createAgent: vi.fn(),
+ updateAgent: vi.fn(),
+ deleteAgent: vi.fn(),
+ getSubagentEnabled: vi.fn(),
+ setSubagentEnabled: vi.fn(),
+ resetSubagentEnabled: vi.fn(),
+ listSkills: vi.fn(),
+ toggleSkill: vi.fn(),
+ listHooks: vi.fn(),
+ createHook: vi.fn(),
+ updateHook: vi.fn(),
+ deleteHook: vi.fn(),
+ listAutomations: vi.fn(),
+ createAutomation: vi.fn(),
+ updateAutomation: vi.fn(),
+ runAutomationOnce: vi.fn(),
+}));
+
+function resetStores({ rootPath = '/test/cwd' }: { rootPath?: string } = {}) {
+ useAgentStore.setState({
+ currentThreadId: null,
+ threads: {},
+ approvalPolicy: 'ask-all',
+ model: '',
+ models: [],
+ contextUsage: null,
+ todoByThreadId: {},
+ pendingInput: null,
+ usageByThreadId: {},
+ isCompressing: false,
+ automations: [],
+ });
+ useWorkspaceStore.setState({
+ rootPath,
+ name: 'test',
+ projects: [],
+ currentProjectId: '',
+ git: { branch: 'main', isDirty: false, staged: [], unstaged: [] },
+ });
+}
+
+beforeEach(() => {
+ deleteSessionMock.mockReset();
+ listSessionsMock.mockReset();
+ deleteSessionMock.mockResolvedValue(undefined);
+ listSessionsMock.mockResolvedValue([]);
+ resetStores();
+});
+
+afterEach(() => {
+ vi.restoreAllMocks();
+});
+
+describe('useAgentRollback().deleteThread', () => {
+ it('calls deleteSession with the workspace rootPath as cwd', async () => {
+ const { result } = renderHook(() => useAgentRollback());
+ await act(async () => {
+ await result.current.deleteThread('thread-1');
+ });
+ expect(deleteSessionMock).toHaveBeenCalledTimes(1);
+ expect(deleteSessionMock).toHaveBeenCalledWith('thread-1', '/test/cwd');
+ });
+
+ it('refreshes the thread list after a successful delete', async () => {
+ listSessionsMock.mockResolvedValue([
+ {
+ sessionId: 's-1',
+ title: 'one',
+ cwd: '/test/cwd',
+ createdAt: '2024-01-01',
+ updatedAt: '2024-01-02',
+ },
+ ]);
+ const { result } = renderHook(() => useAgentRollback());
+ await act(async () => {
+ await result.current.deleteThread('thread-1');
+ });
+ expect(listSessionsMock).toHaveBeenCalledWith('/test/cwd');
+ const threads = useAgentStore.getState().threads;
+ expect(Object.keys(threads)).toEqual(['s-1']);
+ expect(threads['s-1']?.title).toBe('one');
+ });
+
+ it('does not call listSessions when rootPath is empty', async () => {
+ resetStores({ rootPath: '' });
+ const { result } = renderHook(() => useAgentRollback());
+ await act(async () => {
+ await result.current.deleteThread('thread-1');
+ });
+ expect(deleteSessionMock).toHaveBeenCalledWith('thread-1', '');
+ expect(listSessionsMock).not.toHaveBeenCalled();
+ });
+
+ it('clears currentThreadId when the deleted thread is the current one', async () => {
+ act(() => {
+ useAgentStore.setState((s) => {
+ s.currentThreadId = 'thread-current';
+ });
+ });
+ const { result } = renderHook(() => useAgentRollback());
+ await act(async () => {
+ await result.current.deleteThread('thread-current');
+ });
+ expect(useAgentStore.getState().currentThreadId).toBeNull();
+ });
+
+ it('leaves currentThreadId unchanged when deleting a non-current thread', async () => {
+ act(() => {
+ useAgentStore.setState((s) => {
+ s.currentThreadId = 'thread-keep';
+ });
+ });
+ const { result } = renderHook(() => useAgentRollback());
+ await act(async () => {
+ await result.current.deleteThread('thread-other');
+ });
+ expect(useAgentStore.getState().currentThreadId).toBe('thread-keep');
+ });
+
+ it('still clears currentThreadId even if the server delete fails', async () => {
+ deleteSessionMock.mockRejectedValueOnce(new Error('network down'));
+ act(() => {
+ useAgentStore.setState((s) => {
+ s.currentThreadId = 'thread-current';
+ });
+ });
+ const errSpy = vi.spyOn(console, 'error').mockImplementation(() => {});
+ const { result } = renderHook(() => useAgentRollback());
+ await act(async () => {
+ await result.current.deleteThread('thread-current');
+ });
+ expect(useAgentStore.getState().currentThreadId).toBeNull();
+ expect(errSpy).toHaveBeenCalled();
+ errSpy.mockRestore();
+ });
+
+ it('still refreshes the list even if the server delete fails', async () => {
+ deleteSessionMock.mockRejectedValueOnce(new Error('network down'));
+ listSessionsMock.mockResolvedValue([
+ {
+ sessionId: 's-2',
+ title: 'two',
+ cwd: '/test/cwd',
+ createdAt: '2024-01-01',
+ updatedAt: '2024-01-02',
+ },
+ ]);
+ const errSpy = vi.spyOn(console, 'error').mockImplementation(() => {});
+ const { result } = renderHook(() => useAgentRollback());
+ await act(async () => {
+ await result.current.deleteThread('thread-1');
+ });
+ expect(listSessionsMock).toHaveBeenCalled();
+ expect(Object.keys(useAgentStore.getState().threads)).toEqual(['s-2']);
+ errSpy.mockRestore();
+ });
+});
diff --git a/packages/infra/src/config.ts b/packages/infra/src/config.ts
index 38bc7a89..32f7a58f 100644
--- a/packages/infra/src/config.ts
+++ b/packages/infra/src/config.ts
@@ -24,6 +24,7 @@ export interface MemoryConfig {
model: string;
extraTypes: MemoryTypeConfig[];
disabledTypes: string[];
+ promptMaxBytes: number;
}
export interface ActiveModelConfig {
@@ -57,6 +58,7 @@ export const DEFAULT_MEMORY: MemoryConfig = {
model: '',
extraTypes: [],
disabledTypes: [],
+ promptMaxBytes: 8192,
};
export const DEFAULT_CONFIG: AppConfig = {
diff --git a/packages/tui/src/utils.ts b/packages/tui/src/utils.ts
index 09d6b3fb..b8774433 100644
--- a/packages/tui/src/utils.ts
+++ b/packages/tui/src/utils.ts
@@ -3,13 +3,12 @@ import type { UIMessage } from './types.js';
type SessionEvent = {
type: string;
- uuid: string;
+ turnId?: number;
content?: string;
output?: string;
- timestamp: string;
model?: string;
toolName?: string;
- toolCalls?: any[];
+ toolCallId?: string;
};
export function generateId(): string {
@@ -21,36 +20,54 @@ export function formatTime(ts: number): string {
return `${d.getHours()}:${String(d.getMinutes()).padStart(2, '0')}`;
}
+function createTurnScopedIdGenerator() {
+ const counters = new Map();
+ return (prefix: string, turnId: number): string => {
+ const key = `${prefix}:${turnId}`;
+ const next = (counters.get(key) ?? 0) + 1;
+ counters.set(key, next);
+ return `${prefix}-${turnId}-${next}`;
+ };
+}
+
export function historyToUIMessages(history: SessionEvent[]): UIMessage[] {
const messages: UIMessage[] = [];
+ const nextId = createTurnScopedIdGenerator();
+
for (const event of history) {
switch (event.type) {
- case 'user':
+ case 'user': {
+ if (event.turnId === undefined) break;
messages.push({
- id: event.uuid,
- timestamp: new Date(event.timestamp).getTime(),
+ id: nextId('user', event.turnId),
+ timestamp: Date.now(),
role: 'user',
content: event.content ?? '',
});
break;
- case 'assistant':
+ }
+ case 'assistant': {
+ if (event.turnId === undefined) break;
messages.push({
- id: event.uuid,
- timestamp: new Date(event.timestamp).getTime(),
+ id: nextId('assistant', event.turnId),
+ timestamp: Date.now(),
role: 'assistant',
content: event.content ?? '',
model: event.model,
});
break;
- case 'tool_result':
+ }
+ case 'tool_result': {
+ if (event.toolCallId === undefined) break;
messages.push({
- id: event.uuid,
- timestamp: new Date(event.timestamp).getTime(),
+ id: `result-${event.toolCallId}`,
+ timestamp: Date.now(),
role: 'tool',
content: event.output ?? '',
toolName: event.toolName,
});
break;
+ }
}
}
return messages;
diff --git a/packages/tui/test/utils.test.ts b/packages/tui/test/utils.test.ts
index 81154b44..43d976b4 100644
--- a/packages/tui/test/utils.test.ts
+++ b/packages/tui/test/utils.test.ts
@@ -7,56 +7,53 @@ describe('historyToUIMessages', () => {
});
it('should convert user events to UIMessage', () => {
- const history = [
- { type: 'user', uuid: 'u1', content: 'hello', timestamp: '2025-01-01T00:00:00.000Z' },
- ];
+ const history = [{ type: 'user', turnId: 1, content: 'hello' }];
const result = historyToUIMessages(history);
expect(result).toHaveLength(1);
expect(result[0]).toMatchObject({
- id: 'u1',
+ id: 'user-1-1',
role: 'user',
content: 'hello',
});
+ expect(typeof result[0].timestamp).toBe('number');
});
it('should convert assistant events to UIMessage', () => {
- const history = [
- { type: 'assistant', uuid: 'a1', content: 'hi there', timestamp: '2025-01-01T00:00:00.000Z' },
- ];
+ const history = [{ type: 'assistant', turnId: 1, content: 'hi there' }];
const result = historyToUIMessages(history);
expect(result).toHaveLength(1);
expect(result[0]).toMatchObject({
- id: 'a1',
+ id: 'assistant-1-1',
role: 'assistant',
content: 'hi there',
});
+ expect(typeof result[0].timestamp).toBe('number');
});
it('should convert tool_result events to UIMessage with toolName', () => {
const history = [
{
type: 'tool_result',
- uuid: 't1',
+ toolCallId: 'tc1',
output: 'result',
- timestamp: '2025-01-01T00:00:00.000Z',
toolName: 'read',
},
];
const result = historyToUIMessages(history);
expect(result).toHaveLength(1);
expect(result[0]).toMatchObject({
- id: 't1',
+ id: 'result-tc1',
role: 'tool',
content: 'result',
toolName: 'read',
});
+ expect(typeof result[0].timestamp).toBe('number');
});
it('should skip session_meta, role_switch, and compact_boundary events', () => {
const history = [
{
type: 'session_meta',
- uuid: 'm1',
sessionId: 's1',
projectSlug: 'test',
cwd: '/',
@@ -65,16 +62,14 @@ describe('historyToUIMessages', () => {
createdAt: '',
version: '1',
},
- { type: 'user', uuid: 'u1', content: 'hello', timestamp: '2025-01-01T00:00:00.000Z' },
- { type: 'role_switch', uuid: 'r1', fromRole: 'a', toRole: 'b', timestamp: '' },
- { type: 'assistant', uuid: 'a1', content: 'hi', timestamp: '2025-01-01T00:00:00.000Z' },
+ { type: 'user', turnId: 1, content: 'hello' },
+ { type: 'role_switch', fromRole: 'a', toRole: 'b' },
+ { type: 'assistant', turnId: 2, content: 'hi' },
{
type: 'compact_boundary',
- uuid: 'c1',
summary: '...',
replacedRange: [0, 1],
messageCount: 1,
- timestamp: '',
},
];
const result = historyToUIMessages(history);
@@ -85,30 +80,25 @@ describe('historyToUIMessages', () => {
it('should handle conversation with interleaved tool calls', () => {
const history = [
- { type: 'user', uuid: 'u1', content: 'read file', timestamp: '2025-01-01T00:00:01.000Z' },
+ { type: 'user', turnId: 1, content: 'read file' },
{
type: 'assistant',
- uuid: 'a1',
+ turnId: 2,
content: 'let me read that',
- timestamp: '2025-01-01T00:00:02.000Z',
model: 'test-model',
toolCalls: [{ id: 'tc1', name: 'read', arguments: '{}' }],
},
{
type: 'tool_result',
- uuid: 't1',
content: undefined,
output: 'file contents here',
- timestamp: '2025-01-01T00:00:03.000Z',
toolName: 'read',
- parentUuid: 'a1',
toolCallId: 'tc1',
},
{
type: 'assistant',
- uuid: 'a2',
+ turnId: 3,
content: 'the file contains...',
- timestamp: '2025-01-01T00:00:04.000Z',
model: 'test-model',
toolCalls: [],
},
@@ -116,21 +106,41 @@ describe('historyToUIMessages', () => {
const result = historyToUIMessages(history);
expect(result).toHaveLength(4);
expect(result[0].role).toBe('user');
+ expect(result[0].id).toBe('user-1-1');
expect(result[1].role).toBe('assistant');
+ expect(result[1].id).toBe('assistant-2-1');
expect(result[1].model).toBe('test-model');
expect(result[2].role).toBe('tool');
+ expect(result[2].id).toBe('result-tc1');
expect(result[2].toolName).toBe('read');
expect(result[2].content).toBe('file contents here');
expect(result[3].role).toBe('assistant');
+ expect(result[3].id).toBe('assistant-3-1');
});
it('should preserve message order from history', () => {
const history = [
- { type: 'user', uuid: 'u1', content: 'msg1', timestamp: '2025-01-01T00:00:01.000Z' },
- { type: 'user', uuid: 'u2', content: 'msg2', timestamp: '2025-01-01T00:00:02.000Z' },
- { type: 'user', uuid: 'u3', content: 'msg3', timestamp: '2025-01-01T00:00:03.000Z' },
+ { type: 'user', turnId: 1, content: 'msg1' },
+ { type: 'user', turnId: 2, content: 'msg2' },
+ { type: 'user', turnId: 3, content: 'msg3' },
+ ];
+ const result = historyToUIMessages(history);
+ expect(result.map((m) => m.id)).toEqual(['user-1-1', 'user-2-1', 'user-3-1']);
+ });
+
+ it('should scope per-turn ids independently for same turn', () => {
+ const history = [
+ { type: 'user', turnId: 1, content: 'msg1' },
+ { type: 'user', turnId: 1, content: 'msg2' },
+ { type: 'assistant', turnId: 1, content: 'msg3' },
+ { type: 'assistant', turnId: 1, content: 'msg4' },
];
const result = historyToUIMessages(history);
- expect(result.map((m) => m.id)).toEqual(['u1', 'u2', 'u3']);
+ expect(result.map((m) => m.id)).toEqual([
+ 'user-1-1',
+ 'user-1-2',
+ 'assistant-1-1',
+ 'assistant-1-2',
+ ]);
});
});