A mimic of the Task Tool from Claude Code.

task-agent.ts

import Anthropic from '@anthropic-ai/sdk';
import { z } from 'zod';
import { join } from 'path';
import { existsSync, mkdirSync, readFileSync, writeFileSync } from 'fs';

// ============================================================================
// TYPE DEFINITIONS
// ============================================================================

export const TodoStatus = z.enum(['pending', 'in_progress', 'completed']);
export const TodoPriority = z.enum(['high', 'medium', 'low']);

export const TodoItem = z.object({
    content: z.string().min(1, 'Content cannot be empty'),
    status: TodoStatus,
    priority: TodoPriority,
    id: z.string(),
});

export type TodoItemType = z.infer<typeof TodoItem>;
export const TodoList = z.array(TodoItem);

export interface AppConfig {
    maxSteps: number;
    stepDelay: number;
    maxTodos: number;
    todosDir?: string;
}

export interface SessionInfo {
    sessionId: string;
    agentId: string;
    startTime: Date;
    userInstructions: string;
}

export interface ExecutionResult {
    success: boolean;
    totalSteps: number;
    finalTodos: TodoItemType[];
    error?: string;
}

export interface ReportData {
    sessionInfo: SessionInfo;
    executionResult: ExecutionResult;
    endTime: Date;
    todosDir?: string;
}

// ============================================================================
// FILE OPERATIONS
// ============================================================================

export interface FileOperations {
    exists(path: string): boolean;
    mkdir(path: string, options?: { recursive: boolean }): void;
    readFile(path: string, encoding: BufferEncoding): string;
    writeFile(path: string, data: string, encoding?: BufferEncoding): void;
}

export class RealFileOperations implements FileOperations {
    exists(path: string): boolean {
        return existsSync(path);
    }

    mkdir(path: string, options?: { recursive: boolean }): void {
        mkdirSync(path, options);
    }

    readFile(path: string, encoding: BufferEncoding): string {
        return readFileSync(path, { encoding });
    }

    writeFile(path: string, data: string, encoding: BufferEncoding = 'utf-8'): void {
        writeFileSync(path, data, encoding);
    }
}

export class MockFileOperations implements FileOperations {
    private files: Map<string, string> = new Map();
    private directories: Set<string> = new Set();

    exists(path: string): boolean {
        return this.files.has(path) || this.directories.has(path);
    }

    mkdir(path: string, options?: { recursive: boolean }): void {
        this.directories.add(path);
        if (options?.recursive) {
            const parts = path.split('/');
            let current = '';
            for (const part of parts) {
                current = current ? `${current}/${part}` : part;
                this.directories.add(current);
            }
        }
    }

    readFile(path: string, encoding: BufferEncoding): string {
        const content = this.files.get(path);
        if (content === undefined) {
            throw new Error(`File not found: ${path}`);
        }
        return content;
    }

    writeFile(path: string, data: string, encoding?: BufferEncoding): void {
        this.files.set(path, data);
    }

    // Test helper methods
    setFile(path: string, content: string): void {
        this.files.set(path, content);
    }

    getFile(path: string): string | undefined {
        return this.files.get(path);
    }

    clear(): void {
        this.files.clear();
        this.directories.clear();
    }
}

// ============================================================================
// FILE UTILITY FUNCTIONS
// ============================================================================

export function getTodoFilePath(sessionId: string, agentId: string, todosDir = './todos'): string {
    const filename = `${sessionId}-agent-${agentId}.json`;
    return join(todosDir, filename);
}

export function ensureDirectoryExists(path: string, fileOps: FileOperations): void {
    if (!fileOps.exists(path)) {
        fileOps.mkdir(path, { recursive: true });
    }
}

export function readTodosFromFile(
    sessionId: string,
    agentId: string,
    fileOps: FileOperations,
    todosDir = './todos'
): TodoItemType[] {
    const filePath = getTodoFilePath(sessionId, agentId, todosDir);

    try {
        const content = fileOps.readFile(filePath, 'utf-8');
        const data = JSON.parse(content);
        return TodoList.parse(data);
    } catch (error) {
        return [];
    }
}

