Skip to content

feat(ffi): repair_tool_call for every binding via handles in opts_json - #186

Draft
cunninghamcard-bit wants to merge 19 commits into
arcships:masterfrom
cunninghamcard-bit:feat/ffi-tool-call-repair
Draft

cunninghamcard-bit wants to merge 19 commits into
arcships:masterfrom
cunninghamcard-bit:feat/ffi-tool-call-repair

Conversation

@cunninghamcard-bit

@cunninghamcard-bit cunninghamcard-bit commented Sep 16, 2026

Copy link
Copy Markdown
Contributor

Gives every binding the AI SDK repairToolCall hook that #165 landed for Rust callers only.

Problem

GenerateTextOptions.repair_tool_call is #[serde(skip)], so the C ABI's opts_json could not carry it and the seven bindings only ever saw the result of repair (ToolCallRepair, C code 17) — none could supply a repair function. The same was true of abort_signal, which is why aimux_stream_text_with_abort exists as a separate entry point.

Design

opts_json is the whole GenerateTextOptions object, the shape the AI SDK uses. Fields that hold live objects travel as handles and are resolved in parse_opts_arg before the rest deserializes:

{"tools": [...], "abort_signal": 7, "repair_tool_call": 42}
  • aimux_tool_call_repair_new(fn, user_data) -> handle registers a host function with the repairToolCall contract: it receives {tool_call, error, input_schema, tools, messages, instructions} as JSON and returns the repaired RawToolCall JSON, {"error": "<message>"} to record a failure, or NULL to keep the original error; strings come from the new aimux_string_new and aimux frees them. Core parses and validates a returned call from scratch and records an error reply as ToolCallRepair { original_error, cause } — exactly what a Rust closure that returns Err gets, so all seven bindings have one error semantics. parse_tool_call is untouched.
  • The host function runs synchronously on the thread that entered the aimux_* call, inside the existing re-entrancy guard: an aimux_* call from inside it fails with AIMUX_E_FFI_REENTRANT_CALL (204) instead of deadlocking — same rule as the stream callbacks, covered by a test.
  • aimux_tool_call_repair_drop disarms clones held by in-flight calls (they behave as if the function returned NULL), so the host may release user_data after it returns. Handle 0 means "none", matching aimux_tool_call_repair_new(NULL).
  • No existing entry-point signature changes. The explicit abort_handle of the *_with_abort variants still wins when both are given. Future callbacks are one more aimux_*_new constructor plus a JSON field.
  • The RawToolCall / context / reply wire code lives once, in aimux-core (RawToolCall derives serde, ToolCallRepairContext::to_wire_json, parse_repair_reply); the C, Node and Python bridges all use it.

