feat(agents): forward tools the screens' page adds to the agents port

src/buzz-mcp/app-client.ts

@@ -201,13 +201,13 @@ export async function readFile(
   return fromBase64(output.replace(/\s/g, ""));
 }
 
-function toBase64(data: Uint8Array) {
+export function toBase64(data: Uint8Array) {
   let binary = "";
   for (let i = 0; i < data.length; i += 0x8000)
     binary += String.fromCharCode(...data.subarray(i, i + 0x8000));
   return btoa(binary);
 }
-function fromBase64(text: string) {
+export function fromBase64(text: string) {
   const binary = atob(text);
   const data = new Uint8Array(binary.length);
   for (let i = 0; i < binary.length; i++) data[i] = binary.charCodeAt(i);

src/features/agents/port.test.ts

@@ -1,4 +1,6 @@
 import { afterEach, expect, it, vi } from "vitest";
+import type { BuzzClient } from "../../buzz-mcp/client";
+import type { AddedTool } from "../../buzz-mcp/session-tools";
 import type { Events, Frame } from "./event-port";
 import { eventPort } from "./event-port";
 import { keyedRead } from "./handles";
@@ -86,9 +88,10 @@ function connected(agents: ReturnType<typeof fakeAgents>) {
   return client;
 }
 
-it("serves screens in another webview, whichever page loads first", async () => {
+/** Tauri events between webviews, each delivered later and as JSON. */
+function jsonEvents(): Events {
   const handlers = new Set<(event: { payload: Frame }) => void>();
-  const events: Events = {
+  return {
     emitTo: async (_webview, _event, frame) => {
       const payload = JSON.parse(JSON.stringify(frame)) as Frame;
       setTimeout(() => {
@@ -100,6 +103,10 @@ it("serves screens in another webview, whichever page loads first", async () =>
       return () => handlers.delete(handler);
     },
   };
+}
+
+it("serves screens in another webview, whichever page loads first", async () => {
+  const events = jsonEvents();
   const agents = fakeAgents();
   const client = connectAgents(eventPort("main", "agents", events));
   await new Promise((resolve) => setTimeout(resolve, 10));
@@ -127,6 +134,96 @@ it("serves screens in another webview, whichever page loads first", async () =>
   );
 });
 
+it("carries a tool's bytes as JSON, and lists the tools again with a reloaded service", async () => {
+  const events = jsonEvents();
+  const agents = fakeAgents() as unknown as Agents;
+  const added = new Map<string, AddedTool>();
+  const addTool = (tool: AddedTool) => {
+    added.set(tool.name, tool);
+    return () => added.delete(tool.name);
+  };
+  const client = connectAgents(eventPort("main", "agents", events));
+  client.addTool(
+    {
+      name: "files_copy",
+      description: "Copy a file.",
+      inputSchema: { type: "object" },
+      call: async (agent) => {
+        const data = await agent.read("/tmp/a");
+        return (await agent.upload(data, "text/plain")).url;
+      },
+    },
+    { id: "example", revision: "1" },
+  );
+  const first = serveAgents(
+    agents,
+    eventPort("agents", "main", events),
+    addTool,
+  );
+  await vi.waitFor(() => expect(added.has("files_copy")).toBe(true));
+  const upload = vi.fn(async (data: Uint8Array) => ({
+    url: `https://media.example/${data.length}`,
+  }));
+  const agent = {
+    pubkey: "p",
+    key: "k",
+    read: async () => new Uint8Array([1, 2, 3]),
+    upload,
+  } as unknown as BuzzClient;
+  await expect(added.get("files_copy")?.call(agent, {}, {})).resolves.toBe(
+    "https://media.example/3",
+  );
+  expect(upload).toHaveBeenCalledWith(new Uint8Array([1, 2, 3]), "text/plain");
+  // The agents page reloads: its tools go with it, and come back listed.
+  first();
+  expect(added.size).toBe(0);
+  const second = serveAgents(
+    agents,
+    eventPort("agents", "main", events),
+    addTool,
+  );
+  stops.push(() => {
+    client.close();
+    second();
+  });
+  await vi.waitFor(() => expect(added.has("files_copy")).toBe(true));
+});
+
+it("ends a tool's call when the screens' page that answers it reloads", async () => {
+  const events = jsonEvents();
+  const added = new Map<string, AddedTool>();
+  const stop = serveAgents(
+    fakeAgents() as unknown as Agents,
+    eventPort("agents", "main", events),
+    (tool) => {
+      added.set(tool.name, tool);
+      return () => added.delete(tool.name);
+    },
+  );
+  const tool = {
+    name: "notes_wait",
+    description: "Wait.",
+    inputSchema: { type: "object" },
+    call: () => new Promise<string>(() => {}),
+  } as const;
+  const owner = { id: "example", revision: "1" };
+  const before = connectAgents(eventPort("main", "agents", events));
+  before.addTool(tool, owner);
+  await vi.waitFor(() => expect(added.has("notes_wait")).toBe(true));
+  const waiting = added
+    .get("notes_wait")
+    ?.call({ pubkey: "p", key: "k" } as unknown as BuzzClient, {}, {});
+  await new Promise((resolve) => setTimeout(resolve, 10));
+  before.close();
+  const after = connectAgents(eventPort("main", "agents", events));
+  after.addTool(tool, owner);
+  stops.push(() => {
+    after.close();
+    stop();
+  });
+  await expect(waiting).rejects.toThrow(/reloaded/);
+});
+
 it("reads as loading until the service's first push, then as the service", async () => {
   const client = connected(fakeAgents());
   expect(client.snapshot().status).toBe("loading");

src/features/agents/port.ts

@@ -12,6 +12,15 @@ import {
   type SessionPlace,
   type SessionState,
 } from "./service";
+import type { AddedTool } from "../../buzz-mcp/session-tools";
+import type { Reply } from "../../buzz-mcp/tools";
+import {
+  serveTools,
+  type ToolCall,
+  type ToolListing,
+  type ToolOwner,
+  toolsClient,
+} from "./tool-port";
 import {
   serveTypes,
   type TypeCall,
@@ -64,9 +73,13 @@ type State = Readonly<{
   types: TypesState;
   /** Each session state a screen asked for, by `placeKey`. */
   sessions: Readonly<Record<string, SessionState>>;
+  /** This serving of the service. A new one has not heard what the screens'
+   * page asked of the last. */
+  instance: string;
 }>;
 type ToClient =
   | Readonly<{ kind: "state"; state: State }>
+  | Readonly<{ kind: "tool"; call: ToolCall }>
   | Readonly<{ kind: "result"; id: number; value: unknown }>
   | Readonly<{ kind: "error"; id: number; message: string }>;
 type ToService =
@@ -76,7 +89,21 @@ type ToService =
   | Readonly<{ kind: "type"; id: number; call: TypeCall; args: unknown[] }>
   | Readonly<{ kind: "watch"; watch: TypeWatch }>
   | Readonly<{ kind: "session"; key: string; place: SessionPlace }>
-  | Readonly<{ kind: "abort"; signal: number }>;
+  | Readonly<{ kind: "abort"; signal: number }>
+  | Readonly<{ kind: "list"; tool: ToolListing }>
+  | Readonly<{ kind: "withdraw"; name: string }>
+  | Readonly<{
+      kind: "settle";
+      id: number;
+      outcome: { reply: Reply } | { error: string };
+    }>
+  | Readonly<{
+      kind: "client";
+      id: number;
+      client: number;
+      method: string;
+      args: unknown[];
+    }>;
 
 /** A `MessagePort`, or a link to a page in another webview. It delivers each
  * message once, in the order sent. */
@@ -89,7 +116,12 @@ export type AgentsPort = {
 /** `AgentsClient`, and `present` for what this page shows of each type (by
  * type key). */
 export type AgentsPortClient = AgentsClient &
-  Readonly<{ present: ReturnType<typeof typesClient>["present"] }>;
+  Readonly<{
+    present: ReturnType<typeof typesClient>["present"];
+    /** Adds `owner`'s `tool` for the service's agents, answered in this
+     * page, until the returned withdraw. */
+    addTool(tool: AddedTool, owner: ToolOwner): () => void;
+  }>;
 
 const placeKey = (key: string, { channelId, root }: SessionPlace) =>
   JSON.stringify([key, channelId, root ?? ""]);
@@ -98,6 +130,7 @@ function stateOf(
   agents: Agents,
   types: TypesState,
   places: ReadonlyMap<string, readonly [string, SessionPlace]>,
+  instance: string,
 ): State {
   const snapshot = agents.snapshot();
   const memberships: Record<string, ReturnType<Agents["memberships"]>> = {};
@@ -131,13 +164,26 @@ function stateOf(
         agents.sessionState(key, place),
       ]),
     ),
+    instance,
   };
 }
 
 /** Serves `agents` on `port` until the returned stop. It pushes state before
- * each answer, so a call's result never arrives before what it changed. */
-export function serveAgents(agents: Agents, port: AgentsPort) {
+ * each answer, so a call's result never arrives before what it changed. A
+ * tool the screens' page lists is added with `addTool` until it withdraws
+ * it; without `addTool`, its tools are not added. */
+export function serveAgents(
+  agents: Agents,
+  port: AgentsPort,
+  addTool?: (tool: AddedTool, owner: ToolOwner) => () => void,
+) {
   const types = serveTypes(agents);
+  const instance = crypto.randomUUID();
+  const tools = addTool
+    ? serveTools(addTool, (call) =>
+        port.postMessage({ kind: "tool", call } satisfies ToClient),
+      )
+    : undefined;
   const places = new Map<string, readonly [string, SessionPlace]>();
   // The types' own stores, which change without the service's notice.
   const stores = new Map<(listener: () => void) => () => void, () => void>();
@@ -155,7 +201,7 @@ export function serveAgents(agents: Agents, port: AgentsPort) {
       if (!stores.has(subscribe)) stores.set(subscribe, subscribe(changed));
     port.postMessage({
       kind: "state",
-      state: stateOf(agents, types.state(), places),
+      state: stateOf(agents, types.state(), places, instance),
     } satisfies ToClient);
   };
   // Several changes in one task are one push.
@@ -192,6 +238,26 @@ export function serveAgents(agents: Agents, port: AgentsPort) {
       types.abort(data.signal);
       return;
     }
+    if (data.kind === "list") {
+      tools?.list(data.tool);
+      return;
+    }
+    if (data.kind === "withdraw") {
+      tools?.withdraw(data.name);
+      return;
+    }
+    if (data.kind === "settle") {
+      tools?.settle(data.id, data.outcome);
+      return;
+    }
+    if (data.kind === "client") {
+      const { client, method, args } = data;
+      answer(data.id, () => {
+        if (!tools) throw new Error("That tool call has finished");
+        return tools.client(client, method, args);
+      });
+      return;
+    }
     if (data.kind === "type") {
       answer(data.id, () => types.call(data.call, data.args));
       return;
@@ -220,6 +286,7 @@ export function serveAgents(agents: Agents, port: AgentsPort) {
   flush();
   return () => {
     unsubscribe();
+    tools?.stop();
     for (const stop of stores.values()) stop();
     stores.clear();
     port.onmessage = null;
@@ -249,6 +316,7 @@ const STARTING: State = Object.freeze({
     scopes: Object.freeze({}),
   }),
   sessions: Object.freeze({}),
+  instance: "",
 });
 
 /** `next`, keeping each part of `prior` it equals, so screens that compare by
@@ -312,9 +380,18 @@ export function connectAgents(port: AgentsPort) {
     if (data.kind === "state") {
       const prior = state;
       state = shared(prior, data.state);
+      if (prior !== STARTING && state.instance !== prior.instance) {
+        // A new serving, as when its page reloaded: tell it again.
+        asked.clear();
+        tools.relist();
+      }
       if (state !== prior) for (const listener of listeners) listener();
       return;
     }
+    if (data.kind === "tool") {
+      void tools.answer(data.call);
+      return;
+    }
     const pending = waiting.get(data.id);
     if (!pending) return;
     waiting.delete(data.id);
@@ -388,13 +465,36 @@ export function connectAgents(port: AgentsPort) {
   const sends = Object.fromEntries(
     SENDS.map((method) => [method, send(method)]),
   );
+  const tools = toolsClient({
+    list: (tool) => {
+      if (!closed) post({ kind: "list", tool });
+    },
+    withdraw: (name) => {
+      if (!closed) post({ kind: "withdraw", name });
+    },
+    settle: (id, outcome) => {
+      if (!closed) post({ kind: "settle", id, outcome });
+    },
+    call: (client, method, args) =>
+      new Promise((resolve, reject) => {
+        const id = ++next;
+        waiting.set(id, { resolve, reject });
+        try {
+          post({ kind: "client", id, client, method, args });
+        } catch (error) {
+          waiting.delete(id);
+          reject(error);
+        }
+      }),
+  });
   const client = {
     ...calls,
     ...sends,
     snapshot: () => state.snapshot,
     subscribe,
     types: types.types,
     present: types.present,
+    addTool: tools.add,
     find: (pubkey: string) =>
       state.snapshot.agents.find((agent) => agent.pubkey === pubkey),
     classifier: () => state.snapshot.classifier,