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

Clean room implementation of the chat process (#5079)

Browse Source 
Co-authored-by: Douwe Osinga <douwe@squareup.com>
Co-authored-by: Zane <75694352+zanesq@users.noreply.github.com>
This commit is contained in:
Douwe Osinga
2025-10-11 10:45:16 -04:00
committed by GitHub
parent 73f109237e
commit 0c2127124f
36 changed files with 1024 additions and 541 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ui/desktop/src/components/context_management/CompactionMarker.tsx
+4 -3
View File
@@ -1,5 +1,5 @@
import React from 'react';
import { Message, SummarizationRequestedContent } from '../../types/message';
import { Message, SummarizationRequested } from '../../api';
interface CompactionMarkerProps {
message: Message;
@@ -7,8 +7,9 @@ interface CompactionMarkerProps {
export const CompactionMarker: React.FC<CompactionMarkerProps> = ({ message }) => {
const compactionContent = message.content.find(
(content) => content.type === 'summarizationRequested'
) as SummarizationRequestedContent | undefined;
(content): content is SummarizationRequested & { type: 'summarizationRequested' } =>
content.type === 'summarizationRequested'
);
const markerText = compactionContent?.msg || 'Conversation compacted';
ui/desktop/src/components/context_management/ContextManager.tsx
+4 -11
View File
@@ -1,6 +1,6 @@
import React, { createContext, useContext, useState, useCallback } from 'react';
import { Message } from '../../types/message';
import { manageContextFromBackend, convertApiMessageToFrontendMessage } from './index';
import { manageContextFromBackend } from './index';
import { Message } from '../../api';
// Define the context management interface
interface ContextManagerState {
@@ -53,21 +53,14 @@ export const ContextManagerProvider: React.FC<{ children: React.ReactNode }> = (
sessionId: sessionId,
});
// Convert API messages to frontend messages
// The server now handles all visibility - we just display what we receive
const convertedMessages = summaryResponse.messages.map((apiMessage) =>
convertApiMessageToFrontendMessage(apiMessage)
);
// Replace messages with the server-provided messages
setMessages(convertedMessages);
setMessages(summaryResponse.messages);
// Only automatically submit the continuation message for auto-compaction (context limit reached)
// Manual compaction should just compact without continuing the conversation
if (!isManual) {
// Automatically submit the continuation message to continue the conversation
// This should be the third message (index 2) which contains the "I ran into a context length exceeded error..." text
const continuationMessage = convertedMessages[2];
const continuationMessage = summaryResponse.messages[2];
if (continuationMessage) {
setTimeout(() => {
append(continuationMessage);
ui/desktop/src/components/context_management/__tests__/CompactionMarker.test.tsx
+1 -1
View File
@@ -1,7 +1,7 @@
import { describe, it, expect } from 'vitest';
import { render, screen } from '@testing-library/react';
import { CompactionMarker } from '../CompactionMarker';
import { Message } from '../../../types/message';
import { Message } from '../../../api';
describe('CompactionMarker', () => {
it('should render default message when no summarizationRequested content found', () => {
ui/desktop/src/components/context_management/__tests__/ContextManager.test.tsx
+88 -149
View File
@@ -1,20 +1,15 @@
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { renderHook, act } from '@testing-library/react';
import { ContextManagerProvider, useContextManager } from '../ContextManager';
import { Message } from '../../../types/message';
import * as contextManagement from '../index';
import { ContextManageResponse } from '../../../api';
import { ContextManageResponse, Message } from '../../../api';
// Mock the context management functions
vi.mock('../index', () => ({
manageContextFromBackend: vi.fn(),
convertApiMessageToFrontendMessage: vi.fn(),
}));
const mockManageContextFromBackend = vi.mocked(contextManagement.manageContextFromBackend);
const mockConvertApiMessageToFrontendMessage = vi.mocked(
contextManagement.convertApiMessageToFrontendMessage
);
describe('ContextManager', () => {
const mockMessages: Message[] = [
@@ -32,13 +27,6 @@ describe('ContextManager', () => {
},
];
const mockSummaryMessage: Message = {
id: 'summary-1',
role: 'assistant',
created: 3000,
content: [{ type: 'text', text: 'This is a summary of the conversation.' }],
};
const mockSetMessages = vi.fn();
const mockAppend = vi.fn();
@@ -113,6 +101,7 @@ 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: [
{
@@ -120,11 +109,11 @@ describe('ContextManager', () => {
content: [
{ type: 'summarizationRequested', msg: 'Conversation compacted and summarized' },
],
},
} as Message,
{
role: 'assistant',
content: [{ type: 'text', text: 'Summary content' }],
},
} as Message,
{
role: 'assistant',
content: [
@@ -133,36 +122,11 @@ describe('ContextManager', () => {
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 mockCompactionMarker: Message = {
id: 'marker-1',
role: 'assistant',
created: 3000,
content: [{ type: 'summarizationRequested', msg: 'Conversation compacted and summarized' }],
};
const mockContinuationMessage: Message = {
id: 'continuation-1',
role: 'assistant',
created: 3000,
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',
},
],
};
// Mock the conversion function to return different messages based on call order
mockConvertApiMessageToFrontendMessage
.mockReturnValueOnce(mockCompactionMarker) // First call - compaction marker
.mockReturnValueOnce(mockSummaryMessage) // Second call - summary
.mockReturnValueOnce(mockContinuationMessage); // Third call - continuation
const { result } = renderContextManager();
await act(async () => {
@@ -180,39 +144,28 @@ describe('ContextManager', () => {
sessionId: 'test-session-id',
});
// Verify conversion calls with correct parameters
expect(mockConvertApiMessageToFrontendMessage).toHaveBeenNthCalledWith(
1,
expect.objectContaining({
content: [
{ type: 'summarizationRequested', msg: 'Conversation compacted and summarized' },
],
})
);
expect(mockConvertApiMessageToFrontendMessage).toHaveBeenNthCalledWith(
2,
expect.objectContaining({
content: [{ type: 'text', text: 'Summary content' }],
})
);
expect(mockConvertApiMessageToFrontendMessage).toHaveBeenNthCalledWith(
3,
expect.objectContaining({
content: [
{
type: 'text',
text: expect.stringContaining('The previous message contains a summary'),
},
],
})
);
// Expect setMessages to be called with all 3 converted messages
expect(mockSetMessages).toHaveBeenCalledWith([
mockCompactionMarker,
mockSummaryMessage,
mockContinuationMessage,
]);
// 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(() => {
@@ -221,7 +174,16 @@ describe('ContextManager', () => {
// Should append the continuation message (index 2) for auto-compaction
expect(mockAppend).toHaveBeenCalledTimes(1);
expect(mockAppend).toHaveBeenCalledWith(mockContinuationMessage);
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 () => {
@@ -290,8 +252,6 @@ describe('ContextManager', () => {
tokenCounts: [100, 50],
});
mockConvertApiMessageToFrontendMessage.mockReturnValue(mockSummaryMessage);
await act(async () => {
await promise;
});
@@ -363,30 +323,6 @@ describe('ContextManager', () => {
tokenCounts: [8, 100, 50],
});
const mockCompactionMarker: Message = {
id: 'marker-1',
role: 'assistant',
created: 3000,
content: [{ type: 'summarizationRequested', msg: 'Conversation compacted and summarized' }],
};
const mockContinuationMessage: Message = {
id: 'continuation-1',
role: 'assistant',
created: 3000,
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',
},
],
};
mockConvertApiMessageToFrontendMessage
.mockReturnValueOnce(mockCompactionMarker)
.mockReturnValueOnce(mockSummaryMessage)
.mockReturnValueOnce(mockContinuationMessage);
const { result } = renderContextManager();
await act(async () => {
@@ -405,11 +341,26 @@ describe('ContextManager', () => {
});
// Verify all three messages are set
expect(mockSetMessages).toHaveBeenCalledWith([
mockCompactionMarker,
mockSummaryMessage,
mockContinuationMessage,
]);
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: 'Manual 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 check if append would be called
act(() => {
@@ -431,8 +382,6 @@ describe('ContextManager', () => {
tokenCounts: [100, 50],
});
mockConvertApiMessageToFrontendMessage.mockReturnValue(mockSummaryMessage);
const { result } = renderContextManager();
await act(async () => {
@@ -481,30 +430,6 @@ describe('ContextManager', () => {
tokenCounts: [8, 100, 50],
});
const mockCompactionMarker: Message = {
id: 'marker-1',
role: 'assistant',
created: 3000,
content: [{ type: 'summarizationRequested', msg: 'Conversation compacted and summarized' }],
};
const mockContinuationMessage: Message = {
id: 'continuation-1',
role: 'assistant',
created: 3000,
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',
},
],
};
mockConvertApiMessageToFrontendMessage
.mockReturnValueOnce(mockCompactionMarker)
.mockReturnValueOnce(mockSummaryMessage)
.mockReturnValueOnce(mockContinuationMessage);
const { result } = renderContextManager();
await act(async () => {
@@ -517,11 +442,26 @@ describe('ContextManager', () => {
});
// Verify all three messages are set
expect(mockSetMessages).toHaveBeenCalledWith([
mockCompactionMarker,
mockSummaryMessage,
mockContinuationMessage,
]);
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: 'Manual 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 check if append would be called
act(() => {
@@ -559,20 +499,11 @@ describe('ContextManager', () => {
content: [
{ type: 'toolResponse', id: 'test', toolResult: { content: 'Not text content' } },
],
},
} as Message,
],
tokenCounts: [100, 50],
});
const mockMessageWithoutText: Message = {
id: 'summary-1',
role: 'assistant',
created: 3000,
content: [{ type: 'toolResponse', id: 'test', toolResult: { status: 'success' } }],
};
mockConvertApiMessageToFrontendMessage.mockReturnValue(mockMessageWithoutText);
const { result } = renderContextManager();
await act(async () => {
@@ -588,8 +519,16 @@ describe('ContextManager', () => {
expect(result.current.isCompacting).toBe(false);
expect(result.current.compactionError).toBe(null);
// Should still set messages with the converted message
expect(mockSetMessages).toHaveBeenCalledWith([mockMessageWithoutText]);
// 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' } },
],
});
});
});
ui/desktop/src/components/context_management/index.ts
+15 -135
View File
@@ -1,149 +1,29 @@
import {
Message as FrontendMessage,
Content as FrontendContent,
MessageContent as FrontendMessageContent,
ToolCallResult,
ToolCall,
Role,
} from '../../types/message';
import {
ContextManageRequest,
ContextManageResponse,
manageContext,
Message as ApiMessage,
MessageContent as ApiMessageContent,
} from '../../api';
import { generateId } from 'ai';
import { ContextManageRequest, ContextManageResponse, manageContext, Message } from '../../api';
export async function manageContextFromBackend({
messages,
manageAction,
sessionId,
}: {
messages: FrontendMessage[];
messages: Message[];
manageAction: 'truncation' | 'summarize';
sessionId: string;
}): Promise<ContextManageResponse> {
try {
const contextManagementRequest = { manageAction, messages, sessionId };
const contextManagementRequest = { manageAction, messages, sessionId };
// Cast to the API-expected type
const result = await manageContext({
body: contextManagementRequest as unknown as ContextManageRequest,
});
// Cast to the API-expected type
const result = await manageContext({
body: contextManagementRequest as unknown as ContextManageRequest,
});
// Check for errors in the result
if (result.error) {
throw new Error(`Context management failed: ${result.error}`);
}
// Extract the actual data from the result
if (!result.data) {
throw new Error('Context management returned no data');
}
return result.data;
} catch (error) {
console.error(`Context management failed: ${error || 'Unknown error'}`);
throw new Error(
`Context management failed: ${error || 'Unknown error'}\n\nStart a new session.`
);
}
}
// Function to convert API Message to frontend Message
export function convertApiMessageToFrontendMessage(apiMessage: ApiMessage): FrontendMessage {
return {
id: generateId(),
role: apiMessage.role as Role,
created: apiMessage.created ?? Math.floor(Date.now() / 1000),
content: apiMessage.content
.map((apiContent) => mapApiContentToFrontendMessageContent(apiContent))
.filter((content): content is FrontendMessageContent => content !== null),
};
}
// Function to convert API MessageContent to frontend MessageContent
function mapApiContentToFrontendMessageContent(
apiContent: ApiMessageContent
): FrontendMessageContent | null {
// Handle each content type specifically based on its "type" property
if (apiContent.type === 'text') {
return {
type: 'text',
text: apiContent.text,
annotations: apiContent.annotations as Record<string, unknown> | undefined,
};
} else if (apiContent.type === 'image') {
return {
type: 'image',
data: apiContent.data,
mimeType: apiContent.mimeType,
annotations: apiContent.annotations as Record<string, unknown> | undefined,
};
} else if (apiContent.type === 'toolRequest') {
// Ensure the toolCall has the correct type structure
const toolCall = apiContent.toolCall as unknown as ToolCallResult<ToolCall>;
return {
type: 'toolRequest',
id: apiContent.id,
toolCall: toolCall,
};
} else if (apiContent.type === 'toolResponse') {
// Ensure the toolResult has the correct type structure
const toolResult = apiContent.toolResult as unknown as ToolCallResult<FrontendContent[]>;
return {
type: 'toolResponse',
id: apiContent.id,
toolResult: toolResult,
};
} else if (apiContent.type === 'toolConfirmationRequest') {
return {
type: 'toolConfirmationRequest',
id: apiContent.id,
toolName: apiContent.toolName,
arguments: apiContent.arguments as Record<string, unknown>,
prompt: apiContent.prompt === null ? undefined : apiContent.prompt,
};
} else if (apiContent.type === 'contextLengthExceeded') {
return {
type: 'contextLengthExceeded',
msg: apiContent.msg,
};
} else if (apiContent.type === 'summarizationRequested') {
return {
type: 'summarizationRequested',
msg: apiContent.msg,
};
// Check for errors in the result
if (result.error) {
throw new Error(`Context management failed: ${result.error}`);
}
// For types that exist in API but not in frontend, either skip or convert
console.warn(`Skipping unsupported content type: ${apiContent.type}`);
return null;
}
export function createSummarizationRequestMessage(
messages: FrontendMessage[],
requestMessage: string
): FrontendMessage {
// Get the last message
const lastMessage = messages[messages.length - 1];
// Determine the next role (opposite of the last message)
const nextRole: Role = lastMessage.role === 'user' ? 'assistant' : 'user';
// Create the new message with SummarizationRequestedContent
return {
id: generateId(),
role: nextRole,
created: Math.floor(Date.now() / 1000),
content: [
{
type: 'summarizationRequested',
msg: requestMessage,
},
],
};
if (!result.data) {
throw new Error('Context management returned no data');
}
return result.data;
}