fix: sanitize resolved auxiliary chat requests before dispatch

Extend the destination-boundary approach from PR #87840 to synchronous, asynchronous and streaming auxiliary dispatch, including prepared MoA requests. Keep native adapter replay intact.

Co-authored-by: siyoon <siyoon@friendli.ai>
This commit is contained in:
Teknium
2026-09-07 02:12:26 -07:00
parent e24c8499f2
commit ebe4e7bb44
3 changed files with 78 additions and 0 deletions
+9
View File
@@ -2430,6 +2430,9 @@ def _relay_sync_completion(
client: Any, kwargs: dict[str, Any], *, provider: str | None = None,
api_mode: str | None = None, create: Callable[[dict[str, Any]], Any] | None = None,
) -> Any:
from agent.auxiliary_wire import prepare_chat_messages
kwargs = prepare_chat_messages(client, kwargs)
callback = create or (lambda request: client.chat.completions.create(**request))
route = _relay_auxiliary_metadata(provider=provider, api_mode=api_mode)
# Isolate only the provider callback so the owning thread can unwind its lease/DB
@@ -2449,6 +2452,9 @@ async def _relay_async_completion(
client: Any, kwargs: dict[str, Any], *, provider: str | None = None,
api_mode: str | None = None, create: Callable[[dict[str, Any]], Any] | None = None,
) -> Any:
from agent.auxiliary_wire import prepare_chat_messages
kwargs = prepare_chat_messages(client, kwargs)
callback = create or (lambda request: client.chat.completions.create(**request))
route = _relay_auxiliary_metadata(provider=provider, api_mode=api_mode)
if route is None:
@@ -2464,6 +2470,9 @@ async def _relay_async_completion(
def _relay_sync_stream(
client: Any, kwargs: dict[str, Any], *, provider: str | None = None, api_mode: str | None = None
) -> Any:
from agent.auxiliary_wire import prepare_chat_messages
kwargs = prepare_chat_messages(client, kwargs)
route = _relay_auxiliary_metadata(provider=provider, api_mode=api_mode)
if route is None:
return client.chat.completions.create(**kwargs)
+20
View File
@@ -0,0 +1,20 @@
"""Message hygiene at the resolved auxiliary client boundary."""
from openai import AsyncOpenAI, OpenAI
from agent.transports.chat_completions import ChatCompletionsTransport
def prepare_chat_messages(client, kwargs: dict) -> dict:
"""Sanitize actual Chat Completions SDK requests, not native adapter replay.
Auxiliary and MoA callers can retain a prepared request before the virtual
transport sanitizes its copy. The resolved SDK client identifies the wire;
native Messages/Responses adapters must retain their reasoning sidecars.
"""
if not isinstance(client, (OpenAI, AsyncOpenAI)) or "messages" not in kwargs:
return kwargs
messages = ChatCompletionsTransport().convert_messages(
kwargs["messages"], model=kwargs.get("model")
)
return {**kwargs, "messages": messages}
@@ -0,0 +1,49 @@
"""The final auxiliary destination owns message sanitation, not virtual MoA."""
import asyncio
import copy
from types import SimpleNamespace
import pytest
from openai import AsyncOpenAI, OpenAI
from agent.auxiliary_client import _relay_async_completion, _relay_sync_completion
@pytest.mark.parametrize("async_mode", [False, True])
def test_auxiliary_chat_wire_sanitizes_without_mutating_history(async_mode):
history = [{"role": "assistant", "content": "answer", "_db_persisted": True,
"timestamp": 12, "tool_calls": [], "reasoning": "private"}]
original = copy.deepcopy(history)
kwargs = {"model": "fixture-model", "messages": history}
captured = []
if async_mode:
async def run():
async with AsyncOpenAI(api_key="fixture", base_url="http://127.0.0.1:1/v1") as client:
async def send(request):
captured.append(request)
await _relay_async_completion(client, kwargs, create=send)
asyncio.run(run())
else:
with OpenAI(api_key="fixture", base_url="http://127.0.0.1:1/v1") as client:
_relay_sync_completion(client, kwargs, create=captured.append)
assert captured[0]["messages"] == [{"role": "assistant", "content": "answer"}]
assert kwargs["messages"] == original
@pytest.mark.parametrize("async_mode", [False, True])
def test_auxiliary_native_adapters_keep_replay_and_tool_fields(async_mode):
history = [{"role": "assistant", "content": "", "_db_persisted": True,
"codex_reasoning_items": [{"type": "reasoning", "id": "rs_fixture"}],
"thinking_blocks": [{"type": "thinking", "thinking": "native", "signature": "sig"}],
"tool_calls": [{"id": "call_fixture", "type": "function",
"function": {"name": "fixture", "arguments": "{}"}}]}]
kwargs = {"model": "native", "messages": history}
captured = []
if async_mode:
async def send(request):
captured.append(request)
asyncio.run(_relay_async_completion(SimpleNamespace(), kwargs, create=send))
else:
_relay_sync_completion(SimpleNamespace(), kwargs, create=captured.append)
assert captured[0] is kwargs
assert captured[0]["messages"] is history