shockbot/agents/ollama.ts

import { Client } from "@modelcontextprotocol/sdk/client/index.js";
import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
import { Ollama, type Message, type ToolCall } from "ollama";
import { log } from "../utils/cli.ts";
import { retry } from "../utils/retry.ts";
import { agent, type AgentResult, type AgentRunContext } from "./shared.ts";

const DEFAULT_MODEL = "qwen3.6:35b";
const MAX_ITERATIONS = 100;

interface OllamaTool {
    type: "function";
    function: {
        name: string;
        description: string;
        parameters: Record<string, unknown>;
    };
}

async function buildMcpClient(mcpServerUrl: string): Promise<Client> {
    const client = new Client(
        { name: "shockbot-agent", version: "0.1.0" },
        { capabilities: {} },
    );
    const transport = new StreamableHTTPClientTransport(new URL(mcpServerUrl));
    await client.connect(transport);
    return client;
}

async function getOllamaTools(mcpClient: Client): Promise<OllamaTool[]> {
    const { tools } = await mcpClient.listTools();
    return tools.map((t) => ({
        type: "function" as const,
        function: {
            name: t.name,
            description: t.description ?? "",
            parameters: (t.inputSchema as Record<string, unknown>) ?? {
                type: "object",
                properties: {},
            },
        },
    }));
}

async function callMcpTool(
    mcpClient: Client,
    toolName: string,
    args: Record<string, unknown>,
): Promise<string> {
    try {
        const result = await mcpClient.callTool({
            name: toolName,
            arguments: args,
        });
        const content = result.content as
            | Array<{ type: string; text?: string }>
            | undefined;
        if (!content || content.length === 0)
            return JSON.stringify({ success: true });
        const text = content
            .map((c) => (c.type === "text" ? (c.text ?? "") : ""))
            .filter(Boolean)
            .join("\n");
        return text || JSON.stringify(result);
    } catch (err) {
        const msg = err instanceof Error ? err.message : String(err);
        log.debug(`Tool ${toolName} error: ${msg}`);
        return JSON.stringify({ error: msg });
    }
}

/**
 * When context approaches the limit, truncate the content of old tool-result
 * messages to free space. Keeps the most recent N tool results intact so the
 * model still has fresh context; replaces earlier ones with a size notice.
 * Never touches system/user/assistant messages — only tool messages.
 */
function pruneToolMessages(messages: Message[], keepRecent = 6): Message[] {
    const toolIndices: number[] = [];
    for (let i = 0; i < messages.length; i++) {
        if (messages[i].role === "tool") toolIndices.push(i);
    }
    const pruneCount = Math.max(0, toolIndices.length - keepRecent);
    if (pruneCount === 0) return messages;

    const toPrune = new Set(toolIndices.slice(0, pruneCount));
    let pruned = 0;
    const result = messages.map((msg, i) => {
        if (!toPrune.has(i)) return msg;
        const originalLen =
            typeof msg.content === "string" ? msg.content.length : 0;
        pruned++;
        return {
            ...msg,
            content: `[pruned: was ${originalLen} chars — context limit approached]`,
        };
    });
    log.info(`» pruned ${pruned} old tool message(s) to reduce context`);
    return result;
}

async function unloadModel(ollama: Ollama, model: string): Promise<void> {
    try {
        await ollama.generate({ model, keep_alive: 0, prompt: "" });
        log.info(`» unloaded model ${model} from Ollama`);
    } catch (err) {
        const msg = err instanceof Error ? err.message : String(err);
        log.warning(`» failed to unload model ${model}: ${msg}`);
    }
}

