feat(llm-builtin): improve cell link serialization with @link format and any-schema handling (commontoolsinc#2092)

feat(llm-builtin): improve cell link serialization with @link format and any-schema handling

- Use `@link` format instead of `/` for cell link serialization to LLM tools
- Add `@arrayEntry` decorator for array elements that reference cells
- Output schema with the `read` tool
- Handle `any` schema types by converting to cell links when appropriate
- Add base64 validation to prevent false positive link detection
- Add timeout handling for tool calls (30s timeout)
- Refactor chatbot.tsx to use patternTool helper for listMentionable
- Fix schema.ts to properly copy values before annotation to avoid mutations

Author: seefeldb

packages/patterns/chatbot.tsx

@@ -9,6 +9,7 @@ import {
   handler,
   llmDialog,
   NAME,
+  patternTool,
   recipe,
   Stream,
   UI,
@@ -126,19 +127,16 @@ export const TitleGenerator = recipe<
   return title;
 });
 
-const listMentionable = handler<
-  {
-    /** A cell to store the result text */
-    result: Cell<string>;
-  },
-  { mentionable: Cell<MentionableCharm>[] }
+const listMentionable = recipe<
+  { mentionable: Array<MentionableCharm> },
+  { result: Array<{ label: string; cell: Cell<unknown> }> }
 >(
-  (args, state) => {
-    const namesList = state.mentionable.map((charm) => ({
-      label: charm.get()[NAME],
+  ({ mentionable }) => {
+    const result = mentionable.map((charm) => ({
+      label: charm[NAME]!,
       cell: charm,
     }));
-    args.result.set(JSON.stringify(namesList));
+    return { result };
   },
 );
 
@@ -166,11 +164,7 @@ export default recipe<ChatInput, ChatOutput>(
     const recentCharms = schemaifyWish<MentionableCharm[]>("#recent");
 
     const assistantTools = {
-      listMentionable: {
-        description:
-          "List all mentionable items in the space, read() the result.",
-        handler: listMentionable({ mentionable }),
-      },
+      listMentionable: patternTool(listMentionable, { mentionable }),
       listRecent: {
         description:
           "List all recently viewed charms in the space, read() the result.",

packages/runner/src/builtins/llm-dialog.ts

@@ -29,6 +29,7 @@ import {
   getCellOrThrow,
   isCellResultForDereferencing,
 } from "../query-result-proxy.ts";
+import { ContextualFlowControl } from "../cfc.ts";
 
 // Avoid importing from @commontools/charm to prevent circular deps in tests
 
@@ -39,6 +40,7 @@ const logger = getLogger("llm-dialog", {
 
 const client = new LLMClient();
 const REQUEST_TIMEOUT = 1000 * 60 * 5; // 5 minutes
+const TOOL_CALL_TIMEOUT = 1000 * 30 * 1; // 30 seconds
 
 /**
  * Remove the injected `result` field from a JSON schema so tools don't
@@ -201,20 +203,42 @@ function createLLMFriendlyLink(link: NormalizedFullLink): string {
  */
 function traverseAndSerialize(
   value: unknown,
+  schema: JSONSchema | undefined,
+  rootSchema: JSONSchema | undefined = schema,
   seen: Set<unknown> = new Set(),
 ): unknown {
   if (!isRecord(value)) return value;
 
+  // If we encounter an `any` schema, turn value into a cell link
+  if (
+    seen.size > 0 && schema !== undefined &&
+    ContextualFlowControl.isTrueSchema(schema) &&
+    isCellResultForDereferencing(value)
+  ) {
+    // Next step will turn this into a link
+    value = getCellOrThrow(value);
+  }
+
+  // Turn cells into a link, unless they are data: URIs and traverse instead
   if (isCell(value)) {
-    const link = value.getAsNormalizedFullLink();
-    return { "/": encodeJsonPointer(["", link.id, ...link.path]) };
+    const link = value.resolveAsCell().getAsNormalizedFullLink();
+    if (link.id.startsWith("data:")) {
+      return traverseAndSerialize(value.get(), schema, rootSchema, seen);
+    } else {
+      return { "@link": encodeJsonPointer(["", link.id, ...link.path]) };
+    }
   }
 
   // If we've already seen this and it can be mapped to a cell, serialized as
   // cell link, otherwise throw (this should never happen in our cases)
   if (seen.has(value)) {
     if (isCellResultForDereferencing(value)) {
-      return traverseAndSerialize(getCellOrThrow(value), seen);
+      return traverseAndSerialize(
+        getCellOrThrow(value),
+        schema,
+        rootSchema,
+        seen,
+      );
     } else {
       throw new Error(
         "Cannot serialize a value that has already been seen and cannot be mapped to a cell.",
@@ -223,13 +247,43 @@ function traverseAndSerialize(
   }
   seen.add(value);
 
+  const cfc = new ContextualFlowControl();
+
   if (Array.isArray(value)) {
-    return value.map((v) => traverseAndSerialize(v, seen));
+    return value.map((v, index) => {
+      const linkSchema = schema !== undefined
+        ? cfc.schemaAtPath(schema, [index.toString()], rootSchema)
+        : undefined;
+      let result = traverseAndSerialize(v, linkSchema, rootSchema, seen);
+      // Decorate array entries with links that point to underlying cells, if
+      // any. Ignores data: URIs, since they're not useful as links for the LLM.
+      if (isRecord(result) && isCellResultForDereferencing(v)) {
+        const link = getCellOrThrow(v).resolveAsCell()
+          .getAsNormalizedFullLink();
+        if (!link.id.startsWith("data:")) {
+          result = {
+            "@arrayEntry": encodeJsonPointer(["", link.id, ...link.path]),
+            ...result,
+          };
+        }
+      }
+      return result;
+    });
   } else {
     return Object.fromEntries(
-      Object.entries(value).map((
+      Object.entries(value as Record<string, unknown>).map((
         [key, value],
-      ) => [key, traverseAndSerialize(value, seen)]),
+      ) => [
+        key,
+        traverseAndSerialize(
+          value,
+          schema !== undefined
+            ? cfc.schemaAtPath(schema, [key], rootSchema)
+            : undefined,
+          rootSchema,
+          seen,
+        ),
+      ]),
     );
   }
 }
@@ -254,10 +308,10 @@ function traverseAndCellify(
   // - it's a record with a single key "/"
   // - the value of the "/" key is a string that matches the URI pattern
   if (
-    isRecord(value) && typeof value["/"] === "string" &&
-    Object.keys(value).length === 1 && matchLLMFriendlyLink.test(value["/"])
+    isRecord(value) && typeof value["@link"] === "string" &&
+    Object.keys(value).length === 1 && matchLLMFriendlyLink.test(value["@link"])
   ) {
-    const link = parseLLMFriendlyLink(value["/"], space);
+    const link = parseLLMFriendlyLink(value["@link"], space);
     return runtime.getCellFromLink(link);
   }
   if (Array.isArray(value)) {
@@ -1363,15 +1417,17 @@ function handleSchema(
  */
 function handleRead(
   resolved: ResolvedToolCall & { type: "read" },
-): { type: string; value: any } {
-  const serialized = traverseAndSerialize(resolved.cellRef.get());
+): { type: string; value: unknown } {
+  let cell = resolved.cellRef;
+  if (!cell.schema) {
+    cell = cell.asSchema(getCellSchema(cell));
+  }
 
-  // Handle undefined values gracefully - return null for undefined/null
-  const value = serialized === undefined || serialized === null
-    ? null
-    : JSON.parse(JSON.stringify(serialized));
+  const schema = cell.schema;
+  const serialized = traverseAndSerialize(cell.get(), schema);
 
-  return { type: "json", value };
+  // Handle undefined by returning null (valid JSON) instead
+  return { type: "json", value: serialized ?? null, ...(schema && { schema }) };
 }
 
 /**
@@ -1476,8 +1532,22 @@ async function handleRun(
   const cancel = result.sink((r) => {
     r !== undefined && resolve(r);
   });
-  await promise;
-  cancel();
+
+  let timeout;
+  const timeoutPromise = new Promise((_, reject) => {
+    timeout = setTimeout(() => {
+      reject(new Error("Tool call timed out"));
+    }, TOOL_CALL_TIMEOUT);
+  }).then(() => {
+    throw new Error("Tool call timed out");
+  });
+
+  try {
+    await Promise.race([promise, timeoutPromise]);
+  } finally {
+    clearTimeout(timeout);
+    cancel();
+  }
 
   // Get the actual entity ID from the result cell
   const resultLink = createLLMFriendlyLink(result.getAsNormalizedFullLink());

packages/runner/src/schema.ts

@@ -835,8 +835,21 @@ export function validateAndTransform(
   }
 
   // Add the current value to seen before returning
-  seen.push([seenKey, value]);
-  return annotateWithBackToCellSymbols(value, runtime, link, tx);
+  if (isRecord(value)) {
+    const cloned = isObject(value)
+      ? { ...(value as Record<string, unknown>) }
+      : [...(value as unknown[])];
+    seen.push([seenKey, cloned]);
+    return annotateWithBackToCellSymbols(
+      cloned,
+      runtime,
+      link,
+      tx,
+    );
+  } else {
+    seen.push([seenKey, value]);
+    return value;
+  }
 }
 
 /**