export function writeTodosToFile(
    todos: TodoItemType[],
    sessionId: string,
    agentId: string,
    fileOps: FileOperations,
    todosDir = './todos'
): void {
    ensureDirectoryExists(todosDir, fileOps);
    const filePath = getTodoFilePath(sessionId, agentId, todosDir);
    fileOps.writeFile(filePath, JSON.stringify(todos, null, 2));
}

export function saveReportToFile(
    report: string,
    sessionId: string,
    fileOps: FileOperations,
    todosDir = './todos'
): string {
    const reportsPath = join(todosDir, 'reports');
    ensureDirectoryExists(reportsPath, fileOps);
    
    const timestamp = new Date().toISOString().replace(/:/g, '-').split('.')[0];
    const filename = `report-${sessionId}-${timestamp}.txt`;
    const filePath = join(reportsPath, filename);
    
    fileOps.writeFile(filePath, report);
    return filePath;
}

// ============================================================================
// TODO MANAGER
// ============================================================================

export class TodoManager {
    constructor(
        private sessionId: string,
        private agentId: string,
        private fileOps: FileOperations,
        private todosDir = './todos'
    ) {}

    readTodos(): TodoItemType[] {
        return readTodosFromFile(this.sessionId, this.agentId, this.fileOps, this.todosDir);
    }

    writeTodos(todos: TodoItemType[]): void {
        writeTodosToFile(todos, this.sessionId, this.agentId, this.fileOps, this.todosDir);
    }

    updateTodos(todos: TodoItemType[]): { oldTodos: TodoItemType[]; newTodos: TodoItemType[] } {
        const oldTodos = this.readTodos();
        this.writeTodos(todos);
        return { oldTodos, newTodos: todos };
    }

    getTodosByStatus(status: z.infer<typeof TodoStatus>): TodoItemType[] {
        const todos = this.readTodos();
        return todos.filter(todo => todo.status === status);
    }

    getPendingTodos(): TodoItemType[] {
        return this.getTodosByStatus('pending');
    }

    getInProgressTodos(): TodoItemType[] {
        return this.getTodosByStatus('in_progress');
    }

    getCompletedTodos(): TodoItemType[] {
        return this.getTodosByStatus('completed');
    }

    areAllTodosCompleted(): boolean {
        const todos = this.readTodos();
        return todos.length > 0 && todos.every(todo => todo.status === 'completed');
    }

    hasIncompleteTodos(): boolean {
        const todos = this.readTodos();
        return todos.some(todo => todo.status === 'pending' || todo.status === 'in_progress');
    }

    isEmpty(): boolean {
        const todos = this.readTodos();
        return todos.length === 0;
    }

    markTodoAsInProgress(todoId: string): boolean {
        const todos = this.readTodos();
        const todoIndex = todos.findIndex(todo => todo.id === todoId);
        
        if (todoIndex === -1) {
            return false;
        }

        // Check if there are other in_progress tasks
        const hasInProgress = todos.some(todo => todo.status === 'in_progress');
        if (hasInProgress) {
            return false;
        }

        todos[todoIndex].status = 'in_progress';
        this.writeTodos(todos);
        return true;
    }

    markTodoAsCompleted(todoId: string): boolean {
        const todos = this.readTodos();
        const todoIndex = todos.findIndex(todo => todo.id === todoId);
        
        if (todoIndex === -1) {
            return false;
        }

        todos[todoIndex].status = 'completed';
        this.writeTodos(todos);
        return true;
    }

    getNextPendingTodo(): TodoItemType | null {
        const pendingTodos = this.getPendingTodos();
        return pendingTodos.length > 0 ? pendingTodos[0] : null;
    }

    getCurrentInProgressTodo(): TodoItemType | null {
        const inProgressTodos = this.getInProgressTodos();
        return inProgressTodos.length > 0 ? inProgressTodos[0] : null;
    }

