Python: Fix structured_output propagation in ClaudeAgent (#4137)

* Fix structured_output propagation in ClaudeAgent

Capture structured_output from ResultMessage in _get_stream() and
propagate it to AgentResponse.value via a custom finalizer. Previously
structured_output was silently discarded, making output_format unusable.

Fixes #4095

* Address review feedback: use value parameter instead of private properties

- Extend AgentResponse.from_updates() to accept optional value parameter
- Remove structured_output yield from _get_stream()
- Update _finalize_response() to pass value via public API
- Update streaming test to use get_final_response()

* Fix mypy errors: add value parameter to from_updates overloads

Add value parameter to both @overload signatures of
AgentResponse.from_updates() so mypy recognizes the argument.

---------

Co-authored-by: Amit Mukherjee <amimukherjee@microsoft.com>
Co-authored-by: Dmytro Struk <13853051+dmytrostruk@users.noreply.github.com>
This commit is contained in:
Amit Mukherjee
2026-02-23 12:45:02 -06:00
committed by GitHub
Unverified
parent 69eabcd1fc
commit 11628c3166
3 changed files with 183 additions and 2 deletions
@@ -618,12 +618,24 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
"""
response = ResponseStream(
self._get_stream(messages, session=session, options=options, **kwargs),
finalizer=AgentResponse.from_updates,
finalizer=self._finalize_response,
)
if stream:
return response
return response.get_final_response()
def _finalize_response(self, updates: Sequence[AgentResponseUpdate]) -> AgentResponse[Any]:
"""Build AgentResponse and propagate structured_output as value.
Args:
updates: The collected stream updates.
Returns:
An AgentResponse with structured_output set as value if present.
"""
structured_output = getattr(self, "_structured_output", None)
return AgentResponse.from_updates(updates, value=structured_output)
async def _get_stream(
self,
messages: AgentRunInputs | None = None,
@@ -647,6 +659,7 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
await self._apply_runtime_options(dict(options) if options else None)
session_id: str | None = None
structured_output: Any = None
await self._client.query(prompt)
async for message in self._client.receive_response():
@@ -700,7 +713,11 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
error_msg = message.result or "Unknown error from Claude API"
raise AgentException(f"Claude API error: {error_msg}")
session_id = message.session_id
structured_output = message.structured_output
# Update session with session ID
if session_id:
session.service_session_id = session_id
# Store structured output for the finalizer
self._structured_output = structured_output