feat: attach the offered tool definitions to every generation

src/index.ts

@@ -196,6 +196,37 @@ export function extractToolCalls(content: unknown): Array<{ id: string; name: st
     .map((p) => ({ id: p.id, name: p.name }));
 }
 
+export interface ToolDefinitionInput {
+  name: string;
+  description?: string;
+  parameters?: unknown;
+}
+
+export function activeToolDefinitions(
+  pi: Pick<ExtensionAPI, "getAllTools" | "getActiveTools">,
+): ToolDefinitionInput[] {
+  try {
+    const registry = new Map((pi.getAllTools?.() ?? []).map((tool) => [tool.name, tool]));
+    const out: ToolDefinitionInput[] = [];
+    const seen = new Set<string>();
+    for (const name of pi.getActiveTools?.() ?? []) {
+      const tool = registry.get(name);
+      if (!tool || seen.has(name)) continue;
+      seen.add(name);
+      out.push({ name: tool.name, description: tool.description, parameters: tool.parameters });
+    }
+    return out;
+  } catch {
+    return [];
+  }
+}
+
+export function attachToolDefinitions(input: unknown, tools: ToolDefinitionInput[]): unknown {
+  if (!tools.length || !Array.isArray(input)) return input;
+  const [first, ...rest] = input;
+  return first && typeof first === "object" ? [{ ...first, tools }, ...rest] : input;
+}
+
 export interface PiImagePart {
   type: "image";
   data: string;
@@ -768,24 +799,23 @@ export default function (pi: ExtensionAPI) {
     }
     const index = ++state.generationCount;
     const history = lastContextHistory;
-    const baseInput: unknown =
+    const baseInput: unknown[] =
       history ??
       (index === 1
-        ? { role: "user", content: lastPromptText }
+        ? [{ role: "user", content: lastPromptText }]
         : state.pendingToolResults.length
-          ? { role: "tool", tool_results: state.pendingToolResults }
-          : undefined);
+          ? [{ role: "tool", tool_results: state.pendingToolResults }]
+          : []);
     const generationInput = state.systemPrompt