    validateTodoList(todos: TodoItemType[]): { isValid: boolean; errors: string[] } {
        const errors: string[] = [];

        // Maximum TODO count check
        if (todos.length > 3) {
            errors.push('Maximum of 3 TODOs can be created');
        }

        // Check for multiple in_progress tasks
        const inProgressCount = todos.filter(todo => todo.status === 'in_progress').length;
        if (inProgressCount > 1) {
            errors.push('Cannot have multiple tasks in progress simultaneously');
        }

        // Duplicate ID check
        const ids = todos.map(todo => todo.id);
        const uniqueIds = new Set(ids);
        if (ids.length !== uniqueIds.size) {
            errors.push('Duplicate IDs exist');
        }

        return {
            isValid: errors.length === 0,
            errors
        };
    }
}

// ============================================================================
// SYSTEM PROMPT
// ============================================================================

export function buildSystemPrompt(userInstructions: string, config: AppConfig): string {
    return `You are a simple agent that manages and executes TODO lists.\n${userInstructions ? `User instructions: "${userInstructions}"` : ''}\n- Maximum ${config.maxTodos} TODOs\n- Use 'TodoRead' to check the list\n- Use 'TodoWrite' to update\nPlease first call 'TodoRead' and then update with 'TodoWrite' based on the results.`;
}

// ============================================================================
// TOOL SCHEMAS AND DEFINITIONS
// ============================================================================

export const TodoReadInputSchema = z.object({}).passthrough();
export const TodoWriteInputSchema = z.object({
    todos: TodoList.describe('The updated todo list'),
});

export const TodoReadTool: Anthropic.Tool = {
    name: 'TodoRead',
    description: 'Read the current todo list for the session',
    input_schema: {
        type: 'object',
        properties: {},
        additionalProperties: false,
    },
};

export const TodoWriteTool: Anthropic.Tool = {
    name: 'TodoWrite',
    description: 'Update the todo list for the current session. To be used proactively and often to track progress and pending tasks.',
    input_schema: {
        type: 'object',
        properties: {
            todos: {
                type: 'array',
                description: 'The updated todo list',
                items: {
                    type: 'object',
                    properties: {
                        content: { type: 'string', minLength: 1 },
                        status: { type: 'string', enum: ['pending', 'in_progress', 'completed'] },
                        priority: { type: 'string', enum: ['high', 'medium', 'low'] },
                        id: { type: 'string' },
                    },
                    required: ['content', 'status', 'priority', 'id'],
                    additionalProperties: false,
                },
            },
        },
        required: ['todos'],
        additionalProperties: false,
    },
};

// ============================================================================
// AGENT EXECUTOR
// ============================================================================

export class AgentExecutor {
    private todoManager: TodoManager;

    constructor(
        private anthropicClient: Anthropic,
        private sessionInfo: SessionInfo,
        private config: AppConfig,
        fileOps: FileOperations
    ) {
        this.todoManager = new TodoManager(
            sessionInfo.sessionId,
            sessionInfo.agentId,
            fileOps,
            config.todosDir
        );
    }

