[api][plan][runtime][examples] Add pluggable in-chat model routing (MODEL_ROUTER)

[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:<requestId>:<router>"), 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)

Existing agents are unaffected: a request naming a plain chat model takes the unchanged path, including the legacy durable call id.

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.

[weiqingy]

resolveRouter (781-863) has no try/catch and runs before chat(), so a strategy exception escapes error-handling-strategy: a job set to ignore still dies.

It's reachable from the built-in. RuleBasedRoutingStrategy.java:61-66 throws on a rule key that isn't a candidate, and Builder validates names in describe() (ModelRouter.java:176-185) but not in strategy(), so Strategies.rules(Map.of("bg", "...")) builds fine and throws per record.

Could rule keys be checked in build(), where both lists are already in hand? And should a strategy failure honor error-handling-strategy, or is fail-fast the contract?

[purushah]

Great catch — and you found two holes in one comment (the per-record throw AND the policy bypass). Fixed both: build() validates rule keys so the typo fails at registration, and resolveRouter failures now honor error-handling-strategy — deliberately without retries, since v1 strategies are CPU-only. Tests: builderRejectsRuleKeyThatIsNotACandidate, strategyFailureIsIgnoredUnderIgnorePolicy, strategyFailurePropagatesUnderDefaultPolicy.

[weiqingy]

lastError = e.error overwrites and 598 throws only the last, so with [A, B, C] all failing you see C's. The debug line at 585 logs the model and messages but not the error, so A's and B's are gone, including the routed model's own failure. The second ModelRoutingEvent only fires on success (536-549), so exhaustion records nothing about what was tried.

Something like this before the overwrite, if it helps:

if (lastError != null && lastError != e.error) {
    e.error.addSuppressed(lastError);
}

fallbackExhaustedRethrows can't tell today: both scripted errors contain "down" (ChatModelActionRoutingTest.java:408-410). Would distinct messages help?

[purushah]

Really appreciate you noticing the test itself couldn't detect this (both scripted errors shared "down") — that's a level of review rigor we rarely get. Applied your suggestion: errors chain via addSuppressed, exhaustion WARNs with the tried list, and the test now uses distinct markers so it fails against the old behavior.

[weiqingy]

This lookup sits outside the try/catch that produces ChatAttemptFailed (644-676), and the call site catches only that (581). So an unresolvable candidate B propagates straight out of chat(): A's real failure is discarded, the caller sees "resource not found", and the candidate loop is bypassed, including under IGNORE.

Nothing catches the typo upstream either, since ModelRouter's constructor (57-89) never touches the registry.

Could the lookup move inside the conversion, so an unresolvable candidate just counts as that candidate failing? And is there room to cross-validate candidate names at plan construction, the way Agent.java:135-139 argues for the clash check?

[purushah]

This was the sharpest catch of the review — the failure mode (real error discarded, confusing "resource not found" for a model the user never picked, both fallback and IGNORE bypassed) is exactly the worst debugging experience. Fixed: the lookup moved inside the attempt conversion, so it counts as a failed attempt and flows through fallback/IGNORE with prior errors preserved. (Side effect: a plain chat with a missing model now honors IGNORE too — we think that's more consistent.) Test: unresolvableCandidateCountsAsFailedAttemptAndFallsBack.

[weiqingy]

The collections are unmodifiable (55-58, 59-62) but the copy is shallow, and ChatMessage is mutable on main. These are the same instances that go to the model: event.getMessages() is passed into resolveRouter (758) and into chat(...) (764). So a setContent(...) inside a strategy rewrites the prompt that's actually sent, with no event and no trace.

The strategy's reach to other resources is already closed by the type. This is the same thing one level down, and narrowing it is free before v1 ships and breaking after.

firstUserMessage() and lastUserMessage() (100, 114) are what a selection strategy needs. Deep-copy the list, or drop getMessages() until something needs more than the text?

[purushah]

Good catch — fixed with a true deep copy, including the toolCalls list and inner maps (ChatMessage's ctor stores toolCalls by reference; re-verifying your comment caught that second hole too, so this one paid off twice). Tests: routingContextMessagesAreDeepCopied, routingContextToolCallsAreDeepCopiedToo.

[weiqingy]

After a rebase, ModelRoutingEvent won't be in main's ALL_CONSTANTS (EventType.java:45-54), and ConditionExpressionValidator.java:125 rejects any EventType.X that isn't a key of it. So @Action("type == EventType.ModelRoutingEvent && ...") compiles, then fails plan validation. Plain @Action(EventType.ModelRoutingEvent) is unaffected.

Worth adding the entry when you rebase?

[purushah]

Right — will register ModelRoutingEvent (plus a condition-expression test) during the rebase onto current main.

[weiqingy]

String.valueOf(entry.getValue()) returns "null" for a null value, never null, so the regex != null guard below is dead and the entry compiles to Pattern.compile("null", CASE_INSENSITIVE). Any message containing "null" then routes there. Non-String values coerce silently too, and rules arrives as Map<String,Object> off the plan JSON.

Should a null or non-String value be rejected instead?

[purushah]

Nice spot on a classic String.valueOf(null) trap. Fixed: null/non-String rule values are rejected with a descriptive error. Tests: ruleStrategyRejectsNullRuleValue, ruleStrategyRejectsNonStringRuleValue.

[weiqingy]

The guarantee reads unconditional here (69-71) and on RoutingDecision.java:93-95, but the journal is off by default: actionStateStoreBackend defaults to null (AgentConfigOptions.java:65-66), and setupDurableExecutionContext returns early when the store is null (DurableExecutionManager.java:289-292).

The default path is still self-consistent, since the decision and the chat call re-run together. It's the javadoc that promises more than a default deployment gets. Worth a qualifying clause?

[purushah]

Fair — the guarantee is conditional and now reads that way: replay requires a configured action-state store (actionStateStoreBackend); without one the decision re-executes. Doc-only.

[weiqingy]

Retry metrics moved: origin/main:433-434 recorded once on the final response with cumulative totals; this records per attempt (681-698), including on the failure path. Totals over a completed request are unchanged, but requests that ultimately throw now increment where they didn't.

AGENTS.md asks for compatibility impact in the description. Worth a line?

[purushah]

Good eye for catching the semantic shift — it's intentional (per-attempt recording makes fallback attribution land on the model that served the attempt), and it's now documented in a compat note on recordAttemptRetryStats plus the PR description.

[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.

[wenjin272]

Do we need to carry routingMetadata through every tool-request context? The concrete model used for subsequent tool rounds is already stored separately in MODEL; routingMetadata does not participate in model selection and is only used to attach model_routing to the eventual ChatResponseEvent.

Since its lifetime is the whole ReAct loop, could we store it once in an initial-request context keyed by initialRequestId, then retrieve it only when producing the final response? This would avoid repeatedly copying it across tool rounds, remove the special RoutingSelection.carried(...) state, and keep observability metadata out of intermediate ChatMessage.extraArgs. The context could then be cleaned up when the loop completes.

[wenjin272]

ChatModelAction has grown to roughly 1,000 lines and now owns several distinct responsibilities: route resolution, concrete-model invocation/retry, fallback orchestration, routing metadata propagation, and tool-loop state management.

Could we extract the routing-specific pieces into a package-private component such as ModelRoutingResolver / ResolvedModelRoute? RoutingSelection, resolveRouter, candidate ordering, routed durable-call IDs, and routing metadata construction form a cohesive unit. Similarly, chatWithRetries together with ChatAttemptResult / ChatAttemptFailed could become a concrete-model invoker. This would leave ChatModelAction responsible primarily for event orchestration while keeping the routing implementation easier to review and extend.