kaxil opened a new pull request, #76916: URL: https://github.com/apache/airflow/pull/76916
> This builds on #75412 (parked HITL review), which is still open. Until it merges, the diff here includes it: the first commit is a merge of that branch, and the second commit is this change. `AgentOperator(durable=True, enable_hitl_review=True)` raised `ValueError: durable=True and enable_hitl_review=True cannot be used together.` Durable is about to become the default for agent tasks that retry, so the review needs to work with it. On the code before this change, with only that check removed, the new tests fail three ways: - **Regenerations land in another journal.** `regenerate_with_feedback` built a durable agent but ran it outside the operator's `journal_scope`, so its steps went to `current_journal()`'s per-attempt fallback journal over the same storage. That journal numbers runs from zero again and deletes each run as soon as it succeeds, so a retry had no regeneration to replay. - **The journal outlives the task.** The review branch of `_complete_run` returned before the journal cleanup, so an approved task left its journal in the task state store. - **A rejection retries into the rejected answer.** `HITLRejectException` and `HITLMaxIterationsError` fail the task, and the retry replayed the first run in full, so the reviewer was shown the output they had just rejected. With this PR the first output and every regeneration are runs in the task's journal, numbered by review round, and the journal is deleted when the reviewer approves. On Airflow 3.3+ each request for changes is also recorded in the task state store, so when a regeneration fails or the review times out, the retry replays the first output, applies the recorded feedback to it and to each replayed regeneration without asking, and parks on the first round nobody has answered. If the review ends in Reject or `max_hitl_iterations` instead, the retry drops the recorded feedback and asks the model again for the first output, replaying the tool calls before that answer. ## Design rationale **Reviewer answers are kept in the task state store.** Tool approval under durable (the PR below this one) records the reviewer's answer as a journal step because the approval happens inside an agent run. A HITL review round happens between runs, so it has no journal position. Each request for changes is kept under `common_ai_hitl_review_answers` instead, with its round and a sha256 of the output text the reviewer read. The answer applies only to that exact output: if a replayed run produces something else (a tool set to run on every attempt returned new data, or the Dag changed), the reviewer is asked again from that round. A new answer for a round drops the answers for later rounds and deletes the journal runs that followed the old answer. **A reject starts the review over.** I considered replaying rounds 1..n-1 and asking the model again only for the rejected output, but a max-iterations retry would then start on its last allowed round, and a non-durable retry today starts a new review from the first output anyway. `DurableJournal.ask_model_again(0)` marks the first run so a retry replays it up to its last model request and asks again from there, and deletes the regeneration runs after it. **Runs are numbered by review round.** On Airflow 3.3+ a review round resumes in a new process, where the journal's run counter starts at 0 and would collide with the first output's run. The regeneration after round *n* is run *n* in every flow, set with `DurableJournal(..., first_run=n)` on resume. A regeneration also runs under the first run's id with `-review-<n>` appended, so capabilities that key state on the run id (pydantic-ai-harness `SpendLimits`, `StepPersistence`) see the same id on a retry. **A run that succeeds clears an earlier attempt's failure note.** Before this PR a run's journal was deleted right after it succeeded, so a failure note left by an earlier attempt never mattered. A review keeps the journal across processes. Without clearing the note, a later retry would re-run the steps after the old failure point and produce an output the recorded feedback no longer matches. `DurableRun.complete()` rewrites the run's entry when a note exists, the run diverged, or a step it ran is missing from the journal (a live-only tool or a refused write). On a clean first attempt it writes nothing. **A regeneration that ends on a tool needing approval fails the task, like the first run does.** Under `durable`, `DeferredToolRequests` is always an accepted output type, and tool approval is unavailable with `enable_hitl_review`. Without the check, approving such a round returned the pending-approval object as the task's result. ## Airflow 3.1 and 3.2 There the review polls instead of parking and there is no task state store, so feedback is not recorded. A retry replays the first output, asks the reviewer about every round again, and runs each regeneration again, model and tools included. Reject handling and journal cleanup work the same as on 3.3+. ## End to end Real model (`openai:gpt-5.4-mini` through a gateway), Airflow main in breeze, run on this branch's code (the final commit changed one docstring word after the run). The Dag has three `AgentOperator(durable=True, enable_hitl_review=True, retries=1)` tasks that fetch sales figures with a tool and write a summary. All answers went through the HITL Review plugin's REST API, the same calls the chat tab makes. - **Request changes, then approve:** two rounds, then approval. The task state store is empty after approval. - **A regeneration fails:** the reviewer requested changes twice. The second regeneration called a style-guide tool that failed once, so the task retried. Try 2 replayed the first run (5 steps) and the first regeneration (3 steps), applied both recorded answers, ran only the failed regeneration's tail live, and parked on round 3:   - **Reject:** try 1 logs that the review starts over; try 2 replays the fetch and asks the model again for the final answer, which came back reworded:   `fetch_sales` ran once per task across all tries. ## Known issues - Not run on released Airflow 3.3 or 3.4, or on 3.1/3.2: the shared breeze image was built from main. The 3.1/3.2 polling path is covered by unit tests that run without a task state store. - After two faults and a reject, the retry can replay the rejected answer: a run that ends earlier than an earlier attempt, without diverging, keeps that attempt's entries past its end, and if its entry also lost the last-model note (a refused write, or an attempt killed before the run completed), the reject's scan picks the stale model entry. - ToolLoggingCapability prints a tool-call group for a replayed tool call (visible in the reject screenshot), as it does for any durable replay. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
