diff --git a/src/nooa/_llm_state.py b/src/nooa/_llm_state.py index 4bc22571a..3893716ce 100644 --- a/src/nooa/_llm_state.py +++ b/src/nooa/_llm_state.py @@ -10,7 +10,6 @@ from __future__ import annotations -import copy from typing import Any from uuid import uuid4 @@ -32,7 +31,9 @@ def __init__( replay_batch_size: int = 0, ): super().__init__(message) - self.llm_state = copy.deepcopy(llm_state) + # Borrow immutable event state through to the provider adapter. The + # adapter owns serialization and must not mutate caller input. + self.llm_state = llm_state self.reasoning = reasoning self.replay_batch_id = replay_batch_id self.replay_batch_size = replay_batch_size @@ -94,36 +95,17 @@ def demote_reasoning_text(message: dict[str, Any], reasoning: str | None) -> Non message["content"] = reasoning -def demote_chat_reasoning(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: - """Build public Chat messages, withholding opaque state by default.""" - prepared: list[dict[str, Any]] = [] - for original in messages: - state = carried_state(original) - reasoning = carried_reasoning(original) - message = copy.deepcopy(dict(original)) - message.pop(LLM_STATE_KEY, None) - demote_reasoning_text(message, reasoning) - if ( - state is not None - and not reasoning - and message.get("role") == "assistant" - and not message.get("content") - and not message.get("tool_calls") - ): - continue - prepared.append(message) - return prepared - - def demote_responses_batch( batch: list[dict[str, Any]], llm_state: dict[str, Any] | None, reasoning: str | None, ) -> list[dict[str, Any]]: - """Build public Responses items, withholding opaque state by default.""" - clean = [copy.deepcopy(dict(item)) for item in batch] - for item in clean: - item.pop(LLM_STATE_KEY, None) + """Demote reasoning into a cleaned, request-owned Responses batch. + + Replay preparation has already detached public data and removed untrusted + state. Reuse that batch instead of allocating another copy of the history. + """ + clean = batch if not reasoning: if ( diff --git a/src/nooa/context_blocks/models.py b/src/nooa/context_blocks/models.py index 17a34d8c5..fab34f738 100644 --- a/src/nooa/context_blocks/models.py +++ b/src/nooa/context_blocks/models.py @@ -13,7 +13,7 @@ from typing import Annotated, Any, Literal -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, ConfigDict, Field, SkipValidation # Import EventBase here (not forward ref) — possible because events.py imports # Role from roles.py, breaking the circular dependency. @@ -300,10 +300,12 @@ class RenderedMessage(BaseModel): default_factory=tuple, description="Complete ordered tool-call batch on an assistant turn", ) - llm_state: dict[str, Any] | None = Field( + llm_state: SkipValidation[dict[str, Any] | None] = Field( default=None, repr=False, - description="Opaque state carried only to the UnifiedLLM replay boundary", + description=( + "Borrowed immutable opaque state carried through the UnifiedLLM replay boundary" + ), ) reasoning: str | None = Field( default=None, diff --git a/src/nooa/nemo_relay_middleware.py b/src/nooa/nemo_relay_middleware.py index f3847574c..f11630d42 100644 --- a/src/nooa/nemo_relay_middleware.py +++ b/src/nooa/nemo_relay_middleware.py @@ -224,7 +224,12 @@ async def _wrapper(req: Any) -> Any: if hasattr(req, "content") and isinstance(req.content, dict): intercepted = req.content intercepted_msgs = intercepted.get("messages") - if intercepted_msgs is not None: + # Relay crosses a JSON boundary, so even a no-op request returns + # fresh plain dicts. Keep NOOA's original dict subclasses when the + # public payload is unchanged: their non-JSON sidecars carry opaque + # replay state to UnifiedLLM. A real intercept change replaces the + # originals and therefore drops private state fail-closed. + if intercepted_msgs is not None and intercepted_msgs != ctx.messages: ctx.messages = intercepted_msgs # Propagate any supported param changes from the intercept. for key in _PROPAGATABLE_LLM_PARAMS: diff --git a/src/nooa/runtime/middleware.py b/src/nooa/runtime/middleware.py index ea7b35a15..4e78d0f98 100644 --- a/src/nooa/runtime/middleware.py +++ b/src/nooa/runtime/middleware.py @@ -26,7 +26,7 @@ from collections.abc import Awaitable, Callable from typing import Any -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, SkipValidation from nooa.agent import Agent from nooa.events import ExecutionResult @@ -104,7 +104,10 @@ class LLMCallContext(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) - messages: list[dict[str, Any]] + # Preserve in-memory message subclasses used for non-serializable provider + # replay state. Messages are runtime-produced and middleware intentionally + # receives a mutable list, so coercing each dict provides no safety here. + messages: SkipValidation[list[dict[str, Any]]] params: dict[str, Any] = {} agent: Agent | None = None runtime: ActorRuntime | None = None diff --git a/src/nooa/tracing/_litellm_journal.py b/src/nooa/tracing/_litellm_journal.py index 0d68aa620..79933b885 100644 --- a/src/nooa/tracing/_litellm_journal.py +++ b/src/nooa/tracing/_litellm_journal.py @@ -36,6 +36,7 @@ from opentelemetry import trace as otel_trace from nooa.tracing._journal_builder import _encode_image +from nooa.tracing._secret_scrubber import scrub_value from nooa.tracing._session import get_session log = logging.getLogger(__name__) @@ -69,6 +70,12 @@ def _msg_to_dict(msg: Any) -> dict: return {"raw": str(msg)} +def _safe_msg_to_dict(msg: Any) -> dict: + """Normalize one provider message and remove issuer-only opaque state.""" + scrubbed, _ = scrub_value(_msg_to_dict(msg)) + return scrubbed if isinstance(scrubbed, dict) else {} + + def _extract_output_msgs(response_obj: Any) -> list[dict]: """Pull assistant messages out of a litellm completion response. @@ -546,7 +553,7 @@ def log_pre_api_call(self, model: str, messages: list, kwargs: dict) -> None: # No sideband — publish the raw messages as the skeleton # with no block refs. The viewer just uses their content # as-is, matching what the wire shows. - input_skeleton = [_msg_to_dict(m) for m in messages] + input_skeleton = [_safe_msg_to_dict(m) for m in messages] span_id = self._current_span_id() with self._lock: @@ -578,7 +585,9 @@ def log_success_event( # small and re-uses any hash that overlaps with messages the # agent will echo back on subsequent turns. output_blocks: dict[str, str] = {} - output_messages = [_skeleton_dict_message(m, output_blocks) for m in raw_output] + output_messages = [ + _skeleton_dict_message(_safe_msg_to_dict(m), output_blocks) for m in raw_output + ] if output_blocks: self._send_new_blocks(session_id, output_blocks) diff --git a/src/nooa/tracing/_secret_scrubber.py b/src/nooa/tracing/_secret_scrubber.py index 7d9f44038..afaef269f 100644 --- a/src/nooa/tracing/_secret_scrubber.py +++ b/src/nooa/tracing/_secret_scrubber.py @@ -15,6 +15,7 @@ provider.add_span_processor(SecretScrubSpanProcessor(inner_processor)) """ +import json import logging import re import threading @@ -47,6 +48,11 @@ } ) +# Provider-owned replay state is not a user credential, but it has the same +# telemetry rule: it may go back to its issuer and nowhere else. Provider +# adapters add their exact wire keys here as support is introduced. +_OPAQUE_PROVIDER_STATE_KEYS = frozenset({"encrypted_content", "nooa_llm_state"}) + def _is_sensitive_key(key: Any) -> bool: normalized = re.sub(r"[^a-z0-9]+", "_", str(key).lower()).strip("_") @@ -55,6 +61,19 @@ def _is_sensitive_key(key: Any) -> bool: ) +def _is_opaque_provider_state_key(key: Any) -> bool: + normalized = re.sub(r"[^a-z0-9]+", "_", str(key).lower()).strip("_") + return normalized in _OPAQUE_PROVIDER_STATE_KEYS + + +def _redact_key(key: Any) -> str | None: + if _is_sensitive_key(key): + return "sensitive_key" + if _is_opaque_provider_state_key(key): + return "opaque_provider_state" + return None + + # --------------------------------------------------------------------------- # Regex-based secret patterns — high precision, low false positives # --------------------------------------------------------------------------- @@ -214,19 +233,31 @@ def scrub_value(value: Any) -> tuple[Any, int]: number of secrets redacted within this value. """ if isinstance(value, str): - return scrub_string(value) + # OpenInference records LLM inputs as JSON string span attributes. + # Only parse objects/arrays. Matching the prefix avoids allocating a + # stripped copy of a potentially large tool output just to inspect it. + if not re.match(r"\s*[\[{]", value): + return scrub_string(value) + try: + decoded = json.loads(value) + except (json.JSONDecodeError, TypeError): + return scrub_string(value) + decoded, json_count = scrub_value(decoded) + if json_count: + return json.dumps(decoded, separators=(",", ":"), ensure_ascii=False), json_count + return value, 0 if isinstance(value, dict): - scrubbed: dict[Any, Any] = {} + scrubbed_mapping: dict[Any, Any] = {} count = 0 for key, item in value.items(): - if _is_sensitive_key(key): - scrubbed[key] = REDACTED - stats.record("sensitive_key") + if reason := _redact_key(key): + scrubbed_mapping[key] = REDACTED + stats.record(reason) count += 1 else: - scrubbed[key], n = scrub_value(item) + scrubbed_mapping[key], n = scrub_value(item) count += n - return scrubbed, count + return scrubbed_mapping, count if isinstance(value, (list, tuple)): new_items = [] count = 0 @@ -268,9 +299,9 @@ def on_end(self, span: ReadableSpan) -> None: redacted_count = 0 for key, value in span.attributes.items(): - if _is_sensitive_key(key): + if reason := _redact_key(key): new_value, n = REDACTED, 1 - stats.record("sensitive_key") + stats.record(reason) else: new_value, n = scrub_value(value) scrubbed[key] = new_value diff --git a/src/nooa/unifiedllm/fake.py b/src/nooa/unifiedllm/fake.py index 1b242d6f3..8ee613132 100644 --- a/src/nooa/unifiedllm/fake.py +++ b/src/nooa/unifiedllm/fake.py @@ -13,6 +13,8 @@ from nooa.unifiedllm.unifiedllm import LLMResponse, LLMUsage, Tool, ToolCall, UnifiedLLM +from .replay_state import prepare_chat_messages + class FakeLLMClient(UnifiedLLM): """ @@ -81,7 +83,8 @@ async def acall( """ async with self._lock: self.call_count += 1 - self.last_messages = messages + # A non-provider test client must never observe private replay state. + self.last_messages = prepare_chat_messages(messages, None) self.last_tools = tools # Return next response from queue, or empty response if none left @@ -108,7 +111,7 @@ def call( """Synchronous version of acall for UnifiedLLM compatibility.""" # For sync call, we don't need locking since tests are usually single-threaded self.call_count += 1 - self.last_messages = messages + self.last_messages = prepare_chat_messages(messages, None) self.last_tools = tools if self._response_queue: diff --git a/src/nooa/unifiedllm/registry.py b/src/nooa/unifiedllm/registry.py index 7ac724715..ecfb050e6 100644 --- a/src/nooa/unifiedllm/registry.py +++ b/src/nooa/unifiedllm/registry.py @@ -38,6 +38,9 @@ top_p: 1.0 # optional max_tokens: 4096 # optional drop_params: true # optional, defaults to true + store: false # optional Responses API control + include: # optional Responses API output fields + - reasoning.encrypted_content Set a model to ``null`` in a later layer to remove it. """ @@ -385,6 +388,8 @@ def get_llm_client(name: str, *, client_type: str | None = None, **overrides) -> "allowed_openai_params", "additional_drop_params", "extra_body", + "store", + "include", ): if key in config and key not in overrides: params[key] = config[key] diff --git a/src/nooa/unifiedllm/replay_state.py b/src/nooa/unifiedllm/replay_state.py new file mode 100644 index 000000000..0321e5e2b --- /dev/null +++ b/src/nooa/unifiedllm/replay_state.py @@ -0,0 +1,568 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Compatibility-scoped capture and replay of opaque OpenAI reasoning state. + +The event IR treats provider state as an opaque dictionary. This module is the +only code that opens its NOOA envelope or places the payload on provider wire +messages. Unknown providers and compatibility mismatches fail closed. +""" + +from __future__ import annotations + +import copy +import hashlib +import json +import logging +import os +from typing import Any, Literal, cast +from urllib.parse import urlsplit + +import litellm + +from nooa._llm_state import ( + LLM_STATE_KEY, + carried_reasoning, + carried_state, + demote_reasoning_text, + demote_responses_batch, +) + +logger = logging.getLogger(__name__) + +_STATE_VERSION = 2 +_CHAT_FORMAT = "litellm-chat" +_RESPONSES_FORMAT = "openai-responses" +_ENCRYPTED_REASONING_INCLUDE = "reasoning.encrypted_content" + + +def _field(value: Any, key: str, default: Any = None) -> Any: + return value.get(key, default) if isinstance(value, dict) else getattr(value, key, default) + + +def response_item_type(item: Any) -> str | None: + value = _field(item, "type") + return value if isinstance(value, str) else None + + +def unsupported_responses_parts(output: list[Any]) -> list[str]: + """Identify turn parts the canonical response cannot currently project.""" + unsupported: list[str] = [] + for item in output: + item_type = response_item_type(item) + if item_type not in {"reasoning", "function_call", "message"}: + unsupported.append(str(item_type)) + elif item_type == "message": + for block in _field(item, "content", []) or []: + block_type = response_item_type(block) + if block_type != "output_text": + unsupported.append(f"message.{block_type}") + return unsupported + + +def opaque_item(item: Any) -> Any: + """Detach one provider-owned item for durable storage.""" + # Inspect the type: permissive mocks/proxies synthesize arbitrary instance + # attributes and can otherwise recurse forever here. + if callable(getattr(type(item), "model_dump", None)): + return opaque_item(item.model_dump(exclude_none=True)) + if isinstance(item, dict): + return {key: opaque_item(value) for key, value in item.items()} + if isinstance(item, (list, tuple)): + return [opaque_item(value) for value in item] + return copy.deepcopy(item) + + +def _normalized_endpoint(value: Any) -> str: + if not isinstance(value, str) or not value: + return "default" + parsed = urlsplit(value) + if not parsed.scheme or not parsed.netloc: + return value.rstrip("/") + path = parsed.path.rstrip("/") + query = f"?{parsed.query}" if parsed.query else "" + return f"{parsed.scheme.lower()}://{parsed.netloc.lower()}{path}{query}" + + +def _uses_native_openai_endpoint(api_params: dict[str, Any]) -> bool: + endpoint = ( + api_params.get("api_base") + or api_params.get("base_url") + or getattr(litellm, "api_base", None) + or os.getenv("OPENAI_BASE_URL") + or os.getenv("OPENAI_API_BASE") + or "https://api.openai.com/v1" + ) + return _normalized_endpoint(endpoint) == "https://api.openai.com/v1" + + +def replay_scope( + model: str, + api_style: Literal["chat", "responses"], + params: dict[str, Any], +) -> str | None: + """Return a non-secret compatibility key for an opaque provider payload. + + LiteLLM resolves provider and model identity. Provider, API style, and exact + model are intentionally the whole key. Transport routes and authentication + do not change the provider wire format, so gateway or credential changes + must not silently disable capture or replay. + """ + configured_endpoint = params.get("api_base") or params.get("base_url") + try: + resolved_model, provider, _, _ = litellm.get_llm_provider( + model=model, + custom_llm_provider=params.get("custom_llm_provider"), + api_base=configured_endpoint, + ) + except Exception as exc: # noqa: BLE001 - unknown routes fail closed + logger.debug("Could not resolve opaque-state provider for %r: %s", model, exc) + return None + # This PR understands only OpenAI's encrypted reasoning wire formats. + # Other providers may use similarly named fields with different replay + # contracts; they remain fail-closed until their adapters opt in. + if provider not in {"openai", "azure"}: + return None + + digest = hashlib.sha256(resolved_model.encode()).hexdigest() + return f"{api_style}:{provider}:sha256:{digest}" + + +def _envelope(scope: str | None, state_format: str, payload: dict[str, Any]) -> dict | None: + if scope is None or not payload: + return None + return { + "version": _STATE_VERSION, + "scope": scope, + "format": state_format, + "payload": payload, + } + + +def _matching_payload(state: Any, scope: str | None, state_format: str) -> dict | None: + if isinstance(state, dict) and state.get("version") != _STATE_VERSION: + logger.warning( + "Ignoring opaque reasoning state with unsupported or legacy version %r; " + "portable reasoning text will be replayed instead.", + state.get("version"), + ) + return None + if ( + scope is None + or not isinstance(state, dict) + or state.get("version") != _STATE_VERSION + or state.get("scope") != scope + or state.get("format") != state_format + or not isinstance(state.get("payload"), dict) + ): + return None + # Event history owns this payload. Provider adapters may read and serialize + # it, but must not mutate caller input or require a per-request history copy. + return cast(dict[str, Any], state["payload"]) + + +def _fingerprint(value: Any) -> str: + """Hash validation-only public data without retaining a second copy of it.""" + digest = hashlib.sha256() + for chunk in json.JSONEncoder(sort_keys=True, separators=(",", ":")).iterencode(value): + digest.update(chunk.encode()) + return digest.hexdigest() + + +def _valid_fingerprint(value: Any) -> bool: + return ( + isinstance(value, str) and len(value) == 64 and all(c in "0123456789abcdef" for c in value) + ) + + +def _valid_reasoning_items(items: Any) -> bool: + """Validate the OpenAI wire container, not the encrypted contents.""" + return isinstance(items, list) and all( + isinstance(item, dict) + and response_item_type(item) == "reasoning" + and isinstance(item.get("encrypted_content"), str) + and bool(item["encrypted_content"]) + for item in items + ) + + +def _chat_public_carrier(message: Any) -> str | None: + """Fingerprint the public assistant data to which opaque Chat state is bound.""" + if _field(message, "role") != "assistant": + return None + raw_calls = _field(message, "tool_calls") + if raw_calls is None: + raw_calls = [] + if not isinstance(raw_calls, list): + return None + + calls: list[dict[str, str]] = [] + for call in raw_calls: + function = _field(call, "function") + call_id = _field(call, "id") + name = _field(function, "name") + arguments = _field(function, "arguments") + if isinstance(arguments, dict): + arguments = json.dumps(arguments) + if not all(isinstance(value, str) for value in (call_id, name, arguments)): + return None + calls.append({"id": call_id, "name": name, "arguments": arguments}) + + if len({call["id"] for call in calls}) != len(calls): + return None + + content = _field(message, "content") + if content is not None and not isinstance(content, str): + return None + content = (content or None) if calls else (content or "") + return _fingerprint({"content": content, "tool_calls": calls}) + + +def _valid_chat_payload(payload: dict[str, Any]) -> bool: + if set(payload) - {"reasoning_items", "carrier", "state_only"}: + return False + items = payload.get("reasoning_items") + carrier = payload.get("carrier") + if not items or not _valid_reasoning_items(items) or not _valid_fingerprint(carrier): + return False + if "state_only" in payload and payload["state_only"] is not True: + return False + return (payload.get("state_only") is True) == ( + carrier == _fingerprint({"content": "", "tool_calls": []}) + ) + + +def capture_chat_state(message: Any, scope: str | None) -> dict | None: + items = _field(message, "reasoning_items") + if not isinstance(items, list) or not items: + return None + carrier = _chat_public_carrier(message) + if carrier is None: + logger.warning("Discarding OpenAI Chat reasoning state with a malformed carrier.") + return None + payload: dict[str, Any] = { + "reasoning_items": [opaque_item(item) for item in items], + "carrier": carrier, + } + if not _field(message, "content") and not _field(message, "tool_calls"): + payload["state_only"] = True + if not _valid_chat_payload(payload): + logger.warning("Discarding malformed OpenAI Chat reasoning state.") + return None + return _envelope(scope, _CHAT_FORMAT, payload) + + +def responses_message_text(item: Any) -> str: + content = _field(item, "content", []) + if not isinstance(content, list): + return "" + texts = [ + text + for block in content + if _field(block, "type") == "output_text" + and isinstance((text := _field(block, "text")), str) + ] + return "".join(texts) + + +def responses_output_text(output: list[Any]) -> str: + """Match the SDK's separator-free aggregation of assistant text blocks.""" + return "".join( + responses_message_text(item) for item in output if response_item_type(item) == "message" + ) + + +def _responses_call_slot(item: Any) -> dict[str, Any] | None: + # Calls replay from LLMResponse; arguments here only detect edits. Message + # slots still retain text to reconstruct provider boundaries and phases. + call_id = _field(item, "call_id") + name = _field(item, "name") or "" + arguments = _field(item, "arguments") or "" + if isinstance(arguments, dict): + arguments = json.dumps(arguments) + if not all(isinstance(value, str) for value in (call_id, name, arguments)): + return None + return { + "type": "function_call", + "call_id": call_id, + "name": name, + "arguments_sha256": _fingerprint(arguments), + } + + +def _valid_responses_payload(payload: dict[str, Any]) -> bool: + """Validate NOOA-owned structure while leaving encrypted item contents opaque.""" + if set(payload) - {"items", "order", "state_only"}: + return False + items = payload.get("items") + order = payload.get("order") + if not isinstance(items, list) or not isinstance(order, list) or not order: + return False + if not _valid_reasoning_items(items): + return False + + indexes: list[int] = [] + carriers: list[dict[str, Any]] = [] + for slot in order: + if not isinstance(slot, dict): + return False + slot_type = slot.get("type") + if slot_type == "reasoning": + if set(slot) != {"type", "index"}: + return False + index = slot.get("index") + if not isinstance(index, int) or not 0 <= index < len(items): + return False + indexes.append(index) + elif slot_type == "function_call": + if set(slot) != {"type", "call_id", "name", "arguments_sha256"}: + return False + if not all(isinstance(slot.get(key), str) for key in ("call_id", "name")): + return False + if not _valid_fingerprint(slot.get("arguments_sha256")): + return False + carriers.append(slot) + elif slot_type == "message": + if ( + set(slot) - {"type", "content", "phase"} + or not isinstance(slot.get("content"), str) + or ("phase" in slot and slot["phase"] not in ("commentary", "final_answer")) + ): + return False + carriers.append(slot) + else: + return False + + state_only = payload.get("state_only") + if state_only not in (None, True): + return False + return ( + bool(items or carriers != _flatten_responses_carriers(carriers)) + and indexes == list(range(len(items))) + and len({slot["call_id"] for slot in carriers if slot["type"] == "function_call"}) + == sum(slot["type"] == "function_call" for slot in carriers) + and (state_only is True) == (not carriers) + ) + + +def _public_responses_carriers(items: list[dict[str, Any]]) -> list[dict[str, Any]]: + carriers: list[dict[str, Any]] = [] + for item in items: + if response_item_type(item) == "function_call": + slot = _responses_call_slot(item) + if slot is not None: + carriers.append(slot) + elif item.get("role") == "assistant": + content = item.get("content", "") + if isinstance(content, str): + carriers.append({"type": "message", "content": content}) + return carriers + + +def _flatten_responses_carriers(carriers: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Project provider message boundaries into the canonical LLMResponse view.""" + calls = [slot for slot in carriers if slot["type"] == "function_call"] + messages = [slot for slot in carriers if slot["type"] == "message"] + text = "".join(slot["content"] for slot in messages) + return ( + [{"type": "message", "content": text}] if text or (messages and not calls) else [] + ) + calls + + +def capture_responses_state(output: list[Any], scope: str | None) -> dict | None: + unsupported = unsupported_responses_parts(output) + if unsupported: + logger.warning( + "Cannot retain OpenAI Responses state: unsupported turn parts %s would " + "be omitted during replay.", + unsupported, + ) + return None + items: list[Any] = [] + order: list[dict[str, Any]] = [] + has_public_carrier = False + reasoning_items = [item for item in output if response_item_type(item) == "reasoning"] + summary_only = any(_field(item, "encrypted_content") is None for item in reasoning_items) + if summary_only and any( + _field(item, "encrypted_content") is not None for item in reasoning_items + ): + logger.warning( + "Responses reasoning contains items without encrypted content; replaying " + "portable summaries instead of an incomplete opaque reasoning sequence." + ) + for item in output: + item_type = response_item_type(item) + if item_type == "reasoning": + # A summary-only reasoning item has no opaque data to retain. + if summary_only: + continue + order.append({"type": "reasoning", "index": len(items)}) + items.append(opaque_item(item)) + elif item_type == "function_call": + slot = _responses_call_slot(item) + if slot is not None: + order.append(slot) + has_public_carrier = True + elif item_type == "message": + slot = {"type": "message", "content": responses_message_text(item)} + phase = _field(item, "phase") + if phase is not None: + slot["phase"] = phase + order.append(slot) + has_public_carrier = True + carriers = [slot for slot in order if slot["type"] != "reasoning"] + # Keep provider message boundaries/phase when the canonical flat text and + # calls alone cannot reproduce them, even without encrypted reasoning. + if not items and carriers == _flatten_responses_carriers(carriers): + return None + payload: dict[str, Any] = {"items": items, "order": order} + if not has_public_carrier: + payload["state_only"] = True + if not _valid_responses_payload(payload): + logger.warning("Discarding malformed OpenAI Responses reasoning state.") + return None + return _envelope(scope, _RESPONSES_FORMAT, payload) + + +def prepare_chat_messages(messages: list[dict[str, Any]], scope: str | None) -> list[dict]: + """Strip private/raw state and restore only a matching Chat payload.""" + prepared: list[dict[str, Any]] = [] + for original in messages: + state = carried_state(original) + reasoning = carried_reasoning(original) + message = dict(original) + message.pop(LLM_STATE_KEY, None) + message.pop("reasoning_items", None) + message = copy.deepcopy(message) + payload = _matching_payload(state, scope, _CHAT_FORMAT) + restored = False + if payload is not None and not _valid_chat_payload(payload): + logger.warning( + "Opaque Chat reasoning state is malformed; portable reasoning text " + "will be replayed instead." + ) + elif payload is not None and _chat_public_carrier(message) != payload["carrier"]: + logger.warning( + "Opaque Chat reasoning state was not replayed because its public assistant " + "carrier changed; portable reasoning text will be replayed instead." + ) + elif payload is not None: + message["reasoning_items"] = payload["reasoning_items"] + restored = True + if not restored: + demote_reasoning_text(message, reasoning) + if ( + not restored + and state is not None + and not reasoning + and message.get("role") == "assistant" + and not message.get("content") + and not message.get("tool_calls") + ): + continue + prepared.append(message) + return prepared + + +def _clean_responses_batch(batch: Any) -> list[dict[str, Any]]: + if not isinstance(batch, list): + return [] + clean: list[dict[str, Any]] = [] + for original in batch: + if not isinstance(original, dict) or response_item_type(original) == "reasoning": + continue + # Strip private sidecars and rejected wire state before detaching public + # content. Opaque state is borrowed separately after compatibility checks. + item = dict(original) + item.pop(LLM_STATE_KEY, None) + item.pop("reasoning_items", None) + clean.append(copy.deepcopy(item)) + return clean + + +def prepare_responses_batch( + batch: Any, + state: Any, + scope: str | None, + reasoning: str | None = None, +) -> list[dict[str, Any]]: + """Restore a matching Responses payload among its public turn carriers.""" + clean = _clean_responses_batch(batch) + payload = _matching_payload(state, scope, _RESPONSES_FORMAT) + if payload is None: + return demote_responses_batch(clean, state, reasoning) + items = payload.get("items") + order = payload.get("order") + if not _valid_responses_payload(payload): + logger.warning( + "Opaque Responses reasoning state is malformed; portable reasoning text " + "will be replayed instead." + ) + return demote_responses_batch(clean, state, reasoning) + assert isinstance(items, list) + assert isinstance(order, list) + if payload.get("state_only") is True: + if clean != [{"role": "assistant", "content": ""}]: + logger.warning( + "Opaque Responses reasoning state was not replayed because its empty " + "public carrier changed." + ) + return demote_responses_batch(clean, state, reasoning) + return cast(list[dict[str, Any]], items) + + expected_carriers = [ + {key: value for key, value in slot.items() if key != "phase"} + for slot in order + if slot["type"] != "reasoning" + ] + public_carriers = _public_responses_carriers(clean) + if public_carriers not in (expected_carriers, _flatten_responses_carriers(expected_carriers)): + logger.warning( + "Opaque Responses reasoning state was not replayed because its public " + "carriers changed; portable reasoning text will be replayed instead." + ) + return demote_responses_batch(clean, state, reasoning) + + calls = iter(item for item in clean if item.get("type") == "function_call") + messages = iter(item for item in clean if item.get("role") == "assistant") + exact_messages = public_carriers == expected_carriers + replay: list[dict[str, Any]] = [] + for slot in order: + if slot["type"] == "reasoning": + replay.append(items[slot["index"]]) + elif slot["type"] == "function_call": + replay.append(next(calls)) + else: + message = ( + next(messages) + if exact_messages + else {"role": "assistant", "content": slot["content"]} + ) + if "phase" in slot: + message["phase"] = slot["phase"] + replay.append(message) + replay.extend( + item + for item in clean + if item.get("role") != "assistant" and item.get("type") != "function_call" + ) + # Message-phase metadata alone must not suppress portable reasoning text. + return demote_responses_batch(replay, None, reasoning) if not items else replay + + +def add_encrypted_reasoning_include(api_params: dict[str, Any], scope: str | None) -> None: + """Request OpenAI encrypted reasoning only on endpoints known to support it.""" + configured = api_params.get("include") + include = list(configured) if isinstance(configured, (list, tuple, set)) else [] + if configured is not None and not isinstance(configured, (list, tuple, set)): + include.append(configured) + if _ENCRYPTED_REASONING_INCLUDE in include: + api_params["include"] = include + return + if scope and scope.startswith("responses:azure:"): + include.append(_ENCRYPTED_REASONING_INCLUDE) + elif ( + scope and scope.startswith("responses:openai:") and _uses_native_openai_endpoint(api_params) + ): + include.append(_ENCRYPTED_REASONING_INCLUDE) + if include: + api_params["include"] = include diff --git a/src/nooa/unifiedllm/unifiedllm.py b/src/nooa/unifiedllm/unifiedllm.py index 1c9e02288..27d5714e9 100644 --- a/src/nooa/unifiedllm/unifiedllm.py +++ b/src/nooa/unifiedllm/unifiedllm.py @@ -22,11 +22,10 @@ carried_reasoning, carried_replay_batch, carried_state, - demote_chat_reasoning, - demote_responses_batch, ) from nooa.llm_types import LLMResponse, LLMUsage, ToolCall +from . import replay_state from .http_config import HttpConfig from .retry import EmptyContentError, sync_retry, with_retry from .retry_config import RetryConfig @@ -1148,6 +1147,29 @@ def __init__(self, model: str, **config): # concrete subclasses; guarded here so base helpers stay safe. self._http: _ClientHttp | None = None + def _effective_model(self, call_config: dict[str, Any]) -> str: + """Return the model this individual request will actually dispatch.""" + model = call_config.get("model", self.model) + if not isinstance(model, str) or not model: + raise ValueError("model must be a non-empty string") + return model + + @staticmethod + def _validate_request_config(name: str, call_config: dict[str, Any]) -> None: + """Keep provider payloads and routing on their validated top-level paths.""" + if name in call_config: + raise ValueError( + f"{name!r} is managed by UnifiedLLM; pass conversation data through " + "the messages argument" + ) + extra_body = call_config.get("extra_body") + if isinstance(extra_body, dict) and (reserved := {name, "model"} & set(extra_body)): + fields = ", ".join(repr(field) for field in sorted(reserved)) + raise ValueError( + f"extra_body may not override reserved field(s) {fields}; pass model at " + "the top level and conversation data through the messages argument" + ) + def close(self) -> None: """Release this client's sync HTTP resources (its own httpx clients).""" if self._http is not None: @@ -1195,7 +1217,11 @@ def _inject_cache_control_on_content(msg: dict) -> None: msg["cache_control"] = {"type": "ephemeral"} def _inject_cache_control( - self, messages: list[dict[str, Any]], injection_points: list[dict[str, Any]] + self, + messages: list[dict[str, Any]], + injection_points: list[dict[str, Any]], + *, + model: str | None = None, ) -> list[dict[str, Any]]: """Add cache_control to designated messages for prompt caching. @@ -1227,7 +1253,7 @@ def _inject_cache_control( value "last" restricts marking to only the last message of that role. Returns: - A deep copy of messages with cache_control injected at breakpoints. + A copy-on-write view with only breakpoint messages copied. """ if not injection_points: return messages @@ -1246,38 +1272,65 @@ def _inject_cache_control( if not roles_to_cache_all and not roles_to_cache_last: return messages - messages = [copy.deepcopy(msg) for msg in messages] + prepared = messages + copied: set[int] = set() + + def copy_message(index: int, *, copy_last_content_block: bool = False) -> dict[str, Any]: + nonlocal prepared + if index not in copied: + if prepared is messages: + prepared = list(messages) + # A shallow copy is enough here and preserves private replay + # metadata carried by our dict subclass. Nested content is + # detached below before it is changed. + prepared[index] = copy.copy(messages[index]) + copied.add(index) + message = prepared[index] + content = message.get("content") + if copy_last_content_block and isinstance(content, list) and content: + blocks = list(content) + if isinstance(blocks[-1], dict): + blocks[-1] = dict(blocks[-1]) + message["content"] = blocks + return message # Map role names to native Responses API type equivalents _ROLE_TO_TYPE = {"tool": "function_call_output"} - for msg in messages: + for index, original in enumerate(messages): + msg = original role = msg.get("role") if role and role in roles_to_cache_all: + msg = copy_message(index) msg["cache_control"] = {"type": "ephemeral"} elif not role: # Native Responses format: match by type equivalent msg_type = msg.get("type") for r, t in _ROLE_TO_TYPE.items(): if t == msg_type and r in roles_to_cache_all: + msg = copy_message(index) msg["cache_control"] = {"type": "ephemeral"} break # Anthropic needs cache_control on a content block (parts form); other providers # reject a content list on non-user roles, so mark at the message level instead. - anthropic = _is_anthropic_model(self.model) + anthropic = _is_anthropic_model(model or self.model) for role in roles_to_cache_last: # Search for matching messages by role OR by equivalent native type native_type = _ROLE_TO_TYPE.get(role) - for msg in reversed(messages): - if msg.get("role") == role or (native_type and msg.get("type") == native_type): + for index in range(len(messages) - 1, -1, -1): + original = messages[index] + if original.get("role") == role or ( + native_type and original.get("type") == native_type + ): + msg = copy_message(index, copy_last_content_block=anthropic) if anthropic: self._inject_cache_control_on_content(msg) else: msg["cache_control"] = {"type": "ephemeral"} break - return messages + return prepared def count_tokens(self, text: str) -> int: """Count tokens using model-appropriate tokenizer. @@ -1540,51 +1593,10 @@ def _extract_reasoning_and_usage(raw_response: Any) -> tuple[str | None, LLMUsag return reasoning, usage -def _completion_llm_state(message: Any) -> dict[str, Any] | None: - """Capture opaque Chat-Completions reasoning without duplicating the turn. - - LiteLLM 1.97 bridges GPT-5.4+ function-tool requests to the Responses API - and returns the encrypted reasoning state as ``reasoning_items`` on the - chat-shaped message. Public content and tool calls already have canonical - fields on :class:`LLMResponse`; only the opaque provider state belongs here. - """ - reasoning_items = getattr(message, "reasoning_items", None) - if not reasoning_items: - return None - return {"reasoning_items": [_opaque_item(item) for item in reasoning_items]} - - -def _opaque_item(item: Any) -> Any: - """Detach one provider-owned item for durable storage.""" - if hasattr(item, "model_dump"): - return item.model_dump(exclude_none=True) - return copy.deepcopy(item) - - def _item_field(item: Any, name: str) -> Any: return item.get(name) if isinstance(item, dict) else getattr(item, name, None) -def _responses_llm_state(output: list[Any]) -> dict[str, Any] | None: - """Capture opaque Responses items plus lightweight ordering anchors.""" - reasoning_items: list[Any] = [] - order: list[dict[str, Any]] = [] - for item in output: - item_type = _item_field(item, "type") - if item_type == "reasoning": - order.append({"type": "reasoning", "index": len(reasoning_items)}) - reasoning_items.append(_opaque_item(item)) - elif item_type == "function_call": - call_id = _item_field(item, "call_id") - if isinstance(call_id, str): - order.append({"type": "function_call", "call_id": call_id}) - elif item_type == "message": - order.append({"type": "message"}) - if not reasoning_items: - return None - return {"items": reasoning_items, "order": order} - - def _extract_xml_tool_calls(content: str) -> list["ToolCall"]: """Extract tool calls from XML format used by Nemotron/NIM models. @@ -1769,6 +1781,19 @@ def _convert_tool_to_schema(self, tool: Tool) -> dict[str, Any]: }, } + def _completion_http_client(self, call_config: dict[str, Any], *, is_async: bool) -> Any: + """Reuse the owned transport only while its constructor routing still applies.""" + routing_fields = ("api_base", "base_url", "api_key", "custom_llm_provider") + if self._effective_model(call_config) != self.model or any( + call_config.get(key) != self.config.get(key) for key in routing_fields + ): + # LiteLLM uses a supplied OpenAI SDK client's bound URL/key, ignoring + # the corresponding call parameters. Let it build the correct client + # for overrides; these calls use LiteLLM's default HTTP pool settings. + return None + assert self._http is not None + return self._http.async_client if is_async else self._http.sync_client + def call( self, messages: list[dict[str, Any]], @@ -1784,7 +1809,11 @@ def call( If retry_config.retry_on_empty_content is True, will retry when the model returns empty content but has reasoning_content (common with some reasoning models). """ - messages = demote_chat_reasoning(messages) + call_config = {**self.config, **kwargs} + self._validate_request_config("messages", call_config) + effective_model = self._effective_model(call_config) + state_scope = replay_state.replay_scope(effective_model, "chat", call_config) + messages = replay_state.prepare_chat_messages(messages, state_scope) # Inject cache_control at the message level for prompt caching cache_points = ( @@ -1792,13 +1821,15 @@ def call( if cache_control_injection_points is None else cache_control_injection_points ) - prepared_messages = self._inject_cache_control(messages, cache_points) + prepared_messages = self._inject_cache_control( + messages, cache_points, model=effective_model + ) api_params = { "model": self.model, - "messages": prepared_messages, **self.config, **kwargs, + "messages": prepared_messages, } if tools: @@ -1807,13 +1838,13 @@ def call( if output_model is not None: api_params["response_format"] = _maybe_sanitize_response_format( - self.model, output_model + effective_model, output_model ) # Bedrock/Anthropic reject messages with tool_call blocks when tools= is absent. if ( "tools" not in api_params - and _needs_dummy_tool(self.model) + and _needs_dummy_tool(effective_model) and _messages_have_tool_calls(prepared_messages) ): api_params["tools"] = [_DUMMY_TOOL_SCHEMA] @@ -1826,10 +1857,9 @@ def call( retry_on_empty = self.retry_config.retry_on_empty_content if self.retry_config else False - http_client = self._http - assert http_client is not None - if http_client.sync_client is not None: - api_params.setdefault("client", http_client.sync_client) + http_client = self._completion_http_client(call_config, is_async=False) + if http_client is not None: + api_params.setdefault("client", http_client) def _make_call(): raw_response = _collect_sync(litellm.completion(**api_params)) @@ -1843,7 +1873,7 @@ def _make_call(): return raw_response # Track LLM call for debugging (visible via SIGUSR2 if nooa debug handler installed) - with _track_llm_call(model=self.model, endpoint=self.config.get("api_base")): + with _track_llm_call(model=effective_model, endpoint=self.config.get("api_base")): raw_response = ( sync_retry(_make_call, config=self.retry_config) if self.retry_config @@ -1854,7 +1884,7 @@ def _make_call(): if usage: _record_llm_metric("token_usage", usage) _update_token_calibration( - self.model, prepared_messages, usage, tools=api_params.get("tools") + effective_model, prepared_messages, usage, tools=api_params.get("tools") ) raw_tool_calls = cast( list[Any] | None, @@ -1877,7 +1907,7 @@ def _make_call(): ), reasoning=reasoning, usage=usage, - llm_state=_completion_llm_state(response_message), + llm_state=replay_state.capture_chat_state(response_message, state_scope), ) text_content = raw_response.choices[0].message.content or "" # type: ignore[union-attr] @@ -1896,7 +1926,9 @@ def _make_call(): ), reasoning=reasoning, usage=usage, - llm_state=_completion_llm_state(raw_response.choices[0].message), + llm_state=replay_state.capture_chat_state( + raw_response.choices[0].message, state_scope + ), ) if output_model: @@ -1918,7 +1950,9 @@ def _make_call(): finish_reason=_map_completion_finish_reason(raw_response), reasoning=reasoning, usage=usage, - llm_state=_completion_llm_state(raw_response.choices[0].message), + llm_state=replay_state.capture_chat_state( + raw_response.choices[0].message, state_scope + ), ) return LLMResponse( @@ -1928,7 +1962,7 @@ def _make_call(): finish_reason=_map_completion_finish_reason(raw_response), reasoning=reasoning, usage=usage, - llm_state=_completion_llm_state(raw_response.choices[0].message), + llm_state=replay_state.capture_chat_state(raw_response.choices[0].message, state_scope), ) async def acall( @@ -1946,7 +1980,11 @@ async def acall( If retry_config.retry_on_empty_content is True, will retry when the model returns empty content but has reasoning_content (common with some reasoning models). """ - messages = demote_chat_reasoning(messages) + call_config = {**self.config, **kwargs} + self._validate_request_config("messages", call_config) + effective_model = self._effective_model(call_config) + state_scope = replay_state.replay_scope(effective_model, "chat", call_config) + messages = replay_state.prepare_chat_messages(messages, state_scope) # Inject cache_control at the message level for prompt caching cache_points = ( @@ -1954,13 +1992,15 @@ async def acall( if cache_control_injection_points is None else cache_control_injection_points ) - prepared_messages = self._inject_cache_control(messages, cache_points) + prepared_messages = self._inject_cache_control( + messages, cache_points, model=effective_model + ) api_params = { "model": self.model, - "messages": prepared_messages, **self.config, **kwargs, + "messages": prepared_messages, } if tools: @@ -1969,13 +2009,13 @@ async def acall( if output_model is not None: api_params["response_format"] = _maybe_sanitize_response_format( - self.model, output_model + effective_model, output_model ) # Bedrock/Anthropic reject messages with tool_call blocks when tools= is absent. if ( "tools" not in api_params - and _needs_dummy_tool(self.model) + and _needs_dummy_tool(effective_model) and _messages_have_tool_calls(prepared_messages) ): api_params["tools"] = [_DUMMY_TOOL_SCHEMA] @@ -1988,10 +2028,9 @@ async def acall( retry_on_empty = self.retry_config.retry_on_empty_content if self.retry_config else False - http_client = self._http - assert http_client is not None - if http_client.async_client is not None: - api_params.setdefault("client", http_client.async_client) + http_client = self._completion_http_client(call_config, is_async=True) + if http_client is not None: + api_params.setdefault("client", http_client) async def _make_call(): raw_response = await _collect_async(await _litellm_acompletion(api_params)) @@ -2005,7 +2044,7 @@ async def _make_call(): return raw_response # Track LLM call for debugging (visible via SIGUSR2 if nooa debug handler installed) - with _track_llm_call(model=self.model, endpoint=self.config.get("api_base")): + with _track_llm_call(model=effective_model, endpoint=self.config.get("api_base")): raw_response = ( await with_retry(_make_call, config=self.retry_config) if self.retry_config @@ -2016,7 +2055,7 @@ async def _make_call(): if usage: _record_llm_metric("token_usage", usage) _update_token_calibration( - self.model, prepared_messages, usage, tools=api_params.get("tools") + effective_model, prepared_messages, usage, tools=api_params.get("tools") ) raw_tool_calls = cast( list[Any] | None, @@ -2041,7 +2080,7 @@ async def _make_call(): ), reasoning=reasoning, usage=usage, - llm_state=_completion_llm_state(response_message), + llm_state=replay_state.capture_chat_state(response_message, state_scope), ) text_content = raw_response.choices[0].message.content or "" # type: ignore[union-attr] @@ -2060,7 +2099,9 @@ async def _make_call(): ), reasoning=reasoning, usage=usage, - llm_state=_completion_llm_state(raw_response.choices[0].message), + llm_state=replay_state.capture_chat_state( + raw_response.choices[0].message, state_scope + ), ) if output_model: @@ -2082,7 +2123,9 @@ async def _make_call(): finish_reason=_map_completion_finish_reason(raw_response), reasoning=reasoning, usage=usage, - llm_state=_completion_llm_state(raw_response.choices[0].message), + llm_state=replay_state.capture_chat_state( + raw_response.choices[0].message, state_scope + ), ) return LLMResponse( @@ -2092,7 +2135,7 @@ async def _make_call(): finish_reason=_map_completion_finish_reason(raw_response), reasoning=reasoning, usage=usage, - llm_state=_completion_llm_state(raw_response.choices[0].message), + llm_state=replay_state.capture_chat_state(raw_response.choices[0].message, state_scope), ) @@ -2278,23 +2321,29 @@ def call( # OpenAIGPTConfig.remove_cache_control_flag strip — so leaving the marker # on OpenAI/Azure/NIM Responses calls triggers a 400 "Unknown parameter: # input[N].cache_control" at the gateway. - if _is_anthropic_model(self.model): + call_config = {**self.config, **kwargs} + self._validate_request_config("input", call_config) + effective_model = self._effective_model(call_config) + state_scope = replay_state.replay_scope(effective_model, "responses", call_config) + if _is_anthropic_model(effective_model): cache_points = ( self.cache_control_injection_points if cache_control_injection_points is None else cache_control_injection_points ) - prepared_messages = self._inject_cache_control(messages, cache_points) + prepared_messages = self._inject_cache_control( + messages, cache_points, model=effective_model + ) else: prepared_messages = messages - input_messages, instructions = self._transform_messages(prepared_messages) + input_messages, instructions = self._transform_messages(prepared_messages, state_scope) api_params = { "model": self.model, - "input": input_messages, "truncation": "disabled", **self.config, **kwargs, + "input": input_messages, } if instructions: @@ -2311,8 +2360,7 @@ def call( if output_model is not None: api_params.update(_responses_output_params(output_model)) - if reasoning := self.config.get("reasoning"): - api_params["reasoning"] = reasoning + replay_state.add_encrypted_reasoning_include(api_params, state_scope) http_client = self._http assert http_client is not None @@ -2323,7 +2371,7 @@ def _make_call(): return cast("litellm.ResponsesAPIResponse", litellm.responses(**api_params)) # Track LLM call for debugging (visible via SIGUSR2 if nooa debug handler installed) - with _track_llm_call(model=self.model, endpoint=self.config.get("api_base")): + with _track_llm_call(model=effective_model, endpoint=self.config.get("api_base")): raw_response = ( sync_retry(_make_call, config=self.retry_config) if self.retry_config @@ -2332,14 +2380,22 @@ def _make_call(): usage = LLMUsage.from_provider(getattr(raw_response, "usage", None)) if usage: - _update_token_calibration(self.model, messages, usage, tools=api_params.get("tools")) + _update_token_calibration( + effective_model, messages, usage, tools=api_params.get("tools") + ) output: list[Any] = raw_response.output # type: ignore[assignment] - raw_tool_calls = [item for item in output if item.type == "function_call"] + raw_tool_calls = [ + item for item in output if replay_state.response_item_type(item) == "function_call" + ] if raw_tool_calls: tool_calls = [ - ToolCall(id=tc.call_id or "", name=tc.name or "", arguments=tc.arguments or "") + ToolCall( + id=_item_field(tc, "call_id") or "", + name=_item_field(tc, "name") or "", + arguments=_item_field(tc, "arguments") or "", + ) for tc in raw_tool_calls ] @@ -2352,7 +2408,7 @@ def _make_call(): ), reasoning=None, # Responses API doesn't have reasoning usage=usage, - llm_state=_responses_llm_state(output), + llm_state=replay_state.capture_responses_state(output, state_scope), ) text_content = self._extract_text_from_output(raw_response) @@ -2369,7 +2425,7 @@ def _make_call(): finish_reason=_map_responses_finish_reason(raw_response), reasoning=None, usage=usage, - llm_state=_responses_llm_state(output), + llm_state=replay_state.capture_responses_state(output, state_scope), ) return LLMResponse( @@ -2379,7 +2435,7 @@ def _make_call(): finish_reason=_map_responses_finish_reason(raw_response), reasoning=None, usage=usage, - llm_state=_responses_llm_state(output), + llm_state=replay_state.capture_responses_state(output, state_scope), ) async def acall( @@ -2398,23 +2454,29 @@ async def acall( """ # See ResponsesClient.call for why cache_control injection is gated on # Anthropic models only. - if _is_anthropic_model(self.model): + call_config = {**self.config, **kwargs} + self._validate_request_config("input", call_config) + effective_model = self._effective_model(call_config) + state_scope = replay_state.replay_scope(effective_model, "responses", call_config) + if _is_anthropic_model(effective_model): cache_points = ( self.cache_control_injection_points if cache_control_injection_points is None else cache_control_injection_points ) - prepared_messages = self._inject_cache_control(messages, cache_points) + prepared_messages = self._inject_cache_control( + messages, cache_points, model=effective_model + ) else: prepared_messages = messages - input_messages, instructions = self._transform_messages(prepared_messages) + input_messages, instructions = self._transform_messages(prepared_messages, state_scope) api_params = { "model": self.model, - "input": input_messages, "truncation": "disabled", **self.config, **kwargs, + "input": input_messages, } if instructions: @@ -2431,8 +2493,7 @@ async def acall( if output_model is not None: api_params.update(_responses_output_params(output_model)) - if reasoning := self.config.get("reasoning"): - api_params["reasoning"] = reasoning + replay_state.add_encrypted_reasoning_include(api_params, state_scope) http_client = self._http assert http_client is not None @@ -2443,7 +2504,7 @@ async def _make_call(): return cast("litellm.ResponsesAPIResponse", await litellm.aresponses(**api_params)) # Track LLM call for debugging (visible via SIGUSR2 if nooa debug handler installed) - with _track_llm_call(model=self.model, endpoint=self.config.get("api_base")): + with _track_llm_call(model=effective_model, endpoint=self.config.get("api_base")): raw_response = ( await with_retry(_make_call, config=self.retry_config) if self.retry_config @@ -2452,14 +2513,22 @@ async def _make_call(): usage = LLMUsage.from_provider(getattr(raw_response, "usage", None)) if usage: - _update_token_calibration(self.model, messages, usage, tools=api_params.get("tools")) + _update_token_calibration( + effective_model, messages, usage, tools=api_params.get("tools") + ) output: list[Any] = raw_response.output # type: ignore[assignment] - raw_tool_calls = [item for item in output if item.type == "function_call"] + raw_tool_calls = [ + item for item in output if replay_state.response_item_type(item) == "function_call" + ] if raw_tool_calls: tool_calls = [ - ToolCall(id=tc.call_id or "", name=tc.name or "", arguments=tc.arguments or "") + ToolCall( + id=_item_field(tc, "call_id") or "", + name=_item_field(tc, "name") or "", + arguments=_item_field(tc, "arguments") or "", + ) for tc in raw_tool_calls ] @@ -2472,7 +2541,7 @@ async def _make_call(): ), reasoning=None, usage=usage, - llm_state=_responses_llm_state(output), + llm_state=replay_state.capture_responses_state(output, state_scope), ) text_content = self._extract_text_from_output(raw_response) @@ -2489,7 +2558,7 @@ async def _make_call(): finish_reason=_map_responses_finish_reason(raw_response), reasoning=None, usage=usage, - llm_state=_responses_llm_state(output), + llm_state=replay_state.capture_responses_state(output, state_scope), ) return LLMResponse( @@ -2499,11 +2568,13 @@ async def _make_call(): finish_reason=_map_responses_finish_reason(raw_response), reasoning=None, usage=usage, - llm_state=_responses_llm_state(output), + llm_state=replay_state.capture_responses_state(output, state_scope), ) def _transform_messages( - self, messages: list[dict[str, Any]] + self, + messages: list[dict[str, Any]], + state_scope: str | None = None, ) -> tuple[list[dict[str, Any]], str | None]: """Transform messages to Responses API format and extract instructions. @@ -2525,10 +2596,8 @@ def _transform_messages( if skip_batch_items: skip_batch_items -= 1 continue - state = copy.deepcopy(carried_state(original)) + state = carried_state(original) reasoning = carried_reasoning(original) - msg = copy.deepcopy(dict(original)) - msg.pop(LLM_STATE_KEY, None) batch_info = carried_replay_batch(original) if batch_info is not None and (state is not None or reasoning is not None): @@ -2537,15 +2606,27 @@ def _transform_messages( if len(candidates) == batch_size and all( carried_replay_batch(item) == (batch_id, batch_size) for item in candidates ): - batch = [copy.deepcopy(dict(item)) for item in candidates] - transformed.extend(demote_responses_batch(batch, state, reasoning)) + transformed.extend( + replay_state.prepare_responses_batch( + candidates, state, state_scope, reasoning + ) + ) skip_batch_items = batch_size - 1 continue - # Middleware changed the batch. Keep its public items, but do - # not attach private reasoning or state to different neighbors. + # A middleware split or mutated the batch. Keep the public item, + # but fail closed instead of associating state with new neighbors. state = None reasoning = None + # Project without copying nested data yet. Replay preparation owns + # detachment for state-bearing turns; passthrough branches detach + # below, so each public payload is copied only once. + msg = dict(original) + msg.pop(LLM_STATE_KEY, None) + # Provider state supplied outside a valid NOOA envelope is never + # accepted, even if a caller constructs wire dictionaries directly. + msg.pop("reasoning_items", None) + # System messages → extract to instructions if msg.get("role") == "system": content = msg.get("content", "") @@ -2555,10 +2636,14 @@ def _transform_messages( # Already in native Responses format (from ResponsesProviderFormatter) if "type" in msg: + if replay_state.response_item_type(msg) == "reasoning": + continue if state is not None or reasoning is not None: - transformed.extend(demote_responses_batch([msg], state, reasoning)) - else: - transformed.append(msg) + transformed.extend( + replay_state.prepare_responses_batch([msg], state, state_scope, reasoning) + ) + continue + transformed.append(copy.deepcopy(msg)) continue # Legacy OpenAI format: tool result messages @@ -2590,7 +2675,7 @@ def _transform_messages( if isinstance(block, dict) and "cache_control" in block: item["cache_control"] = block["cache_control"] break - transformed.append(item) + transformed.append(copy.deepcopy(item)) continue # Legacy OpenAI format: assistant messages with tool_calls @@ -2609,7 +2694,9 @@ def _transform_messages( "arguments": fn.get("arguments", ""), } ) - transformed.extend(demote_responses_batch(batch, state, reasoning)) + transformed.extend( + replay_state.prepare_responses_batch(batch, state, state_scope, reasoning) + ) continue # User/Assistant text messages → passthrough with cache_control preservation @@ -2621,13 +2708,15 @@ def _transform_messages( if "cache_control" in msg: item["cache_control"] = msg["cache_control"] if (state is not None or reasoning is not None) and msg.get("role") == "assistant": - transformed.extend(demote_responses_batch([item], state, reasoning)) + transformed.extend( + replay_state.prepare_responses_batch([item], state, state_scope, reasoning) + ) else: - transformed.append(item) + transformed.append(copy.deepcopy(item)) continue # Unknown format → passthrough - transformed.append(msg) + transformed.append(copy.deepcopy(msg)) instructions = "\n\n".join(instructions_parts) if instructions_parts else None return transformed, instructions @@ -2636,11 +2725,5 @@ def _extract_text_from_output(self, response: Any) -> str: if hasattr(response, "output_text") and response.output_text: return response.output_text if hasattr(response, "output"): - for item in response.output: - if item.type == "message": - texts = [] - for content_item in item.content: - if hasattr(content_item, "text"): - texts.append(content_item.text) # type: ignore - return "\n".join(texts) + return replay_state.responses_output_text(response.output) return "" diff --git a/tests/context_blocks/test_cached_renderer.py b/tests/context_blocks/test_cached_renderer.py index 689f65b4f..4bbea560c 100644 --- a/tests/context_blocks/test_cached_renderer.py +++ b/tests/context_blocks/test_cached_renderer.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """Tests for the cached renderer (static-prefix / events / dynamic-suffix).""" +from nooa._llm_state import LLM_STATE_KEY, carried_state from nooa.context_blocks.events import ( AssistantEvent, ToolCallEvent, @@ -15,6 +16,7 @@ from nooa.context_blocks.models import BlockMetadata, DynamicContext, ResolvedBlock, Role from nooa.context_blocks.renderer import render_context from nooa.context_blocks.renderers.cached import CachedBlockFormatter +from nooa.events import LLMResponse def _static_block(key: str, content: str, expr: str | None = None) -> ResolvedBlock: @@ -220,6 +222,55 @@ def find_user_event_msg(out: list[dict]) -> dict: # And specifically: no context envelope leaked into the user-event msg. assert "" not in msg1["content"] + def test_opaque_state_is_appended_without_changing_cacheable_prefix(self): + user_event = UserEvent(content="please solve the task", tag="1") + user_block = ResolvedBlock( + key="event_1", + content=user_event.content, + role=Role.USER, + metadata=BlockMetadata(tag="1"), + event=user_event, + ) + first = render_context( + [ + _static_block("sys", "stable instructions"), + user_block, + _dynamic_block("live_state", "version one"), + ], + block_formatter=CachedBlockFormatter(), + provider_formatter=OpenAIProviderFormatter(), + ).output + + state = { + "version": 1, + "scope": "responses:openai:sha256:test", + "format": "openai-responses", + "payload": {"items": [{"type": "reasoning", "encrypted_content": "opaque"}]}, + } + turn = LLMResponse(content="answer", llm_state=state, tag="2") + second = render_context( + [ + _static_block("sys", "stable instructions"), + user_block, + ResolvedBlock( + key="event_2", + content=turn.content, + role=Role.ASSISTANT, + metadata=BlockMetadata(tag="2"), + event=turn, + ), + _dynamic_block("live_state", "version two"), + ], + block_formatter=CachedBlockFormatter(), + provider_formatter=OpenAIProviderFormatter(), + ).output + + assert first[:-1] == second[: len(first) - 1] + assert first[-1] != second[-1] + assert all(LLM_STATE_KEY not in message for message in second) + assert carried_state(second[2]) == state + assert "version two" in second[-1]["content"] + def test_volatile_appended_after_assistant(self): asst_event = AssistantEvent(content="done") asst_event.tag = "2" diff --git a/tests/strategies/test_codeact_text_only_reply.py b/tests/strategies/test_codeact_text_only_reply.py index 90126a8ed..bdffe7741 100644 --- a/tests/strategies/test_codeact_text_only_reply.py +++ b/tests/strategies/test_codeact_text_only_reply.py @@ -91,10 +91,10 @@ async def my_task(self) -> dict: events = agent.event_manager.values() assert events.index(outputs[0]) < events.index(corrections[0]) - # The provider sees the exact assistant text followed by user feedback. + # Provider-visible text reasoning is demoted onto the original assistant turn. assert any( message.get("role") == "assistant" - and message.get("content") == "I think the answer is ready." + and message.get("content") == "I checked the evidence.\n\nI think the answer is ready." for message in fake_llm.last_messages ) assert not any(message.get("tool_calls") for message in fake_llm.last_messages) diff --git a/tests/test_nemo_relay_middleware.py b/tests/test_nemo_relay_middleware.py index 3a3995d31..8f934291a 100644 --- a/tests/test_nemo_relay_middleware.py +++ b/tests/test_nemo_relay_middleware.py @@ -21,6 +21,7 @@ from typing import Any from unittest.mock import MagicMock +from nooa._llm_state import carried_replay_batch, carried_state, carry_replay_batch from nooa.nemo_relay_middleware import ( install_nemo_relay, nemo_relay_agent_call_middleware, @@ -130,6 +131,32 @@ def _make_exec_ctx(code: str = "x = 1", agent: Any = None) -> ExecutePythonConte class TestLLMRequestIntercepts: """Verify that LLM request intercepts (header injection) work end-to-end.""" + @pytest.mark.asyncio + async def test_noop_relay_roundtrip_preserves_private_message_sidecars(self): + """Relay's JSON copy must not erase state when no intercept changed input.""" + messages = carry_replay_batch( + [{"type": "reasoning", "summary": []}, {"type": "function_call", "id": "fc-1"}], + {"scope": "issuer", "payload": {"encrypted_content": "opaque"}}, + None, + ) + ctx = _make_llm_ctx(messages=messages) + seen: list[list[dict[str, Any]]] = [] + + async def nxt(c): + seen.append(c.messages) + c.response = FakeLLMResponse() + return c + + await nemo_relay_llm_middleware(ctx, nxt) + + assert seen[0] is messages + assert carried_state(seen[0][0]) == { + "scope": "issuer", + "payload": {"encrypted_content": "opaque"}, + } + assert carried_replay_batch(seen[0][0]) is not None + assert carried_replay_batch(seen[0][1]) == carried_replay_batch(seen[0][0]) + @pytest.mark.asyncio async def test_request_intercept_injects_header(self): """A registered LLM request intercept can inject HTTP headers.""" diff --git a/tests/tracing/test_journal.py b/tests/tracing/test_journal.py index c98b1e1f8..2e39f8c9d 100644 --- a/tests/tracing/test_journal.py +++ b/tests/tracing/test_journal.py @@ -11,6 +11,7 @@ from __future__ import annotations +import json from types import SimpleNamespace from unittest.mock import patch @@ -22,6 +23,7 @@ ) from nooa.tracing._litellm_journal import ( MessageJournalCallback, + _safe_msg_to_dict, ) @@ -126,6 +128,19 @@ def test_no_sideband_falls_back_to_raw_messages(self, session_ctx): skeleton, _ = cb._call_inputs["cid"] assert skeleton == [{"role": "user", "content": "hello"}] + def test_no_sideband_redacts_opaque_state_from_raw_messages(self, session_ctx): + cb = MessageJournalCallback("http://localhost:5001") + set_journal_payload(None) + + cb.log_pre_api_call( + "m", + [{"type": "reasoning", "encrypted_content": "opaque-openai-state"}], + {"litellm_call_id": "cid"}, + ) + + skeleton, _ = cb._call_inputs["cid"] + assert skeleton == [{"type": "reasoning", "encrypted_content": "[REDACTED]"}] + class TestJournalCallbackSuccessEvent: def test_success_posts_call_record(self, session_ctx): @@ -168,6 +183,50 @@ def test_success_posts_call_record(self, session_ctx): all_blocks = {e["hash"]: e["content"] for post in block_posts for e in post[1]} assert all_blocks[answer_hash] == "answer" + def test_success_redacts_opaque_state_from_provider_output(self, session_ctx): + cb = MessageJournalCallback("http://localhost:5001") + calls, fake_post = _posts() + response = SimpleNamespace( + output=[{"type": "reasoning", "encrypted_content": "opaque-openai-state"}], + usage=None, + ) + + with patch( + "nooa.tracing._litellm_journal._post_json", + side_effect=fake_post, + ): + cb.log_pre_api_call("m", [], {"litellm_call_id": "cid"}) + cb.log_success_event({"litellm_call_id": "cid", "model": "m"}, response, 1.0, 2.0) + + record = next(payload for url, payload in calls if url.endswith("/v1/journal/calls")) + assert record["output_messages"] == [ + {"type": "reasoning", "encrypted_content": "[REDACTED]"} + ] + assert "opaque-openai-state" not in repr(record) + + +def test_safe_msg_to_dict_redacts_json_encoded_opaque_state(): + message = {"content": '{"encrypted_content":"opaque-openai-state"}'} + + safe = _safe_msg_to_dict(message) + + assert "opaque-openai-state" not in safe["content"] + + +@pytest.mark.parametrize("json_encoded", [False, True]) +def test_safe_msg_to_dict_redacts_private_replay_envelope(json_encoded): + from nooa._llm_state import LLM_STATE_KEY + + message = {LLM_STATE_KEY: {"payload": {"future_provider_blob": "opaque-state"}}} + original = {"content": json.dumps(message)} if json_encoded else message + + safe = _safe_msg_to_dict(original) + + decoded = json.loads(safe["content"]) if json_encoded else safe + assert decoded[LLM_STATE_KEY] == "[REDACTED]" + assert "opaque-state" not in repr(safe) + assert message[LLM_STATE_KEY]["payload"]["future_provider_blob"] == "opaque-state" + class TestSentBlocksBounding: """Tests for single-session tracking and deferred hash marking.""" diff --git a/tests/tracing/test_secret_scrubber.py b/tests/tracing/test_secret_scrubber.py index f87cb2a34..fae01cddc 100644 --- a/tests/tracing/test_secret_scrubber.py +++ b/tests/tracing/test_secret_scrubber.py @@ -2,6 +2,9 @@ # SPDX-License-Identifier: Apache-2.0 """Tests for secret scrubbing in telemetry.""" +import json +from unittest.mock import patch + import pytest from nooa.tracing._secret_scrubber import ( @@ -174,6 +177,41 @@ def test_quoted_json_key(self): class TestScrubValue: + @pytest.mark.parametrize("as_array", [False, True]) + def test_json_redaction_preserves_literal_unicode(self, as_array): + payload = {"api_key": "synthetic-secret", "text": "héllo 世界"} + attribute = " \n" + json.dumps([payload] if as_array else payload, ensure_ascii=False) + + result, count = scrub_value(attribute) + + assert "héllo 世界" in result + assert "synthetic-secret" not in result + expected = {"api_key": REDACTED, "text": "héllo 世界"} + assert json.loads(result) == ([expected] if as_array else expected) + assert count == 1 + + @pytest.mark.parametrize("text", ["ordinary code output", " \n héllo 世界", "", "42"]) + def test_plain_text_never_attempts_json_parsing(self, text): + with patch("nooa.tracing._secret_scrubber.json.loads") as parse: + assert scrub_value(text) == (text, 0) + parse.assert_not_called() + + def test_plain_text_still_redacts_without_json_parsing(self): + with patch("nooa.tracing._secret_scrubber.json.loads") as parse: + result, count = scrub_value("api_key=synthetic-secret") + parse.assert_not_called() + assert result == f"api_key={REDACTED}" + assert count == 1 + + def test_invalid_json_falls_back_to_text_scrubbing(self): + result, count = scrub_value("[progress] api_key=synthetic-secret") + assert result == f"[progress] api_key={REDACTED}" + assert count == 1 + + def test_clean_json_keeps_original_formatting(self): + attribute = ' \n{ "text": "héllo 世界" }\n' + assert scrub_value(attribute) == (attribute, 0) + def test_string(self): """A string value is scrubbed and its redaction count returned.""" result, count = scrub_value("key=AKIAIOSFODNN7EXAMPLE") @@ -215,6 +253,25 @@ def test_nested_sensitive_keys(self): assert result == {"safe": {"client_secret": REDACTED, "refresh_token": REDACTED}} assert count == 2 + def test_openai_encrypted_reasoning_is_redacted_from_mapping(self): + result, count = scrub_value( + {"input": [{"type": "reasoning", "encrypted_content": "opaque-openai-state"}]} + ) + + assert result["input"][0]["encrypted_content"] == REDACTED + assert count == 1 + + def test_openai_encrypted_reasoning_is_redacted_from_json_attribute(self): + attribute = json.dumps( + {"input": [{"type": "reasoning", "encrypted_content": "opaque-openai-state"}]} + ) + + result, count = scrub_value(attribute) + + assert json.loads(result)["input"][0]["encrypted_content"] == REDACTED + assert "opaque-openai-state" not in result + assert count == 1 + class TestScrubStats: def test_record_and_snapshot(self): diff --git a/tests/unifiedllm/test_cache_control.py b/tests/unifiedllm/test_cache_control.py index a31e14a3c..d072a6755 100644 --- a/tests/unifiedllm/test_cache_control.py +++ b/tests/unifiedllm/test_cache_control.py @@ -59,6 +59,21 @@ def test_does_not_mutate_original(self, client): assert "cache_control" not in messages[0] assert "cache_control" in result[0] + def test_copies_only_messages_that_receive_a_breakpoint(self, client): + """Stable history is borrowed instead of cloned on every request.""" + messages = [ + {"role": "system", "content": "You are helpful."}, + {"role": "assistant", "content": "Prior answer"}, + {"role": "user", "content": "Continue"}, + ] + + result = client._inject_cache_control(messages, [{"role": "system"}]) + + assert result is not messages + assert result[0] is not messages[0] + assert result[1] is messages[1] + assert result[2] is messages[2] + def test_multiple_roles(self, client): """Can target multiple roles at once.""" messages = [ @@ -93,6 +108,7 @@ def test_no_matching_role(self, client): result = client._inject_cache_control(messages, injection_points) assert "cache_control" not in result[0] + assert result is messages # --------------------------------------------------------------------------- diff --git a/tests/unifiedllm/test_completion_transport_overrides.py b/tests/unifiedllm/test_completion_transport_overrides.py new file mode 100644 index 000000000..e49e5eb48 --- /dev/null +++ b/tests/unifiedllm/test_completion_transport_overrides.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Exercise request routing through LiteLLM and the OpenAI SDK, without network I/O.""" + +import httpx +import pytest +from litellm.llms.openai.openai import OpenAIChatCompletion + +from nooa.unifiedllm import CompletionClient +from nooa.unifiedllm.unifiedllm import _ClientHttp + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize( + "overrides", + [ + {"api_base": "https://override.example/v1"}, + {"api_key": "override-key"}, + {"api_base": "https://override.example/v1", "api_key": "override-key"}, + ], +) +async def test_completion_uses_per_call_route_without_changing_defaults( + monkeypatch, is_async, overrides +): + requests: list[httpx.Request] = [] + + def respond(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response( + 200, + json={ + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 0, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "done"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11}, + }, + ) + + transport = httpx.MockTransport(respond) + monkeypatch.setattr( + _ClientHttp, "_httpx_hardening", staticmethod(lambda: {"transport": transport}) + ) + client = CompletionClient( + "openai/gpt-4o-mini", + api_base="https://original.example/v1", + api_key="original-key", + ) + client_http = client._http + assert client_http is not None + # Intercept LiteLLM's default client too. Its real provider dispatch and SDK + # still choose the URL and Authorization header; only HTTP is replaced. + monkeypatch.setattr( + OpenAIChatCompletion, + "_get_sync_http_client", + staticmethod(lambda: client_http.httpx_sync), + ) + monkeypatch.setattr( + OpenAIChatCompletion, + "_get_async_http_client", + staticmethod(lambda **kwargs: client_http.httpx_async), + ) + try: + owned_client = client_http.async_client if is_async else client_http.sync_client + assert client._completion_http_client(client.config, is_async=is_async) is owned_client + unchanged = {**client.config, "temperature": 0.1} + assert client._completion_http_client(unchanged, is_async=is_async) is owned_client + + for params in ({}, overrides, {}): + messages = [{"role": "user", "content": "Hello"}] + result = ( + await client.acall(messages, **params) + if is_async + else client.call(messages, **params) + ) + assert result.content == "done" + + assert [str(request.url) for request in requests] == [ + "https://original.example/v1/chat/completions", + overrides.get("api_base", "https://original.example/v1") + "/chat/completions", + "https://original.example/v1/chat/completions", + ] + assert [request.headers["Authorization"] for request in requests] == [ + "Bearer original-key", + "Bearer " + overrides.get("api_key", "original-key"), + "Bearer original-key", + ] + finally: + await client.aclose() diff --git a/tests/unifiedllm/test_finish_reason_propagation.py b/tests/unifiedllm/test_finish_reason_propagation.py index 5c1b4a3be..67298ec71 100644 --- a/tests/unifiedllm/test_finish_reason_propagation.py +++ b/tests/unifiedllm/test_finish_reason_propagation.py @@ -236,7 +236,9 @@ def _make_incomplete_responses_tool_response(): def _make_responses_tool_response(text: str): return SimpleNamespace( output=[ - SimpleNamespace(type="message", content=[SimpleNamespace(text=text)]), + SimpleNamespace( + type="message", content=[SimpleNamespace(type="output_text", text=text)] + ), SimpleNamespace( type="function_call", call_id="call_1", @@ -393,11 +395,14 @@ async def my_task(self) -> str: assert await agent_instance.my_task() == "done" events = agent_instance.event_manager.values() - first_output = next(event for event in events if event.event_type == "LLMResponse") + from nooa.context_blocks.events import ToolCallEvent + from nooa.unifiedllm import LLMResponse + + first_output = next(event for event in events if isinstance(event, LLMResponse)) execution = next( event for event in events - if event.event_type == "ToolCallEvent" and event.name == "execute_python" + if isinstance(event, ToolCallEvent) and event.name == "execute_python" ) assert first_output.content == "I will calculate this." assert execution.arguments == {"code": "x = 42"} diff --git a/tests/unifiedllm/test_http_logging.py b/tests/unifiedllm/test_http_logging.py index 9546d2ab6..09315b194 100644 --- a/tests/unifiedllm/test_http_logging.py +++ b/tests/unifiedllm/test_http_logging.py @@ -12,6 +12,29 @@ from nooa.unifiedllm.http_logging import enable_http_request_logging +def test_opaque_reasoning_state_is_redacted_from_http_debug_payloads( + tmp_path, secret_header_server +) -> None: + payload = { + "input": [ + {"encrypted_content": "provider-secret"}, + {"_nooa_llm_state": {"payload": {"items": ["opaque"]}}}, + ] + } + + disable = enable_http_request_logging(output_dir=tmp_path, verbose=False) + try: + httpx.post(f"http://127.0.0.1:{secret_header_server.server_port}/llm", json=payload) + finally: + disable() + redacted = json.loads(next(tmp_path.glob("request_*.json")).read_text()) + + assert redacted["input"] == [ + {"encrypted_content": "[REDACTED]"}, + {"_nooa_llm_state": "[REDACTED]"}, + ] + + class _SecretHeaderHandler(BaseHTTPRequestHandler): status_code = 500 diff --git a/tests/unifiedllm/test_litellm_responses_bridge.py b/tests/unifiedllm/test_litellm_responses_bridge.py index fb1a32109..683af00b2 100644 --- a/tests/unifiedllm/test_litellm_responses_bridge.py +++ b/tests/unifiedllm/test_litellm_responses_bridge.py @@ -95,7 +95,10 @@ async def test_default_reasoning_tool_call_uses_responses_bridge(model: str) -> responses.assert_called_once() assert result.finish_reason == "tool_calls" - assert result.llm_state == {"reasoning_items": [REASONING_ITEM]} + assert result.llm_state is not None + assert result.llm_state["format"] == "litellm-chat" + assert result.llm_state["payload"]["reasoning_items"] == [REASONING_ITEM] + assert len(result.llm_state["payload"]["carrier"]) == 64 finally: await client.aclose() diff --git a/tests/unifiedllm/test_model_registry.py b/tests/unifiedllm/test_model_registry.py index df7d05cad..da6556f4e 100644 --- a/tests/unifiedllm/test_model_registry.py +++ b/tests/unifiedllm/test_model_registry.py @@ -251,6 +251,9 @@ def test_registry_preserves_litellm_pass_through_controls(self, tmp_path): - reasoning_effort extra_body: trace: true + store: false + include: + - reasoning.encrypted_content """, ) reload_registry(path) @@ -261,6 +264,8 @@ def test_registry_preserves_litellm_pass_through_controls(self, tmp_path): assert llm.config["reasoning_effort"] == "high" assert llm.config["allowed_openai_params"] == ["reasoning_effort"] assert llm.config["extra_body"] == {"trace": True} + assert llm.config["store"] is False + assert llm.config["include"] == ["reasoning.encrypted_content"] def test_drop_params_default_true(self): llm = get_llm_client("gpt-4o-mini") diff --git a/tests/unifiedllm/test_reasoning_state_replay.py b/tests/unifiedllm/test_reasoning_state_replay.py new file mode 100644 index 000000000..52a0872ae --- /dev/null +++ b/tests/unifiedllm/test_reasoning_state_replay.py @@ -0,0 +1,1074 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Issuer-gated capture and replay of opaque OpenAI reasoning state.""" + +import json +import sqlite3 +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + +import httpx +import pytest +from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.utils import ModelResponse + +from nooa._llm_state import ( + LLM_STATE_KEY, + ReplayCarryingMessage, + carried_replay_batch, + carried_state, +) +from nooa.context_blocks.events import ToolCallEvent, ToolResult +from nooa.context_blocks.formatter import ( + OpenAIProviderFormatter, + ResponsesProviderFormatter, + XMLBlockFormatter, +) +from nooa.context_blocks.models import ResolvedBlock, Role +from nooa.runtime.middleware import LLMCallContext +from nooa.storage.sqlite import SQLiteEventBackend, _ensure_schema +from nooa.unifiedllm import CompletionClient, LLMResponse, ResponsesClient, Tool +from nooa.unifiedllm.replay_state import ( + prepare_chat_messages, + prepare_responses_batch, + replay_scope, +) +from nooa.unifiedllm.unifiedllm import _ClientHttp + +REASONING = { + "id": "rs_1", + "type": "reasoning", + "encrypted_content": "provider-secret", + "summary": [], +} +REASONING_2 = { + "id": "rs_2", + "type": "reasoning", + "encrypted_content": "provider-secret-2", + "summary": [], +} +MESSAGE = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "done", "annotations": []}], +} +CALL = { + "id": "fc_1", + "type": "function_call", + "call_id": "call_1", + "name": "execute_python", + "arguments": '{"code":"print(1)"}', + "status": "completed", +} +CALL_2 = { + "id": "fc_2", + "type": "function_call", + "call_id": "call_2", + "name": "execute_python", + "arguments": '{"code":"print(2)"}', + "status": "completed", +} + + +def _responses(*items: dict) -> SimpleNamespace: + return SimpleNamespace(output=list(items), output_text="", status="completed", usage=None) + + +def _chat_tool_call(call_id: str = "call_1", code: str = "print(1)") -> dict: + return { + "id": call_id, + "type": "function", + "function": { + "name": "execute_python", + "arguments": json.dumps({"code": code}, separators=(",", ":")), + }, + } + + +def _chat_response( + *, + reasoning_items: list[dict] | None = None, + tool_calls: list[dict] | None = None, +) -> ModelResponse: + return ModelResponse( + model="gpt-5.6", + choices=[ + { + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": "", + "tool_calls": [_chat_tool_call()] if tool_calls is None else tool_calls, + "reasoning_items": reasoning_items, + }, + } + ], + ) + + +def _tool(code: str) -> str: + return code + + +TOOL = Tool(name="execute_python", description="Run code", callable=_tool) + + +def _response_blocks(response: LLMResponse) -> list[ResolvedBlock]: + blocks = [ResolvedBlock(key="turn", content="", role=Role.ASSISTANT, event=response)] + blocks.extend( + ResolvedBlock( + key=f"execution-{call.id}", + content="", + role=Role.ASSISTANT, + event=ToolCallEvent( + tool_call_id=call.id, + name=call.name, + arguments=( + json.loads(call.arguments) + if isinstance(call.arguments, str) + else call.arguments + ), + llm_response_id=response.id, + result=ToolResult(tool_call_id=call.id, content="complete"), + ), + ) + for call in response.tool_calls + ) + return blocks + + +def _render_responses(response: LLMResponse) -> list[dict]: + neutral = XMLBlockFormatter().format(_response_blocks(response)) + return ResponsesProviderFormatter().format(neutral) + + +def _render_chat(response: LLMResponse) -> list[dict]: + neutral = XMLBlockFormatter().format(_response_blocks(response)) + return OpenAIProviderFormatter().format(neutral) + + +def test_responses_text_state_is_captured_and_exactly_replayed() -> None: + client = ResponsesClient(model="openai/gpt-5.6", api_key="account-a") + try: + with patch( + "litellm.responses", + side_effect=[_responses(REASONING, MESSAGE), _responses(MESSAGE)], + ) as call: + first = client.call([{"role": "user", "content": "think"}]) + rendered = _render_responses(first) + assert "provider-secret" not in json.dumps(rendered) + middleware_context = LLMCallContext(messages=rendered) + carrier = next( + message + for message in middleware_context.messages + if carried_state(message) is not None + ) + assert carried_state(carrier) is first.llm_state + client.call(middleware_context.messages + [{"role": "user", "content": "continue"}]) + + assert first.llm_state is not None + assert first.llm_state["payload"]["items"] == [REASONING] + assert "reasoning.encrypted_content" in call.call_args_list[0].kwargs["include"] + replay = call.call_args_list[1].kwargs["input"] + assert replay[:2] == [REASONING, {"role": "assistant", "content": "done"}] + assert LLM_STATE_KEY not in repr(replay) + finally: + client.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("shape", ["text_blocks", "phases", "phase_only", "trailing_message"]) +@pytest.mark.parametrize("persistence", ["json", "sqlite"]) +async def test_real_responses_message_structure_survives_json_resume( + monkeypatch, is_async: bool, shape: str, persistence: str +) -> None: + """Use real LiteLLM/SDK parsing and serialization, including its output_text property.""" + phased_message = {**MESSAGE, "id": "msg_final", "phase": "final_answer"} + if shape == "text_blocks": + output = [ + REASONING, + { + **MESSAGE, + "content": [ + {"type": "output_text", "text": "first", "annotations": []}, + {"type": "output_text", "text": "second", "annotations": []}, + ], + }, + ] + expected = [REASONING, {"role": "assistant", "content": "firstsecond"}] + elif shape == "phases": + output = [ + REASONING, + { + **MESSAGE, + "id": "msg_commentary", + "phase": "commentary", + "content": [{"type": "output_text", "text": "Checking.", "annotations": []}], + }, + REASONING_2, + phased_message, + ] + expected = [ + REASONING, + {"role": "assistant", "content": "Checking.", "phase": "commentary"}, + REASONING_2, + {"role": "assistant", "content": "done", "phase": "final_answer"}, + ] + elif shape == "phase_only": + output = [phased_message] + expected = [{"role": "assistant", "content": "done", "phase": "final_answer"}] + else: + output = [CALL, MESSAGE] + expected = [ + {key: value for key, value in CALL.items() if key not in {"id", "status"}}, + {"role": "assistant", "content": "done"}, + {"type": "function_call_output", "call_id": "call_1", "output": "complete"}, + ] + + raw = ResponsesAPIResponse.model_validate( + { + "id": "resp_test", + "created_at": 0, + "model": "gpt-5.6", + "status": "completed", + "output": output, + } + ) + bodies: list[dict] = [] + + def respond(request: httpx.Request) -> httpx.Response: + bodies.append(json.loads(request.content)) + return httpx.Response(200, json=raw.model_dump(mode="json")) + + transport = httpx.MockTransport(respond) + monkeypatch.setattr( + _ClientHttp, "_httpx_hardening", staticmethod(lambda: {"transport": transport}) + ) + client = ResponsesClient( + model="openai/gpt-5.6", api_key="test", api_base="https://gateway.example/v1" + ) + try: + prompt = [{"role": "user", "content": "think"}] + first = await client.acall(prompt) if is_async else client.call(prompt) + assert first.content == raw.output_text + # The fallback for providers without output_text must use identical text. + assert client._extract_text_from_output(_responses(*output)) == raw.output_text + if persistence == "json": + resumed = LLMResponse.model_validate_json(first.model_dump_json()) + else: + connection = sqlite3.connect(":memory:") + try: + _ensure_schema(connection) + backend = SQLiteEventBackend(connection) + backend.store("response", first) + resumed = backend.get("response") + assert isinstance(resumed, LLMResponse) + finally: + connection.close() + assert resumed.raw_response is None + assert resumed.llm_state is not None + rendered = _render_responses(resumed) + if is_async: + await client.acall(rendered) + else: + client.call(rendered) + assert bodies[1]["input"] == expected + + # A public text edit must still invalidate the saved message structure. + resumed.content = "edited answer" + if is_async: + await client.acall(_render_responses(resumed)) + else: + client.call(_render_responses(resumed)) + assert bodies[2]["input"][0] == {"role": "assistant", "content": "edited answer"} + assert "encrypted_content" not in json.dumps(bodies[2]["input"]) + assert "phase" not in json.dumps(bodies[2]["input"]) + finally: + await client.aclose() + + +def test_direct_multimessage_batch_replays_all_reasoning_in_order() -> None: + from nooa.unifiedllm.replay_state import capture_responses_state + + scope = replay_scope("openai/gpt-5.6", "responses", {}) + second_message = { + **MESSAGE, + "id": "msg_second", + "content": [{"type": "output_text", "text": "second", "annotations": []}], + } + state = capture_responses_state([REASONING, MESSAGE, REASONING_2, second_message], scope) + message = {"role": "assistant", "content": "done"} + second = {"role": "assistant", "content": "second"} + + assert prepare_responses_batch([message, second], state, scope) == [ + REASONING, + message, + REASONING_2, + second, + ] + + +def test_empty_and_summary_only_outputs_do_not_capture_structural_state() -> None: + from nooa.unifiedllm.replay_state import capture_responses_state + + scope = replay_scope("openai/gpt-5.6", "responses", {}) + summary_only = {"type": "reasoning", "summary": [{"type": "summary_text", "text": "why"}]} + assert capture_responses_state([], scope) is None + assert capture_responses_state([summary_only], scope) is None + + +@pytest.mark.parametrize("part", ["refusal", "web_search_call"]) +def test_valid_unprojectable_sdk_turn_does_not_retain_partial_state(part, caplog) -> None: + from openai.types.responses import ( + ResponseFunctionWebSearch, + ResponseOutputMessage, + ResponseReasoningItem, + ) + + from nooa.unifiedllm.replay_state import capture_responses_state + + reasoning = ResponseReasoningItem.model_validate(REASONING) + if part == "refusal": + unsupported = ResponseOutputMessage.model_validate( + {**MESSAGE, "content": [{"type": "refusal", "refusal": "Not permitted."}]} + ) + else: + unsupported = ResponseFunctionWebSearch.model_validate( + { + "id": "ws_1", + "type": "web_search_call", + "status": "completed", + "action": {"type": "search", "query": "weather"}, + } + ) + scope = replay_scope("openai/gpt-5.6", "responses", {}) + assert capture_responses_state([reasoning, unsupported], scope) is None + assert "unsupported turn parts" in caplog.text + assert part in caplog.text + assert "provider-secret" not in caplog.text + assert "Not permitted." not in caplog.text + + +def test_permuted_reasoning_indexes_are_not_replayed(caplog) -> None: + from nooa.unifiedllm.replay_state import capture_responses_state + + scope = replay_scope("openai/gpt-5.6", "responses", {}) + state = capture_responses_state([REASONING, REASONING_2, MESSAGE], scope) + assert state is not None + state["payload"]["order"][0]["index"] = 1 + state["payload"]["order"][1]["index"] = 0 + public = [{"role": "assistant", "content": "done"}] + assert prepare_responses_batch(public, state, scope) == public + assert "is malformed" in caplog.text + + +def test_mixed_summary_only_reasoning_never_replays_a_partial_opaque_sequence(caplog) -> None: + from nooa.unifiedllm.replay_state import capture_responses_state + + scope = replay_scope("openai/gpt-5.6", "responses", {}) + summary_only = {"type": "reasoning", "summary": [{"type": "summary_text", "text": "why"}]} + state = capture_responses_state( + [REASONING, summary_only, {**MESSAGE, "phase": "final_answer"}], scope + ) + assert state is not None + assert state["payload"]["items"] == [] + assert "incomplete opaque reasoning sequence" in caplog.text + assert prepare_responses_batch( + [{"role": "assistant", "content": "done"}], state, scope, "why" + ) == [{"role": "assistant", "content": "why\n\ndone", "phase": "final_answer"}] + + +def test_responses_replay_borrows_stored_payload_without_copying() -> None: + client = ResponsesClient(model="openai/gpt-5.6", api_key="account-a") + try: + with patch("litellm.responses", return_value=_responses(REASONING, MESSAGE)): + first = client.call([{"role": "user", "content": "think"}]) + + observed = None + + def observe_wire_input(**kwargs): + nonlocal observed + reasoning = next(item for item in kwargs["input"] if item.get("type") == "reasoning") + observed = reasoning + return _responses(MESSAGE) + + with patch("litellm.responses", side_effect=observe_wire_input): + client.call(_render_responses(first)) + + assert first.llm_state is not None + assert observed is first.llm_state["payload"]["items"][0] + finally: + client.close() + + +def test_responses_multi_call_state_preserves_provider_order() -> None: + first_raw = _responses(REASONING, CALL, REASONING_2, CALL_2) + client = ResponsesClient(model="openai/gpt-5.6", api_key="account-a") + try: + with patch("litellm.responses", side_effect=[first_raw, _responses(MESSAGE)]) as call: + first = client.call([{"role": "user", "content": "run"}], tools=[TOOL]) + client.call( + _render_responses(first) + [{"role": "user", "content": "continue"}], + tools=[TOOL], + ) + + replay = call.call_args_list[1].kwargs["input"] + assistant_batch = [item for item in replay if item.get("type") != "function_call_output"] + assert [item.get("type") for item in assistant_batch[:4]] == [ + "reasoning", + "function_call", + "reasoning", + "function_call", + ] + assert [ + item.get("call_id") for item in assistant_batch if item.get("type") == "function_call" + ] == [ + "call_1", + "call_2", + ] + finally: + client.close() + + +def test_changed_responses_text_drops_state_and_preserves_edit() -> None: + client = ResponsesClient(model="openai/gpt-5.6", api_key="account-a") + try: + with patch("litellm.responses", return_value=_responses(REASONING, MESSAGE)): + first = client.call([{"role": "user", "content": "think"}]) + rendered = _render_responses(first) + assistant = next(item for item in rendered if carried_state(item) is not None) + assistant["content"] = "replacement text" + + with patch("litellm.responses", return_value=_responses(MESSAGE)) as call: + client.call(rendered) + + replay = call.call_args.kwargs["input"] + assert REASONING not in replay + assert {"role": "assistant", "content": "replacement text"} in replay + finally: + client.close() + + +def test_reordered_responses_calls_drop_state_and_keep_new_order() -> None: + client = ResponsesClient(model="openai/gpt-5.6", api_key="account-a") + try: + with patch("litellm.responses", return_value=_responses(REASONING, MESSAGE, CALL, CALL_2)): + first = client.call([{"role": "user", "content": "run"}], tools=[TOOL]) + rendered = _render_responses(first) + first_call = next( + index for index, item in enumerate(rendered) if item.get("call_id") == "call_1" + ) + second_call = next( + index for index, item in enumerate(rendered) if item.get("call_id") == "call_2" + ) + rendered[first_call], rendered[second_call] = rendered[second_call], rendered[first_call] + + with patch("litellm.responses", return_value=_responses(MESSAGE)) as call: + client.call(rendered, tools=[TOOL]) + + replay = call.call_args.kwargs["input"] + assert REASONING not in replay + assert [item.get("call_id") for item in replay if item.get("type") == "function_call"] == [ + "call_2", + "call_1", + ] + finally: + client.close() + + +def test_reasoning_only_carrier_edit_drops_state_but_keeps_text() -> None: + client = ResponsesClient(model="openai/gpt-5.6", api_key="account-a") + try: + with patch("litellm.responses", return_value=_responses(REASONING)): + first = client.call([{"role": "user", "content": "think"}]) + rendered = _render_responses(first) + assistant = next(item for item in rendered if carried_state(item) is not None) + assistant["content"] = "do not discard me" + + with patch("litellm.responses", return_value=_responses(MESSAGE)) as call: + client.call(rendered) + + assert call.call_args.kwargs["input"] == [ + {"role": "assistant", "content": "do not discard me"} + ] + finally: + client.close() + + +def test_malformed_matching_responses_payload_is_not_forwarded(caplog) -> None: + scope = replay_scope("openai/gpt-5.6", "responses", {}) + state = { + "version": 2, + "scope": scope, + "format": "openai-responses", + "payload": { + "items": [{"type": "reasoning"}], + "order": [ + {"type": "reasoning", "index": 0}, + {"type": "message", "content": "public"}, + ], + }, + } + + assert prepare_responses_batch( + [{"role": "assistant", "content": "public"}], + state, + scope, + "portable reasoning", + ) == [{"role": "assistant", "content": "portable reasoning\n\npublic"}] + assert "is malformed" in caplog.text + + +def test_split_responses_replay_batch_keeps_public_call_but_drops_state() -> None: + """Middleware may edit public items, but cannot reattach state to new neighbors.""" + client = ResponsesClient(model="openai/gpt-5.6", api_key="account-a") + try: + with patch("litellm.responses", return_value=_responses(REASONING, CALL, CALL_2)): + first = client.call([{"role": "user", "content": "run"}], tools=[TOOL]) + rendered = [item for item in _render_responses(first) if carried_replay_batch(item)] + assert len(rendered) == 2 + + with patch("litellm.responses", return_value=_responses(MESSAGE)) as call: + client.call(rendered[:1], tools=[TOOL]) + + replay = call.call_args.kwargs["input"] + assert [item.get("call_id") for item in replay] == ["call_1"] + assert REASONING not in replay + assert "provider-secret" not in repr(replay) + finally: + client.close() + + +def test_reasoning_only_response_replays_without_empty_assistant_message() -> None: + client = ResponsesClient(model="openai/gpt-5.6", api_key="account-a") + try: + with patch( + "litellm.responses", + side_effect=[_responses(REASONING), _responses(MESSAGE)], + ) as call: + first = client.call([{"role": "user", "content": "think"}]) + rendered = _render_responses(first) + client.call(rendered + [{"role": "user", "content": "continue"}]) + + assert first.content == "" + assert first.llm_state is not None + assert first.llm_state["payload"]["state_only"] is True + replay = call.call_args_list[1].kwargs["input"] + assert REASONING in replay + assert {"role": "assistant", "content": ""} not in replay + finally: + client.close() + + +def test_incompatible_reasoning_only_state_drops_its_internal_carrier() -> None: + source = ResponsesClient(model="openai/gpt-5.6", api_key="account-a") + target = ResponsesClient(model="openai/gpt-5.7", api_key="account-a") + try: + with patch("litellm.responses", return_value=_responses(REASONING)): + first = source.call([{"role": "user", "content": "think"}]) + with patch("litellm.responses", return_value=_responses(MESSAGE)) as call: + target.call(_render_responses(first) + [{"role": "user", "content": "continue"}]) + + replay = call.call_args.kwargs["input"] + assert REASONING not in replay + assert {"role": "assistant", "content": ""} not in replay + assert replay == [{"role": "user", "content": "continue"}] + finally: + source.close() + target.close() + + +def test_responses_state_is_hidden_from_a_different_model() -> None: + source = ResponsesClient(model="openai/gpt-5.6", api_key="account-a") + target = ResponsesClient(model="openai/gpt-5.7", api_key="account-a") + try: + with patch("litellm.responses", return_value=_responses(REASONING, MESSAGE)): + first = source.call([{"role": "user", "content": "think"}]) + with patch("litellm.responses", return_value=_responses(MESSAGE)) as call: + target.call(_render_responses(first)) + + replay = call.call_args.kwargs["input"] + assert REASONING not in replay + assert {"role": "assistant", "content": "done"} in replay + assert "provider-secret" not in repr(replay) + finally: + source.close() + target.close() + + +def test_responses_model_override_uses_effective_replay_scope() -> None: + client = ResponsesClient(model="openai/gpt-5.6", api_key="account-a") + try: + with patch("litellm.responses", return_value=_responses(REASONING, MESSAGE)): + first = client.call([{"role": "user", "content": "think"}]) + + with patch("litellm.responses", return_value=_responses(MESSAGE)) as call: + client.call( + _render_responses(first), + model="anthropic/claude-sonnet-4-5", + cache_control_injection_points=[], + ) + + assert call.call_args.kwargs["model"] == "anthropic/claude-sonnet-4-5" + assert REASONING not in call.call_args.kwargs["input"] + assert "include" not in call.call_args.kwargs + finally: + client.close() + + +def test_completion_model_override_uses_effective_replay_scope() -> None: + client = CompletionClient( + model="openai/gpt-5.6", + api_key="account-a", + cache_control_injection_points=[], + ) + try: + with patch("litellm.completion", return_value=_chat_response(reasoning_items=[REASONING])): + first = client.call([{"role": "user", "content": "run"}], tools=[TOOL]) + + with patch("litellm.completion", return_value=_chat_response()) as call: + client.call( + _render_chat(first), + tools=[TOOL], + model="anthropic/claude-sonnet-4-5", + ) + + assert call.call_args.kwargs["model"] == "anthropic/claude-sonnet-4-5" + assert "provider-secret" not in repr(call.call_args.kwargs["messages"]) + finally: + client.close() + + +@pytest.mark.asyncio +async def test_async_clients_use_effective_model_for_replay_scope() -> None: + responses = ResponsesClient(model="openai/gpt-5.6", api_key="account-a") + completion = CompletionClient( + model="openai/gpt-5.6", + api_key="account-a", + cache_control_injection_points=[], + ) + try: + with patch("litellm.responses", return_value=_responses(REASONING, MESSAGE)): + responses_first = responses.call([{"role": "user", "content": "think"}]) + with patch("litellm.completion", return_value=_chat_response(reasoning_items=[REASONING])): + completion_first = completion.call([{"role": "user", "content": "run"}], tools=[TOOL]) + + with ( + patch( + "litellm.aresponses", AsyncMock(return_value=_responses(MESSAGE)) + ) as response_call, + patch("litellm.acompletion", AsyncMock(return_value=_chat_response())) as chat_call, + ): + await responses.acall( + _render_responses(responses_first), + model="anthropic/claude-sonnet-4-5", + cache_control_injection_points=[], + ) + await completion.acall( + _render_chat(completion_first), + tools=[TOOL], + model="anthropic/claude-sonnet-4-5", + ) + + assert REASONING not in response_call.call_args.kwargs["input"] + assert "provider-secret" not in repr(chat_call.call_args.kwargs["messages"]) + finally: + await responses.aclose() + await completion.aclose() + + +def test_azure_responses_state_is_captured_replayed_and_requested() -> None: + client = ResponsesClient( + model="azure/gpt-5.6", + api_key="account-a", + api_base="https://account-a.openai.azure.com", + ) + try: + with patch( + "litellm.responses", + side_effect=[_responses(REASONING, MESSAGE), _responses(MESSAGE)], + ) as call: + first = client.call([{"role": "user", "content": "think"}]) + client.call(_render_responses(first)) + + assert first.llm_state is not None + assert first.llm_state["scope"].startswith("responses:azure:") + assert "reasoning.encrypted_content" in call.call_args_list[0].kwargs["include"] + assert REASONING in call.call_args_list[1].kwargs["input"] + finally: + client.close() + + +def test_responses_state_replays_across_gateways() -> None: + source_client = ResponsesClient( + model="openai/gpt-5.6", + api_key="account-a", + api_base="https://gateway-a.example/v1", + ) + target_client = ResponsesClient( + model="openai/gpt-5.6", + api_key="account-b", + api_base="https://gateway-b.example/v1", + ) + try: + with patch("litellm.responses", return_value=_responses(REASONING, MESSAGE)): + first = source_client.call([{"role": "user", "content": "think"}]) + with patch("litellm.responses", return_value=_responses(MESSAGE)) as target_call: + target_client.call(_render_responses(first)) + + assert REASONING in target_call.call_args.kwargs["input"] + finally: + source_client.close() + target_client.close() + + +def test_scope_is_stable_across_explicit_endpoints() -> None: + source = replay_scope( + "openai/gpt-5.6", + "responses", + {"api_key": "account-a", "api_base": "https://gateway-a.example/v1"}, + ) + target = replay_scope( + "openai/gpt-5.6", + "responses", + {"api_key": "account-a", "api_base": "https://gateway-b.example/v1"}, + ) + + assert source is not None + assert target is not None + assert source == target + + +def test_scope_is_stable_across_environment_selected_endpoints(monkeypatch) -> None: + monkeypatch.setenv("OPENAI_BASE_URL", "https://gateway-a.example/v1") + first = replay_scope("openai/gpt-5.6", "responses", {"api_key": "account-a"}) + monkeypatch.setenv("OPENAI_BASE_URL", "https://gateway-b.example/v1") + second = replay_scope("openai/gpt-5.6", "responses", {"api_key": "account-a"}) + + assert first == second + + +def test_scope_is_stable_across_auth_rotation_and_account_metadata() -> None: + first = replay_scope( + "openai/gpt-5.6", + "responses", + {"api_key": "account-a", "organization": "org-a", "project": "project-a"}, + ) + second = replay_scope( + "openai/gpt-5.6", + "responses", + {"api_key": "account-b", "organization": "org-b", "project": "project-b"}, + ) + without_auth = replay_scope("openai/gpt-5.6", "responses", {}) + + assert first is not None + assert first == second == without_auth + assert "account-a" not in first + + +def test_chat_state_is_captured_replayed_and_api_style_scoped() -> None: + client = CompletionClient( + model="openai/gpt-5.6", + api_key="account-a", + cache_control_injection_points=[], + ) + try: + with patch( + "litellm.completion", + side_effect=[_chat_response(reasoning_items=[REASONING]), _chat_response()], + ) as call: + first = client.call([{"role": "user", "content": "run"}], tools=[TOOL]) + client.call(_render_chat(first), tools=[TOOL]) + + assert first.llm_state is not None + assert first.llm_state["format"] == "litellm-chat" + assert len(first.llm_state["payload"]["carrier"]) == 64 + assert "print(1)" not in json.dumps(first.llm_state) + assistant = next( + item for item in call.call_args_list[1].kwargs["messages"] if item.get("tool_calls") + ) + assert assistant["reasoning_items"] == [REASONING] + assert assistant["reasoning_items"] is first.llm_state["payload"]["reasoning_items"] + + responses = ResponsesClient(model="openai/gpt-5.6", api_key="account-a") + try: + with patch("litellm.responses", return_value=_responses(MESSAGE)) as target_call: + responses.call(_render_chat(first)) + assert REASONING not in target_call.call_args.kwargs["input"] + finally: + responses.close() + finally: + client.close() + + +@pytest.mark.parametrize( + "mutation", ["text", "id", "name", "arguments", "drop", "reorder", "duplicate"] +) +def test_chat_state_is_not_replayed_after_public_carrier_changes( + mutation: str, caplog: pytest.LogCaptureFixture +) -> None: + client = CompletionClient( + model="openai/gpt-5.6", + api_key="account-a", + cache_control_injection_points=[], + ) + try: + response = _chat_response( + reasoning_items=[REASONING], + tool_calls=[_chat_tool_call(), _chat_tool_call("call_2", "print(2)")], + ) + with patch("litellm.completion", return_value=response): + first = client.call([{"role": "user", "content": "run"}], tools=[TOOL]) + + assert first.llm_state is not None + rendered = _render_chat(first) + assistant = next(message for message in rendered if message.get("tool_calls")) + if mutation == "text": + assistant["content"] = "replacement" + elif mutation == "id": + assistant["tool_calls"][0]["id"] = "replacement" + elif mutation == "name": + assistant["tool_calls"][0]["function"]["name"] = "replacement" + elif mutation == "arguments": + assistant["tool_calls"][0]["function"]["arguments"] = '{"code":"changed"}' + elif mutation == "drop": + assistant["tool_calls"].pop() + elif mutation == "reorder": + assistant["tool_calls"].reverse() + else: + assistant["tool_calls"][1]["id"] = assistant["tool_calls"][0]["id"] + + prepared = prepare_chat_messages(rendered, first.llm_state["scope"]) + replayed = next(message for message in prepared if message.get("role") == "assistant") + assert "reasoning_items" not in replayed + assert "provider-secret" not in repr(prepared) + assert "public assistant carrier changed" in caplog.text + finally: + client.close() + + +def test_malformed_chat_carrier_fails_closed(caplog: pytest.LogCaptureFixture) -> None: + client = CompletionClient( + model="openai/gpt-5.6", + api_key="account-a", + cache_control_injection_points=[], + ) + try: + with patch("litellm.completion", return_value=_chat_response(reasoning_items=[REASONING])): + first = client.call([{"role": "user", "content": "run"}], tools=[TOOL]) + + assert first.llm_state is not None + rendered = _render_chat(first) + state = json.loads(json.dumps(first.llm_state)) + del state["payload"]["carrier"] + carrier = next(message for message in rendered if carried_state(message) is not None) + assert isinstance(carrier, ReplayCarryingMessage) + carrier.llm_state = state + prepared = prepare_chat_messages(rendered, first.llm_state["scope"]) + + assert "provider-secret" not in repr(prepared) + assert "Opaque Chat reasoning state is malformed" in caplog.text + finally: + client.close() + + +def test_chat_state_only_carrier_is_bound_to_its_empty_turn() -> None: + client = CompletionClient( + model="openai/gpt-5.6", + api_key="account-a", + cache_control_injection_points=[], + ) + try: + with patch( + "litellm.completion", + return_value=_chat_response(reasoning_items=[REASONING], tool_calls=[]), + ): + first = client.call([{"role": "user", "content": "think"}]) + + assert first.llm_state is not None + assert len(first.llm_state["payload"]["carrier"]) == 64 + assert first.llm_state["payload"]["state_only"] is True + rendered = _render_chat(first) + exact = prepare_chat_messages(rendered, first.llm_state["scope"]) + assert any(message.get("reasoning_items") == [REASONING] for message in exact) + + carrier = next(message for message in rendered if carried_state(message) is not None) + carrier["content"] = "preserve this edit" + edited = prepare_chat_messages(rendered, first.llm_state["scope"]) + assert "provider-secret" not in repr(edited) + assert any(message.get("content") == "preserve this edit" for message in edited) + finally: + client.close() + + +def test_unresolved_route_drops_state_at_capture() -> None: + client = ResponsesClient(model="unknown-route", api_key="account-a") + try: + with ( + patch("litellm.get_llm_provider", side_effect=ValueError("unknown")), + patch("litellm.responses", return_value=_responses(REASONING, MESSAGE)) as call, + ): + response = client.call([{"role": "user", "content": "think"}]) + + assert response.llm_state is None + assert "include" not in call.call_args.kwargs + finally: + client.close() + + +def test_direct_reasoning_items_cannot_bypass_envelope_gate() -> None: + client = ResponsesClient(model="openai/gpt-5.6", api_key="account-a") + try: + with patch("litellm.responses", return_value=_responses(MESSAGE)) as call: + client.call([REASONING, {"role": "user", "content": "continue"}]) + + assert REASONING not in call.call_args.kwargs["input"] + assert "provider-secret" not in repr(call.call_args.kwargs["input"]) + finally: + client.close() + + +def test_responses_input_kwarg_cannot_bypass_replay_gate() -> None: + client = ResponsesClient(model="openai/gpt-5.6", api_key="account-a") + try: + with pytest.raises(ValueError, match="'input' is managed by UnifiedLLM"): + client.call( + [{"role": "user", "content": "continue"}], + input=[REASONING], + ) + finally: + client.close() + + +@pytest.mark.parametrize( + ("client_type", "payload_name"), + [(CompletionClient, "messages"), (ResponsesClient, "input")], +) +def test_constructor_payload_config_cannot_bypass_replay_gate(client_type, payload_name) -> None: + client = client_type( + model="openai/gpt-5.6", + api_key="account-a", + **{payload_name: [REASONING]}, + ) + try: + with pytest.raises(ValueError, match=f"'{payload_name}' is managed by UnifiedLLM"): + client.call([{"role": "user", "content": "continue"}]) + finally: + client.close() + + +@pytest.mark.parametrize( + ("client_type", "payload_name", "nested_field"), + [ + (CompletionClient, "messages", "messages"), + (CompletionClient, "messages", "model"), + (ResponsesClient, "input", "input"), + (ResponsesClient, "input", "model"), + ], +) +def test_extra_body_cannot_override_validated_payload_or_model( + client_type, payload_name, nested_field +) -> None: + client = client_type( + model="openai/gpt-5.6", + api_key="account-a", + extra_body={nested_field: [REASONING]}, + ) + try: + with pytest.raises(ValueError, match="extra_body may not override reserved field"): + client.call([{"role": "user", "content": "continue"}]) + finally: + client.close() + + +@pytest.mark.asyncio +async def test_async_responses_input_kwarg_cannot_bypass_replay_gate() -> None: + client = ResponsesClient(model="openai/gpt-5.6", api_key="account-a") + try: + with pytest.raises(ValueError, match="'input' is managed by UnifiedLLM"): + await client.acall( + [{"role": "user", "content": "continue"}], + input=[REASONING], + ) + finally: + await client.aclose() + + +@pytest.mark.asyncio +async def test_responses_reasoning_uses_per_call_override() -> None: + client = ResponsesClient( + model="openai/gpt-5.6", + api_key="account-a", + reasoning={"effort": "high"}, + ) + try: + with patch("litellm.responses", return_value=_responses(MESSAGE)) as sync_call: + client.call( + [{"role": "user", "content": "continue"}], + reasoning={"effort": "low"}, + ) + with patch("litellm.aresponses", AsyncMock(return_value=_responses(MESSAGE))) as async_call: + await client.acall( + [{"role": "user", "content": "continue"}], + reasoning=None, + ) + + assert sync_call.call_args.kwargs["reasoning"] == {"effort": "low"} + assert async_call.call_args.kwargs["reasoning"] is None + finally: + await client.aclose() + + +@pytest.mark.parametrize("model", ["anthropic/claude-sonnet-4-5", "gemini/gemini-2.5-pro"]) +def test_non_openai_chat_provider_cannot_receive_reasoning_state(model: str) -> None: + assert replay_scope(model, "chat", {"api_key": "account-a"}) is None + client = CompletionClient(model=model, api_key="account-a") + crafted = { + "version": 2, + "scope": f"chat:{model.split('/', 1)[0]}:crafted", + "format": "litellm-chat", + "payload": {"reasoning_items": [REASONING]}, + } + try: + with patch("litellm.completion", return_value=_chat_response()) as call: + client.call( + [ReplayCarryingMessage({"role": "assistant", "content": "public"}, crafted)] + ) + + assert call.call_args.kwargs["messages"] == [{"role": "assistant", "content": "public"}] + assert "provider-secret" not in repr(call.call_args.kwargs) + finally: + client.close() + + +def test_custom_endpoint_does_not_assume_encrypted_include_support() -> None: + client = ResponsesClient( + model="openai/gpt-5.6", + api_key="account-a", + api_base="https://gateway.example/v1", + ) + try: + with patch("litellm.responses", return_value=_responses(MESSAGE)) as call: + client.call([{"role": "user", "content": "hello"}]) + assert "include" not in call.call_args.kwargs + finally: + client.close() + + +@pytest.mark.asyncio +async def test_async_responses_capture_matches_sync() -> None: + client = ResponsesClient(model="openai/gpt-5.6", api_key="account-a") + try: + with patch( + "litellm.aresponses", AsyncMock(return_value=_responses(REASONING, MESSAGE)) + ) as call: + response = await client.acall([{"role": "user", "content": "think"}]) + + assert response.llm_state is not None + assert response.llm_state["payload"]["items"] == [REASONING] + assert "reasoning.encrypted_content" in call.call_args.kwargs["include"] + finally: + await client.aclose() diff --git a/tests/unifiedllm/test_replay_copy_ownership.py b/tests/unifiedllm/test_replay_copy_ownership.py new file mode 100644 index 000000000..3f2d908a3 --- /dev/null +++ b/tests/unifiedllm/test_replay_copy_ownership.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Replay detaches public input once and never copies borrowed/rejected state.""" + +import copy +from typing import Any + +import pytest + +from nooa._llm_state import LLM_STATE_KEY, ReplayCarryingMessage, carry_replay_batch +from nooa.unifiedllm import ResponsesClient +from nooa.unifiedllm.replay_state import prepare_chat_messages, prepare_responses_batch + +SCOPE = "responses:openai:sha256:test" + + +class _NoDeepCopy(dict[str, Any]): + def __deepcopy__(self, memo: dict[int, Any]) -> Any: + raise AssertionError("Opaque state must not be copied during request preparation") + + +class _CountCopies(dict[str, Any]): + def __init__(self, value: dict[str, Any], copies: list[None]): + super().__init__(value) + self.copies = copies + + def __deepcopy__(self, memo: dict[int, Any]) -> "_CountCopies": + self.copies.append(None) + return _CountCopies(copy.deepcopy(dict(self), memo), self.copies) + + +@pytest.mark.parametrize("shape", ["batch", "chat", "native"]) +@pytest.mark.parametrize("compatible", [True, False]) +def test_responses_detaches_public_content_once_without_copying_state( + shape: str, compatible: bool +) -> None: + opaque = _NoDeepCopy( + type="reasoning", id="rs_test", encrypted_content="test-ciphertext", summary=[] + ) + state = { + "version": 2, + "scope": SCOPE, + "format": "openai-responses", + "payload": { + "items": [opaque], + "order": [ + {"type": "reasoning", "index": 0}, + {"type": "message", "content": "hello"}, + ], + }, + } + copies: list[None] = [] + marker = _CountCopies({"type": "ephemeral"}, copies) + message: dict[str, Any] = {"role": "assistant", "content": "hello", "cache_control": marker} + if shape == "native": + message["type"] = "message" + messages: list[dict[str, Any]] = ( + carry_replay_batch([message], state, "why") + if shape == "batch" + else [ReplayCarryingMessage(message, state, "why")] + ) + target = SCOPE if compatible else "responses:openai:sha256:other" + with ResponsesClient(model="openai/gpt-5.6", api_key="test") as client: + prepared, _ = client._transform_messages(messages, target) + + assert len(copies) == 1 + if compatible: + assert prepared[0] is opaque + else: + assert all(item.get("type") != "reasoning" for item in prepared) + assistant = next(item for item in prepared if item.get("role") == "assistant") + assert assistant["content"] == ("hello" if compatible else "why\n\nhello") + assistant["cache_control"]["type"] = "changed" + assert marker == {"type": "ephemeral"} + assert message["content"] == "hello" + + +@pytest.mark.parametrize("path", ["chat", "responses", "batch"]) +@pytest.mark.parametrize("key", [LLM_STATE_KEY, "reasoning_items"]) +def test_rejected_raw_state_is_removed_before_copying(path: str, key: str) -> None: + message = {"role": "assistant", "content": "hello", key: _NoDeepCopy(secret="test")} + if path == "chat": + prepared = prepare_chat_messages([message], None) + elif path == "batch": + prepared = prepare_responses_batch([message], None, None) + else: + with ResponsesClient(model="openai/gpt-5.6", api_key="test") as client: + prepared, _ = client._transform_messages([message]) + + assert prepared == [{"role": "assistant", "content": "hello"}] + assert key in message + + +@pytest.mark.parametrize("role", ["user", "assistant", "tool"]) +def test_responses_passthrough_still_detaches_nested_public_content(role: str) -> None: + marker = {"type": "ephemeral"} + message = {"role": role, "content": "hello", "cache_control": marker} + if role == "tool": + message["tool_call_id"] = "call_test" + with ResponsesClient(model="openai/gpt-5.6", api_key="test") as client: + prepared, _ = client._transform_messages([message]) + + prepared[0]["cache_control"]["type"] = "changed" + assert marker == {"type": "ephemeral"} diff --git a/tests/unifiedllm/test_replay_fingerprints.py b/tests/unifiedllm/test_replay_fingerprints.py new file mode 100644 index 000000000..c173242ad --- /dev/null +++ b/tests/unifiedllm/test_replay_fingerprints.py @@ -0,0 +1,187 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Compact replay bindings survive storage, reject edits, and validate wire items.""" + +import copy +import json +import sqlite3 + +import pytest + +from nooa._llm_state import ReplayCarryingMessage +from nooa.storage.sqlite import SQLiteEventBackend, _ensure_schema +from nooa.unifiedllm import LLMResponse, ToolCall +from nooa.unifiedllm.replay_state import ( + capture_chat_state, + capture_responses_state, + prepare_chat_messages, + prepare_responses_batch, +) + +REASONING = {"type": "reasoning", "id": "rs_test", "encrypted_content": "ciphertext", "summary": []} +INVALID_ITEMS = [ + 123, + {}, + {**REASONING, "type": "message"}, + {**REASONING, "encrypted_content": 123}, + {**REASONING, "encrypted_content": ""}, + {"type": "reasoning", "summary": []}, +] + + +@pytest.mark.parametrize("invalid", INVALID_ITEMS) +def test_chat_capture_rejects_malformed_reasoning_items(invalid, caplog) -> None: + assert ( + capture_chat_state( + {"role": "assistant", "content": "answer", "reasoning_items": [REASONING, invalid]}, + "chat:openai:sha256:test", + ) + is None + ) + assert "malformed" in caplog.text + assert "ciphertext" not in caplog.text + + +@pytest.mark.parametrize("api", ["chat", "responses"]) +@pytest.mark.parametrize("invalid", INVALID_ITEMS) +def test_corrupt_stored_items_never_reach_replay(api, invalid, caplog) -> None: + scope = f"{api}:openai:sha256:test" + public = {"role": "assistant", "content": "answer"} + if api == "chat": + state = capture_chat_state({**public, "reasoning_items": [REASONING]}, scope) + assert state is not None + state["payload"]["reasoning_items"] = [REASONING, invalid] + result = prepare_chat_messages([ReplayCarryingMessage(public, state, "why")], scope) + else: + state = capture_responses_state( + [ + REASONING, + {"type": "message", "content": [{"type": "output_text", "text": "answer"}]}, + ], + scope, + ) + assert state is not None + state["payload"]["items"][0] = invalid + result = prepare_responses_batch([public], state, scope, "why") + assert result == [{"role": "assistant", "content": "why\n\nanswer"}] + assert "malformed" in caplog.text + assert "ciphertext" not in caplog.text + + +@pytest.mark.parametrize("api", ["chat", "responses"]) +@pytest.mark.parametrize("persistence", ["json", "sqlite"]) +def test_large_tool_arguments_are_stored_once_and_replay_exactly(api, persistence) -> None: + arguments = json.dumps({"code": "x" * 100_000 + "é"}) + scope = f"{api}:openai:sha256:test" + call = { + "type": "function_call", + "call_id": "call_test", + "name": "execute_python", + "arguments": arguments, + } + chat = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_test", + "type": "function", + "function": {"name": "execute_python", "arguments": arguments}, + } + ], + } + state = ( + capture_chat_state({**chat, "reasoning_items": [REASONING]}, scope) + if api == "chat" + else capture_responses_state([REASONING, call], scope) + ) + assert state is not None + response = LLMResponse( + tool_calls=[ToolCall(id="call_test", name="execute_python", arguments=arguments)], + llm_state=state, + ) + encoded = response.model_dump_json() + # The envelope adds constant-size bindings, not another 100 KB argument. + assert len(encoded) - len(response.model_dump_json(exclude={"llm_state"})) < 1000 + assert "x" * 100 not in json.dumps(state) + if persistence == "json": + resumed = LLMResponse.model_validate_json(encoded) + else: + connection = sqlite3.connect(":memory:") + try: + _ensure_schema(connection) + backend = SQLiteEventBackend(connection) + backend.store("response", response) + resumed = backend.get("response") + assert isinstance(resumed, LLMResponse) + finally: + connection.close() + assert resumed.llm_state == state + if api == "chat": + result = prepare_chat_messages([ReplayCarryingMessage(chat, resumed.llm_state)], scope) + assert result == [{**chat, "reasoning_items": [REASONING]}] + else: + assert prepare_responses_batch([call], resumed.llm_state, scope) == [REASONING, call] + + +@pytest.mark.parametrize("api", ["chat", "responses"]) +def test_prior_draft_envelope_warns_and_demotes_text(api, caplog) -> None: + scope = f"{api}:openai:sha256:test" + public = {"role": "assistant", "content": "answer"} + state = { + "version": 1, + "scope": scope, + "format": "litellm-chat" if api == "chat" else "openai-responses", + "payload": {}, + } + result = ( + prepare_chat_messages([ReplayCarryingMessage(public, state, "why")], scope) + if api == "chat" + else prepare_responses_batch([public], state, scope, "why") + ) + assert result == [{"role": "assistant", "content": "why\n\nanswer"}] + assert "unsupported or legacy version" in caplog.text + + +@pytest.mark.parametrize("mutation", ["id", "name", "arguments", "drop", "reorder"]) +def test_responses_fingerprints_reject_edited_calls(mutation, caplog) -> None: + scope = "responses:openai:sha256:test" + calls = [ + { + "type": "function_call", + "call_id": f"call_{i}", + "name": "execute_python", + "arguments": '{"code":"print(1)"}', + } + for i in range(2) + ] + state = capture_responses_state([REASONING, *calls], scope) + assert state is not None + edited = copy.deepcopy(calls) + if mutation == "drop": + edited.pop() + elif mutation == "reorder": + edited.reverse() + else: + edited[0]["call_id" if mutation == "id" else mutation] = "changed" + result = prepare_responses_batch(edited, state, scope) + assert result == edited + assert "carriers changed" in caplog.text + + +def test_chat_fingerprint_is_key_order_independent_and_does_not_store_text() -> None: + scope = "chat:openai:sha256:test" + message = { + "role": "assistant", + "content": "large public text" * 1000, + "reasoning_items": [REASONING], + } + state = capture_chat_state(message, scope) + assert state is not None + reordered = dict(reversed(list(message.items()))) + assert capture_chat_state(reordered, scope) == state + assert "large public text" not in json.dumps(state) + public = {"content": message["content"], "role": "assistant"} + assert prepare_chat_messages([ReplayCarryingMessage(public, state)], scope) == [ + {**public, "reasoning_items": [REASONING]} + ] diff --git a/tests/unifiedllm/test_responses_cache_control.py b/tests/unifiedllm/test_responses_cache_control.py index 3ce014a08..aa343735e 100644 --- a/tests/unifiedllm/test_responses_cache_control.py +++ b/tests/unifiedllm/test_responses_cache_control.py @@ -6,6 +6,8 @@ import pytest +from nooa.context_blocks.formatter import ResponsesProviderFormatter +from nooa.context_blocks.models import RenderedMessage, Role, ToolCallInfo from nooa.unifiedllm import ResponsesClient @@ -119,6 +121,28 @@ def test_user_message_cache_control_preserved(self, client): user_msgs = [m for m in input_msgs if m.get("role") == "user"] assert user_msgs[0].get("cache_control") == {"type": "ephemeral"} + def test_stateful_assistant_batch_is_visible_to_cache_injection(self, client): + messages = ResponsesProviderFormatter().format( + [ + RenderedMessage( + role=Role.ASSISTANT, + content="I will run it.", + tool_calls=(ToolCallInfo(id="c1", name="run", arguments="{}"),), + llm_state={"opaque": "provider-only"}, + ) + ] + ) + + prepared = client._inject_cache_control( + messages, [{"role": "assistant", "position": "last"}] + ) + input_messages, _ = client._transform_messages(prepared) + + assistant = next( + message for message in input_messages if message.get("role") == "assistant" + ) + assert assistant["cache_control"] == {"type": "ephemeral"} + class TestToolOutputNotCorrupted: """Ensure tool message output stays a string after position-based injection.""" @@ -167,6 +191,32 @@ class TestResponsesClientEndToEnd: ANTHROPIC_MODEL = "anthropic/claude-haiku-4-5" + def test_assistant_cache_copy_preserves_portable_reasoning(self): + """Copy-on-write cache marking must retain private replay metadata.""" + client = ResponsesClient(model=self.ANTHROPIC_MODEL) + messages = ResponsesProviderFormatter().format( + [ + RenderedMessage( + role=Role.ASSISTANT, + content="public answer", + reasoning="portable reasoning text", + ) + ] + ) + try: + with patch("litellm.responses", return_value=make_mock_responses_response()) as call: + client.call( + messages, + cache_control_injection_points=[{"role": "assistant", "position": "last"}], + ) + + sent = call.call_args.kwargs["input"] + assert sent[0] == {"role": "assistant", "content": "portable reasoning text"} + assert sent[1]["content"][0]["text"] == "public answer" + assert sent[1]["content"][0]["cache_control"] == {"type": "ephemeral"} + finally: + client.close() + @pytest.mark.asyncio async def test_acall_injects_cache_control(self): """acall() injects cache_control on messages before calling litellm.""" diff --git a/tests/unifiedllm/test_responses_client_retry.py b/tests/unifiedllm/test_responses_client_retry.py index 3c0cc4514..4b5053716 100644 --- a/tests/unifiedllm/test_responses_client_retry.py +++ b/tests/unifiedllm/test_responses_client_retry.py @@ -7,6 +7,7 @@ RetryConfig lists 502 as retryable. These tests verify parity with CompletionClient. """ +import hashlib from unittest.mock import AsyncMock, MagicMock, patch import litellm @@ -103,35 +104,52 @@ def test_non_retryable_not_retried(self): def test_reasoning_state_retains_interleaving_without_copying_public_calls(self): """The canonical state has enough anchors for exact ordered replay.""" - reasoning_1 = MagicMock(type="reasoning") - reasoning_1.model_dump.return_value = {"type": "reasoning", "encrypted": "one"} - call_1 = MagicMock(type="function_call", call_id="call-1", arguments="{}") - call_1.name = "one" - reasoning_2 = MagicMock(type="reasoning") - reasoning_2.model_dump.return_value = {"type": "reasoning", "encrypted": "two"} - call_2 = MagicMock(type="function_call", call_id="call-2", arguments="{}") - call_2.name = "two" + reasoning_1 = {"type": "reasoning", "encrypted_content": "one"} + call_1 = { + "type": "function_call", + "call_id": "call-1", + "name": "one", + "arguments": "{}", + } + reasoning_2 = {"type": "reasoning", "encrypted_content": "two"} + call_2 = { + "type": "function_call", + "call_id": "call-2", + "name": "two", + "arguments": "{}", + } raw_response = MagicMock( output=[reasoning_1, call_1, reasoning_2, call_2], output_text="", usage=None, ) - client = ResponsesClient(model="test-model", retry_config=NO_RETRY) + client = ResponsesClient(model="openai/gpt-5.6", api_key="test", retry_config=NO_RETRY) with patch("litellm.responses", return_value=raw_response): response = client.call(messages=[{"role": "user", "content": "hi"}]) assert [call.id for call in response.tool_calls] == ["call-1", "call-2"] - assert response.llm_state == { + assert response.llm_state is not None + assert response.llm_state["payload"] == { "items": [ - {"type": "reasoning", "encrypted": "one"}, - {"type": "reasoning", "encrypted": "two"}, + {"type": "reasoning", "encrypted_content": "one"}, + {"type": "reasoning", "encrypted_content": "two"}, ], "order": [ {"type": "reasoning", "index": 0}, - {"type": "function_call", "call_id": "call-1"}, + { + "type": "function_call", + "call_id": "call-1", + "name": "one", + "arguments_sha256": hashlib.sha256(b'"{}"').hexdigest(), + }, {"type": "reasoning", "index": 1}, - {"type": "function_call", "call_id": "call-2"}, + { + "type": "function_call", + "call_id": "call-2", + "name": "two", + "arguments_sha256": hashlib.sha256(b'"{}"').hexdigest(), + }, ], } diff --git a/tests/unifiedllm/test_responses_formatter.py b/tests/unifiedllm/test_responses_formatter.py index d59bfced8..43d18881f 100644 --- a/tests/unifiedllm/test_responses_formatter.py +++ b/tests/unifiedllm/test_responses_formatter.py @@ -64,6 +64,38 @@ def test_tool_call_format(self): } ] + def test_stateful_tool_batch_remains_valid_public_responses_input(self): + """Replay metadata must not replace public items with a private wrapper key.""" + result = ResponsesProviderFormatter().format( + [ + RenderedMessage( + role=Role.ASSISTANT, + content="I will run it.", + tool_calls=( + ToolCallInfo( + id="call_123", + name="execute_python", + arguments={"code": "print(1)"}, + ), + ), + llm_state={"opaque": "provider-only"}, + ) + ] + ) + + assert result == [ + {"role": "assistant", "content": "I will run it."}, + { + "type": "function_call", + "call_id": "call_123", + "name": "execute_python", + "arguments": json.dumps({"code": "print(1)"}), + }, + ] + assert all("role" in item or "type" in item for item in result) + assert "_batch" not in json.dumps(result) + assert "provider-only" not in json.dumps(result) + def test_tool_result_format(self): """Tool results become function_call_output items.""" messages = [ diff --git a/tests/unifiedllm/test_responses_transport_overrides.py b/tests/unifiedllm/test_responses_transport_overrides.py new file mode 100644 index 000000000..1b58fc238 --- /dev/null +++ b/tests/unifiedllm/test_responses_transport_overrides.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Responses HTTP handlers reuse their pool without binding a URL, key, or model.""" + +import json + +import httpx +import pytest + +from nooa.unifiedllm import ResponsesClient +from nooa.unifiedllm.unifiedllm import _ClientHttp + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize( + "overrides", + [ + {}, + {"api_base": "https://override.example/v1"}, + {"base_url": "https://override.example/v1"}, + {"api_key": "override-key"}, + {"model": "openai/gpt-5.4"}, + {"api_base": "https://override.example/v1", "api_key": "override-key"}, + ], +) +async def test_responses_transport_honors_overrides_without_changing_defaults( + monkeypatch, is_async, overrides +): + requests: list[httpx.Request] = [] + + def respond(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response( + 200, + json={ + "id": "resp_test", + "created_at": 0, + "model": "gpt-5.6", + "status": "completed", + "output": [ + { + "id": "msg_test", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "done", "annotations": []}], + } + ], + }, + ) + + transport = httpx.MockTransport(respond) + monkeypatch.setattr( + _ClientHttp, "_httpx_hardening", staticmethod(lambda: {"transport": transport}) + ) + + # Fail locally if a regression discards our pool and tries the network. + def unexpected_network(*args, **kwargs): + raise AssertionError("Responses must reuse the configured HTTP transport") + + monkeypatch.setattr(httpx.HTTPTransport, "handle_request", unexpected_network) + monkeypatch.setattr(httpx.AsyncHTTPTransport, "handle_async_request", unexpected_network) + client = ResponsesClient( + "openai/gpt-5.6", + api_base="https://original.example/v1", + api_key="original-key", + ) + try: + # Exercise real LiteLLM serialization and dispatch, not mocked call kwargs. + for params in ({}, overrides, {}): + messages = [{"role": "user", "content": "hello"}] + result = ( + await client.acall(messages, **params) + if is_async + else client.call(messages, **params) + ) + assert result.content == "done" + + override_base = overrides.get( + "base_url", overrides.get("api_base", "https://original.example/v1") + ) + assert [str(request.url) for request in requests] == [ + "https://original.example/v1/responses", + override_base + "/responses", + "https://original.example/v1/responses", + ] + assert [request.headers["Authorization"] for request in requests] == [ + "Bearer original-key", + "Bearer " + overrides.get("api_key", "original-key"), + "Bearer original-key", + ] + assert [json.loads(request.content)["model"] for request in requests] == [ + "gpt-5.6", + overrides.get("model", "openai/gpt-5.6").removeprefix("openai/"), + "gpt-5.6", + ] + finally: + await client.aclose()