fix(codemode): run extension before hooks before the host call

packages/codemode/README.md

@@ -91,33 +91,42 @@ runtime.execute(source) // Effect<CodeMode.Result, never, ToolServices>
 The Effect environment is inferred from the supplied tools. `hooks` surround every call the program makes into the
 host: `tool.before`/`tool.after` receive `{ name, input }` with decoded input, `extension.before`/`extension.after`
 receive `{ extension, name, args }`. An `after` hook also receives how the call ended (`success` with its value,
-`failure` with its error, or `interrupted`). A failing `before` hook denies the call, and the program catches the
-failure as a thrown error.
+`failure` with its error, or `interrupted`). A `before` hook runs before the call reaches the host, and a failing one
+denies it: the host function is never called, and the program catches the failure as a thrown error. A denied call
+fails the way the call itself would: a tool call or asynchronous global rejects its promise, and a synchronous global
+throws.
 
 ### `Extension.make`
 
 Extensions are host functions a program calls directly as globals, such as `fetch`. Unlike tools they are not in the
 catalog, not counted against `maxToolCalls`, and not described to the model; the host decides what they mean.
 
 ```ts
-const web = Extension.make({ name: "web", globals: { fetch: (url: string) => globalThis.fetch(url) } })
+const web = Extension.make({ name: "web", globals: { fetch: async (url: string) => globalThis.fetch(url) } })
 const runtime = CodeMode.make({ tools, extensions: [web] })
 ```
 
