fix(policy_engine): keep discarded stream rewrites out of the applied-guardrails header

A streaming step whose rewrite the executor threw away (a tool-call rewrite, a text rewrite
the translation cannot write back, or one the adapter refused) still marked its guardrail as
applied, so the header claimed an output the client never received. The step now returns
right after releasing the original chunks, which leaves the header as the merge base sent it
This commit is contained in:
mateo-berri 2026-09-08 16:48:22 -07:00
parent ea427e33d8
commit ce7cec1a36
2 changed files with 7 additions and 4 deletions

View file

@ -387,7 +387,8 @@ class PipelineExecutor:
yet (a tool-call rewrite, a text rewrite on a translation without write-back, or one
the translation or adapter refused with ``UndeliverableStreamRewrite``) is discarded:
the buffered chunks go back to the originals and the step passes, so the client gets
the stream the merge base sent. The response an earlier step's translation stored under
the stream the merge base sent, and the guardrail stays out of the applied-guardrails
header since its output never reached the client. The response an earlier step's translation stored under
``request_data["response"]`` is dropped first, so this step's hook sees the stream as
the steps before it left it."""
scanner: Final = (
@ -419,9 +420,10 @@ class PipelineExecutor:
)
except UndeliverableStreamRewrite:
_release_original_chunks(step.guardrail, streaming_chunks, originals)
else:
if observer.rewrote_tool_calls or (observer.rewrote_texts and not deliver_rewrites):
_release_original_chunks(step.guardrail, streaming_chunks, originals)
return
if observer.rewrote_tool_calls or (observer.rewrote_texts and not deliver_rewrites):
_release_original_chunks(step.guardrail, streaming_chunks, originals)
return
if not callback.records_own_guardrail_information:
add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=step.guardrail)

View file

@ -1146,6 +1146,7 @@ def _assert_passed_with_discard_warning(result, caplog):
assert result.terminal_action == "allow"
assert [step.outcome for step in result.step_results] == ["pass"]
assert any("'masker'" in record.getMessage() and "discarded" in record.getMessage() for record in caplog.records)
assert "masker" not in ((result.modified_data or {}).get("metadata") or {}).get("applied_guardrails", [])
@pytest.mark.asyncio