fix(orchestrator): apply prompt masking to pass-through request bodies - #2153
Open
cesaryuan wants to merge 4 commits into
Open
fix(orchestrator): apply prompt masking to pass-through request bodies#2153cesaryuan wants to merge 4 commits into
cesaryuan wants to merge 4 commits into
Conversation
-【规则结果】保留提示词保护命中的脱敏规则 -【请求处理】透传请求体沿用统一请求的脱敏结果 -【兼容性】保留旧保护接口并兼容现有调用方 -【数据完整性】避免重建请求体时丢失提供商专有字段 -【测试】补充命中规则记录与请求体脱敏验证
…tection-with-pass-through
Contributor
📝 WalkthroughWalkthroughPrompt protection now returns matched mask rules. Orchestration stores them in request state. Pass-through handling applies those rules to provider-native OpenAI, Anthropic, and Gemini request bodies while preserving other fields and rejecting unsafe layouts. ChangesPrompt protection result capture
Provider-native body patching
Pass-through integration
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant PromptProtection
participant PassThrough
participant Provider
Client->>PromptProtection: submit provider-native request
PromptProtection->>PromptProtection: mask unified request and record matched rules
PromptProtection->>PassThrough: pass original body and PromptProtectionMaskRules
PassThrough->>PassThrough: patch supported text fields
PassThrough->>Provider: send patched provider-native body
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Contributor
Greptile SummaryThe PR preserves prompt-protection masking when provider-native request bodies are passed through, while retaining unknown fields and provider-specific options.
Confidence Score: 5/5The PR appears safe to merge. No blocking failure remains; the previously reported OpenAI Responses Compact omission is fixed by routing its shared
|
| Filename | Overview |
|---|---|
| internal/server/orchestrator/pass_through_prompt_protection.go | Adds provider-format-aware prompt masking for native pass-through bodies, including the previously omitted OpenAI Responses Compact format. |
| internal/server/orchestrator/pass_through.go | Integrates masking into raw-body merging and safely retains the generated protected body when patching fails. |
| internal/server/orchestrator/prompt_protection.go | Captures matched mask rules while retaining compatibility with prompt protectors that only implement the original interface. |
| internal/server/biz/prompt_protection_request.go | Exposes the protected request and matched rules through ProtectWithResult, with Protect retained as a compatibility wrapper. |
| internal/server/orchestrator/state.go | Adds request-scoped storage for matched masking rules shared between inbound and outbound middleware. |
| internal/server/orchestrator/pass_through_prompt_protection_test.go | Covers native prompt layouts, role scoping, Gemini structured tool responses, and OpenAI Responses Compact pass-through preservation. |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart LR
A[Provider-native request] --> B[Unified inbound transformation]
B --> C[Prompt protection]
C --> D[Store matched mask rules]
D --> E[Generate protected outbound body]
E --> F{Pass-through enabled?}
F -- No --> G[Send generated protected body]
F -- Yes --> H[Patch model in original JSON]
H --> I{Raw prompt patch succeeds?}
I -- Yes --> J[Send patched native body]
I -- No --> G
Reviews (2): Last reviewed commit: "fix(orchestrator): preserve masked pass-..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Pull request overview
This PR fixes a security gap where prompt-protection mask transformations could be undone when PassThroughBody replays the original provider-native JSON. It threads matched mask rules through the orchestrator so pass-through can preserve the original JSON structure and provider-specific fields while still patching only the protected prompt text values for supported API formats.
Changes:
- Extend prompt protection to expose matched mask rules (
ProtectWithResult) while keeping the legacyProtectAPI for compatibility. - Persist matched mask rules in orchestrator state and apply them during pass-through body merging for supported provider-native request layouts.
- Add regression tests covering OpenAI Chat/Responses, Anthropic Messages, and Gemini GenerateContent (including structured
functionResponse.response).
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| internal/server/orchestrator/state.go | Stores matched prompt-protection mask rules in request state for later pass-through patching. |
| internal/server/orchestrator/prompt_protection.go | Captures matched rules from prompt protection and writes them into persistence state. |
| internal/server/orchestrator/prompt_protection_test.go | Adds coverage ensuring matched rules are recorded for pass-through. |
| internal/server/orchestrator/prompt_protecter.go | Introduces an optional result-bearing prompt protection interface to keep legacy protectors compatible. |
| internal/server/orchestrator/pass_through.go | Uses the new merge+patch path when pass-through is enabled. |
| internal/server/orchestrator/pass_through_test.go | Adds end-to-end tests validating pass-through patches masked prompt values and fails safely on unsupported layouts. |
| internal/server/orchestrator/pass_through_prompt_protection.go | Implements API-format-specific raw JSON prompt field discovery and rule-based patching. |
| internal/server/orchestrator/pass_through_prompt_protection_test.go | Adds focused tests for raw prompt patching, unsupported formats, and Gemini function response masking. |
| internal/server/biz/prompt_protection_request.go | Adds ProtectWithResult to return both protected request and matched rules; keeps Protect as wrapper. |
| internal/server/biz/prompt_protection_request_test.go | Updates tests to validate matched rules are returned. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
![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: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@internal/server/orchestrator/pass_through_prompt_protection.go`:
- Around line 122-151: Update canonicalPromptProtectionJSON and
replacePassThroughPromptText so JSON numbers retain their exact representation
via json.Decoder.UseNumber, and apply masking recursively only to string leaf
values, never object property names. Preserve all non-matching fields unchanged
when writing the complete function-response object, and add tests covering large
integers and property names that match masking rules.
🪄 Autofix
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: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 46789580-2eea-4570-a83e-052db24b41ee
📒 Files selected for processing (10)
internal/server/biz/prompt_protection_request.gointernal/server/biz/prompt_protection_request_test.gointernal/server/orchestrator/pass_through.gointernal/server/orchestrator/pass_through_prompt_protection.gointernal/server/orchestrator/pass_through_prompt_protection_test.gointernal/server/orchestrator/pass_through_test.gointernal/server/orchestrator/prompt_protecter.gointernal/server/orchestrator/prompt_protection.gointernal/server/orchestrator/prompt_protection_test.gointernal/server/orchestrator/state.go
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
internal/server/orchestrator/pass_through_prompt_protection_test.go (1)
116-144: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a scope-isolation case for Gemini function responses.
This test combines
UserandToolscopes on one rule. It confirms masking happens when both scopes are present, but it does not confirm that aUser-only rule leavesfunctionResponse.responsevalues unmasked. Scope isolation is a stated acceptance criterion for this feature. Add a case withScopes: []objects.PromptProtectionScope{objects.PromptProtectionScopeUser}and assert thatcontents.1.parts.0.functionResponse.response.tokenstays"secret-tool"after patching.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/server/orchestrator/pass_through_prompt_protection_test.go` around lines 116 - 144, Add a scope-isolation test alongside TestPatchPassThroughPromptProtectionMasksGeminiFunctionResponse using a rule scoped only to objects.PromptProtectionScopeUser, then assert the Gemini functionResponse response token remains "secret-tool" after patching while the user content continues to validate user-scope masking.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@internal/server/orchestrator/pass_through_prompt_protection_test.go`:
- Around line 116-144: Add a scope-isolation test alongside
TestPatchPassThroughPromptProtectionMasksGeminiFunctionResponse using a rule
scoped only to objects.PromptProtectionScopeUser, then assert the Gemini
functionResponse response token remains "secret-tool" after patching while the
user content continues to validate user-scope masking.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 4667107c-7c11-45e7-a2c4-f4e2cb3d74cd
📒 Files selected for processing (2)
internal/server/orchestrator/pass_through_prompt_protection.gointernal/server/orchestrator/pass_through_prompt_protection_test.go
🚧 Files skipped from review as they are similar to previous changes (1)
- internal/server/orchestrator/pass_through_prompt_protection.go
This was referenced Aug 9, 2026
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
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
maskresults whenPassThroughBodyreuses the original provider-native JSON request body.PromptProtecterimplementations.functionResponse.responsevalues while retaining every other value in that response object.Why this change matters
Pass-through exists to preserve the request body shape and fields that AxonHub does not own. It should not mean that the original prompt text is immutable after a security policy has changed it.
Before this change, a
maskrule could successfully modify the unified request, but raw-body replay would then replace that request with the original plaintext JSON. The upstream provider therefore received content the user expected to be masked.rejectwas unaffected because it aborts before egress, butmasksilently lost its protection on the pass-through path.Without this fix, users must choose between:
This PR makes both features composable: preserve the native request while applying the minimum security-required string changes.
Implementation details
PromptProtectionRuleService.ProtectWithResultreturns both the protected request and the mask rules that actually matched. The existingProtectmethod remains as a compatibility wrapper.functionResponse.responseis a structured object but is represented as serialized tool content in the unified request. The raw patch follows that same map unmarshal/marshal representation, replaces matching values under thetoolscope, and writes the result back as raw JSON so unrelated response fields survive.Validation
go test ./internal/server/orchestrator -run '^(TestApplyPassThroughBodyMasksPromptProtectedOpenAIContent|TestApplyPassThroughBodyKeepsProtectedBodyForUnsupportedPromptLayout|TestPatchPassThroughPromptProtection|TestProtectPrompts)' -count=1 -vfunctionResponse.response.token; both are masked andprovider_metaremains intact.go test ./internal/server/biz -run '^TestPromptProtectionRuleService_ProtectMask$' -count=1 -vgit diff --check upstream/unstable...HEADThe full
go test ./internal/server/orchestrator -count=1suite currently has pre-existing failures inTestDefaultSelector_SelectModelCandidates_Cache/cache_invalidated_when_channel_updatedandTestChatCompletionOrchestrator_Process_MinuteQuotaExceeded. They are outside this change's execution path; the focused coverage above passes.Closes #2152
Summary by CodeRabbit