Each binding exposes a ToolCallRepair built from a native function and sets it on its GenerateTextOptions (repairToolCall / repair_tool_call), where it serializes as the handle. Registered objects stay alive until close() (Java/Kotlin static registry, Swift passRetained, Go cgo.Handle, Dart NativeCallable + @pragma('vm:isolate-unsendable') so an Isolate.run hand-off fails fast instead of crashing). Node and Python link Core directly and bridge the function natively (napi ThreadsafeFunction; PyO3 on a spawn_blocking thread with the calling thread's GIL released) — in both, the hook may itself call aimux.

Verification

Layer Local Notes
aimux-ffi clippy clean; 129 tests incl. contract repair fixes / gives up / error envelope / re-enters (204) / bad handle (203) / non-integer (5) / handle 0 / dropped clone not invoked / abort_signal via JSON / same on the aimux_stream_text path
Go vet + full suite reference host implementation; repair via StreamText; nested aimux call refused as re-entrant
Java 113 tests, JUnit run directly on JDK 21 gradle needs JDK 17 (CI)
Kotlin 79 tests on JDK 21
Node 91 tests JS repair may call back into aimux (tested); a reply that is neither a call nor {error} is a ToolCallRepair error
Python 89 tests via maturin hook runs on a blocking thread, GIL released on the caller; may call back into aimux (tested); a non-dict reply is a TypeError recorded as the repair cause
Swift library builds; new test expressions type-checked no XCTest here
Flutter not run no dart toolchain here; CI is the first run

Docs: docs/api/{c,go,java,kotlin,swift,flutter,node,python}.md.

Follow-ups tracked separately in #185 (ToolInput raw/parsed enum, unused StreamingToolCallTracker).

… by handle

`opts_json` is now the whole `GenerateTextOptions` in the AI SDK shape. The
two fields Core marks `serde(skip)` because they hold live objects are passed
as handles and resolved in `parse_opts_arg`:

    {"tools": [...], "abort_signal": <handle>, "repair_tool_call": <handle>}

- `aimux_tool_call_repair_new(fn, user_data)` registers a host function with
  the AI SDK `repairToolCall` contract: it receives {tool_call, error,
  input_schema, tools, messages, instructions} as JSON and returns the
  repaired RawToolCall JSON or NULL. It runs synchronously on the calling
  thread inside the re-entrancy guard, like the stream callbacks.
- `aimux_string_new` lets a host hand a string back to aimux; aimux frees it.
- Every existing entry point honors the handle fields with no signature
  change; the explicit `abort_handle` of the *_with_abort variants still wins.
`NewToolCallRepair(fn)` registers a Go `ToolCallRepairFunc` through
`aimux_tool_call_repair_new` and marshals as its handle, so it sits in
`GenerateTextOptions.RepairToolCall` and reaches Core via opts_json like any
other option. The function receives the AI SDK repairToolCall context and
returns the repaired RawToolCall or nil; Go errors and panics stay on the Go
side (`Err()`) and leave the original validation error on the tool call.
… hook

Every Python entry point drives the shared tokio runtime with block_on; a
hook that called back into aimux nested block_on, which tokio rejects with a
panic that PyO3 then resumed on the way out of the enclosing call. Route all
17 block_on sites through one helper that raises RuntimeError when the
current thread is already inside the runtime — the C ABI's re-entrancy rule,
expressed as an ordinary exception that core records as ToolCallRepair.
…t handle 0 as none

- aimux_tool_call_repair_drop clears a shared flag that in-flight clones check
  before calling the host, so releasing user_data after drop is safe; the old
  'in-flight calls keep their clone' wording described a use-after-free.
- The host may reply {"error": "<message>"}; Core records it as
  ToolCallRepair { original_error, cause }, giving every binding the same
  semantics as a Rust closure that returns Err.
- 0 in opts_json means no handle, matching aimux_tool_call_repair_new(NULL).
- opts_json is parsed once into a two-field HandleFields struct plus once into
  GenerateTextOptions instead of a full Value tree; a non-integer handle is
  InvalidArgument (5) like every other schema violation.
- The RawToolCall / context / reply wire code moves into aimux-core
  (RawToolCall derives serde; ToolCallRepairContext::to_wire_json;
  parse_repair_reply) so the Node and Python bridges can drop their copies.
…r; errors through the envelope (unverified locally: no dart toolchain)
…a blocking thread

The calling thread held the GIL for the whole block_on while stream_text's
hook needed it on a runtime worker: with one worker that is a deadlock. Every
GIL-holding block_on now runs under allow_threads and the hook runs in
spawn_blocking, which also lets the hook call back into aimux (tokio's
blocking pool is not an EnterRuntime context), so the re-entrancy guard goes.
Also: SkipJsonSchema on the pydantic field so model_json_schema() works
again, and the wire code moves to core's to_wire_json / parse_repair_reply.
- Go: recover in the exported trampoline itself, so a Close racing an
  in-flight invocation (cgo.Handle.Value panics on a deleted handle)
  yields a NULL reply instead of a Go panic unwinding through Rust.
- Python: a hook that returns anything but a dict or None is a TypeError,
  recorded as the ToolCallRepair cause; before, a returned str dumped to a
  JSON string literal and failed with a message about the wire shape.
- Kotlin: the context's messages are raw JsonElements, as in Go and Java,
  so a message shape the codec does not model cannot fail the repair
  before the user function runs.
- c.md: host exceptions become the {"error"} envelope (or NULL), matching
  the header and every binding; the text said NULL only.
- Tests: repair on the aimux_stream_text path (FFI) and StreamText (Go),
  a Go re-entrancy test, and a malformed reply ({}) in Node and Python.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant