Refactor main work loop

[janole]

Proposal: Refactor the Main Work Loop

This document proposes a new architecture for the main execution loop in src/ai/work.ts to improve clarity, robustness, and maintainability.

The Problem

The current recursive architecture in work.ts, with its interacting work, workInternal, and workTools functions, is a bit clumsy and difficult to follow. The conversation state is managed implicitly by passing the messages array up and down a recursive call stack. This makes the execution flow hard to trace and makes the code brittle and prone to bugs when new features or states are added.

Proposed Architecture: State-Driven Loop

Instead of recursion, we can use a state-driven while loop inside the main work function. This loop will manage the entire conversation turn, continuing as long as the AI's last response was a tool call and exiting only when the AI responds with a final text message.

This makes the execution flow linear, explicit, and much easier to understand.

The flow within the work function would be:

1. Initialize: Set up the LLM, tools, and the initial message list.
2. Start Loop: Enter a while(true) loop that orchestrates the turn.
3. Execute Tools: Process any runnable tools (pending or confirmed).
   • If a tool requires confirmation, exit the function to await user input. The UI will re-trigger work after the user makes a choice.
4. Check for Completion: If the last message is a final text response from the AI (not a tool call), the turn is over. Break the loop.
5. Call LLM: If tool results were just processed, call the LLM with the updated message list.
6. Process Response: Add the AI's response to the message list. If it contains new tool calls, add them as pending.
7. Repeat: The loop continues, ready to execute the newly requested tools.

Benefits of This Architecture

1. Clarity: The while loop makes the sequence of events explicit and easy to follow.
2. Explicit State Management: State is managed within the work function's scope, not implicitly through recursion, making debugging much easier.
3. Single Responsibility: The logic is cleanly separated: work orchestrates, runLLMAndProcessResponse calls the AI, and executeTools runs the tools.
4. Robustness: This pattern is less prone to errors when adding new steps to the process. New logic can be added as a clear step within the loop.

Refactored work.ts Code

Here is a complete, refactored version of work.ts that implements this new architecture.

import { AIMessage } from "@langchain/core/messages";
import { BaseLanguageModel } from "@langchain/core/language_models/base";
import { ToolMessage } from "@langchain/core/tools";
import { ChatSession } from "./chat-session.js";
import { TMessage, ToolProgressMessage } from "./custom-messages.js";
import { getTools, TTool } from "./tools/index.js";
import { chatService } from "./chat-service.js";
import { getStream, tryCatch } from "./utils/index.js";

// --- Main Entry Point ---

export async function work(props: { session: ChatSession, send: (messages: TMessage[]) => void, signal: AbortSignal }) {
    const { session, send, signal } = props;

    // 1. Initialize tools and LLM
    const tools = getTools({ includeDestructiveTools: session.toolMode !== "read-only" });
    const llm = await chatService.getLLM(session).then(llm => llm.withTools(Object.values(tools)));

    let messages = [...session.messages];
    
    // The main loop that drives the conversation turn.
    // It continues as long as the AI is calling tools.
    while (true) {
        // 2. Execute any tools that are ready to run.
        const { messages: postToolMessages, needsConfirmation } = await executeTools({
            messages,
            tools,
            toolMode: session.toolMode,
            workDir: session.workDir,
        });
        
        messages = postToolMessages;
        send(messages); // Update the UI with tool progress

        // If a tool needs confirmation, we stop here. The UI will trigger 'work' again once the user responds.
        if (needsConfirmation) {
            return messages;
        }

        // 3. Check if the last message is a tool result. If not, the AI is done with tools and we can exit.
        const lastMessage = messages[messages.length - 1];
        const isStillProcessingTools = lastMessage?.getType() === 'tool' || (ToolProgressMessage.isTypeOf(lastMessage) && lastMessage.status !== 'pending');
        
        if (!isStillProcessingTools) {
             // If there's nothing more to process, we can assume the last AI message was final.
             const lastAiMessage = messages.slice().reverse().find(m => m.getType() === 'ai') as AIMessage;
             if (!lastAiMessage?.tool_calls?.length) {
                 break;
             }
        }
        
        // 4. Call the LLM with the latest messages (including tool results).
        const { messages: postLlmMessages, hasToolCalls } = await runLLMAndProcessResponse({
            llm,
            messages,
            workDir: session.workDir,
            signal,
            send,
        });

        messages = postLlmMessages;
        send(messages); // Update UI with AI response and pending tool calls

        // If the AI's response has no new tool calls, the conversation turn is complete.
        if (!hasToolCalls) {
            break;
        }
    }

    return messages;
}

