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

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;

  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`);

  const 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}...`);

    let response: Awaited<ReturnType<typeof ollama.chat>>;
    try {
      response = await ollama.chat({
        model,
        messages,
        tools,
        options: {
          think: false,
        } as Record<string, unknown>,
      });
    } catch (err) {
      const lastError = err instanceof Error ? err.message : String(err);
      log.error(`Ollama error: ${lastError}`);
      return { success: false, error: `Ollama request failed: ${lastError}` };
    }

    const assistantMessage = response.message;
    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;
      }

      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;
      messages.push({
        role: "user",
        content:
          "Good. You have selected a mode and received the workflow. " +
          "Do NOT call select_mode again. " +
          "Call the next tool to execute the workflow now (e.g. checkout_pr, get_issue, create_pull_request_review). " +
          "Start immediately.",
      });
    }
  }

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