fix(guardrails): scan reasoning_content in chat completions history

The chat completions guardrail handler only handed message content and
tool calls to guardrails, so an assistant message's reasoning_content
reached the provider exactly as the client sent it. Clients that replay
reasoning, which providers with thinking plus tool calls require, sent any
PII or blocked text in it straight past masking and moderation guardrails

reasoning_content is now extracted as one more text of its message and the
guardrail's rewrite is written back to that field, including on assistant
turns whose content is null
This commit is contained in:
lawrence 2026-10-23 22:52:44 +08:00
parent f2118fc767
commit 49b5a7d005
2 changed files with 55 additions and 4 deletions

View file

@ -19,7 +19,7 @@ import time
import uuid
from collections.abc import Iterator, Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Union, cast
from typing import TYPE_CHECKING, Any, Final, Literal, Union, cast
from typing_extensions import NotRequired, ReadOnly, TypedDict
@ -121,7 +121,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
texts_to_check: Final[list[str]] = []
images_to_check: Final[list[str]] = []
tool_calls_to_check: Final[list[ChatCompletionToolParam]] = []
text_task_mappings: Final[list[tuple[int, int | None]]] = []
text_task_mappings: Final[list[tuple[int, int | Literal["reasoning_content"] | None]]] = []
tool_call_task_mappings: Final[list[tuple[int, int]]] = []
# Step 1: Extract all text content, images, and tool calls
@ -280,7 +280,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
texts_to_check: list[str],
images_to_check: list[str],
tool_calls_to_check: list[ChatCompletionToolParam],
text_task_mappings: list[tuple[int, int | None]],
text_task_mappings: list[tuple[int, int | Literal["reasoning_content"] | None]],
tool_call_task_mappings: list[tuple[int, int]],
skip_system_message: bool = False,
skip_tool_message: bool = False,
@ -325,6 +325,11 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
elif isinstance(image_url, str):
images_to_check.append(image_url)
reasoning_content: Final = message.get("reasoning_content", None)
if isinstance(reasoning_content, str):
texts_to_check.append(reasoning_content)
text_task_mappings.append((msg_idx, "reasoning_content"))
# Extract tool calls (typically in assistant messages)
tool_calls: Final = message.get("tool_calls", None)
if tool_calls is not None and isinstance(tool_calls, list):
@ -338,7 +343,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
self,
messages: list[dict[str, object]],
responses: list[str],
task_mappings: list[tuple[int, int | None]],
task_mappings: list[tuple[int, int | Literal["reasoning_content"] | None]],
) -> None:
"""
Apply guardrail responses back to input message text content.
@ -348,6 +353,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
for task_idx, guardrail_response in enumerate(responses):
mapping = task_mappings[task_idx]
msg_idx = cast(int, mapping[0])
if mapping[1] == "reasoning_content":
message = messages[msg_idx]
message["reasoning_content"] = guardrail_response
continue
content_idx_optional = cast(int | None, mapping[1])
# Handle content

View file

@ -2365,3 +2365,45 @@ class TestStreamingScanKey:
for released in (text_only, finished_tool_call):
ended = handler.released_stream_as_ended(released)
assert len(ended) == len(released) and all(a is b for a, b in zip(ended, released, strict=True))
def _weather_tool_call(location: str) -> dict:
return {
"id": "call_1",
"type": "function",
"function": {"name": "get_weather", "arguments": json.dumps({"location": location})},
}
@pytest.mark.asyncio
@pytest.mark.parametrize("content", [None, "Let me check."])
async def test_history_reasoning_content_is_guardrailed_and_written_back(content: Optional[str]):
guardrail = MockGuardrail()
data = {
"messages": [
{"role": "user", "content": "weather in boston?"},
{
"role": "assistant",
"content": content,
"reasoning_content": "the user wants boston weather",
"tool_calls": [_weather_tool_call("boston")],
},
]
}
await OpenAIChatCompletionsHandler().process_input_messages(data, guardrail)
assert guardrail.last_inputs["texts"] == [
"weather in boston?",
*([] if content is None else [content]),
"the user wants boston weather",
]
assert data["messages"] == [
{"role": "user", "content": "WEATHER IN BOSTON?"},
{
"role": "assistant",
"content": None if content is None else content.upper(),
"reasoning_content": "THE USER WANTS BOSTON WEATHER",
"tool_calls": [_weather_tool_call("BOSTON")],
},
]