-A global may return a value, a Promise, or an Effect. An Effect starts only after the `extension.before` hook admits
-the call, and its requirements join the runtime's environment the same way a tool's do, so the host can provide
-per-execution services such as the calling session:
+A global is synchronous unless it is an `async` function or is marked with `Extension.async`. Mark a plain function
+that returns a Promise or an Effect, so a denied call rejects its promise like any other failure:
+
+```ts
+Extension.make({ name: "web", globals: { fetch: Extension.async((url: string) => globalThis.fetch(url)) } })
+```
+
+A global may return a value, a Promise, or an Effect. An Effect's requirements join the runtime's environment the same
+way a tool's do, so the host can provide per-execution services such as the calling session:
 
 ```ts
 const web = Extension.make({
   name: "web",
   globals: {
-    fetch: (url: string) =>
+    fetch: Extension.async((url: string) =>
       Effect.gen(function* () {
         const session = yield* Session
         yield* session.authorize(url)
         return yield* Effect.promise(() => globalThis.fetch(url))
       }),
+    ),
   },
 })
 const result = yield* CodeMode.make({ extensions: [web] }).execute(code).pipe(Effect.provideService(Session, session))

packages/codemode/src/extension.ts

@@ -5,8 +5,12 @@ import type { Effect } from "effect"
 /**
  * Host functions a program calls directly as globals. Values crossing in either direction are converted, never
  * shared: arguments come in as copies, results go out as copies, and a function inside a result is callable the
- * same way. Extension calls are not tool calls. A global may return a value, a Promise, or an Effect; an Effect runs
- * only after the extension hooks admit the call, and its requirements join the runtime's environment.
+ * same way. Extension calls are not tool calls. A global may return a value, a Promise, or an Effect, and an
+ * Effect's requirements join the runtime's environment.
+ *
+ * A global is synchronous unless it is an `async` function or is marked with `async`. The `extension.before` hook
+ * runs before the call, so a denied call never reaches the host; it fails the way the global would, by rejecting the
+ * promise of an asynchronous global and by throwing from a synchronous one.
  */
 export type Extension<Globals extends Readonly<Record<string, Function>> = Readonly<Record<string, Function>>> = {
   readonly name: string
@@ -32,3 +36,14 @@ export const make = <const Globals extends Readonly<Record<string, Function>>>(
   }
   return { name: options.name, globals: { ...options.globals } }
 }
+
+const asyncGlobals = new WeakSet<Function>()
+const AsyncFunction = (async () => {}).constructor
+
+/** Marks a global that returns a Promise or an Effect, so a denied call rejects its promise instead of throwing. */
+export const async = <F extends Function>(global: F): F => {
+  asyncGlobals.add(global)
+  return global
+}
+
+export const isAsync = (global: Function) => global instanceof AsyncFunction || asyncGlobals.has(global)

packages/codemode/src/interpreter/extensions.ts

@@ -1,11 +1,10 @@
 import { Effect } from "effect"
-import type { Extension } from "../extension.js"
+import { type Extension, isAsync } from "../extension.js"
 import { type ExtensionInvocation, hooked } from "../tool-runtime.js"
 import type { Interpreter } from "./interpreter.js"
 import { createErrorValue, isErrorType } from "./intrinsics.js"
 import { MAX_VALUE_DEPTH } from "./limits.js"
 import { type AstNode, PendingThrow, Throw, typeError } from "./model.js"
-import type { Host } from "./runner.js"
 import { fn } from "./native.js"
 import {
   define,
@@ -134,25 +133,32 @@ export const extensionGlobals = <R>(
     throw typeError(`${label} produced ${describeHost(value)}, which the program cannot hold.`)
   }
 
-  // A host function as a program function. Arguments cross in; a global's call runs inside the host's extension
-  // hooks, which see the host's own error on failure; then whatever came back, or was thrown, crosses out so the
-  // program catches a copy. Functions inside results are part of a value's API and skip the hooks.
+  // A host function as a program function. Arguments cross in; for a global, the `extension.before` hook runs before
+  // the host function is called, so a call it denies never reaches the host. The host's result, or what it threw,
+  // crosses out so the program catches a copy, and `extension.after` sees how the call ended. Functions inside
+  // results are part of a value's API and skip the hooks.
   const wrap = (value: Function, label: string, describe?: (args: ReadonlyArray<unknown>) => ExtensionInvocation) =>
-    fn(builtins, value.name, value.length, (_, values, site) => {
+    fn(builtins, value.name, value.length, function* (_, values, site) {
       const args = values.map((item, index) => toHost(item, `Argument ${index + 1} to ${label}`))
-      const call = callHost<R>(value, args)
+      const cross = (reason: unknown) => new Throw(fromHost(reason, label))
+      const call = describe?.(args)
       const hooks = ctx.tools.hooks
-      const work = (
-        describe === undefined
-          ? call.effect
-          : hooked(describe(args), hooks["extension.before"], hooks["extension.after"], call.effect)
-      ).pipe(
-        Effect.mapBoth({
-          onFailure: (reason) => new Throw(fromHost(reason, label)),
-          onSuccess: (result) => fromHost(result, label),
-        }),
-      )
-      return call.async ? pending(work, site) : blocking(ctx.runner.blocking(work))
+      const before = call && hooks["extension.before"]?.(call).pipe(Effect.mapError(cross))
+      const after = call && hooks["extension.after"]
+      const settle = (effect: Effect.Effect<unknown, unknown, R>) =>
+        (after ? hooked(call, undefined, after, effect) : effect).pipe(
+          Effect.mapBoth({ onFailure: cross, onSuccess: (result) => fromHost(result, label) }),
+        )
+      // A denied call fails the way the global would: an async global rejects its promise, a sync one throws.
+      if (isAsync(value))
+        return pending(
+          Effect.andThen(before ?? Effect.void, settle(Effect.suspend(() => callHost<R>(value, args).effect))),
+          site,
+        )
+      if (before) yield ctx.runner.blocking(before)
+      const host = callHost<R>(value, args)
+      if (host.async) return pending(settle(host.effect), site)
+      return (yield ctx.runner.blocking(settle(host.effect))) as Value
     })
 
   // An asynchronous host call: the program holds a promise that settles when the host's work does.
@@ -176,7 +182,7 @@ export const extensionGlobals = <R>(
 
 /**
  * Calls a host function and describes what the program waits on. A returned Effect or Promise settles later; a
- * returned value or a throw settles now. An Effect has not started yet, so it runs only once the hooks admit it.
+ * returned value or a throw settles now.
  */
 const callHost = <R>(value: Function, args: ReadonlyArray<unknown>) => {
   try {
@@ -191,11 +197,6 @@ const callHost = <R>(value: Function, args: ReadonlyArray<unknown>) => {
   }
 }
 
-// A synchronous host call: its hooks run before any other job, as the call itself did.
-function* blocking(host: Host) {
-  return (yield host) as Value
-}
-
 /**
  * An error crosses as its name, message, `cause`, and own enumerable fields, such as Node's `code`, `errno`,
  * `syscall`, and `path`. `stack` stays on its own side, and no field may shadow an Error method. A field that cannot

packages/codemode/test/codemode.test.ts

@@ -285,6 +285,36 @@ describe("CodeMode call hooks", () => {
     })
   })
 
+  test("a before hook runs before the host function, so a denied call never reaches the host", async () => {
+    const called: Array<string> = []
+    const record = <A>(name: string, result: () => A) => {
+      called.push(name)
+      return result()
+    }
+    const runtime = CodeMode.make({
+      extensions: [
+        Extension.make({
+          name: "host",
+          globals: {
+            value: () => record("value", () => 1),
+            asyncFn: async () => record("asyncFn", () => 1),
+            promise: Extension.async(() => record("promise", () => Promise.resolve(1))),
+            effect: Extension.async(() => Effect.sync(() => record("effect", () => 1))),
+          },
+        }),
+      ],
+      hooks: { "extension.before": (call) => Effect.fail(new Error(`${call.name} is not allowed`)) },
+    })
+    const run = async (code: string) => (await Effect.runPromise(runtime.execute(code))) as { value?: unknown }
+
+    // A denied call fails the way the global would: an async global rejects its promise, a sync one throws.
+    expect(
+      await run(`return (await Promise.allSettled([asyncFn(), promise(), effect()])).map((r) => r.reason.message)`),
+    ).toMatchObject({ value: ["asyncFn is not allowed", "promise is not allowed", "effect is not allowed"] })
+    expect(await run(`try { value() } catch (e) { return e.message }`)).toMatchObject({ value: "value is not allowed" })
+    expect(called).toEqual([])
+  })
+
   test("see a tool defect as the failure the program gets, and an interrupted extension call", async () => {
     const events: Array<string> = []
     const broken = Tool.make({