Skip to content

refactor(agent): support hook-driven runtime model routing - #2252

Draft
gold-silver-copper wants to merge 4 commits into
mainfrom
refactor/runtime-model-swapping
Draft

refactor(agent): support hook-driven runtime model routing#2252
gold-silver-copper wants to merge 4 commits into
mainfrom
refactor/runtime-model-swapping

Conversation

@gold-silver-copper

@gold-silver-copper gold-silver-copper commented Aug 4, 2026

Copy link
Copy Markdown
Contributor

Motivation

Rig's high-level agent types currently carry the concrete completion-model type through every runner, prompt, extractor, and stream. That prevents an application from changing models or providers while retaining the agent's tools, hooks, memory, instructions, context, and output configuration.

This PR makes model behavior a runtime value at the high-level agent boundary while preserving the existing open, typed provider contracts.

Architecture

  • Adds a public opaque, cloneable ModelHandle that erases any CompletionModel once into private unary, streaming, and capability callbacks.
  • Makes Agent, AgentRunner, prompt/stream request handles, agent stream items, and extractor builders independent of the concrete model type. Useful output and typestate generics remain.
  • Keeps CompletionModel, CompletionClient, provider transports, and direct low-level completion/stream APIs typed and externally implementable.
  • Adds synchronous AgentHook::on_model_select routing with borrowed ModelSelection inputs and event-specific ModelSelectionAction results.
  • Composes routing exclusively through the ordered HookStack: each hook sees the prior candidate, the last selection wins, and stop is terminal. Nested stacks preserve the same rules.
  • Keeps one default ModelHandle on AgentRunner; agent replacement and one-run overrides change the initial candidate without suppressing routing hooks.
  • Binds each provider-neutral prepared request to the exact final handle in PreparedModelAttempt, so capability inspection and execution cannot diverge.
  • Normalizes only high-level provider-final stream events to AgentStreamFinal { usage, raw_response }; all other assistant events and correlation IDs pass through unchanged.

Model-call lifecycle

The shared blocking/streaming driver resolves routing exactly once before each AgentRunStep::CallModel. This includes the initial call, post-tool calls, accepted-turn retries, invalid-tool recovery, and structured extraction repair attempts. Completion hooks and capability-sensitive request preparation run only after routing finishes.

Routing hooks run in registration order. Continue preserves the current candidate, Select replaces it for later hooks, and Stop cancels before completion hooks, request preparation, or provider execution. The final handle becomes previous_model for the next model-call boundary and is retained by the prepared attempt, so an in-flight unary future or stream cannot be rebound.

Hooks attached to an agent builder apply to later runners; request/runner hooks apply only to that run. Cloned runners own independent stack snapshots, while explicitly synchronized state captured by a hook follows that hook's clone semantics. Cancellation remains drop-based: no detached task consumes a future or provider stream after its caller is dropped.

ModelHandle, routing hooks, hook stacks, and actions are live runtime behavior and intentionally are not serializable. Applications should persist their own model identifier and resolve it to a handle at runtime. Provider-specific additional_params are not guaranteed to be portable across providers and retain existing validation behavior.

User impact

Applications can keep one concrete Agent, store heterogeneous models in collections, replace defaults between runs, choose a specialist default for one prompt or extraction, or route from a fast research model to a stronger synthesis model after a tool result. The credential-free runtime_model_routing example demonstrates the two-model hook-driven tool round trip.

Verification

Passed locally:

  • cargo fmt --all --check
  • cargo clippy --workspace --all-targets --all-features -- -D warnings
  • cargo test --workspace --all-features
  • cargo test --workspace --all-features --doc
  • cargo test -p rig-agent --all-features --lib agent::hook -- --nocapture (37 passed)
  • cargo test -p rig-agent --test runtime_model_swapping -- --nocapture (16 passed)
  • cargo run -p rig-agent --example runtime_model_routing
  • RUSTDOCFLAGS="-D warnings" cargo doc --package rig-core --package rig-agent --package rig --no-deps --all-features
  • Standalone wasm32-unknown-unknown checks for rig-core, rig-agent, rig, rig-candle, and candle_wasm_chat
  • The native-only rmcp WASM diagnostic check and Candle worker runtime tests
  • Forbidden routing-API identifier search and git diff --check

A fresh independent review covered all 46 tracked files against merge base 6cfae6d829da21f9dc9e775e065fee157b264f7e and found no P0/P1 or lower-severity issues. It separately validated hook composition, stop and retry behavior, attempt binding, cancellation/concurrency, exports, serialization boundaries, and browser-WASM bounds.

No cassette fixture changed. Cassette replay and safety/secret checks passed in the full workspace suite.

Residual risk / unrelated limitation

The full workspace rustdoc command with warnings denied remains blocked by an existing rig-neo4j private intra-doc link at crates/rig-neo4j/src/vector_index.rs:302. The identical failure was reproduced on a detached origin/main worktree. Warning-denied documentation for the affected rig-core, rig-agent, and rig scope passes, so this PR intentionally leaves the unrelated warning unchanged.

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