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:
@@ -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)
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user