    async executeStep(currentStep: number): Promise<{
        toolUsed: boolean;
        message: string;
        shouldContinue: boolean;
        error?: string;
    }> {
        const { userInstructions } = this.sessionInfo;

        const messages: Anthropic.MessageParam[] = [{
            role: "user",
            content: userInstructions ?
                `User instructions: "${userInstructions}"\n\nPlease check the current TODO list and execute incomplete tasks based on these instructions. Use multiple tools efficiently to work effectively.` :
                `Please check the current TODO list and execute incomplete tasks. Use multiple tools efficiently to work effectively.`
        }];

        try {
            // Make all tools available from the beginning
            let currentMessages = [...messages];
            let toolUsed = false;
            let completionMessage = '';
            let maxRounds = 5;
            let roundCount = 0;

            while (roundCount < maxRounds) {
                roundCount++;

                const response = await this.anthropicClient.messages.create({
                    model: "claude-3-5-haiku-latest",
                    max_tokens: 2048,
                    tools: [
                        TodoReadTool,
                        TodoWriteTool,
                        {
                            type: "web_search_20250305",
                            name: "web_search"
                        }
                    ],
                    system: buildSystemPrompt(userInstructions, this.config),
                    messages: currentMessages,
                });

                const toolResults: Anthropic.ToolResultBlockParam[] = [];
                let hasToolsInThisRound = false;

                for (const content of response.content) {
                    if (content.type === 'tool_use') {
                        hasToolsInThisRound = true;
                        toolUsed = true;
                        const toolUse = content;

                        try {
                            if (toolUse.name === 'TodoRead') {
                                const result = await this.handleTodoRead(toolUse.input as z.infer<typeof TodoReadInputSchema>);
                                toolResults.push({
                                    type: 'tool_result',
                                    tool_use_id: toolUse.id,
                                    content: JSON.stringify(result, null, 2)
                                });
                            } else if (toolUse.name === 'TodoWrite') {
                                const result = await this.handleTodoWrite(toolUse.input as z.infer<typeof TodoWriteInputSchema>);
                                toolResults.push({
                                    type: 'tool_result',
                                    tool_use_id: toolUse.id,
                                    content: `TODO update completed: ${result.newTodos.length} tasks updated`
                                });
                            } else if (toolUse.name === 'web_search') {
                                console.log(`🔍 Web search executed for: ${JSON.stringify(toolUse.input)}`);
                            }
                        } catch (error: any) {
                            toolResults.push({
                                type: 'tool_result',
                                tool_use_id: toolUse.id,
                                content: `Error: ${error.message}`,
                                is_error: true
                            });
                        }
                    } else if (content.type === 'text') {
                        completionMessage += content.text;
                    }
                }

                // Exit loop if no tools were used
                if (!hasToolsInThisRound) {
                    break;
                }

                // If there are tool results, pass them to the next round
                if (toolResults.length > 0) {
                    currentMessages.push({
                        role: 'assistant',
                        content: response.content
                    });
                    currentMessages.push({
                        role: 'user',
                        content: toolResults
                    });
                } else {
                    break;
                }
            }

            const shouldContinue = this.shouldContinueExecution(toolUsed);

            return {
                toolUsed,
                message: completionMessage,
                shouldContinue
            };

        } catch (error: any) {
            return {
                toolUsed: false,
                message: '',
                shouldContinue: false,
                error: error.message
            };
        }
    }

    async execute(): Promise<ExecutionResult> {
        let currentStep = 0;
        let lastError: string | undefined;

        // Create initial TODO file
        if (this.todoManager.isEmpty()) {
            this.todoManager.writeTodos([]);
        }

        while (currentStep < this.config.maxSteps) {
            currentStep++;

            const stepResult = await this.executeStep(currentStep);

            if (stepResult.error) {
                lastError = stepResult.error;
                break;
            }

            if (!stepResult.shouldContinue) {
                break;
            }

            // Delay between steps
            if (currentStep < this.config.maxSteps) {
                await new Promise(resolve => setTimeout(resolve, this.config.stepDelay));
            }
        }

        const finalTodos = this.todoManager.readTodos();
        const success = !lastError && this.todoManager.areAllTodosCompleted();

        return {
            success,
            totalSteps: currentStep,
            finalTodos,
            error: lastError
        };
    }

    private async handleTodoRead(input: z.infer<typeof TodoReadInputSchema>): Promise<TodoItemType[]> {
        return this.todoManager.readTodos();
    }

    private async handleTodoWrite(input: z.infer<typeof TodoWriteInputSchema>): Promise<{
        oldTodos: TodoItemType[];
        newTodos: TodoItemType[];
    }> {
        const validation = this.todoManager.validateTodoList(input.todos);
        if (!validation.isValid) {
            throw new Error(`TODO validation failed: ${validation.errors.join(', ')}`);
        }

        return this.todoManager.updateTodos(input.todos);
    }

    private shouldContinueExecution(toolUsed: boolean): boolean {
        if (!toolUsed) {
            return false;
        }

        if (this.todoManager.areAllTodosCompleted()) {
            return false;
        }

        if (this.todoManager.isEmpty() && !this.sessionInfo.userInstructions) {
            return false;
        }

        return true;
    }
}