ci/src/backends/ollama-base.test.ts

import * as fs from "node:fs";
import * as path from "node:path";
import * as os from "node:os";
import { OllamaBaseBackend, OllamaMessage, OllamaChatResponse } from "../backends/ollama-base.js";
import { ToolRegistry } from "../backends/tool-registry.js";
import { BackendRequest } from "../backends/types.js";

class TestableOllamaBaseBackend extends OllamaBaseBackend {
  readonly name = "test-base";
  private mockResponse: OllamaChatResponse;
  private callCount: number;

  constructor(mockResponse: OllamaChatResponse) {
    super(undefined);
    this.mockResponse = mockResponse;
    this.callCount = 0;
  }

  async isAvailable(): Promise<boolean> {
    return true;
  }

  getCallCount(): number {
    return this.callCount;
  }

  protected async callModel(
    messages: OllamaMessage[],
    model: string,
    toolRegistry: ToolRegistry
  ): Promise<OllamaChatResponse> {
    this.callCount++;
    return this.mockResponse;
  }

  protected resolveModel(): string {
    return "test-model";
  }
}

describe("OllamaBaseBackend", () => {
  let tempDir: string;

  beforeEach(() => {
    tempDir = fs.mkdtempSync(path.join(os.tmpdir(), "ciagent-ollama-base-test-"));
  });

  afterEach(() => {
    fs.rmSync(tempDir, { recursive: true, force: true });
  });

  it("returns success when model responds without tool calls", async () => {
    const mockResponse: OllamaChatResponse = {
      choices: [{
        message: {
          content: '{"success": true, "output": "task completed"}',
        },
      }],
      usage: { prompt_tokens: 10, completion_tokens: 20, total_tokens: 30 },
    };

    const backend = new TestableOllamaBaseBackend(mockResponse);
    const request: BackendRequest = {
      persona: "executor",
      workflow: "execute",
      task: "Do something",
      context: {
        project_path: tempDir,
        phase: 1,
        stage: "execute",
        specification: "",
        config_path: "",
      },
      autonomy: "full",
    };

    const result = await backend.execute(request);
    expect(result.success).toBe(true);
    expect(result.output).toContain("task completed");
  });

  it("handles tool calls in response", async () => {
    const writePath = path.join(tempDir, "output.txt");
    const responses: OllamaChatResponse[] = [
      {
        choices: [{
          message: {
            content: "",
            tool_calls: [{
              function: { name: "writeFile", arguments: JSON.stringify({ path: writePath, content: "hello" }) },
            }],
          },
        }],
        usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 },
      },
      {
        choices: [{
          message: {
            content: '{"success": true, "output": "file written"}',
          },
        }],
        usage: { prompt_tokens: 5, completion_tokens: 10, total_tokens: 15 },
      },
    ];

    let callIndex = 0;
    class ToolCallBackend extends OllamaBaseBackend {
      readonly name = "tool-call-test";
      constructor() {
        super(undefined);
      }
      async isAvailable(): Promise<boolean> { return true; }
      protected async callModel(): Promise<OllamaChatResponse> {
        return responses[callIndex++];
      }
      protected resolveModel(): string { return "test-model"; }
    }

    const backend = new ToolCallBackend();
    const request: BackendRequest = {
      persona: "executor",
      workflow: "execute",
      task: "Write a file",
      context: {
        project_path: tempDir,
        phase: 1,
        stage: "execute",
        specification: "",
        config_path: "",
      },
      autonomy: "full",
    };

    const result = await backend.execute(request);
    expect(result.success).toBe(true);
    expect(fs.existsSync(writePath)).toBe(true);
    expect(fs.readFileSync(writePath, "utf-8")).toBe("hello");
    expect(result.artifacts.length).toBe(1);
    expect(result.artifacts[0].path).toBe(writePath);
  });

  it("stops after max tool rounds", async () => {
    const alwaysToolCall: OllamaChatResponse = {
      choices: [{
        message: {
          content: "",
          tool_calls: [{
            function: { name: "readFile", arguments: JSON.stringify({ path: "/etc/hostname" }) },
          }],
        },
      }],
      usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 },
    };

    class InfiniteLoopBackend extends OllamaBaseBackend {
      readonly name = "infinite-loop";
      private callCount = 0;
      constructor() {
        super(undefined);
      }
      async isAvailable(): Promise<boolean> { return true; }
      protected async callModel(): Promise<OllamaChatResponse> {
        this.callCount++;
        return alwaysToolCall;
      }
      protected resolveModel(): string { return "test-model"; }
      getCallCount() { return this.callCount; }
    }

    const backend = new InfiniteLoopBackend();
    const request: BackendRequest = {
      persona: "executor",
      workflow: "execute",
      task: "Infinite loop test",
      context: {
        project_path: tempDir,
        phase: 1,
        stage: "execute",
        specification: "",
        config_path: "",
      },
      autonomy: "full",
    };

    const result = await backend.execute(request);
    expect(result.output).toContain("maximum rounds");
    expect(backend.getCallCount()).toBe(50);
  });

  it("handles error from callModel gracefully", async () => {
    class ErrorBackend extends OllamaBaseBackend {
      readonly name = "error-backend";
      constructor() {
        super(undefined);
      }
      async isAvailable(): Promise<boolean> { return true; }
      protected async callModel(): Promise<OllamaChatResponse> {
        throw new Error("Model connection failed");
      }
      protected resolveModel(): string { return "test-model"; }
    }

    const backend = new ErrorBackend();
    const request: BackendRequest = {
      persona: "executor",
      workflow: "execute",
      task: "Fail test",
      context: {
        project_path: tempDir,
        phase: 1,
        stage: "execute",
        specification: "",
        config_path: "",
      },
      autonomy: "full",
    };

    const result = await backend.execute(request);
    expect(result.success).toBe(false);
    expect(result.error).toContain("Backend execution failed");
  });

  it("modelProfileToModel selects smallest for speed", () => {
    const backend = new TestableOllamaBaseBackend({} as OllamaChatResponse);
    const models = ["llama3.1:70b", "llama3.1:8b", "llama3.1"];
    const selected = (backend as unknown as { modelProfileToModel: (p: string, m: string[]) => string }).modelProfileToModel("speed", models);
    expect(selected).toBe("llama3.1");
  });
});