changes and clean up

This commit is contained in:
dal 2025-10-06 11:06:17 -06:00
parent ce9e51e481
commit 6c77a4f7a3
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
12 changed files with 458 additions and 222 deletions

View File

@ -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

View File

@ -15,6 +15,26 @@ const ConversationSchema = z.object({
export type Conversation = z.infer<typeof ConversationSchema>;
// 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<typeof TodoItemSchema>;
export type TodoList = z.infer<typeof TodoListSchema>;
// 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<TodoList | null> {
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<TodoList> {
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;
}

View File

@ -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({

View File

@ -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<typeof ClarifyingQuestionSchema>;
// 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<typeof AnalyticsEngineerAgentContextSchema>;

View File

@ -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<typeof TodoItemSchema>;
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<Sandbox>(
(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<LanguageModelV2>()
.optional()
@ -57,22 +64,20 @@ const AnalyticsEngineerAgentOptionsSchema = z.object({
.custom<AbortSignal>()
.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<ModelMessage>()).describe('The messages to send to the docs agent'),
});
export type AnalyticsEngineerAgentOptions = z.infer<typeof AnalyticsEngineerAgentOptionsSchema> & {
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<string, any> = {
[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 {

View File

@ -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 } : {}),
};
}

View File

@ -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;

View File

@ -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';

View File

@ -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<TodoWriteToolOutput> {
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}`,
};
}
};
}

View File

@ -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();
});
});

View File

@ -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<typeof TodoWriteToolInputSchema>;
export type TodoWriteToolOutput = z.infer<typeof TodoWriteToolOutputSchema>;
export type TodoWriteToolContext = z.infer<typeof TodoWriteToolContextSchema>;
export function createTodoWriteTool<TAgentContext extends TodoWriteToolContext = TodoWriteToolContext>(
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,
});
}