[purushah]

Linked issue: #897 (design discussion). Supersedes #852.

Purpose of change

Today an agent must name a concrete chat model in each ChatRequestEvent; any routing logic lives in user code and is not first-class in the runtime. Most requests are easy and a small model handles them fine, but a minority really need the strong model — so a fixed choice either overpays on the easy majority or under-serves the hard ones. This PR makes per-request routing a framework capability: an agent sends its request to a router, and the runtime picks one of several candidate models.

My first attempt at this (#852) made the router a special ChatModelSetup that called the chosen backend itself. Review rightly pointed out that this breaks the two things the framework is good at: token metrics collapse onto the router instead of the real backend, and the actual model call disappears from the EventLog. The redesign discussed in #897 fixes this by making the router select instead of delegate: the router only returns a model name, and ChatModelAction runs the normal chat path against that model. A routed request produces the same events and the same per-model metrics as if the agent had named the chosen model directly.

From the user side it stays simple:

env.addResource("router", ResourceType.MODEL_ROUTER,
    ModelRouter.of("small", "big")
        .describe("small", "fast and cheap; chit-chat, simple facts")
        .describe("big", "strong; code, SQL, analysis")
        .strategy(Strategies.rules(Map.of("big", "\\b(code|sql|analyze)\\b")))
        .defaultModel("small")
        .fallback(true)
        .build());

// the agent just sends its ChatRequestEvent to "router"

A few behaviors were worth getting right, and they came out of the #897 discussion:

• The decision is durable. The strategy runs inside its own durable call ("route:<router>" — deterministic across recovery re-processing; per-request uniqueness comes from the store's key/sequence/event scoping), so on recovery the persisted decision is replayed instead of re-running the strategy. Decision latency is measured inside that call, so a replayed run reports the original decision_ms.
• Route once per ReAct loop. Tool-call rounds reuse the already-selected model and carry the routing metadata onto the final response — no re-routing mid-conversation.
• Fallback sits on top of retries, not instead of them. The selected model gets its full retry budget first; only then are the remaining candidates tried in declaration order, and a fallback emits a second ModelRoutingEvent so the event log shows which model actually answered.
• Abstain is not an error, an invalid pick is. A strategy that has no opinion abstains and the router's default model handles the request (decision_source=default); a strategy that returns a non-candidate name fails loudly, because that is a bug, not a runtime condition.
• ModelRoutingEvent is observability-only. It has no built-in consumer and does not drive dispatch — removing every listener does not change which model runs.

Since routers and chat models share the ChatRequestEvent namespace, one name must not be registered as both; this is validated at the addResource call site (with an AgentPlan backstop) so the failure points at the user's own line.

Per the discussion, LLM-as-judge routing is deliberately not in this PR. RoutingContext exposes no chat API, so a strategy cannot make hidden model calls; the framework-managed observable judge (the judge call running on the normal durable/metered chat path) is the agreed follow-up. YAML support is intentionally left out of v1 to keep the first Java API/runtime change small; it and Python support are follow-ups.

Tests

• 20 API unit tests (RoutingTest, RoutingResourceValidationTest): rule matching and abstain over the latest user message, builder and registration validation, candidate descriptions reaching strategies, RoutingDecision invariants and snake_case JSON round-trips, decision_ms surviving replay deserialization, event attribute normalization.
• 9 ChatModelActionRoutingTest integration tests against a scripted fake chat model: routing to the matched candidate, abstain-to-default, invalid candidate failing clearly, non-router requests keeping the legacy "chat" durable id unchanged, fallback across candidates (including exhaustion), and route-once semantics across a tool round with metadata carried to the final response.
• A ResourceCacheTest case covering hasResource visibility of resources inserted directly into the cache (no provider).

Verified locally (JDK 17):

• mvn spotless:check -pl api,plan,runtime,examples
• mvn test -pl api,plan — 312 + 196 tests, 0 failures
• mvn test -pl runtime -Dtest=ResourceCacheTest
• mvn compile -pl examples

API

New public surface, all additive:

• ResourceType.MODEL_ROUTER
• org.apache.flink.agents.api.chat.model.routing: ModelRouter (+ builder), RoutingStrategy, RoutingContext, RoutingDecision, RoutingCandidate, RoutingStrategyDescriptor, Strategies, RuleBasedRoutingStrategy
• ModelRoutingEvent
• RunnerContext.hasResource(name, type) (default method, returns false)

Compatibility impact

A request naming a plain chat model takes the unchanged success path, including the legacy durable call id, and the default error-handling-strategy (FAIL) is unaffected. Three deliberate behavior changes on non-default failure paths:

• retry: retry metrics are recorded per attempt (previously once with cumulative totals on the final response), so requests that ultimately fail now contribute their retry counts and each attempt lands on the model that served it. Totals over a completed request are unchanged.
• retry: the retry WARN log now emits under ChatModelInvoker's logger (was ChatModelAction) — external log alerting keyed to the old logger name needs updating.
• ignore: a request naming a missing chat-model resource is warned and dropped (consistent with other request failures under ignore) instead of failing the job with a raw lookup exception.

Documentation

• doc-needed
• doc-not-needed
• doc-included

The Java examples show usage, but this adds public API, so proper docs are warranted — happy to do a docs follow-up PR once the API settles in review.

[weiqingy]

Thanks for taking this on. A few questions inline.

[purushah]

Thanks @weiqingy for the sharp review — all eight comments are addressed in the follow-up commits with new tests. The EventType registration will come with the rebase onto main.

[wenjin272]

Thanks for pushing this design forward and for addressing the earlier feedback. I have two structural suggestions around tool-round state and the organization of ChatModelAction.

[purushah]

Thanks @wenjin272 — both suggestions implemented in the follow-up commits.

[weiqingy]

Thanks for addressing the comments. While reviewing the reordered code, a few inline comments below.

[purushah]

Thanks @weiqingy for another careful pass — all seven addressed: the Python enum member is in (mixed jobs deserialize; full Python routing stays the follow-up), the three test gaps are covered, the RoutingContext boundary is documented as deliberate, and the PR description now carries the compatibility-impact section and the corrected durable-id format.

[wenjin272]

LGTM from my side. @weiqingy, do you have any remaining comments?

[weiqingy]

Thanks for addressing the comments. Two small things inline, neither blocking. Otherwise, LGTM.

[weiqingy]

The comment here says rule keys are validated at build time so a typo fails at the registration call site instead of throwing per record at runtime. The rule values have the same shape and no check: Pattern.compile runs only in the strategy constructor (RuleBasedRoutingStrategy.java:62), reached from ModelRouter:88 when the resource is provided, and ResourceCache caches only after provide() returns (ResourceCache.java:144). A PatternSyntaxException therefore never caches and re-throws per routed request: under the default FAIL the job dies on the first one, under ignore requests through that router are dropped with a WARN while the job stays RUNNING. Would compiling the patterns in this same block move that to build() too?

[purushah]

Good catch on the caching consequence — fixed: build() now compiles every rule pattern (and rejects null/non-String values) in the same block as the key check, so a malformed regex fails at the registration call site. Two new builder tests cover it.

[weiqingy]

nit: this and selected_model above are both this.selectedModel, and initial_selected_model has no other reference in the PR. Nothing reads it and no test asserts it. The fallback winner is already final_model on the next line. Removing a key from a shipped payload is a compat change later. Is the second key meant to differ from the first on some path, or can it go?

[purushah]

You are right — it duplicated selected_model on every path (fallback winners are final_model). Removed before it ships and becomes a compat constraint.

[weiqingy]

Both addressed at e95f20bd. No remaining comments from me.

[wenjin272]

Thanks for taking this on @purushah. I noticed that this feature has the design discussion (#897) and implementation PR (#964), but no corresponding issue. Could you please create one summarizing the agreed scope, link both the discussion and this PR? This will make the feature easier to track in the project history and release notes.

[purushah]

Thanks @wenjin272 and @weiqingy for the careful reviews and the merge — the feature came out much stronger.

A heads-up on what's next: the Strategies.llm(...) follow-up we agreed to defer in #897 — framework-managed LLM-as-judge, where the engine runs the judge call on its own durable/metered/observable chat path — has been in the works alongside this review cycle for a while now, and it's nearly ready. The shape is exactly what the #897 thread settled on: the judge executes under a durable "judge:<router>" call with engine retries and token attribution, verdicts are constrained to candidate names, judge failures honor the error-handling strategy, and a mistyped judge model fails at plan construction rather than at runtime. I'll open the PR soon.