diff --git a/packages/ai/src/agents/analyst-agent/analyst-agent.ts b/packages/ai/src/agents/analyst-agent/analyst-agent.ts index 81fd5248b..a86f0c797 100644 --- a/packages/ai/src/agents/analyst-agent/analyst-agent.ts +++ b/packages/ai/src/agents/analyst-agent/analyst-agent.ts @@ -123,10 +123,10 @@ export function createAnalystAgent(analystAgentOptions: AnalystAgentOptions) { const docsSystemMessage = docsContent ? ({ - role: 'system', - content: `\n${docsContent}\n`, - providerOptions: DEFAULT_ANTHROPIC_OPTIONS, - } as ModelMessage) + role: 'system', + content: `\n${docsContent}\n`, + providerOptions: DEFAULT_ANTHROPIC_OPTIONS, + } as ModelMessage) : null; async function stream({ messages }: AnalystStreamOptions) { @@ -183,19 +183,19 @@ export function createAnalystAgent(analystAgentOptions: AnalystAgentOptions) { // Create analyst instructions system message with proper escaping const analystInstructionsMessage = analystInstructions ? ({ - role: 'system', - content: `\n${analystInstructions}\n`, - providerOptions: DEFAULT_ANTHROPIC_OPTIONS, - } as ModelMessage) + role: 'system', + content: `\n${analystInstructions}\n`, + providerOptions: DEFAULT_ANTHROPIC_OPTIONS, + } as ModelMessage) : null; // Create user personalization system message const userPersonalizationSystemMessage = userPersonalizationMessageContent ? ({ - role: 'system', - content: userPersonalizationMessageContent, - providerOptions: DEFAULT_ANTHROPIC_OPTIONS, - } as ModelMessage) + role: 'system', + content: userPersonalizationMessageContent, + providerOptions: DEFAULT_ANTHROPIC_OPTIONS, + } as ModelMessage) : null; return wrapTraced( diff --git a/packages/ai/src/agents/analytics-engineer-agent/analytics-engineer-agent.ts b/packages/ai/src/agents/analytics-engineer-agent/analytics-engineer-agent.ts index a184811de..a3a8812bc 100644 --- a/packages/ai/src/agents/analytics-engineer-agent/analytics-engineer-agent.ts +++ b/packages/ai/src/agents/analytics-engineer-agent/analytics-engineer-agent.ts @@ -1,84 +1,19 @@ -import type { LanguageModelV2 } from '@ai-sdk/provider'; -import type { Sandbox } from '@buster/sandbox'; import { type ModelMessage, hasToolCall, stepCountIs, streamText } from 'ai'; -import { wrapTraced } from 'braintrust'; -import z from 'zod'; import { DEFAULT_ANTHROPIC_OPTIONS } from '../../llm/providers/gateway'; import { Sonnet4 } from '../../llm/sonnet-4'; -import { createIdleTool } from '../../tools'; import { IDLE_TOOL_NAME } from '../../tools/communication-tools/idle-tool/idle-tool'; -import { - createEditFileTool, - createLsTool, - createMultiEditFileTool, - createWriteFileTool, -} from '../../tools/file-tools'; -import { createBashTool } from '../../tools/file-tools/bash-tool/bash-tool'; -import { BASH_TOOL_NAME } from '../../tools/file-tools/bash-tool/bash-tool'; -import { EDIT_FILE_TOOL_NAME } from '../../tools/file-tools/edit-file-tool/edit-file-tool'; -import { GREP_TOOL_NAME, createGrepTool } from '../../tools/file-tools/grep-tool/grep-tool'; -import { LS_TOOL_NAME } from '../../tools/file-tools/ls-tool/ls-tool'; -import { MULTI_EDIT_FILE_TOOL_NAME } from '../../tools/file-tools/multi-edit-file-tool/multi-edit-file-tool'; -import { READ_FILE_TOOL_NAME, createReadFileTool } from '../../tools/file-tools/read-file-tool/read-file-tool'; -import { WRITE_FILE_TOOL_NAME } from '../../tools/file-tools/write-file-tool/write-file-tool'; -import { createTaskTool } from '../../tools/task-tools/task-tool/task-tool'; -import { type AgentContext, repairToolCall } from '../../utils/tool-call-repair'; import { createAnalyticsEngineerToolset } from './create-analytics-engineer-toolset'; import { getDocsAgentSystemPrompt as getAnalyticsEngineerAgentSystemPrompt } from './get-analytics-engineer-agent-system-prompt'; +import type { + AnalyticsEngineerAgentOptions, + AnalyticsEngineerAgentStreamOptions, + TodoItem, +} from './types'; export const ANALYST_ENGINEER_AGENT_NAME = 'analyticsEngineerAgent'; const STOP_CONDITIONS = [stepCountIs(100), hasToolCall(IDLE_TOOL_NAME)]; -export const TodoItemSchema = z.object({ - id: z.string().describe('Unique identifier for the todo item. Use existing ID to update, or generate new ID for new items'), - content: z.string().describe('The content/description of the todo'), - status: z.enum(['pending', 'in_progress', 'completed']).describe('Current status of the todo'), - createdAt: z.string().datetime().optional().describe('ISO timestamp when todo was created (optional, will be set automatically for new items)'), - completedAt: z.string().datetime().optional().describe('ISO timestamp when todo was completed (optional)'), -}); - -export type TodoItem = z.infer; - -const AnalyticsEngineerAgentOptionsSchema = z.object({ - folder_structure: z.string().describe('The file structure of the dbt repository'), - userId: z.string(), - chatId: z.string(), - dataSourceId: z.string(), - organizationId: z.string(), - messageId: z.string(), - todosList: - z - .array(TodoItemSchema) - .optional() - .describe('Array of todo items to write/update. Include all todos with their current state.'), - model: z - .custom() - .optional() - .describe('Custom language model to use (defaults to Sonnet4)'), - isSubagent: z - .boolean() - .optional() - .describe('Flag indicating this is a subagent (prevents infinite recursion)'), - abortSignal: z - .custom() - .optional() - .describe('Optional abort signal to cancel agent execution'), -}); - -const AnalyticsEngineerAgentStreamOptionsSchema = z.object({ - messages: z.array(z.custom()).describe('The messages to send to the docs agent'), -}); - -export type AnalyticsEngineerAgentStreamOptions = z.infer< - typeof AnalyticsEngineerAgentStreamOptionsSchema ->; - -export type AnalyticsEngineerAgentOptions = z.infer< - typeof AnalyticsEngineerAgentOptionsSchema ->; - - export function createAnalyticsEngineerAgent( analyticsEngineerAgentOptions: AnalyticsEngineerAgentOptions ) { @@ -89,7 +24,6 @@ export function createAnalyticsEngineerAgent( } as ModelMessage; async function stream({ messages }: AnalyticsEngineerAgentStreamOptions) { - const toolSet = await createAnalyticsEngineerToolset(analyticsEngineerAgentOptions); const streamFn = () => diff --git a/packages/ai/src/agents/analytics-engineer-agent/create-analytics-engineer-toolset.ts b/packages/ai/src/agents/analytics-engineer-agent/create-analytics-engineer-toolset.ts index 511df9045..120402bc8 100644 --- a/packages/ai/src/agents/analytics-engineer-agent/create-analytics-engineer-toolset.ts +++ b/packages/ai/src/agents/analytics-engineer-agent/create-analytics-engineer-toolset.ts @@ -1,11 +1,32 @@ -import type { AnalyticsEngineerAgentOptions } from ".."; -import { createAnalyticsEngineerAgent } from ".."; -import { BASH_TOOL_NAME, EDIT_FILE_TOOL_NAME, GREP_TOOL_NAME, IDLE_TOOL_NAME, LS_TOOL_NAME, MULTI_EDIT_FILE_TOOL_NAME, READ_FILE_TOOL_NAME, WRITE_FILE_TOOL_NAME, createBashTool, createEditFileTool, createGrepTool, createIdleTool, createLsTool, createMultiEditFileTool, createReadFileTool, createTaskTool, createWriteFileTool } from "../../tools"; -import type { AgentFactory } from "../../tools/task-tools/task-tool/task-tool"; +import { + BASH_TOOL_NAME, + EDIT_FILE_TOOL_NAME, + GREP_TOOL_NAME, + IDLE_TOOL_NAME, + LS_TOOL_NAME, + MULTI_EDIT_FILE_TOOL_NAME, + READ_FILE_TOOL_NAME, + TODO_WRITE_TOOL_NAME, + WRITE_FILE_TOOL_NAME, + createBashTool, + createEditFileTool, + createGrepTool, + createIdleTool, + createLsTool, + createMultiEditFileTool, + createReadFileTool, + createTaskTool, + createTodoWriteTool, + createWriteFileTool, +} from '../../tools'; +import type { AgentFactory } from '../../tools/task-tools/task-tool/task-tool'; +import { createAnalyticsEngineerAgent } from './analytics-engineer-agent'; +import type { AnalyticsEngineerAgentOptions } from './types'; -export async function createAnalyticsEngineerToolset(analyticsEngineerAgentOptions: AnalyticsEngineerAgentOptions) { - const idleTool = createIdleTool({ - }); +export async function createAnalyticsEngineerToolset( + analyticsEngineerAgentOptions: AnalyticsEngineerAgentOptions +) { + const idleTool = createIdleTool({}); const writeFileTool = createWriteFileTool({ messageId: analyticsEngineerAgentOptions.messageId, projectDirectory: analyticsEngineerAgentOptions.folder_structure, @@ -34,21 +55,26 @@ export async function createAnalyticsEngineerToolset(analyticsEngineerAgentOptio messageId: analyticsEngineerAgentOptions.messageId, projectDirectory: analyticsEngineerAgentOptions.folder_structure, }); + const todosTool = createTodoWriteTool({ + chatId: analyticsEngineerAgentOptions.chatId, + workingDirectory: analyticsEngineerAgentOptions.folder_structure, + todosList: analyticsEngineerAgentOptions.todosList, + }); // Conditionally create task tool (only for main agent, not for subagents) const taskTool = !analyticsEngineerAgentOptions.isSubagent ? createTaskTool({ - messageId: analyticsEngineerAgentOptions.messageId, - projectDirectory: analyticsEngineerAgentOptions.folder_structure, - // Pass the agent factory function to enable task agent creation - // This needs to match the AgentFactory type signature - createAgent: ((options: AnalyticsEngineerAgentOptions) => { - return createAnalyticsEngineerAgent({ - ...options, - // Inherit model from parent agent if provided - model: analyticsEngineerAgentOptions.model, - }); - }) as unknown as AgentFactory, - }) + messageId: analyticsEngineerAgentOptions.messageId, + projectDirectory: analyticsEngineerAgentOptions.folder_structure, + // Pass the agent factory function to enable task agent creation + // This needs to match the AgentFactory type signature + createAgent: ((options: AnalyticsEngineerAgentOptions) => { + return createAnalyticsEngineerAgent({ + ...options, + // Inherit model from parent agent if provided + model: analyticsEngineerAgentOptions.model, + }); + }) as unknown as AgentFactory, + }) : null; return { @@ -60,6 +86,7 @@ export async function createAnalyticsEngineerToolset(analyticsEngineerAgentOptio [EDIT_FILE_TOOL_NAME]: editFileTool, [MULTI_EDIT_FILE_TOOL_NAME]: multiEditFileTool, [LS_TOOL_NAME]: lsTool, + [TODO_WRITE_TOOL_NAME]: todosTool, ...(taskTool ? { taskTool } : {}), }; -} \ No newline at end of file +} diff --git a/packages/ai/src/agents/analytics-engineer-agent/types.ts b/packages/ai/src/agents/analytics-engineer-agent/types.ts new file mode 100644 index 000000000..50962f4f6 --- /dev/null +++ b/packages/ai/src/agents/analytics-engineer-agent/types.ts @@ -0,0 +1,62 @@ +import type { LanguageModelV2 } from '@ai-sdk/provider'; +import type { ModelMessage } from 'ai'; +import z from 'zod'; + +export const TodoItemSchema = z.object({ + id: z + .string() + .describe( + 'Unique identifier for the todo item. Use existing ID to update, or generate new ID for new items' + ), + content: z.string().describe('The content/description of the todo'), + status: z.enum(['pending', 'in_progress', 'completed']).describe('Current status of the todo'), + createdAt: z + .string() + .datetime() + .optional() + .describe( + 'ISO timestamp when todo was created (optional, will be set automatically for new items)' + ), + completedAt: z + .string() + .datetime() + .optional() + .describe('ISO timestamp when todo was completed (optional)'), +}); + +export type TodoItem = z.infer; + +export const AnalyticsEngineerAgentOptionsSchema = z.object({ + folder_structure: z.string().describe('The file structure of the dbt repository'), + userId: z.string(), + chatId: z.string(), + dataSourceId: z.string(), + organizationId: z.string(), + messageId: z.string(), + todosList: z + .array(TodoItemSchema) + .default([]) + .describe('Array of todo items to write/update. Include all todos with their current state.'), + model: z + .custom() + .optional() + .describe('Custom language model to use (defaults to Sonnet4)'), + isSubagent: z + .boolean() + .optional() + .describe('Flag indicating this is a subagent (prevents infinite recursion)'), + abortSignal: z + .custom() + .optional() + .describe('Optional abort signal to cancel agent execution'), +}); + +export const AnalyticsEngineerAgentStreamOptionsSchema = z.object({ + messages: z.array(z.custom()).describe('The messages to send to the docs agent'), +}); + +export type AnalyticsEngineerAgentStreamOptions = z.infer< + typeof AnalyticsEngineerAgentStreamOptionsSchema +>; + +export type AnalyticsEngineerAgentOptions = z.infer; diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 5d0dd99e6..23d410a0f 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -34,6 +34,3 @@ export type { LsToolInput, LsToolOutput, } from './tools/file-tools/ls-tool/ls-tool'; - -// Export typed tool events for type-safe tool callbacks -export type { ToolEvent, ToolEventCallback } from './agents/analytics-engineer-agent/tool-events'; diff --git a/packages/ai/src/tools/database-tools/super-execute-sql/super-execute-sql.ts b/packages/ai/src/tools/database-tools/super-execute-sql/super-execute-sql.ts index 8e7cecc40..20160586e 100644 --- a/packages/ai/src/tools/database-tools/super-execute-sql/super-execute-sql.ts +++ b/packages/ai/src/tools/database-tools/super-execute-sql/super-execute-sql.ts @@ -1,6 +1,6 @@ import { tool } from 'ai'; import { z } from 'zod'; -import type { AnalyticsEngineerAgentOptions } from '../../../agents/analytics-engineer-agent/analytics-engineer-agent'; +import type { AnalyticsEngineerAgentOptions } from '../../../agents/analytics-engineer-agent/types'; import { createSuperExecuteSqlExecute } from './super-execute-sql-execute'; export const SuperExecuteSqlInputSchema = z.object({ diff --git a/packages/ai/src/tools/index.ts b/packages/ai/src/tools/index.ts index 5e44071ad..c43cb11ed 100644 --- a/packages/ai/src/tools/index.ts +++ b/packages/ai/src/tools/index.ts @@ -1,32 +1,71 @@ // Communication tools export { createDoneTool, DONE_TOOL_NAME } from './communication-tools/done-tool/done-tool'; export { createIdleTool, IDLE_TOOL_NAME } from './communication-tools/idle-tool/idle-tool'; -export { createSubmitThoughtsTool, SUBMIT_THOUGHTS_TOOL_NAME } from './communication-tools/submit-thoughts-tool/submit-thoughts-tool'; +export { + createSubmitThoughtsTool, + SUBMIT_THOUGHTS_TOOL_NAME, +} from './communication-tools/submit-thoughts-tool/submit-thoughts-tool'; // Planning/thinking tools -export { createSequentialThinkingTool, SEQUENTIAL_THINKING_TOOL_NAME } from './planning-thinking-tools/sequential-thinking-tool/sequential-thinking-tool'; +export { + createSequentialThinkingTool, + SEQUENTIAL_THINKING_TOOL_NAME, +} from './planning-thinking-tools/sequential-thinking-tool/sequential-thinking-tool'; // Task tools export { createTaskTool } from './task-tools/task-tool/task-tool'; // Visualization tools -export { createCreateMetricsTool, CREATE_METRICS_TOOL_NAME } from './visualization-tools/metrics/create-metrics-tool/create-metrics-tool'; -export { createModifyMetricsTool, MODIFY_METRICS_TOOL_NAME } from './visualization-tools/metrics/modify-metrics-tool/modify-metrics-tool'; -export { createCreateDashboardsTool, CREATE_DASHBOARDS_TOOL_NAME } from './visualization-tools/dashboards/create-dashboards-tool/create-dashboards-tool'; -export { createModifyDashboardsTool, MODIFY_DASHBOARDS_TOOL_NAME } from './visualization-tools/dashboards/modify-dashboards-tool/modify-dashboards-tool'; -export { createCreateReportsTool, CREATE_REPORTS_TOOL_NAME } from './visualization-tools/reports/create-reports-tool/create-reports-tool'; -export { createModifyReportsTool, MODIFY_REPORTS_TOOL_NAME } from './visualization-tools/reports/modify-reports-tool/modify-reports-tool'; +export { + createCreateMetricsTool, + CREATE_METRICS_TOOL_NAME, +} from './visualization-tools/metrics/create-metrics-tool/create-metrics-tool'; +export { + createModifyMetricsTool, + MODIFY_METRICS_TOOL_NAME, +} from './visualization-tools/metrics/modify-metrics-tool/modify-metrics-tool'; +export { + createCreateDashboardsTool, + CREATE_DASHBOARDS_TOOL_NAME, +} from './visualization-tools/dashboards/create-dashboards-tool/create-dashboards-tool'; +export { + createModifyDashboardsTool, + MODIFY_DASHBOARDS_TOOL_NAME, +} from './visualization-tools/dashboards/modify-dashboards-tool/modify-dashboards-tool'; +export { + createCreateReportsTool, + CREATE_REPORTS_TOOL_NAME, +} from './visualization-tools/reports/create-reports-tool/create-reports-tool'; +export { + createModifyReportsTool, + MODIFY_REPORTS_TOOL_NAME, +} from './visualization-tools/reports/modify-reports-tool/modify-reports-tool'; // Database tools -export { createExecuteSqlTool, EXECUTE_SQL_TOOL_NAME } from './database-tools/execute-sql/execute-sql'; +export { + createExecuteSqlTool, + EXECUTE_SQL_TOOL_NAME, +} from './database-tools/execute-sql/execute-sql'; export { executeSqlDocsAgent } from './database-tools/super-execute-sql/super-execute-sql'; // File tools export { createLsTool, LS_TOOL_NAME } from './file-tools/ls-tool/ls-tool'; -export { createReadFileTool, READ_FILE_TOOL_NAME } from './file-tools/read-file-tool/read-file-tool'; -export { createWriteFileTool, WRITE_FILE_TOOL_NAME } from './file-tools/write-file-tool/write-file-tool'; -export { createEditFileTool, EDIT_FILE_TOOL_NAME } from './file-tools/edit-file-tool/edit-file-tool'; -export { createMultiEditFileTool, MULTI_EDIT_FILE_TOOL_NAME } from './file-tools/multi-edit-file-tool/multi-edit-file-tool'; +export { + createReadFileTool, + READ_FILE_TOOL_NAME, +} from './file-tools/read-file-tool/read-file-tool'; +export { + createWriteFileTool, + WRITE_FILE_TOOL_NAME, +} from './file-tools/write-file-tool/write-file-tool'; +export { + createEditFileTool, + EDIT_FILE_TOOL_NAME, +} from './file-tools/edit-file-tool/edit-file-tool'; +export { + createMultiEditFileTool, + MULTI_EDIT_FILE_TOOL_NAME, +} from './file-tools/multi-edit-file-tool/multi-edit-file-tool'; export { createBashTool, BASH_TOOL_NAME } from './file-tools/bash-tool/bash-tool'; export { createGrepTool, GREP_TOOL_NAME } from './file-tools/grep-tool/grep-tool'; @@ -36,7 +75,10 @@ export { createWebSearchTool } from './web-tools/web-search-tool'; // More planning/thinking tools export { createCheckOffTodoListTool } from './planning-thinking-tools/check-off-todo-list-tool/check-off-todo-list-tool'; export { createUpdateClarificationsFileTool } from './planning-thinking-tools/update-clarifications-file-tool/update-clarifications-file-tool'; -export { createTodoWriteTool, TODO_WRITE_TOOL_NAME } from './planning-thinking-tools/todo-write-tool/todo-write-tool'; +export { + createTodoWriteTool, + TODO_WRITE_TOOL_NAME, +} from './planning-thinking-tools/todo-write-tool/todo-write-tool'; // Legacy exports for backward compatibility (to be deprecated) export { checkOffTodoList } from './planning-thinking-tools/check-off-todo-list-tool/check-off-todo-list-tool'; diff --git a/packages/ai/src/tools/planning-thinking-tools/todo-write-tool/todo-write-tool-execute.ts b/packages/ai/src/tools/planning-thinking-tools/todo-write-tool/todo-write-tool-execute.ts index 56666592d..12bdecb51 100644 --- a/packages/ai/src/tools/planning-thinking-tools/todo-write-tool/todo-write-tool-execute.ts +++ b/packages/ai/src/tools/planning-thinking-tools/todo-write-tool/todo-write-tool-execute.ts @@ -1,14 +1,18 @@ -import type { TodoWriteToolContext, TodoWriteToolInput, TodoWriteToolOutput } from './todo-write-tool'; -import type { TodoItem } from '../../../agents/analytics-engineer-agent/analytics-engineer-agent'; +import type { TodoItem } from '../../../agents/analytics-engineer-agent/types'; +import type { + TodoWriteToolContext, + TodoWriteToolInput, + TodoWriteToolOutput, +} from './todo-write-tool'; /** * Processes todos by setting timestamps and handling status changes */ -function processTodos(inputTodos: TodoItem[], existingTodos: TodoItem[]): TodoItem[] { - const existingById = new Map(existingTodos.map(todo => [todo.id, todo])); +function processTodos(inputTodos: TodoItem[], existingTodos: TodoItem[] = []): TodoItem[] { + const existingById = new Map(existingTodos.map((todo) => [todo.id, todo])); const now = new Date().toISOString(); - return inputTodos.map(todo => { + return inputTodos.map((todo) => { const existing = existingById.get(todo.id); const processed = { ...todo }; @@ -39,38 +43,19 @@ function processTodos(inputTodos: TodoItem[], existingTodos: TodoItem[]): TodoIt */ export function createTodoWriteToolExecute(context: TodoWriteToolContext) { return async function execute(input: TodoWriteToolInput): Promise { - const { chatId, workingDirectory } = context; + const { chatId, todosList = [] } = context; const { todos: inputTodos } = input; console.info(`Writing ${inputTodos.length} todo(s) for chat ${chatId}`); try { - - // Load existing todos from disk - let existingTodos: TodoItem[] = []; - try { - const loaded = await loadTodos(chatId, workingDirectory); - existingTodos = loaded?.todos || []; - } catch (error) { - console.warn('Failed to load existing todos:', error); - } - // Process the todos (set timestamps, handle status changes) - const processedTodos = processTodos(inputTodos, existingTodos); + const processedTodos = processTodos(inputTodos, todosList); - // Save to disk - try { - await saveTodos(chatId, workingDirectory, processedTodos); - } catch (error) { - console.error('Failed to save todos to disk:', error); - return { - success: false, - todos: processedTodos, - message: `Failed to save todos: ${error instanceof Error ? error.message : 'Unknown error'}`, - }; - } + // Update the in-memory todosList by clearing and repopulating + todosList.splice(0, todosList.length, ...processedTodos); - console.info(`Successfully saved ${processedTodos.length} todo(s)`); + console.info(`Successfully updated ${processedTodos.length} todo(s) in memory`); return { success: true, @@ -83,7 +68,7 @@ export function createTodoWriteToolExecute(context: TodoWriteToolContext) { return { success: false, - todos: [], + todos: todosList, message: `Error: ${errorMessage}`, }; } diff --git a/packages/ai/src/tools/planning-thinking-tools/todo-write-tool/todo-write-tool.test.ts b/packages/ai/src/tools/planning-thinking-tools/todo-write-tool/todo-write-tool.test.ts index 358c269a7..0f8c76a1a 100644 --- a/packages/ai/src/tools/planning-thinking-tools/todo-write-tool/todo-write-tool.test.ts +++ b/packages/ai/src/tools/planning-thinking-tools/todo-write-tool/todo-write-tool.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it, vi } from 'vitest'; -import type { TodoItem } from '../../../agents/analytics-engineer-agent/analytics-engineer-agent'; +import type { TodoItem } from '../../../agents/analytics-engineer-agent/types'; import type { TodoWriteToolInput } from './todo-write-tool'; import { createTodoWriteToolExecute } from './todo-write-tool-execute'; @@ -8,13 +8,9 @@ describe('createTodoWriteToolExecute', () => { const workingDirectory = '/test/directory'; it('should create new todos with timestamps', async () => { - // Mock the conversation-history module - vi.doMock('../../../../../apps/cli/src/utils/conversation-history', () => ({ - loadTodos: vi.fn().mockResolvedValue(null), - saveTodos: vi.fn().mockResolvedValue({ chatId, workingDirectory, todos: [], updatedAt: new Date().toISOString() }), - })); + const todosList: TodoItem[] = []; - const execute = createTodoWriteToolExecute({ chatId, workingDirectory }); + const execute = createTodoWriteToolExecute({ chatId, workingDirectory, todosList }); const input: TodoWriteToolInput = { todos: [ { @@ -36,11 +32,12 @@ describe('createTodoWriteToolExecute', () => { expect(result.todos).toHaveLength(2); expect(result.todos[0]?.createdAt).toBeDefined(); expect(result.todos[1]?.createdAt).toBeDefined(); + expect(todosList).toHaveLength(2); }); it('should preserve createdAt for existing todos', async () => { const existingCreatedAt = '2024-01-01T00:00:00.000Z'; - const existingTodos: TodoItem[] = [ + const todosList: TodoItem[] = [ { id: '1', content: 'Existing todo', @@ -49,12 +46,7 @@ describe('createTodoWriteToolExecute', () => { }, ]; - vi.doMock('../../../../../apps/cli/src/utils/conversation-history', () => ({ - loadTodos: vi.fn().mockResolvedValue({ chatId, workingDirectory, todos: existingTodos, updatedAt: new Date().toISOString() }), - saveTodos: vi.fn().mockResolvedValue({ chatId, workingDirectory, todos: existingTodos, updatedAt: new Date().toISOString() }), - })); - - const execute = createTodoWriteToolExecute({ chatId, workingDirectory }); + const execute = createTodoWriteToolExecute({ chatId, workingDirectory, todosList }); const input: TodoWriteToolInput = { todos: [ { @@ -72,7 +64,7 @@ describe('createTodoWriteToolExecute', () => { }); it('should set completedAt when status changes to completed', async () => { - const existingTodos: TodoItem[] = [ + const todosList: TodoItem[] = [ { id: '1', content: 'Todo to complete', @@ -81,12 +73,7 @@ describe('createTodoWriteToolExecute', () => { }, ]; - vi.doMock('../../../../../apps/cli/src/utils/conversation-history', () => ({ - loadTodos: vi.fn().mockResolvedValue({ chatId, workingDirectory, todos: existingTodos, updatedAt: new Date().toISOString() }), - saveTodos: vi.fn().mockResolvedValue({ chatId, workingDirectory, todos: existingTodos, updatedAt: new Date().toISOString() }), - })); - - const execute = createTodoWriteToolExecute({ chatId, workingDirectory }); + const execute = createTodoWriteToolExecute({ chatId, workingDirectory, todosList }); const input: TodoWriteToolInput = { todos: [ { @@ -104,7 +91,7 @@ describe('createTodoWriteToolExecute', () => { }); it('should clear completedAt when status changes from completed', async () => { - const existingTodos: TodoItem[] = [ + const todosList: TodoItem[] = [ { id: '1', content: 'Completed todo', @@ -114,12 +101,7 @@ describe('createTodoWriteToolExecute', () => { }, ]; - vi.doMock('../../../../../apps/cli/src/utils/conversation-history', () => ({ - loadTodos: vi.fn().mockResolvedValue({ chatId, workingDirectory, todos: existingTodos, updatedAt: new Date().toISOString() }), - saveTodos: vi.fn().mockResolvedValue({ chatId, workingDirectory, todos: existingTodos, updatedAt: new Date().toISOString() }), - })); - - const execute = createTodoWriteToolExecute({ chatId, workingDirectory }); + const execute = createTodoWriteToolExecute({ chatId, workingDirectory, todosList }); const input: TodoWriteToolInput = { todos: [ { diff --git a/packages/ai/src/tools/planning-thinking-tools/todo-write-tool/todo-write-tool.ts b/packages/ai/src/tools/planning-thinking-tools/todo-write-tool/todo-write-tool.ts index 4c4f4a58b..a0a548474 100644 --- a/packages/ai/src/tools/planning-thinking-tools/todo-write-tool/todo-write-tool.ts +++ b/packages/ai/src/tools/planning-thinking-tools/todo-write-tool/todo-write-tool.ts @@ -1,6 +1,6 @@ import { tool } from 'ai'; import { z } from 'zod'; -import { TodoItemSchema } from '../../../agents/analytics-engineer-agent/analytics-engineer-agent'; +import { TodoItemSchema } from '../../../agents/analytics-engineer-agent/types'; import { createTodoWriteToolExecute } from './todo-write-tool-execute'; export const TODO_WRITE_TOOL_NAME = 'todoWrite'; @@ -12,21 +12,24 @@ const TodoWriteToolOutputSchema = z.object({ }); const TodoWriteToolInputSchema = z.object({ - todos: z.array(TodoItemSchema).describe('Array of todo items to write/update. Include all todos with their current state.'), + todos: z + .array(TodoItemSchema) + .describe('Array of todo items to write/update. Include all todos with their current state.'), }); const TodoWriteToolContextSchema = z.object({ chatId: z.string().describe('The chat/conversation ID to associate todos with'), workingDirectory: z.string().describe('The working directory for the chat'), + todosList: z.array(TodoItemSchema).default([]).describe('In-memory array of todo items to manipulate'), }); export type TodoWriteToolInput = z.infer; export type TodoWriteToolOutput = z.infer; export type TodoWriteToolContext = z.infer; -export function createTodoWriteTool( - context: TAgentContext -) { +export function createTodoWriteTool< + TAgentContext extends TodoWriteToolContext = TodoWriteToolContext, +>(context: TAgentContext) { const execute = createTodoWriteToolExecute(context); return tool({