Skip to content

Fix output type loss in AgentOperator's human-in-the-loop review - #70132

Open
ColtenOuO wants to merge 1 commit into
apache:mainfrom
ColtenOuO:fix-hitl-review-output-type-loss
Open

Fix output type loss in AgentOperator's human-in-the-loop review#70132
ColtenOuO wants to merge 1 commit into
apache:mainfrom
ColtenOuO:fix-hitl-review-output-type-loss

Conversation

@ColtenOuO

@ColtenOuO ColtenOuO commented Jul 20, 2026

Copy link
Copy Markdown
Contributor

Summary

AgentOperator's human-in-the-loop review path (enable_hitl_review=True) loses the original output type for any output_type that is neither str nor a Pydantic BaseModel (e.g. list[str], bool, dict).

The output is round-tripped through a string while it's shown to a human reviewer: HITLReviewMixin._to_string serializes it before review, and AgentOperator.execute() deserializes it back into output_type after approval. The serialization side used str(output), which produces a Python repr (e.g. "['tag-a', 'tag-b']"), not valid JSON.

The deserialization side then tried json.loads() on that repr, which always raised, and silently fell back to returning the raw repr string instead of the original list/bool/dict — a type change downstream tasks don't expect.

regenerate_with_feedback() (used when a reviewer requests changes) had the identical str(output) bug.

This is the same defect class already fixed for the require_approval=True path in #70075, which switched to TypeAdapter(...).dump_json(...). That fix wasn't applied to the separate enable_hitl_review=True path, so it regressed here.

Changes

  • Both directions of the review round-trip now live next to each other in
    utils/output_type.py: the existing rehydrate_pydantic_output() and a new
    dump_output_to_json().
  • dump_output_to_json() replaces the three copies of the serialization block
    that had drifted apart — HITLReviewMixin._to_string,
    AgentOperator.regenerate_with_feedback and LLMApprovalMixin.defer_for_approval.
    That drift is how the original bug survived Preserve output_type through human approval in LLM operators #70075.
  • AgentOperator.execute()'s HITL-review branch reuses rehydrate_pydantic_output
    instead of a bespoke json.loads/except fallback, so both review paths behave
    consistently.

Guarding an output_type pydantic cannot build a schema for

AgentOperator passes output_type straight through to Agent(...), so pydantic-ai's
[A, B] multi-output lists, ToolOutput/NativeOutput/PromptedOutput markers and
output functions all reach these helpers. TypeAdapter raises
PydanticSchemaGenerationError (a RuntimeError, not a ValidationError) or
AttributeError on those, which neither side's except clause caught:

  • rehydrate_pydantic_output now separates schema build from validation and falls
    back to plain json.loads. Without this an already-approved output is lost, after
    the model call and after the reviewer has signed off.
  • dump_output_to_json falls back to str(output). This one fires before the session
    XCom push, so without it the task dies before the reviewer is ever shown a review
    session. It is reachable through output functions, which only warn
    ("Falling back to unconstrained schema") for a non-schema-able return type and then
    run fine.

Behaviour change

With the default output_type=str, the approved string is no longer parsed as JSON:

Approved string Before After
'42' 42 '42'
'{"text": "hi"}' {'text': 'hi'} '{"text": "hi"}'
'true' True 'true'
'null' None 'null'

This matches the non-HITL path (which returns result.output unparsed) and the
require_approval path. str is the default output_type and enable_hitl_review
has shipped since 0.1.0, so a Dag relying on the implicit parse needs an explicit
output_type (e.g. output_type=int) or its own json.loads() downstream.


Was generative AI tooling used to co-author this PR?
  • Yes — Claude Code (Sonnet 5, Opus 5)

Generated-by: Claude Code (Opus 5) following the guidelines

@ColtenOuO
ColtenOuO requested a review from kaxil July 28, 2026 17:51
Comment thread providers/common/ai/src/airflow/providers/common/ai/mixins/hitl_review.py Outdated
Comment thread providers/common/ai/src/airflow/providers/common/ai/operators/agent.py Outdated
@ColtenOuO
ColtenOuO force-pushed the fix-hitl-review-output-type-loss branch from 5c61485 to a632d04 Compare August 2, 2026 10:04
Add dump_output_to_json() to utils/output_type.py, alongside
rehydrate_pydantic_output(), so both directions of the review round-trip
live in one module. Route all three serialization sites through it:
HITLReviewMixin._to_string, AgentOperator.regenerate_with_feedback and
LLMApprovalMixin.defer_for_approval.

Serialize non-str, non-BaseModel output with TypeAdapter(...).dump_json()
rather than str(output), which produced a Python repr instead of JSON, and
fall back to str(output) when the value's type has no pydantic schema.

Split schema build from validation in rehydrate_pydantic_output(). An
output_type that TypeAdapter cannot build a schema for now falls back to
json.loads() instead of raising PydanticSchemaGenerationError or
AttributeError, neither of which the previous except clause caught.

Use rehydrate_pydantic_output() in AgentOperator.execute()'s HITL branch in
place of the bespoke json.loads()/except fallback. With the default
output_type=str the approved string is no longer parsed as JSON.

Add tests for both schema-build guards, for dump_output_to_json(), and for
output_type=str with a JSON-parseable approved string, and fold the four
TestToString cases into a single delegation test.
@ColtenOuO
ColtenOuO force-pushed the fix-hitl-review-output-type-loss branch from a632d04 to 2adcc5a Compare August 2, 2026 19:03
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants