feat(runner): add secret redaction middleware for AG-UI events#1041
Merged
Gkrumbach07 merged 5 commits intomainfrom Mar 26, 2026
Merged
feat(runner): add secret redaction middleware for AG-UI events#1041Gkrumbach07 merged 5 commits intomainfrom
Gkrumbach07 merged 5 commits intomainfrom
Conversation
Add a new middleware that scrubs secrets from outbound AG-UI events before they reach the frontend. Uses two complementary approaches: 1. Value-based: collects actual secret values from environment variables (ANTHROPIC_API_KEY, GITHUB_TOKEN, etc.) and replaces exact matches. 2. Pattern-based: reuses existing redact_secrets() regex patterns to catch tokens by format (GitHub PATs, Anthropic keys, credential URLs). The middleware is wired into all three bridges (Claude, Gemini CLI, LangGraph) as the innermost layer so tracing middleware (Langfuse) also receives already-redacted events. Also pre-compiles regex patterns in redact_secrets() for better performance since they now run on every text-bearing event. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughAdds a secret_redaction_middleware that collects environment-sourced secrets and pattern matches to redact sensitive strings from async event streams; exposes it in the middleware package, applies it to Claude/GeminiCLI/LangGraph adapter streams before tracing, refactors regex redaction, and adds tests. Changes
Sequence Diagram(s)sequenceDiagram
participant Adapter as Adapter.run()
participant Redaction as secret_redaction_middleware
participant Tracing as tracing_middleware
participant Consumer as Event Consumer
Adapter->>Redaction: async yield BaseEvent
Note over Redaction: collect env secrets\nexact-value replacements\npattern regex substitutions\nrecursive dict/list string redaction
Redaction->>Tracing: async yield redacted BaseEvent
Note over Tracing: attach tracing/span context
Tracing->>Consumer: async yield traced BaseEvent
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In
`@components/runners/ambient-runner/ambient_runner/middleware/secret_redaction.py`:
- Around line 100-141: The code only redacts strings and dicts, so secrets
inside lists (or other nested structures) slip through; implement a recursive
_redact_list(lst, secret_values) that walks lists and redacts
strings/dicts/lists and returns the original list if unchanged, then update
_redact_dict to call _redact_list when a value is a list (and use structural
equality to detect changes, not identity), and update the event handlers
(ToolCallResultEvent -> when content is not str but is dict or list; CustomEvent
-> when value is list) to call the appropriate redaction helper
(_redact_dict/_redact_list/_redact_text) and return model_copy only when the
redacted structure differs from the original. Ensure function names referenced:
_redact_text, _redact_dict, _redact_list, and the event handlers for
ToolCallResultEvent and CustomEvent.
In `@components/runners/ambient-runner/tests/test_secret_redaction.py`:
- Line 36: Replace hardcoded token-shaped literals like the dict entry
{"GITHUB_TOKEN": "ghp_abc123def456ghi789jkl012mno345pqr678st"} with a
runtime-generated test token so secret scanners won't flag the source; in the
test (e.g., inside test functions that set up env dicts / fixtures) create the
value by concatenating a fixed prefix (e.g., "ghp_") with a generated payload
(random choices or repeated characters, or join of safe fragments) and use that
generated string for the GITHUB_TOKEN key and any other token-shaped strings in
the file (including the other occurrences flagged), ensuring the produced token
matches the expected shape/length for the test but is not a hardcoded literal.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 555c2839-a006-44da-8815-f11cf9b955fe
📒 Files selected for processing (7)
components/runners/ambient-runner/ambient_runner/bridges/claude/bridge.pycomponents/runners/ambient-runner/ambient_runner/bridges/gemini_cli/bridge.pycomponents/runners/ambient-runner/ambient_runner/bridges/langgraph/bridge.pycomponents/runners/ambient-runner/ambient_runner/middleware/__init__.pycomponents/runners/ambient-runner/ambient_runner/middleware/secret_redaction.pycomponents/runners/ambient-runner/ambient_runner/platform/utils.pycomponents/runners/ambient-runner/tests/test_secret_redaction.py
components/runners/ambient-runner/ambient_runner/middleware/secret_redaction.py
Outdated
Show resolved
Hide resolved
components/runners/ambient-runner/tests/test_secret_redaction.py
Outdated
Show resolved
Hide resolved
- Fix `is not` → `!=` for recursive dict change detection - Lazy dict allocation: only copy when a change is found Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
♻️ Duplicate comments (2)
components/runners/ambient-runner/ambient_runner/middleware/secret_redaction.py (2)
124-140:⚠️ Potential issue | 🔴 Critical
_redact_dictalso misses list values.Nested lists within dicts won't be redacted. This is part of the same gap flagged above—adopting the unified
_redact_valuehelper would resolve both.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@components/runners/ambient-runner/ambient_runner/middleware/secret_redaction.py` around lines 124 - 140, _redact_dict currently only handles strings and dicts so list values inside dicts are left unredacted; update _redact_dict to delegate handling to a unified helper (e.g., _redact_value) that knows how to process str, dict, list, and other types recursively, then call that helper for each dict value (replace the isinstance(v, str)/dict branches with a single call to _redact_value(v, secret_values)). Ensure you preserve the copy-on-write behavior (only create result = dict(d) when a child was changed) and reference the existing function name _redact_dict and the suggested helper _redact_value so reviewers can find the changes.
100-119:⚠️ Potential issue | 🔴 CriticalSecrets in list payloads bypass redaction.
ToolCallResultEvent.content(lines 100-104) only redacts when content isstr, andCustomEvent.value(lines 111-119) only handlesstranddict. Secrets embedded inlistvalues pass through unredacted—a security gap.🔧 Proposed fix: Add _redact_value helper for recursive list/dict/str handling
+def _redact_value(value: object, secret_values: list[tuple[str, str]]) -> object: + """Recursively redact secrets in str/dict/list structures.""" + if isinstance(value, str): + return _redact_text(value, secret_values) + if isinstance(value, dict): + return _redact_dict(value, secret_values) + if isinstance(value, list): + result = None + for i, item in enumerate(value): + redacted_item = _redact_value(item, secret_values) + if redacted_item is not item: + if result is None: + result = list(value) + result[i] = redacted_item + return result if result is not None else value + return value + + def _redact_event(event: BaseEvent, secret_values: list[tuple[str, str]]) -> BaseEvent:Then update the event handlers:
elif isinstance(event, ToolCallResultEvent): - if isinstance(event.content, str): - redacted = _redact_text(event.content, secret_values) - if redacted != event.content: - return event.model_copy(update={"content": redacted}) + redacted_content = _redact_value(event.content, secret_values) + if redacted_content is not event.content: + return event.model_copy(update={"content": redacted_content})elif isinstance(event, CustomEvent): - if isinstance(event.value, str): - redacted = _redact_text(event.value, secret_values) - if redacted != event.value: - return event.model_copy(update={"value": redacted}) - elif isinstance(event.value, dict): - redacted_dict = _redact_dict(event.value, secret_values) - if redacted_dict != event.value: - return event.model_copy(update={"value": redacted_dict}) + redacted_value = _redact_value(event.value, secret_values) + if redacted_value is not event.value: + return event.model_copy(update={"value": redacted_value})🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@components/runners/ambient-runner/ambient_runner/middleware/secret_redaction.py` around lines 100 - 119, Current handlers only redact strings and dicts, letting secrets inside lists slip through; add a helper named _redact_value that recursively handles str, dict, and list (calling _redact_text for strings and _redact_dict for dicts, and recursing for list elements) and returns a redacted value of the same type; then update the branches handling ToolCallResultEvent.content, RunErrorEvent.message, and CustomEvent.value to call _redact_value(secret_values) (compare returned value to original and return event.model_copy(update={...}) when changed) so lists are properly scanned and redacted.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In
`@components/runners/ambient-runner/ambient_runner/middleware/secret_redaction.py`:
- Around line 124-140: _redact_dict currently only handles strings and dicts so
list values inside dicts are left unredacted; update _redact_dict to delegate
handling to a unified helper (e.g., _redact_value) that knows how to process
str, dict, list, and other types recursively, then call that helper for each
dict value (replace the isinstance(v, str)/dict branches with a single call to
_redact_value(v, secret_values)). Ensure you preserve the copy-on-write behavior
(only create result = dict(d) when a child was changed) and reference the
existing function name _redact_dict and the suggested helper _redact_value so
reviewers can find the changes.
- Around line 100-119: Current handlers only redact strings and dicts, letting
secrets inside lists slip through; add a helper named _redact_value that
recursively handles str, dict, and list (calling _redact_text for strings and
_redact_dict for dicts, and recursing for list elements) and returns a redacted
value of the same type; then update the branches handling
ToolCallResultEvent.content, RunErrorEvent.message, and CustomEvent.value to
call _redact_value(secret_values) (compare returned value to original and return
event.model_copy(update={...}) when changed) so lists are properly scanned and
redacted.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 2e5efb45-dc1e-46fd-a5e9-f068139e9a5e
📒 Files selected for processing (1)
components/runners/ambient-runner/ambient_runner/middleware/secret_redaction.py
Address CodeRabbit review: secrets inside list values (in CustomEvent and nested dicts) were not being redacted. Add _redact_value() helper that recursively handles str, dict, and list types with copy-on-write semantics. Simplifies _redact_dict and event handlers to use it. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In
`@components/runners/ambient-runner/ambient_runner/middleware/secret_redaction.py`:
- Around line 90-98: The two branches handling
TextMessageContentEvent/TextMessageChunkEvent and
ToolCallArgsEvent/ToolCallChunkEvent perform identical redaction on event.delta;
combine them by checking if event is an instance of the union of those four
classes and then call _redact_text(event.delta, secret_values), comparing to
event.delta and returning event.model_copy(update={"delta": redacted}) if
different. Update the conditional to reference the four types
(TextMessageContentEvent, TextMessageChunkEvent, ToolCallArgsEvent,
ToolCallChunkEvent) and keep the call sites to _redact_text and event.model_copy
unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 5cb01cd8-9d92-472e-9ced-cf35f7d64124
📒 Files selected for processing (2)
components/runners/ambient-runner/ambient_runner/middleware/secret_redaction.pycomponents/runners/ambient-runner/tests/test_secret_redaction.py
components/runners/ambient-runner/ambient_runner/middleware/secret_redaction.py
Outdated
Show resolved
Hide resolved
…cret_redaction.py Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In
`@components/runners/ambient-runner/ambient_runner/middleware/secret_redaction.py`:
- Around line 134-143: _react_dict currently only redacts values, leaving
sensitive data in string keys; update _redact_dict to also process keys: for
each k, v call _redact_value on k if isinstance(k, str) (or otherwise leave key
unchanged), call _redact_value on v as before, and only create a shallow copy
(result = dict(d)) when either key or value changes, then assign
result[redacted_key] = redacted_v; ensure you preserve original non-string keys
and maintain the function's return semantics (return original dict if nothing
changed). Use the existing _redact_value helper and the function name
_redact_dict to locate the change.
- Around line 118-119: The current branch calls _redact_text(value,
secret_values) and always returns its result, which can create a new-but-equal
str and break identity-based change detection; update the logic so that after
calling _redact_text you compare the result to the original value and return the
original value object if redacted == value, otherwise return the redacted string
(i.e., preserve identity of the original value when the redaction is a no-op).
Ensure this change is applied at the spot where value is an instance of str in
the same function surrounding the shown snippet and references _redact_text and
value.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: a6058b05-8c8a-467e-9173-3412eaf5528e
📒 Files selected for processing (1)
components/runners/ambient-runner/ambient_runner/middleware/secret_redaction.py
components/runners/ambient-runner/ambient_runner/middleware/secret_redaction.py
Outdated
Show resolved
Hide resolved
- Preserve string identity in _redact_value to maintain zero-copy fast path - Extend _redact_dict to also redact secrets in string keys - Replace hardcoded token-shaped literals in tests with runtime-generated helpers to avoid secret scanner false positives Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In
`@components/runners/ambient-runner/ambient_runner/middleware/secret_redaction.py`:
- Around line 150-167: The middleware secret_redaction_middleware currently only
redacts events that are yielded, letting exceptions raised by event_stream
propagate un-redacted; update secret_redaction_middleware to wrap the async
iteration in a try/except that catches exceptions from event_stream, uses
_collect_secret_values() and _redact_event(...) to sanitize any error payload
(e.g., create or convert to a RunErrorEvent-like object with secrets redacted)
before yielding or re-raising, so observers (obs.cleanup_on_error) never receive
raw exception messages; alternatively adjust tracing.py to redact exceptions
before calling obs.cleanup_on_error (sanitize the exception message or convert
into a redacted RunErrorEvent) and add an integration test that triggers an
adapter/tool error in the upstream stream to assert no secret strings reach
observability.
- Around line 135-147: _redact_dict currently overwrites entries when a redacted
key collides with an existing key or another redacted key; update _redact_dict
to detect collisions and preserve all values losslessly by merging into a list
rather than replacing: when redacted_k != k and result already contains
redacted_k (or when redacted_k equals an existing sibling key in the original
dict), if the existing value at redacted_k is not a list, replace it with a list
[existing_value, redacted_v], otherwise append redacted_v to the list; only
delete the original key (del result[k]) after moving/merging its value, and
ensure the same merge logic applies when multiple original keys redact to the
same redacted_k; add a regression test that feeds {"secret_value_here": 1,
"[REDACTED_MY_KEY]": 2} (and a case with two different keys that redact to the
same placeholder) and asserts both values are preserved (e.g., as a list) under
the redacted key.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 4fc2c080-373c-4071-9589-d81d265635ec
📒 Files selected for processing (2)
components/runners/ambient-runner/ambient_runner/middleware/secret_redaction.pycomponents/runners/ambient-runner/tests/test_secret_redaction.py
Comment on lines
+135
to
+147
| def _redact_dict(d: dict, secret_values: list[tuple[str, str]]) -> dict: | ||
| """Recursively redact keys and values in a dict. Returns original if unchanged.""" | ||
| result: dict | None = None | ||
| for k, v in d.items(): | ||
| redacted_k = _redact_value(k, secret_values) if isinstance(k, str) else k | ||
| redacted_v = _redact_value(v, secret_values) | ||
| if redacted_k is not k or redacted_v is not v: | ||
| if result is None: | ||
| result = dict(d) | ||
| if redacted_k is not k: | ||
| del result[k] | ||
| result[redacted_k] = redacted_v | ||
| return result if result is not None else d |
Contributor
There was a problem hiding this comment.
Handle redacted-key collisions without dropping fields.
Lines 145-146 can silently overwrite data when a redacted key matches an existing sibling key, or when two different keys redact to the same placeholder. A payload like {"secret_value_here": 1, "[REDACTED_MY_KEY]": 2} collapses to one entry. Please detect collisions and choose a lossless fallback instead of mutating into a colliding dict, and add a regression test for that case.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In
`@components/runners/ambient-runner/ambient_runner/middleware/secret_redaction.py`
around lines 135 - 147, _redact_dict currently overwrites entries when a
redacted key collides with an existing key or another redacted key; update
_redact_dict to detect collisions and preserve all values losslessly by merging
into a list rather than replacing: when redacted_k != k and result already
contains redacted_k (or when redacted_k equals an existing sibling key in the
original dict), if the existing value at redacted_k is not a list, replace it
with a list [existing_value, redacted_v], otherwise append redacted_v to the
list; only delete the original key (del result[k]) after moving/merging its
value, and ensure the same merge logic applies when multiple original keys
redact to the same redacted_k; add a regression test that feeds
{"secret_value_here": 1, "[REDACTED_MY_KEY]": 2} (and a case with two different
keys that redact to the same placeholder) and asserts both values are preserved
(e.g., as a list) under the redacted key.
Comment on lines
+150
to
+167
| async def secret_redaction_middleware( | ||
| event_stream: AsyncIterator[BaseEvent], | ||
| ) -> AsyncIterator[BaseEvent]: | ||
| """Wrap an AG-UI event stream with secret redaction. | ||
|
|
||
| Collects secret values from the current environment at invocation time | ||
| and scrubs them from all text-bearing events before yielding. | ||
|
|
||
| Args: | ||
| event_stream: The upstream event stream. | ||
|
|
||
| Yields: | ||
| Events with secrets redacted from text fields. | ||
| """ | ||
| secret_values = _collect_secret_values() | ||
|
|
||
| async for event in event_stream: | ||
| yield _redact_event(event, secret_values) |
Contributor
There was a problem hiding this comment.
Exceptions from the upstream stream still bypass redaction.
This middleware only transforms yielded events. If event_stream raises, components/runners/ambient-runner/ambient_runner/middleware/tracing.py:78-85 sends the original exception to obs.cleanup_on_error(exc) before any RunErrorEvent redaction runs, so secret-bearing adapter/tool errors can still leak to Langfuse. Catch and sanitize upstream exceptions before they reach tracing, or redact them in tracing before handing them to observability, and add an integration test for that failure path.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In
`@components/runners/ambient-runner/ambient_runner/middleware/secret_redaction.py`
around lines 150 - 167, The middleware secret_redaction_middleware currently
only redacts events that are yielded, letting exceptions raised by event_stream
propagate un-redacted; update secret_redaction_middleware to wrap the async
iteration in a try/except that catches exceptions from event_stream, uses
_collect_secret_values() and _redact_event(...) to sanitize any error payload
(e.g., create or convert to a RunErrorEvent-like object with secrets redacted)
before yielding or re-raising, so observers (obs.cleanup_on_error) never receive
raw exception messages; alternatively adjust tracing.py to redact exceptions
before calling obs.cleanup_on_error (sanitize the exception message or convert
into a redacted RunErrorEvent) and add an integration test that triggers an
adapter/tool error in the upstream stream to assert no secret strings reach
observability.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
secret_redaction_middlewarethat scrubs secrets from all text-bearing AG-UI events (assistant messages, tool call args/results, error messages, custom events) before they reach the frontendGITHUB_TOKEN,ANTHROPIC_API_KEY) with pattern-based redaction (regex for known token formats likeghp_*,sk-ant-*)redact_secrets()for better per-event performanceTest plan
🤖 Generated with Claude Code