tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

Compaction overhaul (#5186)

Browse Source 
Co-authored-by: Douwe Osinga <douwe@squareup.com>
Co-authored-by: David Katz <dkatz@squareup.com>
Co-authored-by: Alex Hancock <alexhancock@block.xyz>
This commit is contained in:
Douwe Osinga
2025-10-15 18:39:12 -04:00
committed by GitHub
parent ef9e68f2c6
commit 7c1f5dcb2f
56 changed files with 3069 additions and 4465 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ui/desktop/src/components/context_management/CompactionMarker.tsx
+3 -3
View File
@@ -1,5 +1,5 @@
import React from 'react';
import { Message, SummarizationRequested } from '../../api';
import { Message, ConversationCompacted } from '../../api';
interface CompactionMarkerProps {
message: Message;
@@ -7,8 +7,8 @@ interface CompactionMarkerProps {
export const CompactionMarker: React.FC<CompactionMarkerProps> = ({ message }) => {
const compactionContent = message.content.find(
(content): content is SummarizationRequested & { type: 'summarizationRequested' } =>
content.type === 'summarizationRequested'
(content): content is ConversationCompacted & { type: 'conversationCompacted' } =>
content.type === 'conversationCompacted'
);
const markerText = compactionContent?.msg || 'Conversation compacted';
ui/desktop/src/components/context_management/ContextManager.tsx
+4 -21
View File
@@ -9,12 +9,6 @@ interface ContextManagerState {
}
interface ContextManagerActions {
handleAutoCompaction: (
messages: Message[],
setMessages: (messages: Message[]) => void,
append: (message: Message) => void,
sessionId: string
) => Promise<void>;
handleManualCompaction: (
messages: Message[],
setMessages: (messages: Message[]) => void,
@@ -70,6 +64,7 @@ export const ContextManagerProvider: React.FC<{ children: React.ReactNode }> = (
setIsCompacting(false);
} catch (err) {
// TODO(Douwe): move this to the server
console.error('Error during compaction:', err);
setCompactionError(err instanceof Error ? err.message : 'Unknown error during compaction');
@@ -80,10 +75,11 @@ export const ContextManagerProvider: React.FC<{ children: React.ReactNode }> = (
created: Math.floor(Date.now() / 1000),
content: [
{
type: 'summarizationRequested',
type: 'conversationCompacted',
msg: 'Compaction failed. Please try again or start a new session.',
},
],
metadata: { userVisible: true, agentVisible: true },
};
setMessages([...messages, errorMarker]);
@@ -93,18 +89,6 @@ export const ContextManagerProvider: React.FC<{ children: React.ReactNode }> = (
[]
);
const handleAutoCompaction = useCallback(
async (
messages: Message[],
setMessages: (messages: Message[]) => void,
append: (message: Message) => void,
sessionId: string
) => {
await performCompaction(messages, setMessages, append, sessionId, false);
},
[performCompaction]
);
const handleManualCompaction = useCallback(
async (
messages: Message[],
@@ -118,7 +102,7 @@ export const ContextManagerProvider: React.FC<{ children: React.ReactNode }> = (
);
const hasCompactionMarker = useCallback((message: Message): boolean => {
return message.content.some((content) => content.type === 'summarizationRequested');
return message.content.some((content) => content.type === 'conversationCompacted');
}, []);
const value = {
@@ -127,7 +111,6 @@ export const ContextManagerProvider: React.FC<{ children: React.ReactNode }> = (
compactionError,
// Actions
handleAutoCompaction,
handleManualCompaction,
hasCompactionMarker,
};
ui/desktop/src/components/context_management/__tests__/CompactionMarker.test.tsx
+20 -20
View File
@@ -3,12 +3,20 @@ import { render, screen } from '@testing-library/react';
import { CompactionMarker } from '../CompactionMarker';
import { Message } from '../../../api';
const default_message: Message = {
metadata: {
agentVisible: false,
userVisible: false,
},
id: '1',
role: 'assistant',
created: 1000,content: []
};
describe('CompactionMarker', () => {
it('should render default message when no summarizationRequested content found', () => {
it('should render default message when no conversationCompacted content found', () => {
const message: Message = {
id: '1',
role: 'assistant',
created: 1000,
...default_message,
content: [{ type: 'text', text: 'Regular message' }],
};
@@ -17,14 +25,12 @@ describe('CompactionMarker', () => {
expect(screen.getByText('Conversation compacted')).toBeInTheDocument();
});
it('should render custom message from summarizationRequested content', () => {
it('should render custom message from conversationCompacted content', () => {
const message: Message = {
id: '1',
role: 'assistant',
created: 1000,
...default_message,
content: [
{ type: 'text', text: 'Some other content' },
{ type: 'summarizationRequested', msg: 'Custom compaction message' },
{ type: 'conversationCompacted', msg: 'Custom compaction message' },
],
};
@@ -35,9 +41,7 @@ describe('CompactionMarker', () => {
it('should handle empty message content array', () => {
const message: Message = {
id: '1',
role: 'assistant',
created: 1000,
...default_message,
content: [],
};
@@ -48,10 +52,8 @@ describe('CompactionMarker', () => {
it('should handle summarizationRequested content with empty msg', () => {
const message: Message = {
id: '1',
role: 'assistant',
created: 1000,
content: [{ type: 'summarizationRequested', msg: '' }],
...default_message,
content: [{ type: 'conversationCompacted', msg: '' }],
};
render(<CompactionMarker message={message} />);
@@ -62,11 +64,9 @@ describe('CompactionMarker', () => {
it('should handle summarizationRequested content with undefined msg', () => {
const message: Message = {
id: '1',
role: 'assistant',
created: 1000,
...default_message,
// eslint-disable-next-line @typescript-eslint/no-explicit-any
content: [{ type: 'summarizationRequested' } as any],
content: [{ type: 'conversationCompacted' } as any],
};
render(<CompactionMarker message={message} />);
ui/desktop/src/components/context_management/__tests__/ContextManager.test.tsx
+31 -290
View File
@@ -2,7 +2,18 @@ import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { renderHook, act } from '@testing-library/react';
import { ContextManagerProvider, useContextManager } from '../ContextManager';
import * as contextManagement from '../index';
import { ContextManageResponse, Message } from '../../../api';
import { Message } from '../../../api';
const default_message: Message = {
metadata: {
agentVisible: false,
userVisible: false,
},
id: '1',
role: 'assistant',
created: 1000,
content: [],
};
// Mock the context management functions
vi.mock('../index', () => ({
@@ -14,15 +25,11 @@ const mockManageContextFromBackend = vi.mocked(contextManagement.manageContextFr
describe('ContextManager', () => {
const mockMessages: Message[] = [
{
id: '1',
role: 'user',
created: 1000,
...default_message,
content: [{ type: 'text', text: 'Hello' }],
},
{
id: '2',
role: 'assistant',
created: 2000,
...default_message,
content: [{ type: 'text', text: 'Hi there!' }],
},
];
@@ -51,9 +58,6 @@ describe('ContextManager', () => {
expect(result.current.isCompacting).toBe(false);
expect(result.current.compactionError).toBe(null);
expect(typeof result.current.handleAutoCompaction).toBe('function');
expect(typeof result.current.handleManualCompaction).toBe('function');
expect(typeof result.current.hasCompactionMarker).toBe('function');
});
});
@@ -61,10 +65,8 @@ describe('ContextManager', () => {
it('should return true for messages with summarizationRequested content', () => {
const { result } = renderContextManager();
const messageWithMarker: Message = {
id: '1',
role: 'assistant',
created: 1000,
content: [{ type: 'summarizationRequested', msg: 'Compaction marker' }],
...default_message,
content: [{ type: 'conversationCompacted', msg: 'Compaction marker' }],
};
expect(result.current.hasCompactionMarker(messageWithMarker)).toBe(true);
@@ -73,24 +75,20 @@ describe('ContextManager', () => {
it('should return false for messages without summarizationRequested content', () => {
const { result } = renderContextManager();
const regularMessage: Message = {
id: '1',
role: 'user',
created: 1000,
...default_message,
content: [{ type: 'text', text: 'Hello' }],
};
expect(result.current.hasCompactionMarker(regularMessage)).toBe(false);
});
it('should return true for messages with mixed content including summarizationRequested', () => {
it('should return true for messages with mixed content including conversationCompacted', () => {
const { result } = renderContextManager();
const mixedMessage: Message = {
id: '1',
role: 'assistant',
created: 1000,
...default_message,
content: [
{ type: 'text', text: 'Some text' },
{ type: 'summarizationRequested', msg: 'Compaction marker' },
{ type: 'conversationCompacted', msg: 'Compaction marker' },
],
};
@@ -98,220 +96,22 @@ describe('ContextManager', () => {
});
});
describe('handleAutoCompaction', () => {
it('should successfully perform auto compaction with server-provided messages', async () => {
// Mock the backend response with 3 messages: marker, summary, continuation
// Note: Server messages may not have id/created, which will be added by the code
mockManageContextFromBackend.mockResolvedValue({
messages: [
{
role: 'assistant',
content: [
{ type: 'summarizationRequested', msg: 'Conversation compacted and summarized' },
],
} as Message,
{
role: 'assistant',
content: [{ type: 'text', text: 'Summary content' }],
} as Message,
{
role: 'assistant',
content: [
{
type: 'text',
text: 'The previous message contains a summary that was prepared because a context limit was reached. Do not mention that you read a summary or that conversation summarization occurred Just continue the conversation naturally based on the summarized context',
},
],
} as Message,
],
tokenCounts: [8, 100, 50],
});
const { result } = renderContextManager();
await act(async () => {
await result.current.handleAutoCompaction(
mockMessages,
mockSetMessages,
mockAppend,
'test-session-id'
);
});
expect(mockManageContextFromBackend).toHaveBeenCalledWith({
messages: mockMessages,
manageAction: 'summarize',
sessionId: 'test-session-id',
});
// Expect setMessages to be called with all 3 messages from server
// Note: Server doesn't provide id/created fields, so we don't check for them
expect(mockSetMessages).toHaveBeenCalledTimes(1);
const setMessagesCall = mockSetMessages.mock.calls[0][0];
expect(setMessagesCall).toHaveLength(3);
expect(setMessagesCall[0]).toMatchObject({
role: 'assistant',
content: [{ type: 'summarizationRequested', msg: 'Conversation compacted and summarized' }],
});
expect(setMessagesCall[1]).toMatchObject({
role: 'assistant',
content: [{ type: 'text', text: 'Summary content' }],
});
expect(setMessagesCall[2]).toMatchObject({
role: 'assistant',
content: [
{
type: 'text',
text: 'The previous message contains a summary that was prepared because a context limit was reached. Do not mention that you read a summary or that conversation summarization occurred Just continue the conversation naturally based on the summarized context',
},
],
});
// Fast-forward timers to trigger the append call
act(() => {
vi.advanceTimersByTime(150);
});
// Should append the continuation message (index 2) for auto-compaction
expect(mockAppend).toHaveBeenCalledTimes(1);
const appendedMessage = mockAppend.mock.calls[0][0];
expect(appendedMessage).toMatchObject({
role: 'assistant',
content: [
{
type: 'text',
text: 'The previous message contains a summary that was prepared because a context limit was reached. Do not mention that you read a summary or that conversation summarization occurred Just continue the conversation naturally based on the summarized context',
},
],
});
});
it('should handle compaction errors gracefully', async () => {
const error = new Error('Backend error');
mockManageContextFromBackend.mockRejectedValue(error);
const { result } = renderContextManager();
await act(async () => {
await result.current.handleAutoCompaction(
mockMessages,
mockSetMessages,
mockAppend,
'test-session-id'
);
});
expect(result.current.compactionError).toBe('Backend error');
expect(result.current.isCompacting).toBe(false);
expect(mockSetMessages).toHaveBeenCalledWith([
...mockMessages,
expect.objectContaining({
content: [
{
type: 'summarizationRequested',
msg: 'Compaction failed. Please try again or start a new session.',
},
],
}),
]);
});
it('should set isCompacting state correctly during operation', async () => {
let resolvePromise: (value: ContextManageResponse) => void;
const promise = new Promise<ContextManageResponse>((resolve) => {
resolvePromise = resolve;
});
mockManageContextFromBackend.mockReturnValue(promise);
const { result } = renderContextManager();
// Start compaction
act(() => {
result.current.handleAutoCompaction(
mockMessages,
mockSetMessages,
mockAppend,
'test-session-id'
);
});
// Should be compacting
expect(result.current.isCompacting).toBe(true);
expect(result.current.compactionError).toBe(null);
// Resolve the backend call
resolvePromise!({
messages: [
{
role: 'assistant',
content: [{ type: 'text', text: 'Summary content' }],
},
],
tokenCounts: [100, 50],
});
await act(async () => {
await promise;
});
// Should no longer be compacting
expect(result.current.isCompacting).toBe(false);
});
it('preserves display: false for ancestor messages', async () => {
mockManageContextFromBackend.mockResolvedValue({ messages: [], tokenCounts: [] });
const hiddenMessage: Message = {
id: 'hidden-1',
role: 'user',
created: 1500,
content: [{ type: 'text', text: 'Secret' }],
};
const visibleMessage: Message = {
id: 'visible-1',
role: 'assistant',
created: 1600,
content: [{ type: 'text', text: 'Public' }],
};
const messages: Message[] = [hiddenMessage, visibleMessage];
const { result } = renderContextManager();
await act(async () => {
await result.current.handleAutoCompaction(
messages,
mockSetMessages,
mockAppend,
'test-session-id'
);
});
// No server messages -> setMessages called with empty list
expect(mockSetMessages).toHaveBeenCalledWith([]);
expect(mockAppend).not.toHaveBeenCalled();
});
});
describe('handleManualCompaction', () => {
it('should perform compaction with server-provided messages', async () => {
mockManageContextFromBackend.mockResolvedValue({
messages: [
{
role: 'assistant',
...default_message,
content: [
{ type: 'summarizationRequested', msg: 'Conversation compacted and summarized' },
{ type: 'conversationCompacted', msg: 'Conversation compacted and summarized' },
],
},
{
role: 'assistant',
...default_message,
content: [{ type: 'text', text: 'Manual summary content' }],
},
{
role: 'assistant',
...default_message,
content: [
{
type: 'text',
@@ -346,7 +146,7 @@ describe('ContextManager', () => {
expect(setMessagesCall).toHaveLength(3);
expect(setMessagesCall[0]).toMatchObject({
role: 'assistant',
content: [{ type: 'summarizationRequested', msg: 'Conversation compacted and summarized' }],
content: [{ type: 'conversationCompacted', msg: 'Conversation compacted and summarized' }],
});
expect(setMessagesCall[1]).toMatchObject({
role: 'assistant',
@@ -375,7 +175,7 @@ describe('ContextManager', () => {
mockManageContextFromBackend.mockResolvedValue({
messages: [
{
role: 'assistant',
...default_message,
content: [{ type: 'text', text: 'Manual summary content' }],
},
],
@@ -408,17 +208,17 @@ describe('ContextManager', () => {
mockManageContextFromBackend.mockResolvedValue({
messages: [
{
role: 'assistant',
...default_message,
content: [
{ type: 'summarizationRequested', msg: 'Conversation compacted and summarized' },
{ type: 'conversationCompacted', msg: 'Conversation compacted and summarized' },
],
},
{
role: 'assistant',
...default_message,
content: [{ type: 'text', text: 'Manual summary content' }],
},
{
role: 'assistant',
...default_message,
content: [
{
type: 'text',
@@ -447,7 +247,7 @@ describe('ContextManager', () => {
expect(setMessagesCall).toHaveLength(3);
expect(setMessagesCall[0]).toMatchObject({
role: 'assistant',
content: [{ type: 'summarizationRequested', msg: 'Conversation compacted and summarized' }],
content: [{ type: 'conversationCompacted', msg: 'Conversation compacted and summarized' }],
});
expect(setMessagesCall[1]).toMatchObject({
role: 'assistant',
@@ -473,65 +273,6 @@ describe('ContextManager', () => {
});
});
describe('Error Handling', () => {
it('should handle backend errors with unknown error type', async () => {
mockManageContextFromBackend.mockRejectedValue('String error');
const { result } = renderContextManager();
await act(async () => {
await result.current.handleAutoCompaction(
mockMessages,
mockSetMessages,
mockAppend,
'test-session-id'
);
});
expect(result.current.compactionError).toBe('Unknown error during compaction');
});
it('should handle missing summary content gracefully with server-provided messages', async () => {
mockManageContextFromBackend.mockResolvedValue({
messages: [
{
role: 'assistant',
content: [
{ type: 'toolResponse', id: 'test', toolResult: { content: 'Not text content' } },
],
} as Message,
],
tokenCounts: [100, 50],
});
const { result } = renderContextManager();
await act(async () => {
await result.current.handleAutoCompaction(
mockMessages,
mockSetMessages,
mockAppend,
'test-session-id'
);
});
// Should complete without error even if content is not text
expect(result.current.isCompacting).toBe(false);
expect(result.current.compactionError).toBe(null);
// Should still set messages from server
expect(mockSetMessages).toHaveBeenCalledTimes(1);
const setMessagesCall = mockSetMessages.mock.calls[0][0];
expect(setMessagesCall).toHaveLength(1);
expect(setMessagesCall[0]).toMatchObject({
role: 'assistant',
content: [
{ type: 'toolResponse', id: 'test', toolResult: { content: 'Not text content' } },
],
});
});
});
describe('Context Provider Error', () => {
it('should throw error when useContextManager is used outside provider', () => {
expect(() => {