refactor(agent): support hook-driven runtime model routing - #2252
Draft
gold-silver-copper wants to merge 4 commits into
Draft
refactor(agent): support hook-driven runtime model routing#2252gold-silver-copper wants to merge 4 commits into
gold-silver-copper wants to merge 4 commits into
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
ModelHandlethat erases anyCompletionModelonce into private unary, streaming, and capability callbacks.Agent,AgentRunner, prompt/stream request handles, agent stream items, and extractor builders independent of the concrete model type. Useful output and typestate generics remain.CompletionModel,CompletionClient, provider transports, and direct low-level completion/stream APIs typed and externally implementable.AgentHook::on_model_selectrouting with borrowedModelSelectioninputs and event-specificModelSelectionActionresults.HookStack: each hook sees the prior candidate, the last selection wins, and stop is terminal. Nested stacks preserve the same rules.ModelHandleonAgentRunner; agent replacement and one-run overrides change the initial candidate without suppressing routing hooks.PreparedModelAttempt, so capability inspection and execution cannot diverge.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.
Continuepreserves the current candidate,Selectreplaces it for later hooks, andStopcancels before completion hooks, request preparation, or provider execution. The final handle becomesprevious_modelfor 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-specificadditional_paramsare 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-freeruntime_model_routingexample demonstrates the two-model hook-driven tool round trip.Verification
Passed locally:
cargo fmt --all --checkcargo clippy --workspace --all-targets --all-features -- -D warningscargo test --workspace --all-featurescargo test --workspace --all-features --doccargo 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_routingRUSTDOCFLAGS="-D warnings" cargo doc --package rig-core --package rig-agent --package rig --no-deps --all-featureswasm32-unknown-unknownchecks forrig-core,rig-agent,rig,rig-candle, andcandle_wasm_chatrmcpWASM diagnostic check and Candle worker runtime testsgit diff --checkA fresh independent review covered all 46 tracked files against merge base
6cfae6d829da21f9dc9e775e065fee157b264f7eand 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-neo4jprivate intra-doc link atcrates/rig-neo4j/src/vector_index.rs:302. The identical failure was reproduced on a detachedorigin/mainworktree. Warning-denied documentation for the affectedrig-core,rig-agent, andrigscope passes, so this PR intentionally leaves the unrelated warning unchanged.