diff --git a/apps/cli/src/services/analytics-engineer-handler.ts b/apps/cli/src/services/analytics-engineer-handler.ts index 13338f627..9406f5985 100644 --- a/apps/cli/src/services/analytics-engineer-handler.ts +++ b/apps/cli/src/services/analytics-engineer-handler.ts @@ -59,7 +59,6 @@ export async function runAnalyticsEngineerAgent(params: RunAnalyticsEngineerAgen messageId: randomUUID(), model: proxyModel, abortSignal, - skipTracing: true, // Skip Braintrust tracing in CLI mode }); // Use conversation history - includes user messages, assistant messages, tool calls, and tool results diff --git a/apps/cli/src/utils/conversation-history.ts b/apps/cli/src/utils/conversation-history.ts index ac92dea66..64faf2d84 100644 --- a/apps/cli/src/utils/conversation-history.ts +++ b/apps/cli/src/utils/conversation-history.ts @@ -15,6 +15,26 @@ const ConversationSchema = z.object({ export type Conversation = z.infer; +// Schema for a todo item +const TodoItemSchema = z.object({ + id: z.string().describe('Unique identifier for the todo item'), + 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().describe('ISO timestamp when todo was created'), + completedAt: z.string().datetime().optional().describe('ISO timestamp when todo was completed'), +}); + +// Schema for todos associated with a chat +const TodoListSchema = z.object({ + chatId: z.string().uuid().describe('Unique chat/conversation ID'), + workingDirectory: z.string().describe('Absolute path of the working directory'), + updatedAt: z.string().datetime().describe('ISO timestamp when todos were last updated'), + todos: z.array(TodoItemSchema).describe('Array of todo items'), +}); + +export type TodoItem = z.infer; +export type TodoList = z.infer; + // Base directory for all history const HISTORY_DIR = join(homedir(), '.buster', 'history'); @@ -188,3 +208,52 @@ export async function deleteConversation(chatId: string, workingDirectory: strin const { unlink } = await import('node:fs/promises'); await unlink(filePath); } + +/** + * Gets the file path for todos associated with a chat + */ +function getTodoFilePath(chatId: string, workingDirectory: string): string { + return join(getHistoryDir(workingDirectory), `${chatId}.todos.json`); +} + +/** + * Loads todos for a specific chat + * Returns null if no todos exist for this chat + */ +export async function loadTodos( + chatId: string, + workingDirectory: string +): Promise { + try { + const filePath = getTodoFilePath(chatId, workingDirectory); + const data = await readFile(filePath, 'utf-8'); + const parsed = JSON.parse(data); + return TodoListSchema.parse(parsed); + } catch (error) { + // File doesn't exist or is invalid + return null; + } +} + +/** + * Saves todos for a specific chat + */ +export async function saveTodos( + chatId: string, + workingDirectory: string, + todos: TodoItem[] +): Promise { + await ensureHistoryDir(workingDirectory); + + const todoList: TodoList = { + chatId, + workingDirectory, + updatedAt: new Date().toISOString(), + todos, + }; + + const filePath = getTodoFilePath(chatId, workingDirectory); + await writeFile(filePath, JSON.stringify(todoList, null, 2), { mode: 0o600 }); + + return todoList; +} diff --git a/apps/server/src/api/v2/deploy/POST.ts b/apps/server/src/api/v2/deploy/POST.ts index 7c803144a..c8c14b2c3 100644 --- a/apps/server/src/api/v2/deploy/POST.ts +++ b/apps/server/src/api/v2/deploy/POST.ts @@ -108,7 +108,7 @@ export async function deployHandler( }; // Upsert the dataset - const { updated } = await upsertDataset(tx, datasetParams); + const { updated } = await upsertDataset(datasetParams); if (updated) { modelResult.updated.push({ diff --git a/packages/ai/src/agents/analytics-engineer-agent/analytics-engineer-agent-context.ts b/packages/ai/src/agents/analytics-engineer-agent/analytics-engineer-agent-context.ts deleted file mode 100644 index c9f06ea98..000000000 --- a/packages/ai/src/agents/analytics-engineer-agent/analytics-engineer-agent-context.ts +++ /dev/null @@ -1,29 +0,0 @@ -import type { Sandbox } from '@buster/sandbox'; -import { z } from 'zod'; - -// Best practice: Use const object for keys -export const AnalyticsEngineerAgentContextKeys = { - TodoList: 'todoList', - ClarificationQuestions: 'clarificationQuestions', - DataSourceId: 'dataSourceId', -} as const; - -// Extract type from const object -export type AnalyticsEngineerAgentContextKey = keyof typeof AnalyticsEngineerAgentContextKeys; - -export const ClarifyingQuestionSchema = z.object({ - issue: z.string(), - context: z.string(), - clarificationQuestion: z.string(), -}); - -export type MessageUserClarifyingQuestion = z.infer; - -// Use the const keys in your schema -export const AnalyticsEngineerAgentContextSchema = z.object({ - [AnalyticsEngineerAgentContextKeys.TodoList]: z.string(), - [AnalyticsEngineerAgentContextKeys.ClarificationQuestions]: z.array(ClarifyingQuestionSchema), - [AnalyticsEngineerAgentContextKeys.DataSourceId]: z.string().uuid(), -}); - -export type AnalyticsEngineerAgentContext = z.infer; 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 c517e9eac..a184811de 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 @@ -23,13 +23,23 @@ import { READ_FILE_TOOL_NAME, createReadFileTool } from '../../tools/file-tools/ 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 { ToolEventCallback } from './tool-events'; 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(), @@ -37,14 +47,11 @@ const AnalyticsEngineerAgentOptionsSchema = z.object({ dataSourceId: z.string(), organizationId: z.string(), messageId: z.string(), - sandbox: z - .custom( - (val) => { - return val && typeof val === 'object' && 'id' in val && 'fs' in val; - }, - { message: 'Invalid Sandbox instance' } - ) - .optional(), + todosList: + z + .array(TodoItemSchema) + .optional() + .describe('Array of todo items to write/update. Include all todos with their current state.'), model: z .custom() .optional() @@ -57,22 +64,20 @@ const AnalyticsEngineerAgentOptionsSchema = z.object({ .custom() .optional() .describe('Optional abort signal to cancel agent execution'), - skipTracing: z.boolean().optional().describe('Skip Braintrust tracing (useful for CLI mode)'), }); const AnalyticsEngineerAgentStreamOptionsSchema = z.object({ messages: z.array(z.custom()).describe('The messages to send to the docs agent'), }); -export type AnalyticsEngineerAgentOptions = z.infer & { - onToolEvent?: ToolEventCallback; -}; export type AnalyticsEngineerAgentStreamOptions = z.infer< typeof AnalyticsEngineerAgentStreamOptionsSchema >; -// Extended type for passing to tools (includes sandbox) -export type DocsAgentContextWithSandbox = AnalyticsEngineerAgentOptions & { sandbox: Sandbox }; +export type AnalyticsEngineerAgentOptions = z.infer< + typeof AnalyticsEngineerAgentOptionsSchema +>; + export function createAnalyticsEngineerAgent( analyticsEngineerAgentOptions: AnalyticsEngineerAgentOptions @@ -83,120 +88,23 @@ export function createAnalyticsEngineerAgent( providerOptions: DEFAULT_ANTHROPIC_OPTIONS, } as ModelMessage; - const idleTool = createIdleTool({ - }); - const writeFileTool = createWriteFileTool({ - messageId: analyticsEngineerAgentOptions.messageId, - projectDirectory: analyticsEngineerAgentOptions.folder_structure, - }); - const grepTool = createGrepTool({ - messageId: analyticsEngineerAgentOptions.messageId, - projectDirectory: analyticsEngineerAgentOptions.folder_structure, - }); - const readFileTool = createReadFileTool({ - messageId: analyticsEngineerAgentOptions.messageId, - projectDirectory: analyticsEngineerAgentOptions.folder_structure, - }); - const bashTool = createBashTool({ - messageId: analyticsEngineerAgentOptions.messageId, - projectDirectory: analyticsEngineerAgentOptions.folder_structure, - }); - const editFileTool = createEditFileTool({ - messageId: analyticsEngineerAgentOptions.messageId, - projectDirectory: analyticsEngineerAgentOptions.folder_structure, - }); - const multiEditFileTool = createMultiEditFileTool({ - messageId: analyticsEngineerAgentOptions.messageId, - projectDirectory: analyticsEngineerAgentOptions.folder_structure, - }); - const lsTool = createLsTool({ - messageId: analyticsEngineerAgentOptions.messageId, - projectDirectory: analyticsEngineerAgentOptions.folder_structure, - }); - - // Conditionally create task tool (only for main agent, not for subagents) - const taskTool = !analyticsEngineerAgentOptions.isSubagent - ? createTaskTool({ - messageId: analyticsEngineerAgentOptions.messageId, - projectDirectory: analyticsEngineerAgentOptions.folder_structure, - // Wrap onToolEvent to ensure type compatibility - no optional chaining - ...(analyticsEngineerAgentOptions.onToolEvent && { - // biome-ignore lint/suspicious/noExplicitAny: Bridging between ToolEventType and parent's ToolEvent type - onToolEvent: (event) => analyticsEngineerAgentOptions.onToolEvent?.(event as any), - }), - // Pass the agent factory function to enable task agent creation - // This needs to match the AgentFactory type signature - createAgent: (options) => { - return createAnalyticsEngineerAgent({ - ...options, - // Inherit model from parent agent if provided - model: analyticsEngineerAgentOptions.model, - }); - }, - }) - : null; - - // Create planning tools with simple context async function stream({ messages }: AnalyticsEngineerAgentStreamOptions) { - const agentContext: AgentContext = { - agentName: ANALYST_ENGINEER_AGENT_NAME, - availableTools: [IDLE_TOOL_NAME, GREP_TOOL_NAME, WRITE_FILE_TOOL_NAME, READ_FILE_TOOL_NAME, BASH_TOOL_NAME, EDIT_FILE_TOOL_NAME, MULTI_EDIT_FILE_TOOL_NAME, LS_TOOL_NAME], - }; - // Build tools object conditionally including task tool - // biome-ignore lint/suspicious/noExplicitAny: tools object contains various tool types - const tools: Record = { - [IDLE_TOOL_NAME]: idleTool, - [GREP_TOOL_NAME]: grepTool, - [WRITE_FILE_TOOL_NAME]: writeFileTool, - [READ_FILE_TOOL_NAME]: readFileTool, - [BASH_TOOL_NAME]: bashTool, - [EDIT_FILE_TOOL_NAME]: editFileTool, - [MULTI_EDIT_FILE_TOOL_NAME]: multiEditFileTool, - [LS_TOOL_NAME]: lsTool, - }; - - // Add task tool only if not a subagent (prevent recursion) - if (taskTool) { - tools.taskTool = taskTool; - } + const toolSet = await createAnalyticsEngineerToolset(analyticsEngineerAgentOptions); const streamFn = () => streamText({ model: analyticsEngineerAgentOptions.model || Sonnet4, providerOptions: DEFAULT_ANTHROPIC_OPTIONS, - tools, + tools: toolSet, messages: [systemMessage, ...messages], stopWhen: STOP_CONDITIONS, toolChoice: 'required', maxOutputTokens: 10000, temperature: 0, - experimental_context: analyticsEngineerAgentOptions, - ...(analyticsEngineerAgentOptions.abortSignal && { - abortSignal: analyticsEngineerAgentOptions.abortSignal, - }), - experimental_repairToolCall: async (repairContext) => { - return repairToolCall({ - toolCall: repairContext.toolCall, - tools: repairContext.tools, - error: repairContext.error, - messages: repairContext.messages, - ...(repairContext.system && { system: repairContext.system }), - ...(repairContext.inputSchema && { inputSchema: repairContext.inputSchema }), - agentContext, - }); - }, }); - // Skip tracing for CLI mode - if (analyticsEngineerAgentOptions.skipTracing) { - return streamFn(); - } - - // Use Braintrust tracing for production - return wrapTraced(streamFn, { - name: 'Docs Agent', - })(); + return streamFn(); } return { 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 new file mode 100644 index 000000000..511df9045 --- /dev/null +++ b/packages/ai/src/agents/analytics-engineer-agent/create-analytics-engineer-toolset.ts @@ -0,0 +1,65 @@ +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"; + +export async function createAnalyticsEngineerToolset(analyticsEngineerAgentOptions: AnalyticsEngineerAgentOptions) { + const idleTool = createIdleTool({ + }); + const writeFileTool = createWriteFileTool({ + messageId: analyticsEngineerAgentOptions.messageId, + projectDirectory: analyticsEngineerAgentOptions.folder_structure, + }); + const grepTool = createGrepTool({ + messageId: analyticsEngineerAgentOptions.messageId, + projectDirectory: analyticsEngineerAgentOptions.folder_structure, + }); + const readFileTool = createReadFileTool({ + messageId: analyticsEngineerAgentOptions.messageId, + projectDirectory: analyticsEngineerAgentOptions.folder_structure, + }); + const bashTool = createBashTool({ + messageId: analyticsEngineerAgentOptions.messageId, + projectDirectory: analyticsEngineerAgentOptions.folder_structure, + }); + const editFileTool = createEditFileTool({ + messageId: analyticsEngineerAgentOptions.messageId, + projectDirectory: analyticsEngineerAgentOptions.folder_structure, + }); + const multiEditFileTool = createMultiEditFileTool({ + messageId: analyticsEngineerAgentOptions.messageId, + projectDirectory: analyticsEngineerAgentOptions.folder_structure, + }); + const lsTool = createLsTool({ + messageId: analyticsEngineerAgentOptions.messageId, + projectDirectory: analyticsEngineerAgentOptions.folder_structure, + }); + // 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, + }) + : null; + + return { + [IDLE_TOOL_NAME]: idleTool, + [WRITE_FILE_TOOL_NAME]: writeFileTool, + [GREP_TOOL_NAME]: grepTool, + [READ_FILE_TOOL_NAME]: readFileTool, + [BASH_TOOL_NAME]: bashTool, + [EDIT_FILE_TOOL_NAME]: editFileTool, + [MULTI_EDIT_FILE_TOOL_NAME]: multiEditFileTool, + [LS_TOOL_NAME]: lsTool, + ...(taskTool ? { taskTool } : {}), + }; +} \ No newline at end of file diff --git a/packages/ai/src/agents/analytics-engineer-agent/tool-events.ts b/packages/ai/src/agents/analytics-engineer-agent/tool-events.ts deleted file mode 100644 index 99334fc07..000000000 --- a/packages/ai/src/agents/analytics-engineer-agent/tool-events.ts +++ /dev/null @@ -1,53 +0,0 @@ -import type { IdleInput, IdleOutput } from '../../tools/communication-tools/idle-tool/idle-tool'; -import type { BashToolInput, BashToolOutput } from '../../tools/file-tools/bash-tool/bash-tool'; -import type { - EditFileToolInput, - EditFileToolOutput, -} from '../../tools/file-tools/edit-file-tool/edit-file-tool'; -import type { GrepToolInput, GrepToolOutput } from '../../tools/file-tools/grep-tool/grep-tool'; -import type { LsToolInput, LsToolOutput } from '../../tools/file-tools/ls-tool/ls-tool'; -import type { - ReadFileToolInput, - ReadFileToolOutput, -} from '../../tools/file-tools/read-file-tool/read-file-tool'; -import type { - WriteFileToolInput, - WriteFileToolOutput, -} from '../../tools/file-tools/write-file-tool/write-file-tool'; - -/** - * Discriminated union of all possible tool events in the analytics engineer agent - * This provides full type safety from tools -> CLI display - */ -export type ToolEvent = - // Idle tool events - | { tool: 'idleTool'; event: 'start'; args: IdleInput } - | { tool: 'idleTool'; event: 'complete'; result: IdleOutput; args: IdleInput } - // Bash tool events - | { tool: 'bashTool'; event: 'start'; args: BashToolInput } - | { tool: 'bashTool'; event: 'complete'; result: BashToolOutput; args: BashToolInput } - // Grep tool events - | { tool: 'grepTool'; event: 'start'; args: GrepToolInput } - | { tool: 'grepTool'; event: 'complete'; result: GrepToolOutput; args: GrepToolInput } - // Read file tool events - | { tool: 'readFileTool'; event: 'start'; args: ReadFileToolInput } - | { tool: 'readFileTool'; event: 'complete'; result: ReadFileToolOutput; args: ReadFileToolInput } - // Write file tool events - | { tool: 'writeFileTool'; event: 'start'; args: WriteFileToolInput } - | { - tool: 'writeFileTool'; - event: 'complete'; - result: WriteFileToolOutput; - args: WriteFileToolInput; - } - // Edit file tool events - | { tool: 'editFileTool'; event: 'start'; args: EditFileToolInput } - | { tool: 'editFileTool'; event: 'complete'; result: EditFileToolOutput; args: EditFileToolInput } - // Ls tool events - | { tool: 'lsTool'; event: 'start'; args: LsToolInput } - | { tool: 'lsTool'; event: 'complete'; result: LsToolOutput; args: LsToolInput }; - -/** - * Callback type for tool events - single typed callback for all tools - */ -export type ToolEventCallback = (event: ToolEvent) => void; diff --git a/packages/ai/src/tools/index.ts b/packages/ai/src/tools/index.ts index 21a3c518f..5e44071ad 100644 --- a/packages/ai/src/tools/index.ts +++ b/packages/ai/src/tools/index.ts @@ -1,32 +1,42 @@ -export { createDoneTool } from './communication-tools/done-tool/done-tool'; -export { createIdleTool } from './communication-tools/idle-tool/idle-tool'; -export { createSubmitThoughtsTool } from './communication-tools/submit-thoughts-tool/submit-thoughts-tool'; -export { createSequentialThinkingTool } from './planning-thinking-tools/sequential-thinking-tool/sequential-thinking-tool'; -// Task tools - factory functions +// 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'; + +// Planning/thinking tools +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 - factory functions -export { createCreateMetricsTool } from './visualization-tools/metrics/create-metrics-tool/create-metrics-tool'; -export { createModifyMetricsTool } from './visualization-tools/metrics/modify-metrics-tool/modify-metrics-tool'; -export { createCreateDashboardsTool } from './visualization-tools/dashboards/create-dashboards-tool/create-dashboards-tool'; -export { createModifyDashboardsTool } from './visualization-tools/dashboards/modify-dashboards-tool/modify-dashboards-tool'; -export { createCreateReportsTool } from './visualization-tools/reports/create-reports-tool/create-reports-tool'; -export { createModifyReportsTool } from './visualization-tools/reports/modify-reports-tool/modify-reports-tool'; -export { createExecuteSqlTool } from './database-tools/execute-sql/execute-sql'; + +// 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'; + +// Database tools +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 - factory functions -export { createLsTool } from './file-tools/ls-tool/ls-tool'; -export { createReadFileTool } from './file-tools/read-file-tool/read-file-tool'; -export { createWriteFileTool } from './file-tools/write-file-tool/write-file-tool'; -export { createEditFileTool } from './file-tools/edit-file-tool/edit-file-tool'; -export { createMultiEditFileTool } from './file-tools/multi-edit-file-tool/multi-edit-file-tool'; -export { createBashTool } from './file-tools/bash-tool/bash-tool'; -export { createGrepTool } from './file-tools/grep-tool/grep-tool'; -// Web tools - factory functions + +// 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 { createBashTool, BASH_TOOL_NAME } from './file-tools/bash-tool/bash-tool'; +export { createGrepTool, GREP_TOOL_NAME } from './file-tools/grep-tool/grep-tool'; + +// Web tools export { createWebSearchTool } from './web-tools/web-search-tool'; -// Planning/thinking tools - factory functions +// 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'; // 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-description.txt b/packages/ai/src/tools/planning-thinking-tools/todo-write-tool/todo-write-tool-description.txt new file mode 100644 index 000000000..e69de29bb 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 new file mode 100644 index 000000000..56666592d --- /dev/null +++ b/packages/ai/src/tools/planning-thinking-tools/todo-write-tool/todo-write-tool-execute.ts @@ -0,0 +1,91 @@ +import type { TodoWriteToolContext, TodoWriteToolInput, TodoWriteToolOutput } from './todo-write-tool'; +import type { TodoItem } from '../../../agents/analytics-engineer-agent/analytics-engineer-agent'; + +/** + * 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])); + const now = new Date().toISOString(); + + return inputTodos.map(todo => { + const existing = existingById.get(todo.id); + const processed = { ...todo }; + + // Set createdAt for new todos + if (!existing) { + processed.createdAt = processed.createdAt || now; + } else { + // Keep existing createdAt + processed.createdAt = existing.createdAt; + } + + // Set completedAt when status changes to completed + if (processed.status === 'completed' && (!existing || existing.status !== 'completed')) { + processed.completedAt = now; + } + + // Clear completedAt if status changes from completed to something else + if (processed.status !== 'completed' && existing?.status === 'completed') { + processed.completedAt = undefined; + } + + return processed; + }); +} + +/** + * Creates the execute function for the todo write tool + */ +export function createTodoWriteToolExecute(context: TodoWriteToolContext) { + return async function execute(input: TodoWriteToolInput): Promise { + const { chatId, workingDirectory } = 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); + + // 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'}`, + }; + } + + console.info(`Successfully saved ${processedTodos.length} todo(s)`); + + return { + success: true, + todos: processedTodos, + message: `Successfully updated ${processedTodos.length} todo(s)`, + }; + } catch (error) { + const errorMessage = error instanceof Error ? error.message : 'Unknown error'; + console.error('Error in todo write tool:', errorMessage); + + return { + success: false, + todos: [], + 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 new file mode 100644 index 000000000..358c269a7 --- /dev/null +++ b/packages/ai/src/tools/planning-thinking-tools/todo-write-tool/todo-write-tool.test.ts @@ -0,0 +1,138 @@ +import { describe, expect, it, vi } from 'vitest'; +import type { TodoItem } from '../../../agents/analytics-engineer-agent/analytics-engineer-agent'; +import type { TodoWriteToolInput } from './todo-write-tool'; +import { createTodoWriteToolExecute } from './todo-write-tool-execute'; + +describe('createTodoWriteToolExecute', () => { + const chatId = 'test-chat-id'; + 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 execute = createTodoWriteToolExecute({ chatId, workingDirectory }); + const input: TodoWriteToolInput = { + todos: [ + { + id: '1', + content: 'First todo', + status: 'pending', + }, + { + id: '2', + content: 'Second todo', + status: 'in_progress', + }, + ], + }; + + const result = await execute(input); + + expect(result.success).toBe(true); + expect(result.todos).toHaveLength(2); + expect(result.todos[0]?.createdAt).toBeDefined(); + expect(result.todos[1]?.createdAt).toBeDefined(); + }); + + it('should preserve createdAt for existing todos', async () => { + const existingCreatedAt = '2024-01-01T00:00:00.000Z'; + const existingTodos: TodoItem[] = [ + { + id: '1', + content: 'Existing todo', + status: 'pending', + createdAt: existingCreatedAt, + }, + ]; + + 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 input: TodoWriteToolInput = { + todos: [ + { + id: '1', + content: 'Updated todo', + status: 'completed', + }, + ], + }; + + const result = await execute(input); + + expect(result.success).toBe(true); + expect(result.todos[0]?.createdAt).toBe(existingCreatedAt); + }); + + it('should set completedAt when status changes to completed', async () => { + const existingTodos: TodoItem[] = [ + { + id: '1', + content: 'Todo to complete', + status: 'in_progress', + createdAt: '2024-01-01T00:00:00.000Z', + }, + ]; + + 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 input: TodoWriteToolInput = { + todos: [ + { + id: '1', + content: 'Todo to complete', + status: 'completed', + }, + ], + }; + + const result = await execute(input); + + expect(result.success).toBe(true); + expect(result.todos[0]?.completedAt).toBeDefined(); + }); + + it('should clear completedAt when status changes from completed', async () => { + const existingTodos: TodoItem[] = [ + { + id: '1', + content: 'Completed todo', + status: 'completed', + createdAt: '2024-01-01T00:00:00.000Z', + completedAt: '2024-01-02T00:00:00.000Z', + }, + ]; + + 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 input: TodoWriteToolInput = { + todos: [ + { + id: '1', + content: 'Completed todo', + status: 'in_progress', + }, + ], + }; + + const result = await execute(input); + + expect(result.success).toBe(true); + expect(result.todos[0]?.completedAt).toBeUndefined(); + }); +}); 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 new file mode 100644 index 000000000..4c4f4a58b --- /dev/null +++ b/packages/ai/src/tools/planning-thinking-tools/todo-write-tool/todo-write-tool.ts @@ -0,0 +1,38 @@ +import { tool } from 'ai'; +import { z } from 'zod'; +import { TodoItemSchema } from '../../../agents/analytics-engineer-agent/analytics-engineer-agent'; +import { createTodoWriteToolExecute } from './todo-write-tool-execute'; + +export const TODO_WRITE_TOOL_NAME = 'todoWrite'; + +const TodoWriteToolOutputSchema = z.object({ + success: z.boolean().describe('Whether the operation was successful'), + todos: z.array(TodoItemSchema).describe('The full current state of all todos'), + message: z.string().optional().describe('Optional message about the operation'), +}); + +const TodoWriteToolInputSchema = z.object({ + 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'), +}); + +export type TodoWriteToolInput = z.infer; +export type TodoWriteToolOutput = z.infer; +export type TodoWriteToolContext = z.infer; + +export function createTodoWriteTool( + context: TAgentContext +) { + const execute = createTodoWriteToolExecute(context); + + return tool({ + description: `Write and manage todo items for the current chat session. Accepts an array of todo items with their current state (pending, in_progress, or completed). All todos are persisted to disk and associated with the current chat. Returns the full current state of all todos.`, + inputSchema: TodoWriteToolInputSchema, + outputSchema: TodoWriteToolOutputSchema, + execute, + }); +}