fix(codemode): let search be awaited or chained like a Promise (#59981)

Author: rekram1-node

packages/codemode/README.md

@@ -198,7 +198,8 @@ every visible tool. Hosts render their own model-facing instructions from these
 and `CodeMode.toolExpression(path)` supply the exact callable forms.
 
 The synchronous `search(...)` built-in is always available. It supports exact-path lookup, namespace-scoped search,
-empty-query browsing, and pagination, and returns callable paths with full signatures. Search counts toward
+empty-query browsing, and pagination, and returns callable paths with full signatures. Its result can be used directly,
+awaited, or chained with `then`, `catch`, and `finally`, so programs work whichever way they treat it. Search counts toward
 `maxToolCalls`. Search also matches descriptions from enclosing `Namespace` values.
 
 ## Execution Limits

packages/codemode/interpreter-support.md

@@ -31,7 +31,8 @@ the ultimate source of truth. Upstream test262 files run verbatim from `test/tes
 - [x] Live Date, RegExp, Map, Set, URL, URLSearchParams, Headers, and Uint8Array values inside CodeMode.
 - [x] Tool calls through the host-provided `tools` tree only.
 - [x] The global `search(...)` built-in: synchronous tool discovery that counts as an admitted tool call and is
-      shadowable by program declarations like other globals.
+      shadowable by program declarations like other globals. Its result can also be awaited or chained with
+      `then`, `catch`, and `finally`.
 - [x] Cooperative timeout, an optional total tool-call limit, output bounding, and unrestricted tool-call concurrency.
 - [x] The timeout fires between interpreter steps, so one built-in is bounded in what it may build: strings up to
       2^24 characters (`repeat`, `pad*`, `concat`, `join`, `+`, template literals, `JSON.stringify`), arrays up to

packages/codemode/src/interpreter/globals.ts

@@ -87,8 +87,7 @@ type Factory = <R>(ctx: Interpreter<R>) => Value
 // A table rather than a list so the names are known before any runtime exists.
 const table: Record<string, Factory> = {
   tools: () => new ToolReference([]),
-  search: (ctx) =>
-    native(ctx.builtins, { name: "search", call: (_, args) => ctx.tool(ctx.tools.search, args), callback: false }),
+  search: (ctx) => native(ctx.builtins, { name: "search", call: (_, args) => ctx.search(args), callback: false }),
   undefined: () => undefined,
   NaN: () => NaN,
   Infinity: () => Infinity,

packages/codemode/src/interpreter/interpreter.ts

@@ -118,7 +118,7 @@ import {
 } from "./runner.js"
 import { describeValue, isOpaque, rejectCircularInsertion, typeofValue } from "./references.js"
 import { fn } from "./native.js"
-import { capability } from "./promises.js"
+import { capability, promiseLike } from "./promises.js"
 import { ScopeStack } from "./scope.js"
 import { constructRegExp } from "../stdlib/regexp.js"
 import { enumerableSource } from "../stdlib/object.js"
@@ -433,6 +433,11 @@ export class Interpreter<R> {
     for (const arg of args) json.push(yield* toBoundary(this, arg))
     return fromJson(this, (yield this.runner.blocking(run(json))) as Json | undefined)
   }
+
+  /** The built-in search: its result is usable directly, and can also be awaited or chained like a Promise. */
+  *search(args: Array<Value>) {
+    return promiseLike(this, yield* this.tool(this.tools.search, args))
+  }
 }
 
 /** One activation: the top-level program or a single function call, evaluating against its own scope chain. */

packages/codemode/src/interpreter/promises.ts

@@ -1,5 +1,5 @@
 import { type AstNode, Throw, typeError } from "./model.js"
-import { Callable, define, get, hidden, Arr, Fn, PromiseObj, record, type Value } from "./objects.js"
+import { Callable, define, get, hidden, Arr, Fn, Obj, PromiseObj, record, type Value } from "./objects.js"
 import { constructor, fn, methods, native, receiver, requiresNew } from "./native.js"
 import { createAggregateErrorValue, materialize } from "./errors.js"
 import { describeValue, typeofValue } from "./references.js"
@@ -258,6 +258,41 @@ function* constructPromise<R>(ctx: Interpreter<R>, executor: Value, site: AstNod
   return made.promise
 }
 
+const promiseLikePrototypes = new WeakMap<object, Obj>()
+
+/**
+ * Lets a synchronous object result also be awaited or chained like a Promise. It inherits non-enumerable
+ * then/catch/finally that settle with a plain copy, so data views never show them and awaiting does not recurse.
+ */
+export const promiseLike = <R>(ctx: Interpreter<R>, value: Value): Value => {
+  if (!(value instanceof Obj) || value.proto !== ctx.builtins.Object) return value
+  value.proto = promiseLikePrototype(ctx)
+  return value
+}
+
+const promiseLikePrototype = <R>(ctx: Interpreter<R>): Obj => {
+  const cached = promiseLikePrototypes.get(ctx)
+  if (cached !== undefined) return cached
+  const proto = new Obj(ctx.builtins.Object)
+  const settled = (thisValue: Value, name: string) => {
+    const source = receiver(Obj, thisValue, name)
+    const copy = new Obj(ctx.builtins.Object)
+    for (const [key, slot] of source.props) copy.props.set(key, { ...slot })
+    return ctx.runner.promiseResolve(copy)
+  }
+  methods(
+    ctx.builtins,
+    proto,
+    (["then", "catch", "finally"] as const).map((name) => [
+      name,
+      name === "then" ? 2 : 1,
+      (thisValue, args, site) => instanceMethod(ctx, name, settled(thisValue, name), args, site),
+    ]),
+  )
+  promiseLikePrototypes.set(ctx, proto)
+  return proto
+}
+
 export const promiseGlobal = <R>(ctx: Interpreter<R>) => {
   const builtins = ctx.builtins
   const proto = builtins.Promise

packages/codemode/test/errors.test.ts

@@ -151,11 +151,11 @@ describe("source locations", () => {
   test("names a missing method on a call instead of the previous line", async () => {
     const failure = await error(`// Try search with different namespaces
 for (const ns of ["github", "tools.github", "tools", ""]) {
-  const s = await search({query: "star", namespace: ns, limit: 100}).catch(e=>({items:[],error:String(e)}));
+  const s = await search({query: "star", namespace: ns, limit: 100}).filter((item) => item.path);
   return s
 }`)
     expect(failure.location).toEqual({ line: 3, column: 19 })
-    expect(failure.message).toBe("TypeError: search(...).catch is not a function. (line 3, col 19)")
+    expect(failure.message).toBe("TypeError: search(...).filter is not a function. (line 3, col 19)")
   })
 })
 

packages/codemode/test/parity.test.ts

@@ -997,8 +997,8 @@ describe("coercion parity: unknown static members read as undefined", () => {
     expect(await value(`try { JSON.rawJSON("1") } catch (e) { return e.message }`)).toBe(
       "JSON.rawJSON is not a function.",
     )
-    expect(await value(`try { search({ query: "star" }).catch(() => 1) } catch (e) { return e.message }`)).toBe(
-      "search(...).catch is not a function.",
+    expect(await value(`try { search({ query: "star" }).filter(() => 1) } catch (e) { return e.message }`)).toBe(
+      "search(...).filter is not a function.",
     )
     expect(
       await value(`const foo = () => ({ bar: () => ({}) }); try { foo().bar().baz() } catch (e) { return e.message }`),

packages/codemode/test/tool-paths.test.ts

@@ -355,6 +355,65 @@ describe("tools.search alias", () => {
   })
 })
 
+describe("search results used like a Promise", () => {
+  const runtime = CodeMode.make({ tools: { api: { list: echo("List things", "listed") } } })
+
+  test("search(...) can be read, awaited, or chained", async () => {
+    const call = "search"
+    expect(
+      await value(
+        runtime,
+        `
+        const query = { query: "list" }
+        return [
+          ${call}(query).items[0].path,
+          (await ${call}(query)).items[0].path,
+          await ${call}(query).then((result) => result.items[0].path),
+          (await ${call}(query).catch(() => "unused")).items[0].path,
+          (await ${call}(query).finally(() => "ignored")).items[0].path,
+          (await Promise.all([${call}(query), ${call}({})])).map((result) => result.items.length),
+        ]
+        `,
+      ),
+    ).toStrictEqual(["tools.api.list", "tools.api.list", "tools.api.list", "tools.api.list", "tools.api.list", [1, 1]])
+  })
+
+  test("the Promise methods stay out of data views", async () => {
+    const expected = { items: [{ path: "tools.api.list" }], remaining: 0, next: null }
+    const result = await value(
+      runtime,
+      `
+      const result = search({ query: "list" })
+      const forIn = []
+      for (const key in result) forIn.push(key)
+      return {
+        json: JSON.parse(JSON.stringify(result)),
+        keys: Object.keys(result),
+        spread: Object.keys({ ...result }),
+        forIn,
+        returned: result,
+      }
+      `,
+    )
+    expect(result).toMatchObject({
+      json: expected,
+      keys: ["items", "remaining", "next"],
+      spread: ["items", "remaining", "next"],
+      forIn: ["items", "remaining", "next"],
+      returned: expected,
+    })
+  })
+
+  test("invalid input throws at the call, so try/catch sees it in either style", async () => {
+    expect((await failure(runtime, `return search({ limit: "ten" }).items`)).message).toContain(
+      "Invalid input for tool 'search'",
+    )
+    expect(
+      await value(runtime, `try { await search({ limit: "ten" }) } catch (error) { return error.message }`),
+    ).toContain("Invalid input for tool 'search'")
+  })
+})
+
 describe("tool references under ==", () => {
   test("compare by identity against data objects without converting them", async () => {
     const runtime = CodeMode.make({ tools: { probe: echo("Probe", "ok") } })

packages/core/src/codemode/instructions.ts

@@ -10,7 +10,7 @@ const prompt = (hasMoreTools: boolean) => `# Code Mode
 
 Use the \`execute\` tool to call the tools listed below. They cannot be called directly${hasMoreTools ? ", and neither can \`search\`. Both" : ". They"} only work inside code you pass to \`execute\`.
 
-${hasMoreTools ? `The catalog is partial. Inside \`execute\`, use \`search(...)\` to find a tool, then call it by the \`path\` in the result. \`search\` is synchronous. Call it without \`await\`; it does not return a Promise. Do not guess tool names.
+${hasMoreTools ? `The catalog is partial. Inside \`execute\`, use \`search(...)\` to find a tool, then call it by the \`path\` in the result. Do not guess tool names.
 
 - ${searchSignature}` : "The catalog is complete. Do not guess tool names."}
 

packages/core/test/codemode/catalog.test.ts

@@ -131,7 +131,7 @@ describe("CodeModeInstructions.render", () => {
     expect(partial).toContain("- orders (1 tool, none shown)")
     expect(partial).toContain("They cannot be called directly, and neither can `search`.")
     expect(partial).toContain(
-      "The catalog is partial. Inside `execute`, use `search(...)` to find a tool, then call it by the `path` in the result. `search` is synchronous. Call it without `await`; it does not return a Promise.",
+      "The catalog is partial. Inside `execute`, use `search(...)` to find a tool, then call it by the `path` in the result.",
     )
     expect(partial).toContain("- search({")
     expect(partial).toContain("  /** @integer @exclusiveMinimum 0 */\n  limit?: number,")