fix(guardrails): send only new messages since last assistant turn to CrowdStrike AIDR - #31974
Conversation
…CrowdStrike AIDR Previously, every guardrail request forwarded the full conversation history to CrowdStrike AIDR. In a multi-turn conversation this means every prior message gets re-scanned on every new call, even though those messages were already evaluated in earlier turns. CrowdStrike AIDR internally has a conversation boundary optimization in place for just this scenario (ref. <https://aidr-docs.crowdstrike.com/docs/aidr/apis#messages-array-optional---array-of-message-objects-containing-a-conversation-segment-with-the-ai-system>). However, it is nevertheless wasteful to send so much data to the API when only a subset of it will be processed. It also risks hitting the documented 1 MiB request size limit. So now we filter down to system messages plus either the messages after the last assistant turn, or the last assistant message itself when that is what is being guarded. We also preserve the original, full message history within the guardrail in order to stitch back any transformations.
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
Greptile SummaryThis PR optimizes CrowdStrike AIDR guardrail calls by filtering messages down to system messages plus only the messages since the last assistant turn (rather than forwarding the full conversation history on every request), and strips history entirely from response guardrail calls. It also refactors the API response from raw
Confidence Score: 4/5Safe to merge with the noted caveats; the core filtering logic is well-tested and the fallback paths are sound. The filtering and writeback logic is thoroughly covered by 33 passing tests spanning multi-turn scenarios, skip-system/tool interactions, Anthropic and OpenAI provider paths, and edge cases. The two substantive concerns are: (1) the response guardrail behavioral change from sending full history to assistant-only could affect users with context-sensitive AIDR policies — while the AIDR documentation supports this design, it was not gated by a feature flag; (2) the id()-based writeback silently falls back when object identities do not match, leaving structured_messages untransformed in skip-filter scenarios. The production implementation in crowdstrike_aidr.py, specifically _build_guard_input_for_response, _writeback_messages, and _call_crowdstrike_aidr_guard.
|
| Filename | Overview |
|---|---|
| litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py | Refactors message filtering to send only new messages (since last assistant turn) to CrowdStrike AIDR, adds typed Pydantic response parsing, and introduces writeback logic to stitch transformations back into the full message list. Logic is sound with minor concerns around assert usage and id()-based identity matching. |
| tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py | Updates existing tests to reflect the new message-filtering behavior and adds extensive new test coverage for multi-turn filtering, transformation stitching, skip-system/tool interactions, and Anthropic/OpenAI provider writeback paths. |
Reviews (1): Last reviewed commit: "fix(guardrails): send only new messages ..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| response = await self.async_handler.post(url=endpoint, json=payload, headers=headers) | ||
| assert response is not None | ||
| response.raise_for_status() |
There was a problem hiding this comment.
-O (optimize flag), turning this into a silent no-op. If async_handler.post ever returns None, the subsequent response.raise_for_status() would raise a confusing AttributeError instead of surfacing a meaningful error. A proper if guard is safer here.
| response = await self.async_handler.post(url=endpoint, json=payload, headers=headers) | |
| assert response is not None | |
| response.raise_for_status() | |
| response = await self.async_handler.post(url=endpoint, json=payload, headers=headers) | |
| if response is None: | |
| raise RuntimeError(f"CrowdStrike AIDR Guardrail ({hook_name}): HTTP client returned None for {endpoint}") | |
| response.raise_for_status() |
| def _writeback_messages( | ||
| self, | ||
| guard_output: Mapping[str, Any], | ||
| num_assistant_messages: int, | ||
| ) -> list[str]: | ||
| transformed_messages = guard_output.get("messages", []) | ||
| tail = transformed_messages[-num_assistant_messages:] if num_assistant_messages > 0 else [] | ||
| return [(_extract_text_from_content(msg.get("content")) if isinstance(msg, dict) else "") for msg in tail] | ||
| structured_messages: list[AllMessageValues], | ||
| guard_output: _GuardInput, | ||
| sent_indices: tuple[int, ...], | ||
| request_data: dict, | ||
| ) -> list[AllMessageValues] | None: | ||
| if effective_skip_system_message_for_guardrail(self) or effective_skip_tool_message_for_guardrail(self): | ||
| request_messages = request_data.get("messages") | ||
| full_messages = ( | ||
| cast("list[AllMessageValues]", request_messages) | ||
| if isinstance(request_messages, list) | ||
| else structured_messages | ||
| ) | ||
| else: | ||
| full_messages = structured_messages | ||
| return _redacted_messages(structured_messages, guard_output, sent_indices, full_messages) |
There was a problem hiding this comment.
_redacted_messages uses id() to match structured_messages entries against full_messages. When skip_system_message_in_guardrail or skip_tool_message_in_guardrail is enabled, full_messages comes from request_data["messages"] — a different list. If the guardrail framework ever deep-copies message dicts before building structured_messages, id(structured_messages[idx]) won't match any element in request_data["messages"], the guard not redactions.keys() <= {id(message) for message in full_messages} fires, and _writeback_messages returns None. The caller then falls back to the original, untransformed structured_messages, silently skipping the redaction writeback on the full message list. texts are still correctly transformed (covered by test_apply_guardrail_request_keeps_original_messages_when_skip_filters_differ), but structured_messages would be left unredacted in that scenario.
| def _build_guard_input_for_response(self, inputs: GenericGuardrailAPIInputs) -> _GuardInput: | ||
| output_texts: list[str] = inputs.get("texts", []) | ||
| if len(output_texts) == 0: | ||
| verbose_proxy_logger.warning("CrowdStrike AIDR Guardrail: No text in output response.") | ||
| return None | ||
|
|
||
| input_messages = request_data.get("messages", []) | ||
|
|
||
| return _GuardInput( | ||
| messages=[ | ||
| _Message(role=role, content=content) | ||
| for (role, content) in ( | ||
| (message["role"], _normalize_content(message.get("content"))) for message in input_messages | ||
| ) | ||
| if content is not None and len(content) > 0 | ||
| ] | ||
| + [_Message(role="assistant", content=text) for text in output_texts] | ||
| messages=[_Message(role="assistant", content=text) for text in output_texts], | ||
| tools=inputs.get("tools", []), | ||
| ) |
There was a problem hiding this comment.
Previously the response guardrail forwarded the full request history + assistant messages to CrowdStrike AIDR; now it sends only the assistant messages. The PR description notes that AIDR performs conversation-boundary optimization internally, so detection results should be equivalent. However, per the repo's policy on backwards-incompatible changes, a change that alters what data is sent to a third-party guardrail API should be gated by a user-controlled flag (e.g., send_full_history_on_response=True) so that operators relying on AIDR policies that inspect conversation context can opt out of the new behavior.
Rule Used: What: avoid backwards-incompatible changes without... (source)
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
![veria-ai[bot]](https://avatars.githubusercontent.com/in/1723776?s=60&v=4){kind=small}
| if not messages: | ||
| return _FilteredMessages([], ()) | ||
|
|
||
| if messages[-1]["role"] == "assistant": |
There was a problem hiding this comment.
High: Input guardrail bypass
An authenticated client can hide text from CrowdStrike by putting it in an earlier user message and appending any assistant message; this branch sends only system messages plus the final assistant message to AIDR, while the full client-supplied history is still forwarded to the model. For stateless chat-completions requests, do not assume messages before the last assistant turn were previously vetted unless the proxy tracks that checkpoint server-side; inspect all user/tool content in the current request or gate this pruning on server-owned conversation state.
There was a problem hiding this comment.
I will repost my previous response to this:
A client can send
[{role: "user", content: <blocked prompt>}, {role: "assistant", content: "ok"}, {role: "user", content: "continue"}], causing AIDR to inspect only the final benign user message while the model still receives the blocked prompt in its context.Clients probably should not be keeping blocked messages in the chat history. And even then, the CrowdStrike AIDR API is internally going to apply the publicly-documented conversation boundary optimization anyways, so worrying about it here in the gateway is moot.
PR overviewThis pull request changes the CrowdStrike AIDR guardrail integration so it prunes chat history and sends only messages considered new since the last assistant turn for inspection. The touched code is in the LiteLLM proxy guardrail hook that prepares chat-completion messages for CrowdStrike AIDR. There is one open security issue in the current implementation: a client can place uninspected user content before an assistant message while that same full history is still sent onward to the model. This creates a concrete input-guardrail bypass unless message pruning is tied to server-owned conversation state or all relevant user/tool content in the current request is inspected. No issues have been addressed yet, so the PR still carries a meaningful security risk. Open issues (1)
Fixed/addressed: 0 · PR risk: 7/10 |
3 participants
Relevant issues
Ports #31230 (originally authored by @kenany) onto a
litellm_-prefixed branch offlitellm_internal_stagingso CI runsLinear ticket
Pre-Submission checklist
@greptileaiand received a Confidence Score of at least 4/5 before requesting a maintainer reviewScreenshots / Proof of Fix
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py
Type
🐛 Bug Fix
Changes
Previously, every guardrail request forwarded the full conversation history to CrowdStrike AIDR. In a multi-turn conversation this means every prior message gets re-scanned on every new call, even though those messages were already evaluated in earlier turns.
CrowdStrike AIDR internally has a conversation boundary optimization in place for just this scenario (ref. https://aidr-docs.crowdstrike.com/docs/aidr/apis#messages-array-optional---array-of-message-objects-containing-a-conversation-segment-with-the-ai-system). However, it is nevertheless wasteful to send so much data to the API when only a subset of it will be processed. It also risks hitting the documented 1 MiB request size limit.
So now we filter down to system messages plus either the messages after the last assistant turn, or the last assistant message itself when that is what is being guarded. We also preserve the original, full message history within the guardrail in order to stitch back any transformations