// --- Helper Functions ---

/**
 * Calls the LLM and processes the response, adding pending tool messages.
 */
async function runLLMAndProcessResponse(props: { llm: BaseLanguageModel, messages: TMessage[], workDir: string, signal: AbortSignal, send: (messages: TMessage[]) => void }) {
    const { llm, messages, workDir, signal, send } = props;
    
    const { stream, error } = await getStream(llm, messages, { metadata: { workDir }, signal });
    if (error || !stream) throw (error || new Error("Failed to get stream from LLM."));

    let aiMessage: AIMessage | null = null;
    for await (const chunk of stream) {
        aiMessage = aiMessage ? aiMessage.concat(chunk) : (chunk instanceof AIMessage ? chunk : new AIMessage(chunk.content));
        send([...messages, aiMessage]);
    }
    if (!aiMessage) throw new Error("LLM stream ended without a message.");

    const newMessages = [...messages, aiMessage];
    const hasToolCalls = !!aiMessage.tool_calls?.length;

    // If the AI wants to call tools, add them to the message list in "pending" state.
    if (hasToolCalls) {
        const toolProgressMessages = aiMessage.tool_calls!.map(toolCall => new ToolProgressMessage(toolCall));
        newMessages.push(...toolProgressMessages);
    }
    
    return { messages: newMessages, hasToolCalls };
}

/**
 * Finds and executes all tools that are in a runnable state.
 */
async function executeTools(props: { messages: TMessage[], tools: Record<string, TTool>, toolMode: ChatSession['toolMode'], workDir: string }) {
    const { messages, tools, toolMode, workDir } = props;
    
    let needsConfirmation = false;
    const workingMessages = [...messages];

    for (let i = 0; i < workingMessages.length; i++) {
        const msg = workingMessages[i];
        if (!ToolProgressMessage.isTypeOf(msg)) continue;

        const toolProgress = msg as ToolProgressMessage;
        const toolCall = toolProgress.toolCall;
        
        // Skip tools that are not ready to be processed.
        if (toolProgress.status !== 'pending' && toolProgress.status !== 'confirmed') {
            continue;
        }

        const selectedTool = tools[toolCall.name];

        // Handle denied or invalid tool calls
        if (!selectedTool) {
            toolProgress.status = "error";
            toolProgress.content = "Tool not found";
        } else if (toolProgress.status === "pending" && selectedTool.metadata?.["destructive"]) {
            if (toolMode === "read-only") {
                toolProgress.status = "error";
                toolProgress.content = "Tool call denied by security policy (read-only mode)";
            } else if (toolMode === "confirm") {
                toolProgress.status = "pending-confirmation";
                needsConfirmation = true;
                continue; // Move to the next message without executing this one
            }
        }

        // Execute the tool if it's runnable
        if(toolProgress.status === 'confirmed' || (toolProgress.status === 'pending' && !selectedTool.metadata?.["destructive"])){
            const { result, error } = await tryCatch<ToolMessage>(selectedTool.invoke(toolCall, { metadata: { workDir } }));
            
            if (result) {
                toolProgress.status = result.text.startsWith("ERROR: ") ? "error" : "success";
                toolProgress.content = result.text;
                // Find the original AI message and add the real ToolMessage for the next LLM call
                workingMessages.push(result);
            } else {
                toolProgress.status = "error";
                toolProgress.content = error?.message || "Unknown Error";
            }
        }
    }

    return { messages: workingMessages, needsConfirmation };
}

export function needsToolConfirmation(messages: TMessage[]) {
    return messages.some(m => ToolProgressMessage.isTypeOf(m) && m.status === "pending-confirmation");
}