shockbot/review.ts

import { Ollama } from "ollama";
import type { Message } from "ollama";
import type {
  DiffChunk,
  ChunkReview,
  ReviewResult,
  AggregatedReview,
  Finding,
} from "./types.ts";
import type { ToolServer } from "./mcp.ts";

const ollama = new Ollama({ host: process.env.OLLAMA_HOST });

const SYSTEM_PROMPT = `You are a code reviewer. Your job is to find real bugs, security issues, and significant problems in the changed lines below. Do not comment on style, formatting, or personal preference unless they cause functional problems.

Be concise. Each comment should say what the problem is and why it matters.

You have tools available to read files, list directories, and search for symbols in the repository. Use them when you need more context to understand the code being reviewed.

When you are done gathering context and ready to report your findings, you MUST respond with ONLY a valid JSON object in this exact shape:
{
  "issues": [
    { "line": <line number>, "severity": "<error|warning|nit>", "comment": "<your comment>" }
  ],
  "summary": "<one paragraph overall summary>"
}

Do not include any text outside the JSON object.`;

export function buildReviewPrompt(
  chunk: DiffChunk,
  userPrompt: string,
): string {
  return `${userPrompt}

File: ${chunk.filename} (${chunk.language})
Hunk: ${chunk.context}

${chunk.hunk}`;
}

function parseReviewResult(content: string): ReviewResult | null {
  // Strip markdown code fences the model may have added despite instructions
  const stripped = content
    .replace(/^```(?:json)?\s*/m, "")
    .replace(/\s*```\s*$/m, "")
    .trim();
  try {
    const parsed = JSON.parse(stripped);
    if (!Array.isArray(parsed.issues) || typeof parsed.summary !== "string")
      return null;
    return parsed as ReviewResult;
  } catch {
    return null;
  }
}

export async function reviewChunk(
  chunk: DiffChunk,
  userPrompt: string,
  model: string,
  contextWindow: number,
  mcpServer: ToolServer,
): Promise<ChunkReview> {
  console.log(`Reviewing chunk from ${chunk.filename} with model ${model}...`);
  mcpServer.resetCallCount();

  const messages: Message[] = [
    { role: "system", content: SYSTEM_PROMPT },
    { role: "user", content: buildReviewPrompt(chunk, userPrompt) },
  ];

  let forceConclusion = false;
  let parseFailures = 0;

  while (true) {
    const response = await ollama.chat({
      model,
      messages,
      tools: forceConclusion ? undefined : mcpServer.tools,
      format: "json",
      stream: false,
      keep_alive: -1,
      options: { num_ctx: contextWindow, temperature: 0.2 },
    });

    const msg = response.message;

    // Model called one or more tools
    if (msg.tool_calls && msg.tool_calls.length > 0) {
      messages.push({
        role: "assistant",
        content: msg.content ?? "",
        tool_calls: msg.tool_calls,
      });

      for (const call of msg.tool_calls) {
        const result = await mcpServer.execute(
          call.function.name,
          call.function.arguments as Record<string, unknown>,
        );
        messages.push({
          role: "tool",
          content: result,
          tool_name: call.function.name,
        });
      }

      if (mcpServer.atLimit) {
        forceConclusion = true;
        messages.push({
          role: "user",
          content:
            "You have reached the tool call limit. Respond now with your final JSON review.",
        });
      }
      continue;
    }

    // Model produced content — try to parse as final JSON review
    const content = (msg.content ?? "").trim();

    if (content) {
      const result = parseReviewResult(content);
      if (result) return { chunk, result };

      // Parse failed — give the model one chance to fix it
      parseFailures++;
      if (parseFailures <= 1) {
        messages.push(
          { role: "assistant", content },
          {
            role: "user",
            content:
              "Your response was not valid JSON. Respond with ONLY the JSON review object, no other text.",
          },
        );
        forceConclusion = true;
        continue;
      }

      // Second failure — return raw content as summary rather than crash
      return {
        chunk,
        result: {
          issues: [],
          summary: `[parse error] ${content.slice(0, 500)}`,
        },
      };
    }

    // Empty response — shouldn't happen, but bail rather than loop forever
    return {
      chunk,
      result: { issues: [], summary: "[empty response from model]" },
    };
  }
}

export function aggregateFindings(results: ChunkReview[]): AggregatedReview {
  console.log(`Aggregating findings from ${results.length} reviewed chunks...`);
  const findings: Finding[] = results.flatMap(({ chunk, result }) =>
    result.issues.map((issue) => ({ ...issue, filename: chunk.filename })),
  );
  const summaries = results.map((r) => r.result.summary).filter(Boolean);
  return { findings, summaries };
}