feat(middleware): result handling, model-state isolation, and system-prompt fidelity

[zastrowm]

Description

Follow-up to the internal InvokeModelStage middleware (#2760), aligning the Python port with the TypeScript behavioral spec. Two behavior changes.

1. Result handling. The MiddlewareResult sentinel no longer rides the event stream. Instead the last yielded event is the result, matching the SDK convention (ModelStopReason ends stream_messages(), ToolResultEvent ends tool execution). The chain is now transparent. Output handlers — which transform the result — take/return a MiddlewareResult wrapper so we can extend it later without changing the signature.

# Before: integration site sniffs the stream for the sentinel
async for event in registry.invoke(InvokeModelStage, ctx, terminal):
    if isinstance(event, MiddlewareResult):
        model_result = event.value
    else:
        yield event

# After: events flow straight through; the last one is the result
async for event in registry.invoke(InvokeModelStage, ctx, terminal):
    yield event
stop_reason, message, usage, metrics = event["stop"]

2. Model-state isolation. Model runtime state is snapshotted before the chain and written back only on success, so middleware can't corrupt provider state (e.g. server-side response IDs) via the agent escape hatch.

# Before: terminal streamed against live agent state
model_state=agent._model_state,

# After: stream against a snapshot, write back after the chain succeeds
model_state_snapshot = copy.deepcopy(agent._model_state)
# ... terminal uses model_state_snapshot ...
agent._model_state = model_state_snapshot

Public API Changes

InvokeModelStage is internal, so no public surface changes. For middleware authors, the Output-phase handler shape changes to a MiddlewareResult wrapper:

def output_handler(result):  # result: MiddlewareResult
    stop_reason, message, usage, metrics = result.value["stop"]
    return result.replace(value=ModelStopReason(stop_reason="end_turn", message=message, usage=usage, metrics=metrics))

Wrap and Input handlers are unchanged.

Related Issues

• Middleware - Python #2735 (Middleware – Python tracking issue), builds on feat: add internal middleware system for InvokeModelStage #2760
• Behavioral spec: strands-ts/src/middleware/README.md

Documentation PR

None — internal stage; divergences documented in strands-py/src/strands/_middleware/README.md.

Type of Change

New feature

Testing

• I ran hatch run prepare

Checklist

• I have read the CONTRIBUTING document
• I have reviewed and understand every line of code in this PR, including any generated by AI tools, and I can explain why it works
• My change is focused and reasonably small; I have split unrelated work into separate PRs
• I have added any necessary tests that prove my fix is effective or my feature works
• I have updated the documentation accordingly
• I have added an appropriate example to the documentation to outline the feature, or no new docs are needed
• My changes generate no new warnings
• Any dependent changes have been merged and published

---

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

[github-actions]

Issue: event here is the async for loop variable, so if the middleware chain yields zero events it's never bound and this line raises UnboundLocalError: cannot access local variable 'event'. The previous code guarded this with an explicit if model_result is None: raise RuntimeError("Middleware chain did not produce a MiddlewareResult..."), which gave middleware authors a clear, actionable message. That guard is gone now.

This is reachable today — a Wrap middleware that short-circuits without yielding the result event triggers it. Reproduced:

async def silent(context, next_fn):
    if False:
        yield  # never yields the result event

agent._middleware_registry.add_middleware(InvokeModelStage, silent)
agent("test")  # -> UnboundLocalError at this line

Suggestion: Restore an explicit guard so a misbehaving middleware fails loudly instead of with a cryptic UnboundLocalError. For example, sentinel-initialize before the loop and check after:

last_event = None
async for event in agent._middleware_registry.invoke(...):
    last_event = event
    yield event

if last_event is None:
    raise RuntimeError(
        "Middleware chain did not yield a result event. "
        "Ensure middleware forwards events from next()."
    )
agent._model_state = model_state_snapshot
stop_reason, message, usage, metrics = last_event["stop"]

[github-actions]

Issue: If an Output handler forgets to return (returns None), this line raises AttributeError: 'NoneType' object has no attribute 'value', which doesn't point the author at the real mistake. The contract (handler must return a MiddlewareResult) is documented, but the failure mode is opaque.

Suggestion: Consider a targeted check after awaiting transformed to produce a clearer message, e.g.:

if not isinstance(transformed, MiddlewareResult):
    raise TypeError(
        f"Output handler must return a MiddlewareResult, got {type(transformed).__name__}"
    )
yield transformed.value

Optional/low-priority since this is an internal stage, but it's a cheap guardrail given the wrapper contract is new for middleware authors.

[github-actions]

Assessment: Comment

Clean, well-tested port. The result-encoding simplification (last event is the result) and the model-state snapshot/writeback isolation are both clear improvements over the sentinel approach, and the system-prompt content-block fidelity fix is a nice correctness catch. 97 unit tests pass locally and they assert on observable behavior (full messages, model-received state) rather than internal plumbing.

Review themes

• Error ergonomics: The previous explicit RuntimeError guard for "chain produced no result" was dropped; a misbehaving middleware now surfaces as an UnboundLocalError/AttributeError. Both are reachable and reproduced — see inline comments. Restoring loud, actionable failures keeps the "obvious path is the happy path" tenet for middleware authors.
• Model-state isolation: Snapshot-before / writeback-on-success is correct and the discard-on-error semantics match the TS spec. Tests cover before-next, after-next, deletions, and writeback well.
• Docs: README updates (both Python and TS spec) accurately describe the new wrapper contract and isolation behavior.

Nice work keeping the chain transparent while still giving Output handlers room to evolve via the MiddlewareResult wrapper.

[codecov]

Codecov Report

❌ Patch coverage is 96.00000% with 1 line in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
strands-py/src/strands/_middleware/registry.py	90.90%	0 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!