done tool shift in params

This commit is contained in:
dal 2025-08-14 14:39:51 -06:00
parent 5ec0af4273
commit 597f1b56a8
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
8 changed files with 104 additions and 109 deletions

View File

@ -4,12 +4,7 @@ import {
OptimisticJsonParser,
getOptimisticValue,
} from '../../../utils/streaming/optimistic-json-parser';
import {
type DoneToolContext,
type DoneToolInput,
DoneToolInputSchema,
type DoneToolState,
} from './done-tool';
import type { DoneToolContext, DoneToolInput, DoneToolState } from './done-tool';
import {
createDoneToolRawLlmMessageEntry,
createDoneToolResponseMessage,
@ -17,7 +12,7 @@ import {
// Type-safe key extraction from the schema - will cause compile error if field name changes
// Using keyof with the inferred type ensures we're using the actual schema keys
const FINAL_RESPONSE_KEY = 'final_response' as const satisfies keyof DoneToolInput;
const FINAL_RESPONSE_KEY = 'finalResponse' as const satisfies keyof DoneToolInput;
export function createDoneToolDelta(doneToolState: DoneToolState, context: DoneToolContext) {
return async function doneToolDelta(
@ -37,7 +32,7 @@ export function createDoneToolDelta(doneToolState: DoneToolState, context: DoneT
if (finalResponse !== undefined && finalResponse !== '') {
// Update the state with the extracted final_response
doneToolState.final_response = finalResponse;
doneToolState.finalResponse = finalResponse;
// Create the response entries with the current state
const doneToolResponseEntry = createDoneToolResponseMessage(

View File

@ -10,7 +10,7 @@ export function createDoneToolFinish(doneToolState: DoneToolState, context: Done
return async function doneToolFinish(
options: { input: DoneToolInput } & ToolCallOptions
): Promise<void> {
doneToolState.entry_id = options.toolCallId;
doneToolState.toolCallId = options.toolCallId;
const doneToolResponseEntry = createDoneToolResponseMessage(doneToolState, options.toolCallId);
const doneToolMessage = createDoneToolRawLlmMessageEntry(doneToolState, options.toolCallId);

View File

@ -14,7 +14,7 @@ import {
// Factory function that creates a type-safe callback for the specific agent context
export function createDoneToolStart(doneToolState: DoneToolState, context: DoneToolContext) {
return async function doneToolStart(options: ToolCallOptions): Promise<void> {
doneToolState.entry_id = options.toolCallId;
doneToolState.toolCallId = options.toolCallId;
// Extract files from the tool call responses in messages
if (options.messages) {

View File

@ -18,9 +18,9 @@ describe('Done Tool Streaming Tests', () => {
describe('createDoneToolStart', () => {
test('should initialize state with entry_id on start', async () => {
const state: DoneToolState = {
entry_id: undefined,
toolCallId: undefined,
args: undefined,
final_response: undefined,
finalResponse: undefined,
};
const startHandler = createDoneToolStart(state, mockContext);
@ -31,14 +31,14 @@ describe('Done Tool Streaming Tests', () => {
await startHandler(options);
expect(state.entry_id).toBe('tool-call-123');
expect(state.toolCallId).toBe('tool-call-123');
});
test('should handle start with messages containing file tool calls', async () => {
const state: DoneToolState = {
entry_id: undefined,
toolCallId: undefined,
args: undefined,
final_response: undefined,
finalResponse: undefined,
};
const startHandler = createDoneToolStart(state, mockContext);
@ -87,14 +87,14 @@ describe('Done Tool Streaming Tests', () => {
await startHandler(options);
expect(state.entry_id).toBe('tool-call-123');
expect(state.toolCallId).toBe('tool-call-123');
});
test('should handle start without messages', async () => {
const state: DoneToolState = {
entry_id: undefined,
toolCallId: undefined,
args: undefined,
final_response: undefined,
finalResponse: undefined,
};
const startHandler = createDoneToolStart(state, mockContext);
@ -105,7 +105,7 @@ describe('Done Tool Streaming Tests', () => {
await startHandler(options);
expect(state.entry_id).toBe('tool-call-456');
expect(state.toolCallId).toBe('tool-call-456');
});
test('should handle context without messageId', async () => {
@ -114,9 +114,9 @@ describe('Done Tool Streaming Tests', () => {
workflowStartTime: Date.now(),
};
const state: DoneToolState = {
entry_id: undefined,
toolCallId: undefined,
args: undefined,
final_response: undefined,
finalResponse: undefined,
};
const startHandler = createDoneToolStart(state, contextWithoutMessageId);
@ -126,16 +126,16 @@ describe('Done Tool Streaming Tests', () => {
};
await expect(startHandler(options)).resolves.not.toThrow();
expect(state.entry_id).toBe('tool-call-789');
expect(state.toolCallId).toBe('tool-call-789');
});
});
describe('createDoneToolDelta', () => {
test('should accumulate text deltas to args', async () => {
const state: DoneToolState = {
entry_id: 'test-entry',
toolCallId: 'test-entry',
args: '',
final_response: undefined,
finalResponse: undefined,
};
const deltaHandler = createDoneToolDelta(state, mockContext);
@ -159,9 +159,9 @@ describe('Done Tool Streaming Tests', () => {
test('should extract partial final_response from incomplete JSON', async () => {
const state: DoneToolState = {
entry_id: 'test-entry',
toolCallId: 'test-entry',
args: '',
final_response: undefined,
finalResponse: undefined,
};
const deltaHandler = createDoneToolDelta(state, mockContext);
@ -173,14 +173,14 @@ describe('Done Tool Streaming Tests', () => {
});
expect(state.args).toBe('{"final_response": "This is a partial response that is still being');
expect(state.final_response).toBe('This is a partial response that is still being');
expect(state.finalResponse).toBe('This is a partial response that is still being');
});
test('should handle complete JSON in delta', async () => {
const state: DoneToolState = {
entry_id: 'test-entry',
toolCallId: 'test-entry',
args: '',
final_response: undefined,
finalResponse: undefined,
};
const deltaHandler = createDoneToolDelta(state, mockContext);
@ -192,14 +192,14 @@ describe('Done Tool Streaming Tests', () => {
});
expect(state.args).toBe('{"final_response": "Complete response message"}');
expect(state.final_response).toBe('Complete response message');
expect(state.finalResponse).toBe('Complete response message');
});
test('should handle markdown content in final_response', async () => {
const state: DoneToolState = {
entry_id: 'test-entry',
toolCallId: 'test-entry',
args: '',
final_response: undefined,
finalResponse: undefined,
};
const deltaHandler = createDoneToolDelta(state, mockContext);
@ -217,14 +217,14 @@ describe('Done Tool Streaming Tests', () => {
messages: [],
});
expect(state.final_response).toBe(markdownContent);
expect(state.finalResponse).toBe(markdownContent);
});
test('should handle escaped characters in JSON', async () => {
const state: DoneToolState = {
entry_id: 'test-entry',
toolCallId: 'test-entry',
args: '',
final_response: undefined,
finalResponse: undefined,
};
const deltaHandler = createDoneToolDelta(state, mockContext);
@ -235,14 +235,14 @@ describe('Done Tool Streaming Tests', () => {
messages: [],
});
expect(state.final_response).toBe('Line 1\nLine 2\n"Quoted text"');
expect(state.finalResponse).toBe('Line 1\nLine 2\n"Quoted text"');
});
test('should not update state when no final_response is extracted', async () => {
const state: DoneToolState = {
entry_id: 'test-entry',
toolCallId: 'test-entry',
args: '',
final_response: undefined,
finalResponse: undefined,
};
const deltaHandler = createDoneToolDelta(state, mockContext);
@ -254,14 +254,14 @@ describe('Done Tool Streaming Tests', () => {
});
expect(state.args).toBe('{"other_field": "value"}');
expect(state.final_response).toBeUndefined();
expect(state.finalResponse).toBeUndefined();
});
test('should handle empty final_response gracefully', async () => {
const state: DoneToolState = {
entry_id: 'test-entry',
toolCallId: 'test-entry',
args: '',
final_response: undefined,
finalResponse: undefined,
};
const deltaHandler = createDoneToolDelta(state, mockContext);
@ -273,22 +273,22 @@ describe('Done Tool Streaming Tests', () => {
});
expect(state.args).toBe('{"final_response": ""}');
expect(state.final_response).toBeUndefined();
expect(state.finalResponse).toBeUndefined();
});
});
describe('createDoneToolFinish', () => {
test('should update state with final input on finish', async () => {
const state: DoneToolState = {
entry_id: undefined,
toolCallId: undefined,
args: '{"final_response": "Final message"}',
final_response: 'Final message',
finalResponse: 'Final message',
};
const finishHandler = createDoneToolFinish(state, mockContext);
const input: DoneToolInput = {
final_response: 'This is the final response message',
finalResponse: 'This is the final response message',
};
await finishHandler({
@ -297,20 +297,20 @@ describe('Done Tool Streaming Tests', () => {
messages: [],
});
expect(state.entry_id).toBe('tool-call-123');
expect(state.toolCallId).toBe('tool-call-123');
});
test('should handle finish without prior entry_id', async () => {
const state: DoneToolState = {
entry_id: undefined,
toolCallId: undefined,
args: undefined,
final_response: undefined,
finalResponse: undefined,
};
const finishHandler = createDoneToolFinish(state, mockContext);
const input: DoneToolInput = {
final_response: 'Response without prior start',
finalResponse: 'Response without prior start',
};
await finishHandler({
@ -319,14 +319,14 @@ describe('Done Tool Streaming Tests', () => {
messages: [],
});
expect(state.entry_id).toBe('tool-call-456');
expect(state.toolCallId).toBe('tool-call-456');
});
test('should handle markdown formatted final response', async () => {
const state: DoneToolState = {
entry_id: undefined,
toolCallId: undefined,
args: undefined,
final_response: undefined,
finalResponse: undefined,
};
const finishHandler = createDoneToolFinish(state, mockContext);
@ -346,7 +346,7 @@ The following items were processed:
`;
const input: DoneToolInput = {
final_response: markdownResponse,
finalResponse: markdownResponse,
};
await finishHandler({
@ -355,7 +355,7 @@ The following items were processed:
messages: [],
});
expect(state.entry_id).toBe('tool-call-789');
expect(state.toolCallId).toBe('tool-call-789');
});
});
@ -373,9 +373,9 @@ The following items were processed:
};
const state: DoneToolState = {
entry_id: undefined,
toolCallId: undefined,
args: undefined,
final_response: undefined,
finalResponse: undefined,
};
const handler1 = createDoneToolStart(state, validContext);
@ -387,9 +387,9 @@ The following items were processed:
test('should maintain state type consistency through streaming lifecycle', async () => {
const state: DoneToolState = {
entry_id: undefined,
toolCallId: undefined,
args: undefined,
final_response: undefined,
finalResponse: undefined,
};
const startHandler = createDoneToolStart(state, mockContext);
@ -397,7 +397,7 @@ The following items were processed:
const finishHandler = createDoneToolFinish(state, mockContext);
await startHandler({ toolCallId: 'test-123', messages: [] });
expect(state.entry_id).toBeTypeOf('string');
expect(state.toolCallId).toBeTypeOf('string');
await deltaHandler({
inputTextDelta: '{"final_response": "Testing"}',
@ -405,22 +405,22 @@ The following items were processed:
messages: [],
});
expect(state.args).toBeTypeOf('string');
expect(state.final_response).toBeTypeOf('string');
expect(state.finalResponse).toBeTypeOf('string');
const input: DoneToolInput = {
final_response: 'Final test',
finalResponse: 'Final test',
};
await finishHandler({ input, toolCallId: 'test-123', messages: [] });
expect(state.entry_id).toBeTypeOf('string');
expect(state.toolCallId).toBeTypeOf('string');
});
});
describe('Streaming Flow Integration', () => {
test('should handle complete streaming flow from start to finish', async () => {
const state: DoneToolState = {
entry_id: undefined,
toolCallId: undefined,
args: undefined,
final_response: undefined,
finalResponse: undefined,
};
const startHandler = createDoneToolStart(state, mockContext);
@ -430,7 +430,7 @@ The following items were processed:
const toolCallId = 'streaming-test-123';
await startHandler({ toolCallId, messages: [] });
expect(state.entry_id).toBe(toolCallId);
expect(state.toolCallId).toBe(toolCallId);
const chunks = [
'{"final_',
@ -452,23 +452,23 @@ The following items were processed:
expect(state.args).toBe(
'{"final_response": "This is a streaming response that comes in multiple chunks"}'
);
expect(state.final_response).toBe(
expect(state.finalResponse).toBe(
'This is a streaming response that comes in multiple chunks'
);
const input: DoneToolInput = {
final_response: 'This is a streaming response that comes in multiple chunks',
finalResponse: 'This is a streaming response that comes in multiple chunks',
};
await finishHandler({ input, toolCallId, messages: [] });
expect(state.entry_id).toBe(toolCallId);
expect(state.toolCallId).toBe(toolCallId);
});
test('should handle streaming with special characters and formatting', async () => {
const state: DoneToolState = {
entry_id: undefined,
toolCallId: undefined,
args: undefined,
final_response: undefined,
finalResponse: undefined,
};
const deltaHandler = createDoneToolDelta(state, mockContext);
@ -492,7 +492,7 @@ The following items were processed:
});
}
expect(state.final_response).toBe(
expect(state.finalResponse).toBe(
'## Results\n\n- Success: 90%\n- Failed: 10%\n\n**Note:** Review failed items'
);
});

View File

@ -44,9 +44,9 @@ describe('Done Tool Integration Tests', () => {
describe('Database Message Updates', () => {
test('should create initial entries when done tool starts', async () => {
const state: DoneToolState = {
entry_id: undefined,
toolCallId: undefined,
args: undefined,
final_response: undefined,
finalResponse: undefined,
};
const startHandler = createDoneToolStart(state, mockContext);
@ -66,9 +66,9 @@ describe('Done Tool Integration Tests', () => {
test('should update entries during streaming delta', async () => {
const state: DoneToolState = {
entry_id: undefined,
toolCallId: undefined,
args: '',
final_response: undefined,
finalResponse: undefined,
};
const startHandler = createDoneToolStart(state, mockContext);
@ -89,14 +89,14 @@ describe('Done Tool Integration Tests', () => {
.where(and(eq(messages.id, testMessageId), isNull(messages.deletedAt)));
expect(message?.rawLlmMessages).toBeDefined();
expect(state.final_response).toBe('Partial response');
expect(state.finalResponse).toBe('Partial response');
});
test('should finalize entries when done tool finishes', async () => {
const state: DoneToolState = {
entry_id: undefined,
toolCallId: undefined,
args: '',
final_response: undefined,
finalResponse: undefined,
};
const startHandler = createDoneToolStart(state, mockContext);
@ -106,7 +106,7 @@ describe('Done Tool Integration Tests', () => {
await startHandler({ toolCallId, messages: [] });
const input: DoneToolInput = {
final_response: 'This is the complete final response',
finalResponse: 'This is the complete final response',
};
await finishHandler({ input, toolCallId, messages: [] });
@ -124,9 +124,9 @@ describe('Done Tool Integration Tests', () => {
describe('Complete Streaming Flow', () => {
test('should handle full streaming lifecycle with database updates', async () => {
const state: DoneToolState = {
entry_id: undefined,
toolCallId: undefined,
args: '',
final_response: undefined,
finalResponse: undefined,
};
const startHandler = createDoneToolStart(state, mockContext);
@ -163,10 +163,10 @@ The following tasks have been completed:
All operations completed successfully.`;
expect(state.final_response).toBe(expectedResponse);
expect(state.finalResponse).toBe(expectedResponse);
const input: DoneToolInput = {
final_response: expectedResponse,
finalResponse: expectedResponse,
};
await finishHandler({ input, toolCallId, messages: [] });
@ -183,15 +183,15 @@ All operations completed successfully.`;
test('should handle multiple done tool invocations in sequence', async () => {
const state1: DoneToolState = {
entry_id: undefined,
toolCallId: undefined,
args: '',
final_response: undefined,
finalResponse: undefined,
};
const state2: DoneToolState = {
entry_id: undefined,
toolCallId: undefined,
args: '',
final_response: undefined,
finalResponse: undefined,
};
const startHandler1 = createDoneToolStart(state1, mockContext);
@ -205,14 +205,14 @@ All operations completed successfully.`;
await startHandler1({ toolCallId: toolCallId1, messages: [] });
await finishHandler1({
input: { final_response: 'First response' },
input: { finalResponse: 'First response' },
toolCallId: toolCallId1,
messages: [],
});
await startHandler2({ toolCallId: toolCallId2, messages: [] });
await finishHandler2({
input: { final_response: 'Second response' },
input: { finalResponse: 'Second response' },
toolCallId: toolCallId2,
messages: [],
});
@ -237,23 +237,23 @@ All operations completed successfully.`;
};
const state: DoneToolState = {
entry_id: undefined,
toolCallId: undefined,
args: '',
final_response: undefined,
finalResponse: undefined,
};
const startHandler = createDoneToolStart(state, invalidContext);
const toolCallId = randomUUID();
await expect(startHandler({ toolCallId, messages: [] })).resolves.not.toThrow();
expect(state.entry_id).toBe(toolCallId);
expect(state.toolCallId).toBe(toolCallId);
});
test('should continue processing even if database update fails', async () => {
const state: DoneToolState = {
entry_id: undefined,
toolCallId: undefined,
args: '',
final_response: undefined,
finalResponse: undefined,
};
const invalidContext: DoneToolContext = {
@ -272,16 +272,16 @@ All operations completed successfully.`;
})
).resolves.not.toThrow();
expect(state.final_response).toBe('Test');
expect(state.finalResponse).toBe('Test');
});
});
describe('Message Entry Modes', () => {
test('should use append mode for start operations', async () => {
const state: DoneToolState = {
entry_id: undefined,
toolCallId: undefined,
args: undefined,
final_response: undefined,
finalResponse: undefined,
};
const startHandler = createDoneToolStart(state, mockContext);
@ -313,9 +313,9 @@ All operations completed successfully.`;
test('should use update mode for delta and finish operations', async () => {
const state: DoneToolState = {
entry_id: undefined,
toolCallId: undefined,
args: '',
final_response: undefined,
finalResponse: undefined,
};
await updateMessageEntries({

View File

@ -8,7 +8,7 @@ import { createDoneToolStart } from './done-tool-start';
export const DONE_TOOL_NAME = 'doneTool';
export const DoneToolInputSchema = z.object({
final_response: z
finalResponse: z
.string()
.min(1, 'Final response is required')
.describe(
@ -26,14 +26,14 @@ const DoneToolContextSchema = z.object({
});
const DoneToolStateSchema = z.object({
entry_id: z
toolCallId: z
.string()
.optional()
.describe(
'The entry ID of the entry that triggered the done tool. This is optional and will be set by the tool start'
),
args: z.string().optional().describe('The arguments of the done tool'),
final_response: z
finalResponse: z
.string()
.optional()
.describe(
@ -48,9 +48,9 @@ export type DoneToolState = z.infer<typeof DoneToolStateSchema>;
export function createDoneTool(context: DoneToolContext) {
const state: DoneToolState = {
entry_id: undefined,
toolCallId: undefined,
args: undefined,
final_response: undefined,
finalResponse: undefined,
};
const execute = createDoneToolExecute();

View File

@ -7,7 +7,7 @@ export function createDoneToolResponseMessage(
toolCallId?: string
): ChatMessageResponseMessage_Text | null {
// Use entry_id from state or fallback to provided toolCallId
const id = doneToolState.entry_id || toolCallId;
const id = doneToolState.toolCallId || toolCallId;
if (!id) {
return null;
@ -16,7 +16,7 @@ export function createDoneToolResponseMessage(
return {
id,
type: 'text',
message: doneToolState.final_response || '',
message: doneToolState.finalResponse || '',
is_final_message: true,
};
}
@ -25,7 +25,7 @@ export function createDoneToolRawLlmMessageEntry(
doneToolState: DoneToolState,
toolCallId?: string
): ModelMessage | undefined {
const id = doneToolState.entry_id || toolCallId;
const id = doneToolState.toolCallId || toolCallId;
if (!id) {
return undefined;
@ -41,8 +41,8 @@ export function createDoneToolRawLlmMessageEntry(
input: {},
},
// Optionally include any accumulated text content
...(doneToolState.final_response
? [{ type: 'text' as const, text: doneToolState.final_response }]
...(doneToolState.finalResponse
? [{ type: 'text' as const, text: doneToolState.finalResponse }]
: []),
],
};