diff --git a/src/core/brain_ollama.ts b/src/core/brain_ollama.ts index 1449065..ec4e56c 100644 --- a/src/core/brain_ollama.ts +++ b/src/core/brain_ollama.ts @@ -15,7 +15,8 @@ import type { Brain, TaskCommand } from "./brain.js"; import { EventQueue } from "./brain.js"; import type { BrainEvent } from "./brain_protocol.js"; -import { TOOLS } from "./brain_protocol.js"; +import { TOOL_DEFINITIONS } from "./tool_registry.js"; +import { TOOLS, type ToolName } from "./brain_protocol.js"; import type { ToolResult } from "./tool_executor.js"; import { ollamaChat, @@ -47,7 +48,7 @@ const DEFAULT_MAX_TURNS = 24; // The 8 canonical tools, advertised to the model as OpenAI function schemas. The // ONE implementation lives host-side in tool_executor.ts; this only describes // them so the model can request them. Names are pinned by TOOLS (protocol v3). -const TOOL_SCHEMAS: readonly ToolSchema[] = buildToolSchemas(); +const TOOL_SCHEMAS: readonly ToolSchema[] = ollamaToolSchemas(); const SYSTEM_PERSONA = "You are Aether Code, an autonomous coding agent running locally. You work in " + @@ -60,9 +61,11 @@ export class OllamaBrain implements Brain { private readonly queue = new EventQueue(); private readonly opts: OllamaBrainOptions; private readonly chat: OllamaChatFn; - // Pending tool result: the loop awaits this promise after emitting a - // tool_call; sendToolResult resolves it (mirrors LocalBrain's stdin reply). - private pending: ((r: ToolResult) => void) | null = null; + // Outstanding tool calls, keyed by the id the loop emitted. Keyed rather than + // anonymous: a single resolver accepted any result for whatever call happened + // to be waiting, so a duplicate vanished silently and a result carrying the + // wrong id was indistinguishable from the right one. + private readonly pending = new Map void>(); private aborted = false; constructor(opts: OllamaBrainOptions = {}) { @@ -77,34 +80,53 @@ export class OllamaBrain implements Brain { return this.queue.drain(); } - sendToolResult(_id: string, result: ToolResult): void { - if (this.pending) { - const resolve = this.pending; - this.pending = null; - resolve(result); + sendToolResult(id: string, result: ToolResult): void { + const resolve = this.pending.get(id); + if (!resolve) { + // Unknown id, or a second result for a call already settled. Either way + // the host and the brain disagree about what is in flight; say so rather + // than letting the loop advance on a result it never asked for. + this.queue.push({ type: "error", msg: `tool result for unknown or already-settled call ${id}` }); + return; } + this.pending.delete(id); + resolve(result); } - control(_action: "pause" | "resume" | "steer", _note?: string): void { - // Steering is not yet wired for the local-Ollama brain (single-pass loop). - // The seam exists so the host can call it uniformly; it is a safe no-op. + control(action: "pause" | "resume" | "steer", note?: string): void { + // This brain runs a single-pass loop with no interruption point, so it + // cannot honour pause, resume or steer. Silently returning made the host + // believe the instruction landed; a dropped steer then reads as the model + // ignoring the user. Report it instead of accepting it. + this.queue.push({ + type: "monologue", + text: + `[${action} is not supported by the local Ollama brain — the instruction was not applied` + + (note ? `: "${note}"` : "") + + "]", + depth: 0, + }); } close(): void { this.aborted = true; - // Unblock a loop parked on a tool result so it can observe the abort. - if (this.pending) { - const resolve = this.pending; - this.pending = null; + // Unblock every loop parked on a tool result so it can observe the abort. + // Drained as a set so no waiter can be stranded by an early return. + for (const [id, resolve] of this.pending) { + this.pending.delete(id); resolve({ output: "[aborted]", exitCode: 130 }); } this.queue.end(); } - /** Await the host's reply to the tool_call we just emitted. */ - private waitForTool(): Promise { + /** + * Register a waiter for a tool call BEFORE the event is emitted. Registering + * afterwards only worked because the consumer resumes on a microtask; a host + * that replied synchronously would have found no waiter and been rejected. + */ + private waitForTool(id: string): Promise { return new Promise((resolve) => { - this.pending = resolve; + this.pending.set(id, resolve); }); } @@ -158,8 +180,9 @@ export class OllamaBrain implements Brain { for (const call of calls) { if (this.aborted) break; const args = parseArgs(call.function.arguments); + const waiting = this.waitForTool(call.id); this.queue.push({ type: "tool_call", id: call.id, name: call.function.name, args }); - const toolResult = await this.waitForTool(); + const toolResult = await waiting; messages.push({ role: "tool", tool_call_id: call.id, @@ -223,24 +246,56 @@ function parseArgs(raw: string): Record { } } -/** Minimal OpenAI function schema per canonical tool (parameters left open). */ -function buildToolSchemas(): readonly ToolSchema[] { - const descriptions: Record = { - read_file: "Read a workspace file. args: {path}", - write_file: "Write/overwrite a workspace file. args: {path, content}", - run_shell: "Run a shell command in the workspace. args: {command}", - run_tests: "Run the project's test command. args: {command?}", - repo_search: "Grep the repository for a string. args: {query}", - git_commit: "Stage all changes and commit. args: {message}", - web_search: "Search the web. args: {query, limit?}", - web_fetch: "Fetch a web page as readable text. args: {url}", +/** + * Model-facing tool schemas, GENERATED from TOOL_DEFINITIONS. + * + * These were hand-written, with `parameters: { properties: {}, additionalProperties: true }` + * — so the model was told nothing about argument names, types, required-ness or + * bounds, and the only description of the real shapes lived in free text that + * nothing kept in step with the validator. Generating them means the host's + * validator and the model's contract cannot drift. + */ +export function ollamaToolSchemas(): readonly ToolSchema[] { + const summaries: Readonly> = { + read_file: "Read a workspace file.", + write_file: "Write or overwrite a workspace file.", + run_shell: "Run a shell command in the workspace.", + run_tests: "Run the project's test command.", + repo_search: "Grep the repository for a string.", + git_commit: "Stage this run's changes and commit them.", + web_search: "Search the web.", + web_fetch: "Fetch a web page as readable text.", }; - return TOOLS.map((name) => ({ - type: "function" as const, - function: { - name, - description: descriptions[name] ?? name, - parameters: { type: "object", properties: {}, additionalProperties: true }, - }, - })); + + return TOOLS.map((name) => { + const definition = TOOL_DEFINITIONS[name]; + const properties: Record> = {}; + const required: string[] = []; + + for (const [argument, spec] of Object.entries(definition.args)) { + properties[argument] = + spec.type === "integer" + ? { type: "integer", minimum: spec.min, maximum: spec.max } + // maxBytes is a byte budget; as maxLength it is a ceiling in + // characters, which is conservative for any multi-byte input. + : { type: "string", maxLength: spec.maxBytes }; + if (spec.required !== false) required.push(argument); + } + + return { + type: "function" as const, + function: { + name, + description: `${summaries[name]} side effect: ${definition.sideEffect}.`, + parameters: { + type: "object", + properties, + required, + // The host rejects unknown arguments, so advertising them as allowed + // only invites the model to send something that will be refused. + additionalProperties: false, + }, + }, + }; + }); } diff --git a/test/brain_ollama.test.ts b/test/brain_ollama.test.ts index 33dd0d0..4d8b92e 100644 --- a/test/brain_ollama.test.ts +++ b/test/brain_ollama.test.ts @@ -1,6 +1,8 @@ import { test } from "node:test"; import assert from "node:assert/strict"; -import { OllamaBrain } from "../src/core/brain_ollama.js"; +import { OllamaBrain, ollamaToolSchemas } from "../src/core/brain_ollama.js"; +import { TOOLS, type ToolName } from "../src/core/brain_protocol.js"; +import { TOOL_DEFINITIONS } from "../src/core/tool_registry.js"; import type { Brain, TaskCommand } from "../src/core/brain.js"; import type { BrainEvent } from "../src/core/brain_protocol.js"; import type { ChatMessage, ChatReply } from "../src/core/ollama.js"; @@ -128,3 +130,131 @@ test("OllamaBrain stops after maxTurns without a final answer", async () => { assert.ok(done, "loop terminates with a done even on a runaway model"); assert.equal(done?.type === "done" ? done.reason : "", "max-turns"); }); + +// ── tool-call correlation (SC-A4) ─────────────────────────────────────────── +// The pending tool result used to be a single anonymous resolver and +// sendToolResult ignored its `id` parameter entirely. Any result satisfied +// whatever call happened to be waiting, a duplicate vanished silently, and a +// result for an unknown id was indistinguishable from the real one. + +function toolThenAnswer(): readonly ChatReply[] { + return [ + { + role: "assistant", + content: "", + tool_calls: [{ id: "call-1", type: "function", function: { name: "read_file", arguments: '{"path":"a.ts"}' } }], + }, + { role: "assistant", content: "done" }, + ]; +} + +test("a tool result for an unknown id is rejected, not applied to the waiting call", async () => { + const { chat } = fakeChat(toolThenAnswer()); + const brain = new OllamaBrain({ chat }); + const events: BrainEvent[] = []; + for await (const ev of brain.run(task)) { + events.push(ev); + if (ev.type === "tool_call") { + brain.sendToolResult("not-the-right-id", { output: "WRONG", exitCode: 0 }); + brain.sendToolResult(ev.id, { output: "right", exitCode: 0 }); + } + } + brain.close(); + const errors = events.filter((ev) => ev.type === "error"); + assert.equal(errors.length >= 1, true, "an unrecognised tool-call id must be surfaced, not swallowed"); + assert.equal(events.some((ev) => ev.type === "done"), true, "the run still terminates"); +}); + +test("a duplicate tool result does not advance the loop twice", async () => { + const { chat, calls } = fakeChat(toolThenAnswer()); + const brain = new OllamaBrain({ chat }); + for await (const ev of brain.run(task)) { + if (ev.type === "tool_call") { + brain.sendToolResult(ev.id, { output: "first", exitCode: 0 }); + brain.sendToolResult(ev.id, { output: "second", exitCode: 0 }); + } + } + brain.close(); + const toolReplies = calls() + .flat() + .filter((message) => message.role === "tool"); + const outputs = toolReplies.map((message) => message.content); + assert.equal(outputs.includes("second"), false, "the second result must not reach the model"); + assert.equal(outputs.filter((text) => text === "first").length >= 1, true); +}); + +test("close resolves an outstanding waiter so the run cannot hang", async () => { + const { chat } = fakeChat(toolThenAnswer()); + const brain = new OllamaBrain({ chat }); + const events: BrainEvent[] = []; + for await (const ev of brain.run(task)) { + events.push(ev); + if (ev.type === "tool_call") brain.close(); // never send a result + } + assert.equal(events.some((ev) => ev.type === "tool_call"), true); +}); + +// ── control honesty ───────────────────────────────────────────────────────── + +test("control() reports that steering is unsupported instead of silently succeeding", async () => { + const { chat } = fakeChat(toolThenAnswer()); + const brain = new OllamaBrain({ chat }); + const events: BrainEvent[] = []; + for await (const ev of brain.run(task)) { + events.push(ev); + if (ev.type === "tool_call") { + brain.control("steer", "actually, do something else"); + brain.sendToolResult(ev.id, { output: "ok", exitCode: 0 }); + } + } + brain.close(); + const said = events + .filter((ev): ev is Extract => ev.type === "monologue") + .map((ev) => ev.text) + .join(" "); + assert.match(said, /steer/i, "a steer this brain cannot honour must be visible to the user"); + assert.match(said, /not supported|unsupported|cannot/i); +}); + +// ── advertised tool schemas ───────────────────────────────────────────────── + +test("advertised tool schemas are generated from TOOL_DEFINITIONS, not hand-written", () => { + const schemas = ollamaToolSchemas(); + assert.deepEqual( + schemas.map((schema) => schema.function.name).sort(), + [...TOOLS].sort(), + "every protocol tool is advertised, and nothing else", + ); + + for (const schema of schemas) { + const name = schema.function.name as ToolName; + const definition = TOOL_DEFINITIONS[name]; + const parameters = schema.function.parameters as { + properties: Record; + required?: string[]; + additionalProperties?: boolean; + }; + + assert.equal(parameters.additionalProperties, false, `${name} must not accept arbitrary extra arguments`); + assert.deepEqual( + Object.keys(parameters.properties).sort(), + Object.keys(definition.args).sort(), + `${name} advertises exactly the arguments the host validates`, + ); + assert.deepEqual( + (parameters.required ?? []).sort(), + Object.entries(definition.args) + .filter(([, argument]) => argument.required !== false) + .map(([key]) => key) + .sort(), + `${name} required set matches the host validator`, + ); + } + + // Spot-check that bounds actually crossed over rather than being dropped. + const search = schemas.find((schema) => schema.function.name === "web_search"); + const limit = (search!.function.parameters as { properties: Record }) + .properties["limit"]; + assert.equal(limit?.minimum, 1); + assert.equal(limit?.maximum, 10); +});