feat(model-routing): thread per-call model through InvokeModelStage

strands-py/src/strands/_middleware/stages.py

@@ -13,6 +13,7 @@
     from ..agent.agent import Agent
     from ..experimental.bidi import BidiAgent
     from ..interrupt import _InterruptState
+    from ..models.model import Model
     from ..types._events import ModelStopReason, ToolResultEvent, TypedEvent
     from ..types.content import Messages, SystemPrompt
     from ..types.tools import AgentTool, ToolChoice, ToolSpec, ToolUse
@@ -25,6 +26,9 @@ class InvokeModelContext:
     All collection fields (messages, system_prompt, tool_specs, tool_choice) are
     defensive copies — middleware cannot accidentally mutate agent state.
     invocation_state is shared by reference (hooks and tools write to it during streaming).
+
+    ``model`` is the model this call invokes; it starts as ``agent.model`` and middleware
+    may replace it per call.
     """
 
     agent: Agent
@@ -33,6 +37,7 @@ class InvokeModelContext:
     tool_specs: list[ToolSpec]
     tool_choice: ToolChoice | None
     invocation_state: dict[str, Any]
+    model: Model
     projected_input_tokens: int | None = None
 
 

strands-py/src/strands/event_loop/event_loop.py

@@ -537,6 +537,7 @@ async def _handle_model_execution(
                 tool_specs=copy.deepcopy(tool_specs),
                 tool_choice=copy.deepcopy(structured_output_context.tool_choice),
                 invocation_state=invocation_state,
+                model=agent.model,
                 projected_input_tokens=projected_input_tokens,
             )
 
@@ -559,8 +560,7 @@ async def _handle_model_execution(
 
             if last_event is None:
                 raise RuntimeError(
-                    "Middleware chain did not yield a result event. "
-                    "Ensure middleware forwards events from next()."
+                    "Middleware chain did not yield a result event. Ensure middleware forwards events from next()."
                 )
 
             # Write the post-stream model state back to the agent. Skipped on error
@@ -663,7 +663,7 @@ def _make_invoke_model_terminal(
     async def terminal(ctx: InvokeModelContext) -> AsyncGenerator[Any, None]:
         system_prompt_str, system_prompt_content = split_system_prompt(ctx.system_prompt)
 
-        model_id = agent.model.config.get("model_id") if hasattr(agent.model, "config") else None
+        model_id = ctx.model.config.get("model_id") if hasattr(ctx.model, "config") else None
         model_invoke_span = tracer.start_model_invoke_span(
             messages=ctx.messages,
             parent_span=cycle_span,
@@ -675,7 +675,7 @@ async def terminal(ctx: InvokeModelContext) -> AsyncGenerator[Any, None]:
         with trace_api.use_span(model_invoke_span, end_on_exit=False):
             try:
                 async for event in stream_messages(
-                    agent.model,
+                    ctx.model,
                     system_prompt_str,
                     ctx.messages,
                     ctx.tool_specs,

strands-py/tests/strands/agent/conversation_manager/test_token_usage_middleware.py

@@ -25,6 +25,7 @@ def make_context(**overrides) -> InvokeModelContext:
         tool_specs=[],
         tool_choice=None,
         invocation_state={},
+        model=Mock(),
     )
     defaults.update(overrides)
     return InvokeModelContext(**defaults)

strands-py/tests/strands/injection/test_message_injection.py

@@ -51,6 +51,7 @@ def invoke_ctx(messages: list[dict], agent: Any = None) -> InvokeModelContext:
         tool_specs=[],
         tool_choice=None,
         invocation_state={},
+        model=MagicMock(),
     )
 
 

strands-py/tests/strands/memory/test_memory_manager.py

@@ -1192,6 +1192,7 @@ def _invoke_ctx(messages: list[dict], agent: Any) -> Any:
         tool_specs=[],
         tool_choice=None,
         invocation_state={},
+        model=getattr(agent, "model", None),
     )
 
 

strands-py/tests/strands/middleware/test_agent_middleware.py

@@ -720,3 +720,38 @@ async def retry_middleware(context, next_fn):
     result = agent("test")
     assert result.message["content"][0]["text"] == "Success!"
     assert call_count == 3
+
+
+# --- per-call model on context (routing plumbing) ---
+
+
+def test_invoke_model_context_exposes_agent_model(agent):
+    """InvokeModelContext.model defaults to agent.model."""
+    captured: list = []
+
+    async def capture(context, next_fn):
+        captured.append(context.model)
+        async for event in next_fn(context):
+            yield event
+
+    agent._middleware_registry.add_middleware(InvokeModelStage, capture)
+    agent("test")
+
+    assert captured == [agent.model]
+
+
+def test_terminal_streams_context_model_override():
+    """The terminal streams the model set on the context, not agent.model."""
+    model_a = MockedModelProvider([{"role": "assistant", "content": [{"text": "A"}]}])
+    model_b = MockedModelProvider([{"role": "assistant", "content": [{"text": "B"}]}])
+    agent = Agent(model=model_a, callback_handler=None)
+    model_a.stream = AsyncMock(wraps=model_a.stream)
+
+    def route_to_b(context):
+        return replace(context, model=model_b)
+
+    agent._middleware_registry.add_middleware(InvokeModelStage.Input, route_to_b)
+    result = agent("test")
+
+    assert result.message["content"][0]["text"] == "B"
+    model_a.stream.assert_not_called()

strands-py/tests/strands/vended_plugins/context_injector/test_plugin.py

@@ -40,6 +40,7 @@ def invoke_ctx(messages: list[dict], agent: Any) -> InvokeModelContext:
         tool_specs=[],
         tool_choice=None,
         invocation_state={},
+        model=MagicMock(),
     )