-      ? [
-          { role: "system", content: state.systemPrompt },
-          ...(Array.isArray(baseInput) ? baseInput : baseInput ? [baseInput] : []),
-        ]
-      : baseInput;
+      ? [{ role: "system", content: state.systemPrompt }, ...baseInput]
+      : baseInput.length
+        ? baseInput
+        : undefined;
 
     const obs = state.root.startObservation(
       GENERATION_PREFIX,
       {
-        input: generationInput,
+        input: attachToolDefinitions(generationInput, activeToolDefinitions(pi)),
         model: ctx.model?.id,
         metadata: {
           assistant_index: index - 1,

test/tool-definitions.test.ts

@@ -0,0 +1,127 @@
+import assert from "node:assert/strict";
+import { after, before, describe, it } from "node:test";
+import { activeToolDefinitions, attachToolDefinitions, type ToolDefinitionInput } from "../src/index.ts";
+import {
+  type Capture,
+  type CapturedSpan,
+  createSandbox,
+  runPi,
+  startCaptureServer,
+  startMockProvider,
+  waitForRequests,
+} from "./helpers.ts";
+
+const TOOLS: ToolDefinitionInput[] = [{ name: "read", description: "Read a file", parameters: { type: "object" } }];
+
+describe("attachToolDefinitions", () => {
+  it("adds the tools to the first message without mutating the caller's data", () => {
+    const input = [
+      { role: "system", content: "be brief" },
+      { role: "user", content: "hi" },
+    ];
+    const out = attachToolDefinitions(input, TOOLS) as Array<Record<string, unknown>>;
+    assert.deepEqual(out, [{ role: "system", content: "be brief", tools: TOOLS }, { role: "user", content: "hi" }]);
+    assert.equal("tools" in input[0]!, false);
+  });
+
+  it("passes the input through when there is nothing to attach or nothing to attach to", () => {
+    const messages = [{ role: "user", content: "hi" }];
+    assert.equal(attachToolDefinitions(messages, []), messages);
+    assert.equal(attachToolDefinitions(undefined, TOOLS), undefined);
+    const bare = { role: "user", content: "hi" };
+    assert.equal(attachToolDefinitions(bare, TOOLS), bare);
+  });
+});
+
+describe("activeToolDefinitions", () => {
+  const registry = [
+    { name: "read", description: "Read a file", parameters: { type: "object" }, promptGuidelines: ["x"], sourceInfo: {} },
+    { name: "bash", description: "Run a command", parameters: { type: "object" }, sourceInfo: {} },
+    { name: "grep", description: "Search files", parameters: { type: "object" }, sourceInfo: {} },
+  ];
+
+  it("keeps only the active tools, in the active order, with only the ChatML fields", () => {
+    const pi = { getAllTools: () => registry as never, getActiveTools: () => ["bash", "read"] };
+    assert.deepEqual(activeToolDefinitions(pi), [
+      { name: "bash", description: "Run a command", parameters: { type: "object" } },
+      { name: "read", description: "Read a file", parameters: { type: "object" } },
+    ]);
+  });
+
+  it("degrades to nothing on a pi without the accessors or with throwing ones", () => {
+    assert.deepEqual(activeToolDefinitions({} as never), []);
+    const throwing = {
+      getAllTools: () => {
+        throw new Error("not ready");
+      },
+      getActiveTools: () => ["read"],
+    };
+    assert.deepEqual(activeToolDefinitions(throwing as never), []);
+  });
+});
+
+function langfuseEnv(capture: Capture): Record<string, string> {
+  return {
+    LANGFUSE_PUBLIC_KEY: "pk-lf-test",
+    LANGFUSE_SECRET_KEY: "sk-lf-test",
+    LANGFUSE_BASE_URL: `http://127.0.0.1:${capture.port}`,
+  };
+}
+
+function generationInputs(spans: CapturedSpan[]): unknown[] {
+  return spans
+    .filter((s) => s.name === "LLM Call")
+    .sort((a, b) => (a.startNs < b.startNs ? -1 : 1))
+    .map((s) => JSON.parse(String(s.attrs["langfuse.observation.input"])));
+}
+
+describe("integration: tool definitions", () => {
+  let mock: { port: number; close: () => void };
+
+  before(async () => {
+    mock = await startMockProvider();
+  });
+  after(() => mock.close());
+
+  it("attaches the active tools to every generation, as a ChatML message array", async () => {
+    const capture = await startCaptureServer();
+    try {
+      const sandbox = createSandbox(mock.port);
+      const result = await runPi(sandbox, "Explore this project and summarize it", { env: langfuseEnv(capture) });
+      assert.equal(result.status, 0, `pi failed: ${result.stderr}`);
+      await waitForRequests(capture, 1);
+
+      const inputs = generationInputs(capture.spans());
+      assert.equal(inputs.length, 3, "the mock script runs three generations");
+      for (const [i, input] of inputs.entries()) {
+        assert.ok(Array.isArray(input), `generation ${i} input must be an array of messages, got ${typeof input}`);
+        const [head, ...rest] = input as Array<Record<string, unknown>>;
+        assert.equal(head?.role, "system", `generation ${i} must start with the system message`);
+        assert.ok(rest.length >= 1, `generation ${i} carries at least one base message`);
+        for (const [j, message] of rest.entries()) {
+          assert.ok(
+            message && typeof message === "object" && !Array.isArray(message),
+            `generation ${i} message ${j} must be a flat message object, not a nested array`,
+          );
+          assert.equal(typeof message.role, "string", `generation ${i} message ${j} must carry a role`);
+        }
+        // Since #33 the base is the whole conversation history, so it grows per
+        // generation; the newest message is still the delta this call adds.
+        assert.equal(rest[0]!.role, "user", `generation ${i} history opens on the user turn`);
+        assert.equal(rest.at(-1)!.role, i === 0 ? "user" : "tool", `generation ${i} ends on its delta`);
+        const tools = head!.tools as ToolDefinitionInput[];
+        assert.ok(Array.isArray(tools), `generation ${i} must carry tools on its first message`);
+        assert.deepEqual(
+          tools.map((t) => t.name),
+          ["read", "bash", "edit", "write"],
+        );
+        for (const tool of tools) {
+          assert.equal(typeof tool.description, "string");
+          assert.equal((tool.parameters as { type?: unknown })?.type, "object", `${tool.name} keeps its JSON schema`);
+        }
+      }
+    } finally {
+      capture.close();
+    }
+  });
+});