async function runOllamaLoop(ctx: AgentRunContext): Promise<AgentResult> {
    const ollamaHost = process.env.OLLAMA_HOST ?? "";
    if (!ollamaHost) {
        const errorMsg =
            "OLLAMA_HOST environment variable is not set. Please set it to the URL of your Ollama instance.";
        log.error(errorMsg);
        return { success: false, error: errorMsg };
    }

    const model = ctx.model ?? process.env.OLLAMA_MODEL ?? DEFAULT_MODEL;
    const numCtx = ctx.payload.contextWindow ?? 262144;

    log.info(`» connecting to Ollama at ${ollamaHost}, model ${model}`);

    const ollama = new Ollama({ host: ollamaHost });
    const mcpClient = await buildMcpClient(ctx.mcpServerUrl);

    log.info("» fetching MCP tool list...");
    const tools = await getOllamaTools(mcpClient);
    log.info(`» ${tools.length} tools available`);

    let messages: Message[] = [
        {
            role: "user",
            content: ctx.instructions.full,
        },
    ];

    // Tools that signal the agent has produced its final output.
    const OUTPUT_TOOLS = new Set([
        "create_pull_request_review",
        "report_progress",
        "set_output",
    ]);

    let iterations = 0;
    let pendingModeNudge = false;
    let calledOutputTool = false;
    let addedContinueNudge = false;

    while (iterations < MAX_ITERATIONS) {
        iterations++;
        log.info(`» Ollama turn ${iterations}/${MAX_ITERATIONS}...`);

        // Non-streaming with a heartbeat timer so the activity monitor stays alive
        // during long prefill. Streaming was tried but Ollama only emits one tool
        // call per chunk — batched tool calls collapse to one-per-turn, turning a
        // 7-turn run into 26 turns. The heartbeat fires every 60s to prevent the
        // 300s activity timeout from triggering during large-context prefill.
        let response: Awaited<ReturnType<typeof ollama.chat>>;
        const turnStart = Date.now();
        const heartbeat = setInterval(() => {
            log.info(`» still waiting for model... (${Math.round((Date.now() - turnStart) / 1000)}s)`);
        }, 60_000);
        try {
            response = await retry(
                () => ollama.chat({
                    model,
                    messages,
                    tools,
                    keep_alive: -1,
                    think: false,
                    options: { num_ctx: numCtx, temperature: 0.1 },
                }),
                {
                    delaysMs: [3_000, 8_000, 15_000, 30_000, 600_000], // up to 6 attempts over ~16 minutes
                    shouldRetry: (err) => {
                        const msg = err instanceof Error ? err.message : String(err);
                        return /unexpected EOF|XML syntax error|ECONNRESET|ETIMEDOUT|fetch failed/i.test(msg);
                    },
                    label: `Ollama turn ${iterations}`,
                },
            );
        } catch (err) {
            clearInterval(heartbeat);
            await unloadModel(ollama, model);
            const lastError = err instanceof Error ? err.message : String(err);
            log.error(`Ollama error: ${lastError}`);
            return { success: false, error: `Ollama request failed: ${lastError}` };
        }
        clearInterval(heartbeat);

        const promptTokens = response.prompt_eval_count;
        const evalTokens = response.eval_count;
        const assistantMessage = response.message;

        if (promptTokens !== undefined) {
            const total = promptTokens + (evalTokens ?? 0);
            const pct = Math.round((total / numCtx) * 100);
            log.info(
                `  context: ${promptTokens} prompt + ${evalTokens ?? 0} eval = ${total} tokens (${pct}% of ${numCtx} limit)`,
            );
            if (promptTokens > numCtx * 0.77) {
                messages = pruneToolMessages(messages);
            }
        }

        messages.push(assistantMessage);

        const toolCalls: ToolCall[] | undefined = assistantMessage.tool_calls;

        if (!toolCalls || toolCalls.length === 0) {
            log.debug(`  model text: ${assistantMessage.content?.slice(0, 500)}`);

            // If the model stopped before ever calling an output tool and we haven't
            // nudged yet, give it one more push to continue the workflow — regardless
            // of whether the mode nudge is still pending (the model may have stopped
            // right after select_mode before acting on the guidance).
            if (!calledOutputTool && !addedContinueNudge) {
                log.info(
                    "» model stopped before completing task — nudging to continue",
                );
                addedContinueNudge = true;
                messages.push({
                    role: "user",
                    content:
                        "Your task is not complete yet. Continue executing the workflow — " +
                        "call the next required tool to finish. " +
                        "Do not stop until you have submitted a review (create_pull_request_review) " +
                        "or called report_progress with a final summary.",
                });
                continue;
            }

            await unloadModel(ollama, model);

            if (pendingModeNudge) {
                log.info("» agent finished after mode nudge (no tool calls)");
            } else {
                log.info("» agent finished (no tool calls)");
            }

            return {
                success: true,
                output: assistantMessage.content || undefined,
            };
        }

        pendingModeNudge = false;

        const calledSelectMode = toolCalls.some(
            (tc) => tc.function.name === "select_mode",
        );

        for (const toolCall of toolCalls) {
            const toolName = toolCall.function.name;
            const toolArgs = toolCall.function.arguments;

            log.info(`» calling tool: ${toolName}`);
            log.debug(`  args: ${JSON.stringify(toolArgs)}`);

            if (OUTPUT_TOOLS.has(toolName)) {
                calledOutputTool = true;
            }

            if (ctx.onToolUse) {
                ctx.onToolUse({ toolName, input: toolArgs });
            }

            const result = await callMcpTool(
                mcpClient,
                toolName,
                toolArgs as Record<string, unknown>,
            );
            log.debug(`  result: ${result.slice(0, 200)}`);

            messages.push({
                role: "tool",
                content: result,
            });
        }

        // After the FIRST select_mode call, nudge the model to act on the guidance.
        // Only nudge once — repeated nudging causes a loop where the model keeps
        // re-calling select_mode instead of executing the workflow.
        if (calledSelectMode && !pendingModeNudge) {
            pendingModeNudge = true;

            // Parse the selected mode name from the tool result so we can give a
            // more specific first-step instruction.
            let selectedMode = "";
            try {
                const lastToolMsg = messages[messages.length - 1];
                const parsed = JSON.parse(
                    typeof lastToolMsg.content === "string" ? lastToolMsg.content : "",
                );
                if (typeof parsed?.modeName === "string")
                    selectedMode = parsed.modeName;
            } catch {
                // best-effort
            }

            const firstStep =
                selectedMode === "Review" || selectedMode === "IncrementalReview"
                    ? "Your first tool call must be checkout_pr with the PR number from the event context."
                    : "Call the first tool required by the workflow now.";

            messages.push({
                role: "user",
                content:
                    `Good. You have selected ${selectedMode || "a"} mode and received the workflow. ` +
                    "Do NOT call select_mode again. " +
                    `${firstStep} ` +
                    "Execute the complete workflow step by step until you call create_pull_request_review or report_progress.",
            });
        }
    }

    await unloadModel(ollama, model);

    log.warning(`» agent hit max iterations (${MAX_ITERATIONS})`);

    return {
        success: false,
        error: `Agent exceeded maximum iterations (${MAX_ITERATIONS})`,
    };
}

export const ollamaAgent = agent({
    name: "ollama",
    run: async (ctx: AgentRunContext): Promise<AgentResult> => {
        return runOllamaLoop(ctx);
    },
});