From 346898319acd188bd19224ef3a9cda089048ed10 Mon Sep 17 00:00:00 2001 From: Paul Furgale Date: Wed, 9 Sep 2026 17:14:56 +0000 Subject: [PATCH 01/13] feat(llm): retain cross-provider reasoning state Signed-off-by: Paul Furgale --- skills/nooa-context-and-state/SKILL.md | 2 + src/nooa/unifiedllm/http_logging.py | 15 +- src/nooa/unifiedllm/replay_state.py | 269 +++++++---- src/nooa/unifiedllm/unifiedllm.py | 24 +- .../test_cross_provider_reasoning.py | 423 ++++++++++++++++++ tests/unifiedllm/test_http_logging.py | 6 + .../unifiedllm/test_reasoning_state_replay.py | 7 +- 7 files changed, 653 insertions(+), 93 deletions(-) create mode 100644 tests/unifiedllm/test_cross_provider_reasoning.py diff --git a/skills/nooa-context-and-state/SKILL.md b/skills/nooa-context-and-state/SKILL.md index 5364d7e4e..5a91c808a 100644 --- a/skills/nooa-context-and-state/SKILL.md +++ b/skills/nooa-context-and-state/SKILL.md @@ -87,6 +87,8 @@ async def solve(self, problem: str) -> str: Event history is what fills the LLM's conversation window. Key model-visible event types (names have no "Event" suffix): `Task`, `Message`, `Reasoning`, `Error`, `Feedback`, `LLMResponse`, `PythonOutput`, `Summary`, `Notification`. `LLMResponse` is both the canonical assistant turn and the home of its token/cost metadata; renderers expose only its conversational fields. Runtime-only events (never shown to the LLM) include `BeforeAgentCall`/`AfterAgentCall` and `LLMCallStart`/`LLMCallEnd`. +`LLMResponse.reasoning` is provider-exposed text and remains useful across model switches, where UnifiedLLM replays it as ordinary assistant text. `LLMResponse.llm_state` is opaque provider state; it is persisted for resume but is replayed only through a matching provider/API/issuer gate. + ```python # Query (AND semantics; chronological; limit keeps most recent) recent = agent.events.query(limit=20) diff --git a/src/nooa/unifiedllm/http_logging.py b/src/nooa/unifiedllm/http_logging.py index 1349df488..fa27ab430 100644 --- a/src/nooa/unifiedllm/http_logging.py +++ b/src/nooa/unifiedllm/http_logging.py @@ -8,13 +8,24 @@ from nooa._llm_state import LLM_STATE_KEY from nooa.tracing._secret_scrubber import REDACTED, _is_sensitive_key, scrub_value +_OPAQUE_STATE_KEYS = { + LLM_STATE_KEY, + "encrypted_content", + "signature", + "thought_signature", + "thought_signatures", + "thoughtSignature", + "thoughtSignatures", +} + def _redact_opaque_state(value): - """Remove replay envelopes and encrypted reasoning from debug logs.""" + """Remove provider-owned reasoning state from debug logs.""" if isinstance(value, dict): + redacted_thinking = value.get("type") == "redacted_thinking" return { key: REDACTED - if key in {LLM_STATE_KEY, "encrypted_content"} + if key in _OPAQUE_STATE_KEYS or (redacted_thinking and key == "data") else _redact_opaque_state(item) for key, item in value.items() } diff --git a/src/nooa/unifiedllm/replay_state.py b/src/nooa/unifiedllm/replay_state.py index 314c590bb..ac390d03f 100644 --- a/src/nooa/unifiedllm/replay_state.py +++ b/src/nooa/unifiedllm/replay_state.py @@ -1,6 +1,6 @@ # 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. +"""Compatibility-scoped capture and replay of closed-provider 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 @@ -33,6 +33,13 @@ _CHAT_FORMAT = "litellm-chat" _RESPONSES_FORMAT = "openai-responses" _ENCRYPTED_REASONING_INCLUDE = "reasoning.encrypted_content" +_INLINE_THOUGHT_SIGNATURE_SEPARATOR = "__thought__" +_SUPPORTED_PROVIDERS = { + "openai", + "azure", + "anthropic", + "gemini", +} def _field(value: Any, key: str, default: Any = None) -> Any: @@ -102,10 +109,9 @@ def replay_scope( 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"}: + if provider not in _SUPPORTED_PROVIDERS or ( + api_style == "responses" and provider not in {"openai", "azure"} + ): return None digest = hashlib.sha256(resolved_model.encode()).hexdigest() @@ -138,73 +144,121 @@ def _matching_payload(state: Any, scope: str | None, state_format: str) -> dict return cast(dict[str, Any], state["payload"]) -def _chat_public_carrier(message: Any) -> dict[str, Any] | None: - """Project 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 +def _is_state_only(state: Any) -> bool: + return ( + isinstance(state, dict) + and state.get("version") == _STATE_VERSION + and isinstance(state.get("payload"), dict) + and state["payload"].get("state_only") is True + ) + - 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}) - - content = _field(message, "content") - if content is not None and not isinstance(content, str): +def _scope_provider(scope: str | None) -> str | None: + if not isinstance(scope, str): return None - content = (content or None) if calls else (content or "") - return {"content": content, "tool_calls": calls} + parts = scope.split(":", 2) + return parts[1] if len(parts) == 3 else None + + +def _sanitize_chat_payload(payload: dict[str, Any], scope: str | None) -> dict[str, Any]: + """Keep only replay fields whose LiteLLM Chat semantics NOOA knows.""" + provider = _scope_provider(scope) + clean: dict[str, Any] = {} + fields_by_provider = { + "openai": ("reasoning_items", "thinking_blocks"), + "azure": ("reasoning_items", "thinking_blocks"), + "anthropic": ("thinking_blocks",), + "gemini": ("thinking_blocks",), + } + for key in fields_by_provider.get(provider or "", ()): + value = payload.get(key) + if isinstance(value, list) and value: + clean[key] = opaque_item(value) + + provider_fields = ( + payload.get("provider_specific_fields") + if provider in {"openai", "azure", "gemini"} + else None + ) + signatures = ( + provider_fields.get("thought_signatures") if isinstance(provider_fields, dict) else None + ) + if ( + isinstance(signatures, list) + and signatures + and all(isinstance(signature, str) and signature for signature in signatures) + ): + clean["provider_specific_fields"] = {"thought_signatures": copy.deepcopy(signatures)} + tool_state = ( + payload.get("tool_calls") + if provider in {"openai", "azure", "gemini"} + else None + ) + if isinstance(tool_state, list): + calls: list[dict[str, Any] | None] = [] + for item in tool_state: + fields = item.get("provider_specific_fields") if isinstance(item, dict) else None + signature = fields.get("thought_signature") if isinstance(fields, dict) else None + calls.append( + {"provider_specific_fields": {"thought_signature": signature}} + if isinstance(signature, str) and signature + else None + ) + if any(item is not None for item in calls): + clean["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 isinstance(items, list) or not items or not isinstance(carrier, dict): - return False - if set(carrier) != {"content", "tool_calls"}: - return False - if carrier.get("content") is not None and not isinstance(carrier.get("content"), str): - return False - calls = carrier.get("tool_calls") - if not isinstance(calls, list) or not all( - isinstance(call, dict) - and set(call) == {"id", "name", "arguments"} - and all(isinstance(call.get(key), str) for key in ("id", "name", "arguments")) - for call in calls + if clean and payload.get("state_only") is True: + clean["state_only"] = True + return clean + + +def _tool_call_state(tool_call: Any) -> dict[str, Any] | None: + dumped = opaque_item(tool_call) + if not isinstance(dumped, dict): + return None + fields = dumped.get("provider_specific_fields") + signature = fields.get("thought_signature") if isinstance(fields, dict) else None + call_id = dumped.get("id") + if ( + not signature + and isinstance(call_id, str) + and _INLINE_THOUGHT_SIGNATURE_SEPARATOR in call_id ): - return False - if len({call["id"] for call in calls}) != len(calls): - return False - if "state_only" in payload and payload["state_only"] is not True: - return False - return (payload.get("state_only") is True) == (carrier == {"content": "", "tool_calls": []}) + signature = call_id.split(_INLINE_THOUGHT_SIGNATURE_SEPARATOR, 1)[1] + if not isinstance(signature, str) or not signature: + return None + return {"provider_specific_fields": {"thought_signature": signature}} + + +def public_tool_call_id(value: Any, scope: str | None) -> str: + """Return an application call id without LiteLLM's inline Gemini state.""" + call_id = _field(value, "id", "") + if not isinstance(call_id, str): + return str(call_id or "") + if _scope_provider(scope) != "gemini": + return call_id + return call_id.split(_INLINE_THOUGHT_SIGNATURE_SEPARATOR, 1)[0] 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.") + payload: dict[str, Any] = {} + for key in ("reasoning_items", "thinking_blocks"): + value = _field(message, key) + if isinstance(value, list) and value: + payload[key] = opaque_item(value) + + provider_fields = _field(message, "provider_specific_fields") + if isinstance(provider_fields, dict): + payload["provider_specific_fields"] = provider_fields + + tool_state = [_tool_call_state(call) for call in (_field(message, "tool_calls") or [])] + if any(item is not None for item in tool_state): + payload["tool_calls"] = tool_state + + payload = _sanitize_chat_payload(payload, scope) + if not payload: 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): @@ -305,7 +359,52 @@ def _public_responses_carriers(items: list[dict[str, Any]]) -> list[dict[str, An return carriers +def _strip_inline_signature(value: Any) -> Any: + if isinstance(value, str) and _INLINE_THOUGHT_SIGNATURE_SEPARATOR in value: + return value.split(_INLINE_THOUGHT_SIGNATURE_SEPARATOR, 1)[0] + return value + + +def _strip_chat_state(message: dict[str, Any], *, strip_inline_signatures: bool) -> None: + for key in ("reasoning_items", "thinking_blocks", "provider_specific_fields"): + message.pop(key, None) + tool_calls = message.get("tool_calls") + if isinstance(tool_calls, list): + for call in tool_calls: + if not isinstance(call, dict): + continue + if strip_inline_signatures: + call["id"] = _strip_inline_signature(call.get("id")) + call.pop("provider_specific_fields", None) + function = call.get("function") + if isinstance(function, dict): + function.pop("provider_specific_fields", None) + if strip_inline_signatures and "tool_call_id" in message: + message["tool_call_id"] = _strip_inline_signature(message["tool_call_id"]) + + +def _restore_chat_state(message: dict[str, Any], payload: dict[str, Any]) -> None: + for key in ("reasoning_items", "thinking_blocks", "provider_specific_fields"): + if key in payload: + message[key] = copy.deepcopy(payload[key]) + tool_calls = message.get("tool_calls") + tool_state = payload.get("tool_calls") + if not isinstance(tool_calls, list) or not isinstance(tool_state, list): + return + for call, state in zip(tool_calls, tool_state, strict=False): + if not isinstance(call, dict) or not isinstance(state, dict): + continue + fields = state.get("provider_specific_fields") + if isinstance(fields, dict): + call["provider_specific_fields"] = copy.deepcopy(fields) + + def capture_responses_state(output: list[Any], scope: str | None) -> dict | None: + # NOOA only knows the OpenAI/Azure Responses item contract. Other providers + # may expose a similarly shaped API through a gateway, but that is not + # evidence that their opaque state is wire-compatible. + if _scope_provider(scope) not in {"openai", "azure"}: + return None items: list[Any] = [] order: list[dict[str, Any]] = [] has_public_carrier = False @@ -333,6 +432,19 @@ def capture_responses_state(output: list[Any], scope: str | None) -> dict | None return _envelope(scope, _RESPONSES_FORMAT, payload) +def responses_reasoning_text(output: list[Any]) -> str | None: + """Return provider-visible Responses reasoning summaries as plain text.""" + texts: list[str] = [] + for item in output: + if response_item_type(item) != "reasoning": + continue + for summary in _field(item, "summary", []) or []: + text = _field(summary, "text") + if isinstance(text, str) and text: + texts.append(text) + return "\n".join(texts) or None + + 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]] = [] @@ -341,23 +453,16 @@ def prepare_chat_messages(messages: list[dict[str, Any]], scope: str | None) -> reasoning = carried_reasoning(original) message = copy.deepcopy(dict(original)) message.pop(LLM_STATE_KEY, None) - message.pop("reasoning_items", None) + source_scope = state.get("scope") if isinstance(state, dict) else None + _strip_chat_state( + message, strip_inline_signatures=_scope_provider(source_scope) == "gemini" + ) 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: + if payload is not None: + payload = _sanitize_chat_payload(payload, scope) or None + if payload: + _restore_chat_state(message, payload) + else: demote_reasoning_text(message, reasoning) if ( not restored @@ -394,7 +499,11 @@ def prepare_responses_batch( ) -> 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) + payload = ( + _matching_payload(state, scope, _RESPONSES_FORMAT) + if _scope_provider(scope) in {"openai", "azure"} + else None + ) if payload is None: return demote_responses_batch(clean, state, reasoning) items = payload.get("items") diff --git a/src/nooa/unifiedllm/unifiedllm.py b/src/nooa/unifiedllm/unifiedllm.py index d3d824967..77f114623 100644 --- a/src/nooa/unifiedllm/unifiedllm.py +++ b/src/nooa/unifiedllm/unifiedllm.py @@ -1894,7 +1894,11 @@ def _make_call(): if raw_tool_calls: response_message = raw_response.choices[0].message tool_calls = [ - ToolCall(id=tc.id, name=tc.function.name or "", arguments=tc.function.arguments) + ToolCall( + id=replay_state.public_tool_call_id(tc, state_scope), + name=tc.function.name or "", + arguments=tc.function.arguments, + ) for tc in raw_tool_calls ] @@ -2066,7 +2070,9 @@ async def _make_call(): response_message = raw_response.choices[0].message tool_calls = [ ToolCall( - id=tc.id, name=tc.function.name or "", arguments=tc.function.arguments or "" + id=replay_state.public_tool_call_id(tc, state_scope), + name=tc.function.name or "", + arguments=tc.function.arguments or "", ) # type: ignore[union-attr] for tc in raw_tool_calls ] @@ -2385,6 +2391,7 @@ def _make_call(): ) output: list[Any] = raw_response.output # type: ignore[assignment] + reasoning = replay_state.responses_reasoning_text(output) raw_tool_calls = [ item for item in output if replay_state.response_item_type(item) == "function_call" ] @@ -2406,7 +2413,7 @@ def _make_call(): finish_reason=_finish_reason_for_tool_calls( _map_responses_finish_reason(raw_response) ), - reasoning=None, # Responses API doesn't have reasoning + reasoning=reasoning, usage=usage, llm_state=replay_state.capture_responses_state(output, state_scope), ) @@ -2423,7 +2430,7 @@ def _make_call(): parsed=parsed_content, tool_calls=[], finish_reason=_map_responses_finish_reason(raw_response), - reasoning=None, + reasoning=reasoning, usage=usage, llm_state=replay_state.capture_responses_state(output, state_scope), ) @@ -2433,7 +2440,7 @@ def _make_call(): content=text_content, tool_calls=[], finish_reason=_map_responses_finish_reason(raw_response), - reasoning=None, + reasoning=reasoning, usage=usage, llm_state=replay_state.capture_responses_state(output, state_scope), ) @@ -2518,6 +2525,7 @@ async def _make_call(): ) output: list[Any] = raw_response.output # type: ignore[assignment] + reasoning = replay_state.responses_reasoning_text(output) raw_tool_calls = [ item for item in output if replay_state.response_item_type(item) == "function_call" ] @@ -2539,7 +2547,7 @@ async def _make_call(): finish_reason=_finish_reason_for_tool_calls( _map_responses_finish_reason(raw_response) ), - reasoning=None, + reasoning=reasoning, usage=usage, llm_state=replay_state.capture_responses_state(output, state_scope), ) @@ -2556,7 +2564,7 @@ async def _make_call(): parsed=parsed_content, tool_calls=[], finish_reason=_map_responses_finish_reason(raw_response), - reasoning=None, + reasoning=reasoning, usage=usage, llm_state=replay_state.capture_responses_state(output, state_scope), ) @@ -2566,7 +2574,7 @@ async def _make_call(): content=text_content, tool_calls=[], finish_reason=_map_responses_finish_reason(raw_response), - reasoning=None, + reasoning=reasoning, usage=usage, llm_state=replay_state.capture_responses_state(output, state_scope), ) diff --git a/tests/unifiedllm/test_cross_provider_reasoning.py b/tests/unifiedllm/test_cross_provider_reasoning.py new file mode 100644 index 000000000..7299d74f5 --- /dev/null +++ b/tests/unifiedllm/test_cross_provider_reasoning.py @@ -0,0 +1,423 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Closed-provider opaque state and provider-independent text reasoning replay.""" + +import json +from types import SimpleNamespace +from typing import Any, Literal, cast +from unittest.mock import AsyncMock, patch + +import pytest +from litellm.types.utils import Choices, Message, ModelResponse + +from nooa._llm_state import ReplayCarryingMessage +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.unifiedllm import CompletionClient, LLMResponse, ResponsesClient, Tool +from nooa.unifiedllm.replay_state import ( + capture_responses_state, + prepare_chat_messages, + prepare_responses_batch, + replay_scope, +) + +ANTHROPIC_THINKING = [ + {"type": "thinking", "thinking": "Check the inputs.", "signature": "anthropic-sig"}, + {"type": "redacted_thinking", "data": "anthropic-redacted"}, +] +GEMINI_SIGNATURE = "Z2VtaW5pLXNpZ25hdHVyZQ==" +GEMINI_SIGNATURE_2 = "c2Vjb25kLXNpZ25hdHVyZQ==" + + +def _execute_python(code: str) -> str: + return code + + +TOOL = Tool(name="execute_python", description="Run code", callable=_execute_python) + + +def _tool_call(call_id: str = "call_1", provider_specific_fields: dict | None = None) -> dict: + call = { + "id": call_id, + "type": "function", + "function": {"name": "execute_python", "arguments": '{"code":"print(1)"}'}, + } + if provider_specific_fields: + call["provider_specific_fields"] = provider_specific_fields + return call + + +def _chat_response(message: Message, finish_reason: str = "tool_calls") -> ModelResponse: + return ModelResponse( + model="test-model", + choices=[Choices(message=message, finish_reason=finish_reason)], + ) + + +def _anthropic_response() -> ModelResponse: + return _chat_response( + Message( + role="assistant", + content=None, + tool_calls=[_tool_call()], + thinking_blocks=cast(Any, ANTHROPIC_THINKING), + reasoning_content="Check the inputs.", + ) + ) + + +def _gemini_response() -> ModelResponse: + return _chat_response( + Message( + role="assistant", + content=None, + tool_calls=[ + _tool_call( + f"call_1__thought__{GEMINI_SIGNATURE}", + {"thought_signature": GEMINI_SIGNATURE, "private": "discard-me"}, + ), + _tool_call( + f"call_2__thought__{GEMINI_SIGNATURE_2}", + {"thought_signature": GEMINI_SIGNATURE_2}, + ), + ], + provider_specific_fields={ + "thought_signatures": [GEMINI_SIGNATURE, GEMINI_SIGNATURE_2], + "private": "discard-me", + }, + reasoning_content="Inspect the value.", + ) + ) + + +def _render(response: LLMResponse, *, responses: bool = False) -> list[dict]: + 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 + ) + neutral = XMLBlockFormatter().format(blocks) + formatter = ResponsesProviderFormatter() if responses else OpenAIProviderFormatter() + return formatter.format(neutral) + + +def test_anthropic_thinking_blocks_round_trip_exactly() -> None: + client = CompletionClient(model="anthropic/claude-sonnet-4", api_key="account-a") + try: + with patch( + "litellm.completion", side_effect=[_anthropic_response(), _anthropic_response()] + ) as completion: + first = client.call([{"role": "user", "content": "run"}], tools=[TOOL]) + client.call(_render(first), tools=[TOOL]) + + assert first.reasoning == "Check the inputs." + assert first.llm_state is not None + assert first.llm_state["payload"]["thinking_blocks"] == ANTHROPIC_THINKING + assistant = next( + message + for message in completion.call_args_list[1].kwargs["messages"] + if message.get("role") == "assistant" + ) + assert assistant["thinking_blocks"] == ANTHROPIC_THINKING + assert assistant["content"] is None + finally: + client.close() + + +@pytest.mark.asyncio +async def test_async_anthropic_capture_matches_sync() -> None: + client = CompletionClient(model="anthropic/claude-sonnet-4", api_key="account-a") + try: + with patch("litellm.acompletion", AsyncMock(return_value=_anthropic_response())): + response = await client.acall([{"role": "user", "content": "run"}], tools=[TOOL]) + assert response.llm_state is not None + assert response.llm_state["payload"]["thinking_blocks"] == ANTHROPIC_THINKING + finally: + await client.aclose() + + +def test_gemini_signatures_round_trip_without_becoming_public_call_ids() -> None: + client = CompletionClient(model="gemini/gemini-2.5-pro", api_key="account-a") + try: + with patch( + "litellm.completion", side_effect=[_gemini_response(), _gemini_response()] + ) as completion: + first = client.call([{"role": "user", "content": "run"}], tools=[TOOL]) + client.call(_render(first), tools=[TOOL]) + + assert [call.id for call in first.tool_calls] == ["call_1", "call_2"] + assert "discard-me" not in json.dumps(first.llm_state) + assistant = next( + message + for message in completion.call_args_list[1].kwargs["messages"] + if message.get("role") == "assistant" + ) + assert assistant["provider_specific_fields"] == { + "thought_signatures": [GEMINI_SIGNATURE, GEMINI_SIGNATURE_2] + } + assert [call["id"] for call in assistant["tool_calls"]] == ["call_1", "call_2"] + assert [ + call["provider_specific_fields"]["thought_signature"] + for call in assistant["tool_calls"] + ] == [GEMINI_SIGNATURE, GEMINI_SIGNATURE_2] + finally: + client.close() + + +def test_non_gemini_tool_call_id_with_thought_substring_is_unchanged() -> None: + call_id = "call_business__thought__phase" + response = _chat_response( + Message(role="assistant", content=None, tool_calls=[_tool_call(call_id)]) + ) + client = CompletionClient(model="openai/gpt-4o", api_key="account-a") + try: + with patch("litellm.completion", return_value=response): + first = client.call([{"role": "user", "content": "run"}], tools=[TOOL]) + assert first.tool_calls[0].id == call_id + + rendered = _render(first) + assistant = next(message for message in rendered if message.get("tool_calls")) + assert assistant["tool_calls"][0]["id"] == call_id + prepared = prepare_chat_messages(rendered, None) + assistant = next(message for message in prepared if message.get("tool_calls")) + assert assistant["tool_calls"][0]["id"] == call_id + finally: + client.close() + + +def test_cross_provider_replay_hides_opaque_state_but_keeps_reasoning_text() -> None: + source = CompletionClient(model="gemini/gemini-2.5-pro", api_key="account-a") + target = CompletionClient(model="anthropic/claude-sonnet-4", api_key="account-a") + try: + with patch("litellm.completion", return_value=_gemini_response()): + first = source.call([{"role": "user", "content": "run"}], tools=[TOOL]) + with patch( + "litellm.completion", + return_value=_chat_response(Message(role="assistant", content="done"), "stop"), + ) as completion: + target.call(_render(first), tools=[TOOL]) + + replayed = completion.call_args.kwargs["messages"] + assistant = next(message for message in replayed if message.get("role") == "assistant") + assert assistant["content"] == "Inspect the value." + assert [call["id"] for call in assistant["tool_calls"]] == ["call_1", "call_2"] + assert GEMINI_SIGNATURE not in json.dumps(replayed) + assert GEMINI_SIGNATURE_2 not in json.dumps(replayed) + finally: + source.close() + target.close() + + +def test_plain_reasoning_replays_as_ordinary_text_for_every_model() -> None: + source = CompletionClient(model="deepseek/deepseek-reasoner", api_key="account-a") + target = CompletionClient(model="openai/gpt-4o", api_key="account-a") + response = _chat_response( + Message( + role="assistant", + content="Visible answer.", + reasoning_content="Plain reasoning.", + ), + "stop", + ) + try: + with patch("litellm.completion", return_value=response): + first = source.call([{"role": "user", "content": "think"}]) + with patch("litellm.completion", return_value=response) as completion: + target.call(_render(first)) + + assert first.llm_state is None + assistant = next( + message + for message in completion.call_args.kwargs["messages"] + if message.get("role") == "assistant" + ) + assert assistant == { + "role": "assistant", + "content": "Plain reasoning.\n\nVisible answer.", + } + finally: + source.close() + target.close() + + +def _responses_output(*items: dict) -> SimpleNamespace: + return SimpleNamespace(output=list(items), output_text="", status="completed", usage=None) + + +RESPONSES_REASONING = { + "id": "rs_1", + "type": "reasoning", + "encrypted_content": "openai-secret", + "summary": [{"type": "summary_text", "text": "Check the evidence."}], +} +RESPONSES_MESSAGE = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "Answer.", "annotations": []}], +} + + +def test_responses_summary_stays_exact_on_match_and_demotes_on_model_change() -> 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_output(RESPONSES_REASONING, RESPONSES_MESSAGE), + ): + first = source.call([{"role": "user", "content": "think"}]) + + assert first.reasoning == "Check the evidence." + rendered = _render(first, responses=True) + with patch( + "litellm.responses", return_value=_responses_output(RESPONSES_MESSAGE) + ) as matching: + source.call(rendered) + assert matching.call_args.kwargs["input"][:2] == [ + RESPONSES_REASONING, + {"role": "assistant", "content": "Answer."}, + ] + + with patch( + "litellm.responses", return_value=_responses_output(RESPONSES_MESSAGE) + ) as changed: + target.call(rendered) + assert changed.call_args.kwargs["input"] == [ + {"role": "assistant", "content": "Check the evidence.\n\nAnswer."} + ] + assert "openai-secret" not in json.dumps(changed.call_args.kwargs["input"]) + finally: + source.close() + target.close() + + +def test_reasoning_only_responses_turn_demotes_without_an_empty_message() -> 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_output(RESPONSES_REASONING)): + first = source.call([{"role": "user", "content": "think"}]) + with patch( + "litellm.responses", return_value=_responses_output(RESPONSES_MESSAGE) + ) as changed: + target.call(_render(first, responses=True)) + + assert changed.call_args.kwargs["input"] == [ + {"role": "assistant", "content": "Check the evidence."} + ] + finally: + source.close() + target.close() + + +def test_state_only_turn_drops_empty_carrier_across_api_styles() -> None: + responses_state = { + "version": 1, + "scope": "responses:openai:sha256:source", + "format": "openai-responses", + "payload": {"items": [RESPONSES_REASONING], "order": [], "state_only": True}, + } + chat_carrier = ReplayCarryingMessage( + {"role": "assistant", "content": ""}, responses_state + ) + assert prepare_chat_messages([chat_carrier], "chat:openai:sha256:target") == [] + + chat_state = { + "version": 1, + "scope": "chat:openai:sha256:source", + "format": "litellm-chat", + "payload": {"reasoning_items": [{"type": "reasoning"}], "state_only": True}, + } + assert prepare_responses_batch( + [{"role": "assistant", "content": ""}], + chat_state, + "responses:openai:sha256:target", + ) == [] + + +def test_responses_demotion_keeps_native_output_content_valid() -> None: + native_message = { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Answer."}], + } + + assert prepare_responses_batch([native_message], None, None, "Check the evidence.") == [ + {"role": "assistant", "content": "Check the evidence."}, + native_message, + ] + + +@pytest.mark.parametrize( + ("model", "api_style"), + [ + ("anthropic/claude-sonnet-4", "responses"), + ("gemini/gemini-2.5-pro", "responses"), + ("vertex_ai/gemini-2.5-pro", "responses"), + ("vertex_ai/gemini-2.5-pro", "chat"), + ], +) +def test_unverified_closed_provider_routes_have_no_opaque_replay_scope( + model: str, api_style: Literal["chat", "responses"] +) -> None: + assert replay_scope(model, api_style, {"api_key": "account-a"}) is None + + +def test_non_openai_responses_scope_cannot_capture_or_restore_opaque_items() -> None: + fabricated_scope = "responses:anthropic:sha256:untrusted" + assert capture_responses_state([RESPONSES_REASONING], fabricated_scope) is None + + state = { + "version": 1, + "scope": fabricated_scope, + "format": "openai-responses", + "payload": {"items": [RESPONSES_REASONING], "order": []}, + } + assert prepare_responses_batch( + [RESPONSES_MESSAGE], state, fabricated_scope, "Check the evidence." + ) == [ + {"role": "assistant", "content": "Check the evidence."}, + RESPONSES_MESSAGE, + ] + + +@pytest.mark.parametrize( + ("model", "environment"), + [ + ("anthropic/claude-sonnet-4", "ANTHROPIC_API_BASE"), + ("gemini/gemini-2.5-pro", "GEMINI_API_BASE"), + ], +) +def test_closed_provider_environment_endpoint_partitions_scope( + monkeypatch, model: str, environment: str +) -> None: + monkeypatch.setenv(environment, "https://issuer-a.example/v1") + first = replay_scope(model, "chat", {"api_key": "account-a"}) + monkeypatch.setenv(environment, "https://issuer-b.example/v1") + second = replay_scope(model, "chat", {"api_key": "account-a"}) + + assert first is not None + assert second is not None + assert first != second diff --git a/tests/unifiedllm/test_http_logging.py b/tests/unifiedllm/test_http_logging.py index 73bc13f09..26e274f29 100644 --- a/tests/unifiedllm/test_http_logging.py +++ b/tests/unifiedllm/test_http_logging.py @@ -16,6 +16,9 @@ def test_opaque_reasoning_state_is_redacted_from_http_debug_payloads() -> None: payload = { "input": [ {"encrypted_content": "provider-secret"}, + {"type": "thinking", "thinking": "visible", "signature": "anthropic-secret"}, + {"type": "redacted_thinking", "data": "anthropic-redacted"}, + {"thought_signature": "gemini-secret"}, {"_nooa_llm_state": {"payload": {"items": ["opaque"]}}}, ] } @@ -24,6 +27,9 @@ def test_opaque_reasoning_state_is_redacted_from_http_debug_payloads() -> None: assert redacted["input"] == [ {"encrypted_content": "[REDACTED]"}, + {"type": "thinking", "thinking": "visible", "signature": "[REDACTED]"}, + {"type": "redacted_thinking", "data": "[REDACTED]"}, + {"thought_signature": "[REDACTED]"}, {"_nooa_llm_state": "[REDACTED]"}, ] diff --git a/tests/unifiedllm/test_reasoning_state_replay.py b/tests/unifiedllm/test_reasoning_state_replay.py index aecd1841d..cc6187457 100644 --- a/tests/unifiedllm/test_reasoning_state_replay.py +++ b/tests/unifiedllm/test_reasoning_state_replay.py @@ -823,12 +823,13 @@ async def test_responses_reasoning_uses_per_call_override() -> None: @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 +def test_non_openai_chat_provider_cannot_receive_openai_reasoning_state(model: str) -> None: + scope = replay_scope(model, "chat", {"api_key": "account-a"}) + assert scope is not None client = CompletionClient(model=model, api_key="account-a") crafted = { "version": 1, - "scope": f"chat:{model.split('/', 1)[0]}:crafted", + "scope": scope, "format": "litellm-chat", "payload": {"reasoning_items": [REASONING]}, } From 2113a4f46601d7b5fe8922e4376283b525857c93 Mon Sep 17 00:00:00 2001 From: Paul Furgale Date: Wed, 9 Sep 2026 18:18:04 +0000 Subject: [PATCH 02/13] style(llm): format cross-provider replay Signed-off-by: Paul Furgale --- src/nooa/unifiedllm/replay_state.py | 6 +----- .../unifiedllm/test_cross_provider_reasoning.py | 17 +++++++++-------- 2 files changed, 10 insertions(+), 13 deletions(-) diff --git a/src/nooa/unifiedllm/replay_state.py b/src/nooa/unifiedllm/replay_state.py index ac390d03f..510a48701 100644 --- a/src/nooa/unifiedllm/replay_state.py +++ b/src/nooa/unifiedllm/replay_state.py @@ -190,11 +190,7 @@ def _sanitize_chat_payload(payload: dict[str, Any], scope: str | None) -> dict[s ): clean["provider_specific_fields"] = {"thought_signatures": copy.deepcopy(signatures)} - tool_state = ( - payload.get("tool_calls") - if provider in {"openai", "azure", "gemini"} - else None - ) + tool_state = payload.get("tool_calls") if provider in {"openai", "azure", "gemini"} else None if isinstance(tool_state, list): calls: list[dict[str, Any] | None] = [] for item in tool_state: diff --git a/tests/unifiedllm/test_cross_provider_reasoning.py b/tests/unifiedllm/test_cross_provider_reasoning.py index 7299d74f5..f7352a6ec 100644 --- a/tests/unifiedllm/test_cross_provider_reasoning.py +++ b/tests/unifiedllm/test_cross_provider_reasoning.py @@ -339,9 +339,7 @@ def test_state_only_turn_drops_empty_carrier_across_api_styles() -> None: "format": "openai-responses", "payload": {"items": [RESPONSES_REASONING], "order": [], "state_only": True}, } - chat_carrier = ReplayCarryingMessage( - {"role": "assistant", "content": ""}, responses_state - ) + chat_carrier = ReplayCarryingMessage({"role": "assistant", "content": ""}, responses_state) assert prepare_chat_messages([chat_carrier], "chat:openai:sha256:target") == [] chat_state = { @@ -350,11 +348,14 @@ def test_state_only_turn_drops_empty_carrier_across_api_styles() -> None: "format": "litellm-chat", "payload": {"reasoning_items": [{"type": "reasoning"}], "state_only": True}, } - assert prepare_responses_batch( - [{"role": "assistant", "content": ""}], - chat_state, - "responses:openai:sha256:target", - ) == [] + assert ( + prepare_responses_batch( + [{"role": "assistant", "content": ""}], + chat_state, + "responses:openai:sha256:target", + ) + == [] + ) def test_responses_demotion_keeps_native_output_content_valid() -> None: From 0d3b58582b11e33f34d07d9a7c82e2b41c07147f Mon Sep 17 00:00:00 2001 From: Paul Furgale Date: Wed, 9 Sep 2026 18:47:16 +0000 Subject: [PATCH 03/13] fix(llm): harden cross-provider response edges Signed-off-by: Paul Furgale --- src/nooa/unifiedllm/unifiedllm.py | 2 +- tests/unifiedllm/test_finish_reason_propagation.py | 10 ++++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/src/nooa/unifiedllm/unifiedllm.py b/src/nooa/unifiedllm/unifiedllm.py index 77f114623..793bcb475 100644 --- a/src/nooa/unifiedllm/unifiedllm.py +++ b/src/nooa/unifiedllm/unifiedllm.py @@ -1897,7 +1897,7 @@ def _make_call(): ToolCall( id=replay_state.public_tool_call_id(tc, state_scope), name=tc.function.name or "", - arguments=tc.function.arguments, + arguments=tc.function.arguments or "", ) for tc in raw_tool_calls ] diff --git a/tests/unifiedllm/test_finish_reason_propagation.py b/tests/unifiedllm/test_finish_reason_propagation.py index 5c1b4a3be..2f24507bf 100644 --- a/tests/unifiedllm/test_finish_reason_propagation.py +++ b/tests/unifiedllm/test_finish_reason_propagation.py @@ -197,6 +197,16 @@ def test_sync_tool_call_preserves_accompanying_text(self, client): out = client.call([{"role": "user", "content": "Hi"}]) assert out.content == "I will use the tool." + def test_sync_tool_call_normalizes_missing_arguments(self, client): + tc = make_tool_call("call_1", "do_thing", "{}") + tc.function.arguments = None # type: ignore[assignment] + response = make_mock_response(content=None, tool_calls=[tc]) + + with patch("litellm.completion", return_value=response): + out = client.call([{"role": "user", "content": "Hi"}]) + + assert out.tool_calls[0].arguments == "" + @pytest.mark.asyncio async def test_async_tool_call_preserves_accompanying_text(self, client): tc = make_tool_call("call_1", "do_thing", "{}") From 875a946858d2016987ec17690d7ea62f46edbfa3 Mon Sep 17 00:00:00 2001 From: Paul Furgale Date: Wed, 9 Sep 2026 19:22:38 +0000 Subject: [PATCH 04/13] fix(llm): redact cross-provider replay state from traces Signed-off-by: Paul Furgale --- src/nooa/tracing/_secret_scrubber.py | 13 +++++++++++-- tests/tracing/test_journal.py | 16 +++++++++++++++ tests/tracing/test_secret_scrubber.py | 28 +++++++++++++++++++++++++++ 3 files changed, 55 insertions(+), 2 deletions(-) diff --git a/src/nooa/tracing/_secret_scrubber.py b/src/nooa/tracing/_secret_scrubber.py index 864760cb6..49c84488a 100644 --- a/src/nooa/tracing/_secret_scrubber.py +++ b/src/nooa/tracing/_secret_scrubber.py @@ -51,7 +51,9 @@ # 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"}) +_OPAQUE_PROVIDER_STATE_KEYS = frozenset( + {"encrypted_content", "nooa_llm_state", "thought_signature", "thought_signatures"} +) def _is_sensitive_key(key: Any) -> bool: @@ -247,8 +249,15 @@ def scrub_value(value: Any) -> tuple[Any, int]: if isinstance(value, dict): scrubbed_mapping: dict[Any, Any] = {} count = 0 + is_thinking = value.get("type") == "thinking" + is_redacted_thinking = value.get("type") == "redacted_thinking" for key, item in value.items(): - if reason := _redact_key(key): + reason = _redact_key(key) + if reason is None and is_thinking and key == "signature": + reason = "opaque_provider_state" + if reason is None and is_redacted_thinking and key == "data": + reason = "opaque_provider_state" + if reason is not None: scrubbed_mapping[key] = REDACTED stats.record(reason) count += 1 diff --git a/tests/tracing/test_journal.py b/tests/tracing/test_journal.py index 2e39f8c9d..4723ed0ae 100644 --- a/tests/tracing/test_journal.py +++ b/tests/tracing/test_journal.py @@ -228,6 +228,22 @@ def test_safe_msg_to_dict_redacts_private_replay_envelope(json_encoded): assert message[LLM_STATE_KEY]["payload"]["future_provider_blob"] == "opaque-state" +@pytest.mark.parametrize( + "message", + [ + {"thinking_blocks": [{"type": "thinking", "signature": "anthropic-sig"}]}, + {"thinking_blocks": [{"type": "redacted_thinking", "data": "opaque-data"}]}, + {"provider_specific_fields": {"thought_signature": "gemini-sig"}}, + ], +) +def test_safe_msg_to_dict_redacts_cross_provider_state(message): + safe = _safe_msg_to_dict(message) + + assert "anthropic-sig" not in repr(safe) + assert "opaque-data" not in repr(safe) + assert "gemini-sig" not in repr(safe) + + 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 2a9e29177..3af883bbf 100644 --- a/tests/tracing/test_secret_scrubber.py +++ b/tests/tracing/test_secret_scrubber.py @@ -236,6 +236,34 @@ def test_openai_encrypted_reasoning_is_redacted_from_json_attribute(self): assert "opaque-openai-state" not in result assert count == 1 + @pytest.mark.parametrize( + ("provider_state", "secret"), + [ + ( + {"thinking_blocks": [{"type": "thinking", "signature": "anthropic-sig"}]}, + "anthropic-sig", + ), + ( + {"thinking_blocks": [{"type": "redacted_thinking", "data": "opaque-data"}]}, + "opaque-data", + ), + ( + {"provider_specific_fields": {"thought_signature": "gemini-sig"}}, + "gemini-sig", + ), + ( + {"provider_specific_fields": {"thought_signatures": ["gemini-sig"]}}, + "gemini-sig", + ), + ], + ) + def test_cross_provider_opaque_state_is_redacted(self, provider_state, secret): + result, count = scrub_value(provider_state) + + assert secret not in repr(result) + assert REDACTED in repr(result) + assert count == 1 + class TestScrubStats: def test_record_and_snapshot(self): From af47a3a13fe4f7e4f179fdd895de1c5b2455c095 Mon Sep 17 00:00:00 2001 From: Paul Furgale Date: Wed, 9 Sep 2026 19:31:24 +0000 Subject: [PATCH 05/13] fix(llm): scrub Gemini inline replay signatures Signed-off-by: Paul Furgale --- src/nooa/tracing/_secret_scrubber.py | 14 +++++++++++++- tests/tracing/test_journal.py | 4 ++++ tests/tracing/test_secret_scrubber.py | 21 +++++++++++++++++++++ 3 files changed, 38 insertions(+), 1 deletion(-) diff --git a/src/nooa/tracing/_secret_scrubber.py b/src/nooa/tracing/_secret_scrubber.py index 49c84488a..0a7abb92d 100644 --- a/src/nooa/tracing/_secret_scrubber.py +++ b/src/nooa/tracing/_secret_scrubber.py @@ -52,7 +52,14 @@ # 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", "thought_signature", "thought_signatures"} + { + "encrypted_content", + "nooa_llm_state", + "thought_signature", + "thought_signatures", + "thoughtsignature", + "thoughtsignatures", + } ) @@ -81,6 +88,11 @@ def _redact_key(key: Any) -> str | None: # --------------------------------------------------------------------------- _SECRET_PATTERNS: list[tuple[str, re.Pattern[str]]] = [ + # LiteLLM may append Gemini thought signatures to tool-call IDs. + ( + "gemini_inline_thought_signature", + re.compile(r"__thought__(?P[A-Za-z0-9+/=_-]+)"), + ), # AWS Access Key IDs (AKIA, ASIA, AIDA, AROA + 16 alphanumeric) ("aws_access_key", re.compile(r"(?P(?:AKIA|ASIA|AIDA|AROA)[A-Z0-9]{16})")), # AWS Secret Access Keys (40-char base64 after known prefix) diff --git a/tests/tracing/test_journal.py b/tests/tracing/test_journal.py index 4723ed0ae..92fdeab90 100644 --- a/tests/tracing/test_journal.py +++ b/tests/tracing/test_journal.py @@ -234,6 +234,8 @@ def test_safe_msg_to_dict_redacts_private_replay_envelope(json_encoded): {"thinking_blocks": [{"type": "thinking", "signature": "anthropic-sig"}]}, {"thinking_blocks": [{"type": "redacted_thinking", "data": "opaque-data"}]}, {"provider_specific_fields": {"thought_signature": "gemini-sig"}}, + {"providerSpecificFields": {"thoughtSignature": "gemini-camel-sig"}}, + {"tool_calls": [{"id": "call_1__thought__Z2VtaW5pLXNpZw=="}]}, ], ) def test_safe_msg_to_dict_redacts_cross_provider_state(message): @@ -242,6 +244,8 @@ def test_safe_msg_to_dict_redacts_cross_provider_state(message): assert "anthropic-sig" not in repr(safe) assert "opaque-data" not in repr(safe) assert "gemini-sig" not in repr(safe) + assert "gemini-camel-sig" not in repr(safe) + assert "Z2VtaW5pLXNpZw==" not in repr(safe) class TestSentBlocksBounding: diff --git a/tests/tracing/test_secret_scrubber.py b/tests/tracing/test_secret_scrubber.py index 3af883bbf..044a986a6 100644 --- a/tests/tracing/test_secret_scrubber.py +++ b/tests/tracing/test_secret_scrubber.py @@ -255,6 +255,14 @@ def test_openai_encrypted_reasoning_is_redacted_from_json_attribute(self): {"provider_specific_fields": {"thought_signatures": ["gemini-sig"]}}, "gemini-sig", ), + ( + {"providerSpecificFields": {"thoughtSignature": "gemini-camel-sig"}}, + "gemini-camel-sig", + ), + ( + {"providerSpecificFields": {"thoughtSignatures": ["gemini-camel-sig"]}}, + "gemini-camel-sig", + ), ], ) def test_cross_provider_opaque_state_is_redacted(self, provider_state, secret): @@ -264,6 +272,19 @@ def test_cross_provider_opaque_state_is_redacted(self, provider_state, secret): assert REDACTED in repr(result) assert count == 1 + def test_gemini_inline_tool_call_signature_is_redacted_from_flat_string(self): + result, count = scrub_value("call_1__thought__Z2VtaW5pLXNpZw==") + + assert result == f"call_1__thought__{REDACTED}" + assert count == 1 + + def test_gemini_inline_tool_call_signature_is_redacted_from_json_attribute(self): + result, count = scrub_value('{"tool_calls":[{"id":"call_1__thought__Z2VtaW5pLXNpZw=="}]}') + + assert "Z2VtaW5pLXNpZw==" not in result + assert REDACTED in result + assert count == 1 + class TestScrubStats: def test_record_and_snapshot(self): From cfcaa0e228d0723092a7765521b27b92f1fd22bc Mon Sep 17 00:00:00 2001 From: Paul Furgale Date: Wed, 9 Sep 2026 20:20:38 +0000 Subject: [PATCH 06/13] fix(llm): bind Gemini state to tool identities Signed-off-by: Paul Furgale --- src/nooa/unifiedllm/replay_state.py | 35 +++++++++++++++++-- .../test_cross_provider_reasoning.py | 29 +++++++++++++++ .../test_litellm_responses_bridge.py | 11 +----- 3 files changed, 63 insertions(+), 12 deletions(-) diff --git a/src/nooa/unifiedllm/replay_state.py b/src/nooa/unifiedllm/replay_state.py index 510a48701..0323f9fa0 100644 --- a/src/nooa/unifiedllm/replay_state.py +++ b/src/nooa/unifiedllm/replay_state.py @@ -204,6 +204,15 @@ def _sanitize_chat_payload(payload: dict[str, Any], scope: str | None) -> dict[s if any(item is not None for item in calls): clean["tool_calls"] = calls + tool_call_ids = payload.get("tool_call_ids") + if ( + clean + and isinstance(tool_call_ids, list) + and tool_call_ids + and all(isinstance(call_id, str) and call_id for call_id in tool_call_ids) + ): + clean["tool_call_ids"] = copy.deepcopy(tool_call_ids) + if clean and payload.get("state_only") is True: clean["state_only"] = True return clean @@ -248,10 +257,17 @@ def capture_chat_state(message: Any, scope: str | None) -> dict | None: if isinstance(provider_fields, dict): payload["provider_specific_fields"] = provider_fields - tool_state = [_tool_call_state(call) for call in (_field(message, "tool_calls") or [])] + raw_tool_calls = list(_field(message, "tool_calls") or []) + tool_state = [_tool_call_state(call) for call in raw_tool_calls] if any(item is not None for item in tool_state): payload["tool_calls"] = tool_state + if payload and raw_tool_calls: + tool_call_ids = [public_tool_call_id(call, scope) for call in raw_tool_calls] + if not all(tool_call_ids) or len(set(tool_call_ids)) != len(tool_call_ids): + return None + payload["tool_call_ids"] = tool_call_ids + payload = _sanitize_chat_payload(payload, scope) if not payload: return None @@ -380,6 +396,19 @@ def _strip_chat_state(message: dict[str, Any], *, strip_inline_signatures: bool) def _restore_chat_state(message: dict[str, Any], payload: dict[str, Any]) -> None: + expected_call_ids = payload.get("tool_call_ids") + if expected_call_ids is not None: + tool_calls = message.get("tool_calls") + if ( + not isinstance(tool_calls, list) + or [call.get("id") if isinstance(call, dict) else None for call in tool_calls] + != expected_call_ids + ): + return + elif isinstance(payload.get("tool_calls"), list): + # Positional per-call state cannot safely survive public call mutation. + return + for key in ("reasoning_items", "thinking_blocks", "provider_specific_fields"): if key in payload: message[key] = copy.deepcopy(payload[key]) @@ -387,7 +416,9 @@ def _restore_chat_state(message: dict[str, Any], payload: dict[str, Any]) -> Non tool_state = payload.get("tool_calls") if not isinstance(tool_calls, list) or not isinstance(tool_state, list): return - for call, state in zip(tool_calls, tool_state, strict=False): + if len(tool_calls) != len(tool_state): + return + for call, state in zip(tool_calls, tool_state, strict=True): if not isinstance(call, dict) or not isinstance(state, dict): continue fields = state.get("provider_specific_fields") diff --git a/tests/unifiedllm/test_cross_provider_reasoning.py b/tests/unifiedllm/test_cross_provider_reasoning.py index f7352a6ec..260133cd5 100644 --- a/tests/unifiedllm/test_cross_provider_reasoning.py +++ b/tests/unifiedllm/test_cross_provider_reasoning.py @@ -166,7 +166,9 @@ def test_gemini_signatures_round_trip_without_becoming_public_call_ids() -> None client.call(_render(first), tools=[TOOL]) assert [call.id for call in first.tool_calls] == ["call_1", "call_2"] + assert first.llm_state is not None assert "discard-me" not in json.dumps(first.llm_state) + assert first.llm_state["payload"]["tool_call_ids"] == ["call_1", "call_2"] assistant = next( message for message in completion.call_args_list[1].kwargs["messages"] @@ -184,6 +186,33 @@ def test_gemini_signatures_round_trip_without_becoming_public_call_ids() -> None client.close() +@pytest.mark.parametrize("mutation", ["drop", "reorder", "duplicate"]) +def test_gemini_tool_state_fails_closed_when_public_calls_change(mutation: str) -> None: + client = CompletionClient(model="gemini/gemini-2.5-pro", api_key="account-a") + try: + with patch("litellm.completion", return_value=_gemini_response()): + first = client.call([{"role": "user", "content": "run"}], tools=[TOOL]) + + assert first.llm_state is not None + rendered = _render(first) + assistant = next(message for message in rendered if message.get("tool_calls")) + if 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("tool_calls")) + assert "provider_specific_fields" not in replayed + assert all("provider_specific_fields" not in call for call in replayed["tool_calls"]) + assert GEMINI_SIGNATURE not in json.dumps(prepared) + assert GEMINI_SIGNATURE_2 not in json.dumps(prepared) + finally: + client.close() + + def test_non_gemini_tool_call_id_with_thought_substring_is_unchanged() -> None: call_id = "call_business__thought__phase" response = _chat_response( diff --git a/tests/unifiedllm/test_litellm_responses_bridge.py b/tests/unifiedllm/test_litellm_responses_bridge.py index 3437f03f2..8c1cf44d2 100644 --- a/tests/unifiedllm/test_litellm_responses_bridge.py +++ b/tests/unifiedllm/test_litellm_responses_bridge.py @@ -99,16 +99,7 @@ async def test_default_reasoning_tool_call_uses_responses_bridge(model: str) -> assert result.llm_state["format"] == "litellm-chat" assert result.llm_state["payload"] == { "reasoning_items": [REASONING_ITEM], - "carrier": { - "content": None, - "tool_calls": [ - { - "id": "fc_123", - "name": "execute_python", - "arguments": '{"code":"print(1)"}', - } - ], - }, + "tool_call_ids": ["fc_123"], } finally: await client.aclose() From efdb6d46f376bf83286b79e353e64c36959faa20 Mon Sep 17 00:00:00 2001 From: Paul Furgale Date: Thu, 10 Sep 2026 12:24:03 +0000 Subject: [PATCH 07/13] test(llm): permit replay across gateway routes Signed-off-by: Paul Furgale --- skills/nooa-context-and-state/SKILL.md | 2 +- src/nooa/tracing/_litellm_journal.py | 2 +- src/nooa/tracing/_secret_scrubber.py | 4 +- .../test_cross_provider_reasoning.py | 40 ++++++++++++++----- 4 files changed, 34 insertions(+), 14 deletions(-) diff --git a/skills/nooa-context-and-state/SKILL.md b/skills/nooa-context-and-state/SKILL.md index 5a91c808a..4f8d68dde 100644 --- a/skills/nooa-context-and-state/SKILL.md +++ b/skills/nooa-context-and-state/SKILL.md @@ -87,7 +87,7 @@ async def solve(self, problem: str) -> str: Event history is what fills the LLM's conversation window. Key model-visible event types (names have no "Event" suffix): `Task`, `Message`, `Reasoning`, `Error`, `Feedback`, `LLMResponse`, `PythonOutput`, `Summary`, `Notification`. `LLMResponse` is both the canonical assistant turn and the home of its token/cost metadata; renderers expose only its conversational fields. Runtime-only events (never shown to the LLM) include `BeforeAgentCall`/`AfterAgentCall` and `LLMCallStart`/`LLMCallEnd`. -`LLMResponse.reasoning` is provider-exposed text and remains useful across model switches, where UnifiedLLM replays it as ordinary assistant text. `LLMResponse.llm_state` is opaque provider state; it is persisted for resume but is replayed only through a matching provider/API/issuer gate. +`LLMResponse.reasoning` is provider-exposed text and remains useful across model switches, where UnifiedLLM replays it as ordinary assistant text. `LLMResponse.llm_state` is opaque provider state; it is persisted for resume but is replayed only through a matching provider/API/model gate. ```python # Query (AND semantics; chronological; limit keeps most recent) diff --git a/src/nooa/tracing/_litellm_journal.py b/src/nooa/tracing/_litellm_journal.py index 79933b885..4e4756f62 100644 --- a/src/nooa/tracing/_litellm_journal.py +++ b/src/nooa/tracing/_litellm_journal.py @@ -71,7 +71,7 @@ def _msg_to_dict(msg: Any) -> dict: def _safe_msg_to_dict(msg: Any) -> dict: - """Normalize one provider message and remove issuer-only opaque state.""" + """Normalize one provider message and remove provider-only opaque state.""" scrubbed, _ = scrub_value(_msg_to_dict(msg)) return scrubbed if isinstance(scrubbed, dict) else {} diff --git a/src/nooa/tracing/_secret_scrubber.py b/src/nooa/tracing/_secret_scrubber.py index 0a7abb92d..d03f3e0af 100644 --- a/src/nooa/tracing/_secret_scrubber.py +++ b/src/nooa/tracing/_secret_scrubber.py @@ -49,8 +49,8 @@ ) # 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. +# telemetry rule: it may go back to its matching provider and nowhere else. +# Provider adapters add their exact wire keys here as support is introduced. _OPAQUE_PROVIDER_STATE_KEYS = frozenset( { "encrypted_content", diff --git a/tests/unifiedllm/test_cross_provider_reasoning.py b/tests/unifiedllm/test_cross_provider_reasoning.py index 260133cd5..5455657a5 100644 --- a/tests/unifiedllm/test_cross_provider_reasoning.py +++ b/tests/unifiedllm/test_cross_provider_reasoning.py @@ -122,13 +122,22 @@ def _render(response: LLMResponse, *, responses: bool = False) -> list[dict]: def test_anthropic_thinking_blocks_round_trip_exactly() -> None: - client = CompletionClient(model="anthropic/claude-sonnet-4", api_key="account-a") + source = CompletionClient( + model="anthropic/claude-sonnet-4", + api_key="account-a", + api_base="https://gateway-a.example/v1", + ) + target = CompletionClient( + model="anthropic/claude-sonnet-4", + api_key="account-b", + api_base="https://gateway-b.example/v1", + ) try: with patch( "litellm.completion", side_effect=[_anthropic_response(), _anthropic_response()] ) as completion: - first = client.call([{"role": "user", "content": "run"}], tools=[TOOL]) - client.call(_render(first), tools=[TOOL]) + first = source.call([{"role": "user", "content": "run"}], tools=[TOOL]) + target.call(_render(first), tools=[TOOL]) assert first.reasoning == "Check the inputs." assert first.llm_state is not None @@ -141,7 +150,8 @@ def test_anthropic_thinking_blocks_round_trip_exactly() -> None: assert assistant["thinking_blocks"] == ANTHROPIC_THINKING assert assistant["content"] is None finally: - client.close() + source.close() + target.close() @pytest.mark.asyncio @@ -157,13 +167,22 @@ async def test_async_anthropic_capture_matches_sync() -> None: def test_gemini_signatures_round_trip_without_becoming_public_call_ids() -> None: - client = CompletionClient(model="gemini/gemini-2.5-pro", api_key="account-a") + source = CompletionClient( + model="gemini/gemini-2.5-pro", + api_key="account-a", + api_base="https://gateway-a.example/v1", + ) + target = CompletionClient( + model="gemini/gemini-2.5-pro", + api_key="account-b", + api_base="https://gateway-b.example/v1", + ) try: with patch( "litellm.completion", side_effect=[_gemini_response(), _gemini_response()] ) as completion: - first = client.call([{"role": "user", "content": "run"}], tools=[TOOL]) - client.call(_render(first), tools=[TOOL]) + first = source.call([{"role": "user", "content": "run"}], tools=[TOOL]) + target.call(_render(first), tools=[TOOL]) assert [call.id for call in first.tool_calls] == ["call_1", "call_2"] assert first.llm_state is not None @@ -183,7 +202,8 @@ def test_gemini_signatures_round_trip_without_becoming_public_call_ids() -> None for call in assistant["tool_calls"] ] == [GEMINI_SIGNATURE, GEMINI_SIGNATURE_2] finally: - client.close() + source.close() + target.close() @pytest.mark.parametrize("mutation", ["drop", "reorder", "duplicate"]) @@ -440,7 +460,7 @@ def test_non_openai_responses_scope_cannot_capture_or_restore_opaque_items() -> ("gemini/gemini-2.5-pro", "GEMINI_API_BASE"), ], ) -def test_closed_provider_environment_endpoint_partitions_scope( +def test_closed_provider_scope_is_stable_across_environment_endpoints( monkeypatch, model: str, environment: str ) -> None: monkeypatch.setenv(environment, "https://issuer-a.example/v1") @@ -450,4 +470,4 @@ def test_closed_provider_environment_endpoint_partitions_scope( assert first is not None assert second is not None - assert first != second + assert first == second From 5c5b6d35899340c7b9ba6ab8bee178b0abcad029 Mon Sep 17 00:00:00 2001 From: Paul Furgale Date: Thu, 10 Sep 2026 13:08:00 +0000 Subject: [PATCH 08/13] perf(llm): borrow closed-provider replay state Signed-off-by: Paul Furgale --- src/nooa/unifiedllm/replay_state.py | 12 ++++++------ tests/unifiedllm/test_cross_provider_reasoning.py | 5 +++++ 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/src/nooa/unifiedllm/replay_state.py b/src/nooa/unifiedllm/replay_state.py index 0323f9fa0..b0f9cd79f 100644 --- a/src/nooa/unifiedllm/replay_state.py +++ b/src/nooa/unifiedllm/replay_state.py @@ -173,7 +173,7 @@ def _sanitize_chat_payload(payload: dict[str, Any], scope: str | None) -> dict[s for key in fields_by_provider.get(provider or "", ()): value = payload.get(key) if isinstance(value, list) and value: - clean[key] = opaque_item(value) + clean[key] = value provider_fields = ( payload.get("provider_specific_fields") @@ -188,7 +188,7 @@ def _sanitize_chat_payload(payload: dict[str, Any], scope: str | None) -> dict[s and signatures and all(isinstance(signature, str) and signature for signature in signatures) ): - clean["provider_specific_fields"] = {"thought_signatures": copy.deepcopy(signatures)} + clean["provider_specific_fields"] = {"thought_signatures": signatures} tool_state = payload.get("tool_calls") if provider in {"openai", "azure", "gemini"} else None if isinstance(tool_state, list): @@ -211,7 +211,7 @@ def _sanitize_chat_payload(payload: dict[str, Any], scope: str | None) -> dict[s and tool_call_ids and all(isinstance(call_id, str) and call_id for call_id in tool_call_ids) ): - clean["tool_call_ids"] = copy.deepcopy(tool_call_ids) + clean["tool_call_ids"] = tool_call_ids if clean and payload.get("state_only") is True: clean["state_only"] = True @@ -255,7 +255,7 @@ def capture_chat_state(message: Any, scope: str | None) -> dict | None: provider_fields = _field(message, "provider_specific_fields") if isinstance(provider_fields, dict): - payload["provider_specific_fields"] = provider_fields + payload["provider_specific_fields"] = opaque_item(provider_fields) raw_tool_calls = list(_field(message, "tool_calls") or []) tool_state = [_tool_call_state(call) for call in raw_tool_calls] @@ -411,7 +411,7 @@ def _restore_chat_state(message: dict[str, Any], payload: dict[str, Any]) -> Non for key in ("reasoning_items", "thinking_blocks", "provider_specific_fields"): if key in payload: - message[key] = copy.deepcopy(payload[key]) + message[key] = payload[key] tool_calls = message.get("tool_calls") tool_state = payload.get("tool_calls") if not isinstance(tool_calls, list) or not isinstance(tool_state, list): @@ -423,7 +423,7 @@ def _restore_chat_state(message: dict[str, Any], payload: dict[str, Any]) -> Non continue fields = state.get("provider_specific_fields") if isinstance(fields, dict): - call["provider_specific_fields"] = copy.deepcopy(fields) + call["provider_specific_fields"] = fields def capture_responses_state(output: list[Any], scope: str | None) -> dict | None: diff --git a/tests/unifiedllm/test_cross_provider_reasoning.py b/tests/unifiedllm/test_cross_provider_reasoning.py index 5455657a5..5fabe0893 100644 --- a/tests/unifiedllm/test_cross_provider_reasoning.py +++ b/tests/unifiedllm/test_cross_provider_reasoning.py @@ -148,6 +148,7 @@ def test_anthropic_thinking_blocks_round_trip_exactly() -> None: if message.get("role") == "assistant" ) assert assistant["thinking_blocks"] == ANTHROPIC_THINKING + assert assistant["thinking_blocks"] is first.llm_state["payload"]["thinking_blocks"] assert assistant["content"] is None finally: source.close() @@ -196,6 +197,10 @@ def test_gemini_signatures_round_trip_without_becoming_public_call_ids() -> None assert assistant["provider_specific_fields"] == { "thought_signatures": [GEMINI_SIGNATURE, GEMINI_SIGNATURE_2] } + assert ( + assistant["provider_specific_fields"]["thought_signatures"] + is first.llm_state["payload"]["provider_specific_fields"]["thought_signatures"] + ) assert [call["id"] for call in assistant["tool_calls"]] == ["call_1", "call_2"] assert [ call["provider_specific_fields"]["thought_signature"] From ab8ed7c4c3c57519653509148e95adbc48806cda Mon Sep 17 00:00:00 2001 From: Paul Furgale Date: Thu, 10 Sep 2026 14:47:21 +0000 Subject: [PATCH 09/13] fix(llm): surface reasoning replay failures Signed-off-by: Paul Furgale --- skills/nooa-context-and-state/SKILL.md | 2 + src/nooa/strategies/codeact.py | 6 +- src/nooa/strategies/pure_python.py | 5 + src/nooa/unifiedllm/__init__.py | 2 + src/nooa/unifiedllm/replay_state.py | 380 ++++++++++++------ src/nooa/unifiedllm/unifiedllm.py | 17 +- tests/strategies/test_error_recovery_gl106.py | 38 +- .../test_cross_provider_reasoning.py | 138 ++++++- .../unifiedllm/test_reasoning_state_replay.py | 204 ++++------ 9 files changed, 540 insertions(+), 252 deletions(-) diff --git a/skills/nooa-context-and-state/SKILL.md b/skills/nooa-context-and-state/SKILL.md index 4f8d68dde..af0234843 100644 --- a/skills/nooa-context-and-state/SKILL.md +++ b/skills/nooa-context-and-state/SKILL.md @@ -89,6 +89,8 @@ Event history is what fills the LLM's conversation window. Key model-visible eve `LLMResponse.reasoning` is provider-exposed text and remains useful across model switches, where UnifiedLLM replays it as ordinary assistant text. `LLMResponse.llm_state` is opaque provider state; it is persisted for resume but is replayed only through a matching provider/API/model gate. +An intentional model/API mismatch logs a warning and falls back to the portable text. A malformed current-version envelope or provider signature raises `ReasoningReplayError`; do not catch it and silently continue, because it signals archive corruption or an unsupported provider contract change. + ```python # Query (AND semantics; chronological; limit keeps most recent) recent = agent.events.query(limit=20) diff --git a/src/nooa/strategies/codeact.py b/src/nooa/strategies/codeact.py index 5f70d9928..afe20c6e9 100644 --- a/src/nooa/strategies/codeact.py +++ b/src/nooa/strategies/codeact.py @@ -69,7 +69,7 @@ run_postconditions, run_preconditions, ) -from nooa.unifiedllm import LLMResponse, Tool, ToolCall +from nooa.unifiedllm import LLMResponse, ReasoningReplayError, Tool, ToolCall if TYPE_CHECKING: from nooa.config.strategy_config import CodeActConfig @@ -874,6 +874,8 @@ async def _run_generation( try: with _init_hm.timer("time_prefill"): await self._run_prefill(runtime, call, builtins, session) + except ReasoningReplayError: + raise except Exception as e: logger.warning(f"[CODEACT] Prefill error (continuing): {e}") runtime.event_manager.add(Error(content=f"Prefill error: {e}")) @@ -919,6 +921,8 @@ async def _run_generation( tool_choice=tool_choice, **self._build_sampling_kwargs(), ) + except ReasoningReplayError: + raise except BlockSyntaxError as e: self._handle_block_syntax_error(e, session, runtime) continue diff --git a/src/nooa/strategies/pure_python.py b/src/nooa/strategies/pure_python.py index eccdb5af3..55adc6247 100644 --- a/src/nooa/strategies/pure_python.py +++ b/src/nooa/strategies/pure_python.py @@ -49,6 +49,7 @@ ReturnValueValidator, ) from nooa.strategies.template import TemplateStrategy +from nooa.unifiedllm import ReasoningReplayError # Import httpx timeout exceptions if available (used by litellm) try: @@ -268,6 +269,8 @@ async def execute(self, runtime: RuntimeServices, call: "CurrentCall") -> Any: if self.prefill: try: await self._run_prefill(runtime, call, builtins, session) + except ReasoningReplayError: + raise except Exception as e: logger.warning(f"[PURE_PYTHON] Prefill error (continuing): {e}") runtime.event_manager.add(Error(content=f"Prefill error: {e}")) @@ -305,6 +308,8 @@ async def execute(self, runtime: RuntimeServices, call: "CurrentCall") -> Any: generate_event_id: str | None = None try: code, generate_event_id = await self._generate_code(runtime, session) + except ReasoningReplayError: + raise except _HTTPX_TIMEOUT_EXCEPTIONS as e: # Catch httpx timeout exceptions and preserve them session.record_error() diff --git a/src/nooa/unifiedllm/__init__.py b/src/nooa/unifiedllm/__init__.py index 591218d68..dc7899afe 100644 --- a/src/nooa/unifiedllm/__init__.py +++ b/src/nooa/unifiedllm/__init__.py @@ -11,6 +11,7 @@ reload_registry, resolve_api_key_from_config, ) +from nooa.unifiedllm.replay_state import ReasoningReplayError from nooa.unifiedllm.retry import ( EmptyContentError, RetryingWrapper, @@ -48,6 +49,7 @@ # Response types "LLMResponse", "LLMUsage", + "ReasoningReplayError", # HTTP config "HttpConfig", # Retry utilities diff --git a/src/nooa/unifiedllm/replay_state.py b/src/nooa/unifiedllm/replay_state.py index b0f9cd79f..f087be0e7 100644 --- a/src/nooa/unifiedllm/replay_state.py +++ b/src/nooa/unifiedllm/replay_state.py @@ -4,7 +4,8 @@ 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. +messages. Expected incompatibility warns and demotes portable text; malformed +current state raises instead of silently hiding a framework or provider change. """ from __future__ import annotations @@ -42,6 +43,10 @@ } +class ReasoningReplayError(RuntimeError): + """Opaque reasoning state is present but violates NOOA's replay contract.""" + + def _field(value: Any, key: str, default: Any = None) -> Any: return value.get(key, default) if isinstance(value, dict) else getattr(value, key, default) @@ -64,6 +69,17 @@ def opaque_item(item: Any) -> Any: return copy.deepcopy(item) +def _warn_unknown_fields(fields: dict[str, Any], known: set[str], location: str) -> None: + unknown = sorted(key for key, value in fields.items() if key not in known and value is not None) + if unknown: + logger.warning( + "Ignoring unrecognized provider field(s) %s on %s; opaque reasoning retention " + "may need updating.", + ", ".join(unknown), + location, + ) + + def _normalized_endpoint(value: Any) -> str: if not isinstance(value, str) or not value: return "default" @@ -106,8 +122,12 @@ def replay_scope( 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) + except Exception as exc: # noqa: BLE001 - model routing is third-party input + logger.warning( + "Opaque reasoning replay is disabled because LiteLLM could not resolve model %r: %s", + model, + exc, + ) return None if provider not in _SUPPORTED_PROVIDERS or ( api_style == "responses" and provider not in {"openai", "azure"} @@ -119,8 +139,13 @@ def replay_scope( def _envelope(scope: str | None, state_format: str, payload: dict[str, Any]) -> dict | None: - if scope is None or not payload: + if not payload: return None + if scope is None: + raise ReasoningReplayError( + "The provider returned opaque reasoning state, but NOOA could not establish " + "a safe replay scope for it." + ) return { "version": _STATE_VERSION, "scope": scope, @@ -130,46 +155,70 @@ def _envelope(scope: str | None, state_format: str, payload: dict[str, Any]) -> def _matching_payload(state: Any, scope: str | None, state_format: str) -> dict | None: + if state is None: + return None + if not isinstance(state, dict) or state.get("version") != _STATE_VERSION: + version = state.get("version") if isinstance(state, dict) else None + logger.warning( + "Ignoring opaque reasoning state with unsupported or legacy version %r; " + "portable reasoning text will be replayed instead.", + version, + ) + return None + + source_scope = state.get("scope") + source_format = state.get("format") + payload = state.get("payload") 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) + set(state) != {"version", "scope", "format", "payload"} + or not isinstance(source_scope, str) + or not isinstance(source_format, str) + or source_format not in {_CHAT_FORMAT, _RESPONSES_FORMAT} + or not isinstance(payload, dict) ): + raise ReasoningReplayError( + "Malformed version-1 opaque reasoning envelope: expected a known format, " + "a string scope, and a mapping payload." + ) + _scope_provider(source_scope) + if source_scope != scope or source_format != state_format: + logger.warning( + "Opaque reasoning state captured for %s (%s) is incompatible with %s (%s); " + "portable reasoning text will be replayed instead.", + source_scope, + source_format, + scope or "an unsupported target", + state_format, + ) 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 _is_state_only(state: Any) -> bool: - return ( - isinstance(state, dict) - and state.get("version") == _STATE_VERSION - and isinstance(state.get("payload"), dict) - and state["payload"].get("state_only") is True - ) + return cast(dict[str, Any], payload) def _scope_provider(scope: str | None) -> str | None: - if not isinstance(scope, str): + if scope is None: return None + if not isinstance(scope, str): + raise ReasoningReplayError("Malformed opaque reasoning scope: expected a string.") parts = scope.split(":", 2) - return parts[1] if len(parts) == 3 else None + if len(parts) != 3 or not all(parts): + raise ReasoningReplayError( + f"Malformed opaque reasoning scope {scope!r}: expected api:provider:identity." + ) + return parts[1] def _sanitize_chat_payload(payload: dict[str, Any], scope: str | None) -> dict[str, Any]: - """Keep only replay fields whose LiteLLM Chat semantics NOOA knows.""" + """Validate the small provider-specific envelope owned by NOOA.""" provider = _scope_provider(scope) - clean: dict[str, Any] = {} fields_by_provider = { "openai": ("reasoning_items", "thinking_blocks"), "azure": ("reasoning_items", "thinking_blocks"), "anthropic": ("thinking_blocks",), "gemini": ("thinking_blocks",), } + clean: dict[str, Any] = {} for key in fields_by_provider.get(provider or "", ()): value = payload.get(key) if isinstance(value, list) and value: @@ -210,28 +259,53 @@ def _sanitize_chat_payload(payload: dict[str, Any], scope: str | None) -> dict[s and isinstance(tool_call_ids, list) and tool_call_ids and all(isinstance(call_id, str) and call_id for call_id in tool_call_ids) + and len(set(tool_call_ids)) == len(tool_call_ids) ): clean["tool_call_ids"] = tool_call_ids if clean and payload.get("state_only") is True: clean["state_only"] = True + if ( + not clean + or clean != payload + or ( + "tool_calls" in clean + and len(clean["tool_calls"]) != len(clean.get("tool_call_ids", ())) + ) + or (clean.get("state_only") is True and ("tool_calls" in clean or "tool_call_ids" in clean)) + ): + raise ReasoningReplayError( + f"Opaque chat reasoning state is malformed or unsupported for provider {provider!r}." + ) return clean def _tool_call_state(tool_call: Any) -> dict[str, Any] | None: dumped = opaque_item(tool_call) if not isinstance(dumped, dict): - return None + raise ReasoningReplayError("Malformed provider tool call: expected a mapping.") fields = dumped.get("provider_specific_fields") - signature = fields.get("thought_signature") if isinstance(fields, dict) else None - call_id = dumped.get("id") + if fields is not None and not isinstance(fields, dict): + raise ReasoningReplayError("Malformed tool-call provider_specific_fields.") + if fields: + _warn_unknown_fields(fields, {"thought_signature"}, "a provider tool call") + signature = fields.get("thought_signature") if fields else None if ( - not signature - and isinstance(call_id, str) - and _INLINE_THOUGHT_SIGNATURE_SEPARATOR in call_id + fields + and "thought_signature" in fields + and (not isinstance(signature, str) or not signature) ): - signature = call_id.split(_INLINE_THOUGHT_SIGNATURE_SEPARATOR, 1)[1] - if not isinstance(signature, str) or not signature: + raise ReasoningReplayError("Malformed tool-call thought_signature.") + call_id = dumped.get("id") + inline_signature = None + if isinstance(call_id, str) and _INLINE_THOUGHT_SIGNATURE_SEPARATOR in call_id: + inline_signature = call_id.split(_INLINE_THOUGHT_SIGNATURE_SEPARATOR, 1)[1] + if not inline_signature: + raise ReasoningReplayError("Malformed inline tool-call thought signature.") + if signature and inline_signature and signature != inline_signature: + raise ReasoningReplayError("Conflicting thought signatures on one provider tool call.") + signature = signature or inline_signature + if signature is None: return None return {"provider_specific_fields": {"thought_signature": signature}} @@ -239,8 +313,8 @@ def _tool_call_state(tool_call: Any) -> dict[str, Any] | None: def public_tool_call_id(value: Any, scope: str | None) -> str: """Return an application call id without LiteLLM's inline Gemini state.""" call_id = _field(value, "id", "") - if not isinstance(call_id, str): - return str(call_id or "") + if not isinstance(call_id, str) or not call_id: + raise ReasoningReplayError("Malformed provider tool call: expected a non-empty id.") if _scope_provider(scope) != "gemini": return call_id return call_id.split(_INLINE_THOUGHT_SIGNATURE_SEPARATOR, 1)[0] @@ -250,27 +324,43 @@ def capture_chat_state(message: Any, scope: str | None) -> dict | None: payload: dict[str, Any] = {} for key in ("reasoning_items", "thinking_blocks"): value = _field(message, key) - if isinstance(value, list) and value: - payload[key] = opaque_item(value) + if value is None or value == []: + continue + if not isinstance(value, list): + raise ReasoningReplayError(f"Malformed provider response field {key!r}.") + payload[key] = opaque_item(value) provider_fields = _field(message, "provider_specific_fields") - if isinstance(provider_fields, dict): - payload["provider_specific_fields"] = opaque_item(provider_fields) + if provider_fields is not None and not isinstance(provider_fields, dict): + raise ReasoningReplayError("Malformed provider_specific_fields in provider response.") + if provider_fields: + _warn_unknown_fields(provider_fields, {"thought_signatures"}, "a provider message") + # LiteLLM also puts benign fields such as `refusal` here. Only a known + # thought-signature field belongs in the opaque replay envelope. + if isinstance(provider_fields, dict) and "thought_signatures" in provider_fields: + payload["provider_specific_fields"] = opaque_item( + {"thought_signatures": provider_fields["thought_signatures"]} + ) - raw_tool_calls = list(_field(message, "tool_calls") or []) + raw_tool_calls_value = _field(message, "tool_calls") + if raw_tool_calls_value is not None and not isinstance(raw_tool_calls_value, list): + raise ReasoningReplayError("Malformed tool_calls in provider response.") + raw_tool_calls = raw_tool_calls_value or [] tool_state = [_tool_call_state(call) for call in raw_tool_calls] if any(item is not None for item in tool_state): payload["tool_calls"] = tool_state if payload and raw_tool_calls: tool_call_ids = [public_tool_call_id(call, scope) for call in raw_tool_calls] - if not all(tool_call_ids) or len(set(tool_call_ids)) != len(tool_call_ids): - return None + if len(set(tool_call_ids)) != len(tool_call_ids): + raise ReasoningReplayError( + "Cannot retain opaque reasoning state for duplicate provider tool-call ids." + ) payload["tool_call_ids"] = tool_call_ids - payload = _sanitize_chat_payload(payload, scope) if not payload: return None + payload = _sanitize_chat_payload(payload, scope) if not _field(message, "content") and not _field(message, "tool_calls"): payload["state_only"] = True if not _valid_chat_payload(payload): @@ -378,7 +468,9 @@ def _strip_inline_signature(value: Any) -> Any: def _strip_chat_state(message: dict[str, Any], *, strip_inline_signatures: bool) -> None: + removed = False for key in ("reasoning_items", "thinking_blocks", "provider_specific_fields"): + removed = key in message or removed message.pop(key, None) tool_calls = message.get("tool_calls") if isinstance(tool_calls, list): @@ -387,15 +479,29 @@ def _strip_chat_state(message: dict[str, Any], *, strip_inline_signatures: bool) continue if strip_inline_signatures: call["id"] = _strip_inline_signature(call.get("id")) + removed = "provider_specific_fields" in call or removed call.pop("provider_specific_fields", None) function = call.get("function") if isinstance(function, dict): + removed = "provider_specific_fields" in function or removed function.pop("provider_specific_fields", None) if strip_inline_signatures and "tool_call_id" in message: message["tool_call_id"] = _strip_inline_signature(message["tool_call_id"]) + if removed: + logger.warning( + "Removed untrusted provider reasoning fields from a public chat message; " + "replay opaque state through a persisted LLMResponse instead." + ) -def _restore_chat_state(message: dict[str, Any], payload: dict[str, Any]) -> None: +def _restore_chat_state(message: dict[str, Any], payload: dict[str, Any]) -> bool: + if payload.get("state_only") is True and message != {"role": "assistant", "content": ""}: + logger.warning( + "Opaque reasoning state was not replayed because its empty public carrier " + "changed; portable reasoning text will be replayed instead." + ) + return False + expected_call_ids = payload.get("tool_call_ids") if expected_call_ids is not None: tool_calls = message.get("tool_calls") @@ -404,26 +510,32 @@ def _restore_chat_state(message: dict[str, Any], payload: dict[str, Any]) -> Non or [call.get("id") if isinstance(call, dict) else None for call in tool_calls] != expected_call_ids ): - return - elif isinstance(payload.get("tool_calls"), list): - # Positional per-call state cannot safely survive public call mutation. - return + logger.warning( + "Opaque reasoning state was not replayed because the public tool-call " + "sequence changed; portable reasoning text will be replayed instead." + ) + return False + elif message.get("tool_calls"): + logger.warning( + "Opaque reasoning state was not replayed because tool calls were added to its " + "public carrier; portable reasoning text will be replayed instead." + ) + return False for key in ("reasoning_items", "thinking_blocks", "provider_specific_fields"): if key in payload: message[key] = payload[key] tool_calls = message.get("tool_calls") tool_state = payload.get("tool_calls") - if not isinstance(tool_calls, list) or not isinstance(tool_state, list): - return - if len(tool_calls) != len(tool_state): - return - for call, state in zip(tool_calls, tool_state, strict=True): - if not isinstance(call, dict) or not isinstance(state, dict): - continue - fields = state.get("provider_specific_fields") - if isinstance(fields, dict): - call["provider_specific_fields"] = fields + if tool_state is None: + return True + # The payload validator and identity check above guarantee equal lists. + for call, state in zip(cast(list[Any], tool_calls), cast(list[Any], tool_state), strict=True): + if state is not None: + cast(dict[str, Any], call)["provider_specific_fields"] = cast(dict[str, Any], state)[ + "provider_specific_fields" + ] + return True def capture_responses_state(output: list[Any], scope: str | None) -> dict | None: @@ -431,10 +543,18 @@ def capture_responses_state(output: list[Any], scope: str | None) -> dict | None # may expose a similarly shaped API through a gateway, but that is not # evidence that their opaque state is wire-compatible. if _scope_provider(scope) not in {"openai", "azure"}: + if any(response_item_type(item) == "reasoning" for item in output): + raise ReasoningReplayError( + "The provider returned Responses reasoning state, but NOOA only supports " + "opaque Responses replay for OpenAI and Azure routes." + ) return None items: list[Any] = [] order: list[dict[str, Any]] = [] - has_public_carrier = False + call_ids: list[str] = [] + message_count = 0 + malformed_call_id = False + unknown_output_types: set[str] = set() for item in output: item_type = response_item_type(item) if item_type == "reasoning": @@ -442,20 +562,36 @@ def capture_responses_state(output: list[Any], scope: str | None) -> dict | None 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 + if slot is None: + malformed_call_id = True + continue + order.append(slot) + call_ids.append(cast(str, slot["call_id"])) elif item_type == "message": order.append({"type": "message", "content": _responses_message_text(item)}) - has_public_carrier = True + message_count += 1 + else: + unknown_output_types.add(item_type or "") if not items: return None + if unknown_output_types: + raise ReasoningReplayError( + "Cannot retain Responses reasoning state beside unsupported output type(s): " + + ", ".join(sorted(unknown_output_types)) + ) + if malformed_call_id: + raise ReasoningReplayError( + "Cannot retain Responses reasoning state beside a malformed function call id." + ) + if len(set(call_ids)) != len(call_ids) or message_count > 1: + raise ReasoningReplayError( + "Cannot retain Responses reasoning state with duplicate call ids or multiple messages." + ) payload: dict[str, Any] = {"items": items, "order": order} - if not has_public_carrier: + if not call_ids and not message_count: payload["state_only"] = True if not _valid_responses_payload(payload): - logger.warning("Discarding malformed OpenAI Responses reasoning state.") - return None + raise ReasoningReplayError("Malformed OpenAI Responses reasoning state.") return _envelope(scope, _RESPONSES_FORMAT, payload) @@ -465,9 +601,16 @@ def responses_reasoning_text(output: list[Any]) -> str | None: for item in output: if response_item_type(item) != "reasoning": continue - for summary in _field(item, "summary", []) or []: + summary_items = _field(item, "summary") + if summary_items is None: + summary_items = [] + elif not isinstance(summary_items, list): + raise ReasoningReplayError("Malformed Responses reasoning summary.") + for summary in summary_items: text = _field(summary, "text") - if isinstance(text, str) and text: + if not isinstance(text, str): + raise ReasoningReplayError("Malformed Responses reasoning summary text.") + if text: texts.append(text) return "\n".join(texts) or None @@ -480,16 +623,22 @@ def prepare_chat_messages(messages: list[dict[str, Any]], scope: str | None) -> reasoning = carried_reasoning(original) message = copy.deepcopy(dict(original)) message.pop(LLM_STATE_KEY, None) - source_scope = state.get("scope") if isinstance(state, dict) else None + source_scope = ( + state.get("scope") + if isinstance(state, dict) and state.get("version") == _STATE_VERSION + else None + ) _strip_chat_state( - message, strip_inline_signatures=_scope_provider(source_scope) == "gemini" + message, + strip_inline_signatures=( + isinstance(source_scope, str) and _scope_provider(source_scope) == "gemini" + ), ) payload = _matching_payload(state, scope, _CHAT_FORMAT) if payload is not None: - payload = _sanitize_chat_payload(payload, scope) or None - if payload: - _restore_chat_state(message, payload) - else: + payload = _sanitize_chat_payload(payload, scope) + restored = payload is not None and _restore_chat_state(message, payload) + if not restored: demote_reasoning_text(message, reasoning) if ( not restored @@ -506,14 +655,25 @@ def prepare_chat_messages(messages: list[dict[str, Any]], scope: str | None) -> def _clean_responses_batch(batch: Any) -> list[dict[str, Any]]: if not isinstance(batch, list): - return [] + raise ReasoningReplayError("Malformed Responses replay batch: expected a list.") clean: list[dict[str, Any]] = [] for original in batch: - if not isinstance(original, dict) or response_item_type(original) == "reasoning": + if not isinstance(original, dict): + raise ReasoningReplayError("Malformed Responses replay item: expected a mapping.") + if response_item_type(original) == "reasoning": + logger.warning( + "Removed an untrusted reasoning item from public Responses input; replay " + "opaque state through a persisted LLMResponse instead." + ) continue item = copy.deepcopy(original) item.pop(LLM_STATE_KEY, None) - item.pop("reasoning_items", None) + if "reasoning_items" in item: + logger.warning( + "Removed untrusted reasoning_items from public Responses input; replay " + "opaque state through a persisted LLMResponse instead." + ) + item.pop("reasoning_items") clean.append(item) return clean @@ -526,28 +686,29 @@ def prepare_responses_batch( ) -> 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 _scope_provider(scope) in {"openai", "azure"} - else None - ) + payload = _matching_payload(state, scope, _RESPONSES_FORMAT) if payload is None: return demote_responses_batch(clean, state, reasoning) + if _scope_provider(scope) not in {"openai", "azure"}: + raise ReasoningReplayError( + "Opaque Responses replay is only supported for OpenAI and Azure scopes." + ) + unknown = sorted(set(payload) - {"items", "order", "state_only"}) + if unknown: + raise ReasoningReplayError( + f"Malformed Responses reasoning state: unknown field(s) {', '.join(unknown)}." + ) 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) + raise ReasoningReplayError("Malformed OpenAI Responses reasoning state.") 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." + "Opaque Responses reasoning state was not replayed because its empty public " + "carrier changed; portable reasoning text will be replayed instead." ) return demote_responses_batch(clean, state, reasoning) return cast(list[dict[str, Any]], items) @@ -560,12 +721,10 @@ def prepare_responses_batch( ) return demote_responses_batch(clean, state, reasoning) - calls = { - item.get("call_id"): item - for item in clean - if item.get("type") == "function_call" and isinstance(item.get("call_id"), str) - } - message = next((item for item in clean if item.get("role") == "assistant"), None) + calls = [item for item in clean if item.get("type") == "function_call"] + messages = [item for item in clean if item.get("role") == "assistant"] + calls_by_id = {cast(str, item["call_id"]): item for item in calls} + message = messages[0] if messages else None emitted_calls: set[str] = set() emitted_message = False pending: list[dict[str, Any]] = [] @@ -573,35 +732,22 @@ def prepare_responses_batch( last_carrier_emitted = False for slot in order: - if not isinstance(slot, dict): - continue if slot.get("type") == "reasoning": - index = slot.get("index") - if isinstance(index, int) and 0 <= index < len(items): - item = items[index] - if isinstance(item, dict): - pending.append(item) + pending.append(cast(dict[str, Any], items[cast(int, slot["index"])])) continue if slot.get("type") == "function_call": - call_id = slot.get("call_id") - carrier = calls.get(call_id) if isinstance(call_id, str) else None - if carrier is not None and isinstance(call_id, str): - replay.extend(pending) - replay.append(carrier) - emitted_calls.add(call_id) - last_carrier_emitted = True - else: - last_carrier_emitted = False + call_id = cast(str, slot["call_id"]) + replay.extend(pending) + replay.append(calls_by_id[call_id]) + emitted_calls.add(call_id) + last_carrier_emitted = True pending = [] continue if slot.get("type") == "message": - if message is not None and not emitted_message: - replay.extend(pending) - replay.append(message) - emitted_message = True - last_carrier_emitted = True - else: - last_carrier_emitted = False + replay.extend(pending) + replay.append(cast(dict[str, Any], message)) + emitted_message = True + last_carrier_emitted = True pending = [] if last_carrier_emitted: diff --git a/src/nooa/unifiedllm/unifiedllm.py b/src/nooa/unifiedllm/unifiedllm.py index 793bcb475..9e77cece5 100644 --- a/src/nooa/unifiedllm/unifiedllm.py +++ b/src/nooa/unifiedllm/unifiedllm.py @@ -2623,14 +2623,23 @@ def _transform_messages( continue # A middleware split or mutated the batch. Keep the public item, # but fail closed instead of associating state with new neighbors. + logger.warning( + "Opaque Responses reasoning state was not replayed because middleware " + "split its public carrier batch; portable reasoning text will be " + "replayed on the surviving carrier when available." + ) state = None - reasoning = None msg = copy.deepcopy(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) + if "reasoning_items" in msg: + logger.warning( + "Removed untrusted reasoning_items from public Responses input; replay " + "opaque state through a persisted LLMResponse instead." + ) + msg.pop("reasoning_items") # System messages → extract to instructions if msg.get("role") == "system": @@ -2642,6 +2651,10 @@ def _transform_messages( # Already in native Responses format (from ResponsesProviderFormatter) if "type" in msg: if replay_state.response_item_type(msg) == "reasoning": + logger.warning( + "Removed an untrusted reasoning item from public Responses input; " + "replay opaque state through a persisted LLMResponse instead." + ) continue if state is not None or reasoning is not None: transformed.extend( diff --git a/tests/strategies/test_error_recovery_gl106.py b/tests/strategies/test_error_recovery_gl106.py index 4bc3309a1..805fa28ee 100644 --- a/tests/strategies/test_error_recovery_gl106.py +++ b/tests/strategies/test_error_recovery_gl106.py @@ -23,7 +23,7 @@ from nooa.strategies.codeact import CodeActStrategy from nooa.strategies.predict import PredictStrategy from nooa.strategies.pure_python import PurePythonStrategy -from nooa.unifiedllm import FakeLLMClient, LLMResponse, ToolCall +from nooa.unifiedllm import FakeLLMClient, LLMResponse, ReasoningReplayError, ToolCall # --------------------------------------------------------------------------- # Module-level Pydantic models (required for PredictStrategy type resolution) @@ -98,6 +98,42 @@ async def acall( return await super().acall(messages, tools, output_model, **kwargs) +class ReasoningReplayFailingLLM(FakeLLMClient): + """Raise a framework replay error and count attempts.""" + + def __init__(self): + super().__init__() + self.attempts = 0 + + async def acall( + self, messages: list[dict], tools=None, output_model=None, **kwargs + ) -> LLMResponse: + self.attempts += 1 + raise ReasoningReplayError("malformed retained reasoning state") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("strategy_name", ["codeact", "pure_python", "predict"]) +async def test_reasoning_replay_errors_bypass_strategy_retries(strategy_name: str) -> None: + if strategy_name == "codeact": + selected = CodeActStrategy(config=CodeActConfig(max_retries=3)) + elif strategy_name == "pure_python": + selected = PurePythonStrategy(max_retries=3) + else: + selected = PredictStrategy(config=PredictConfig(max_retries=3)) + + class TestAgent(Agent, llm=_DUMMY_LLM): + @strategy(selected) + async def compute(self, value: int) -> int: + """Return {value}.""" + ... + + failing_llm = ReasoningReplayFailingLLM() + with pytest.raises(ReasoningReplayError, match="malformed retained reasoning state"): + await TestAgent(llm=failing_llm).compute(1) + assert failing_llm.attempts == 1 + + # --------------------------------------------------------------------------- # CodeActStrategy — LLM API error exhaustion (codeact.py lines 668-699) # --------------------------------------------------------------------------- diff --git a/tests/unifiedllm/test_cross_provider_reasoning.py b/tests/unifiedllm/test_cross_provider_reasoning.py index 5fabe0893..dd856657e 100644 --- a/tests/unifiedllm/test_cross_provider_reasoning.py +++ b/tests/unifiedllm/test_cross_provider_reasoning.py @@ -20,6 +20,8 @@ from nooa.context_blocks.models import ResolvedBlock, Role from nooa.unifiedllm import CompletionClient, LLMResponse, ResponsesClient, Tool from nooa.unifiedllm.replay_state import ( + ReasoningReplayError, + capture_chat_state, capture_responses_state, prepare_chat_messages, prepare_responses_batch, @@ -212,7 +214,9 @@ def test_gemini_signatures_round_trip_without_becoming_public_call_ids() -> None @pytest.mark.parametrize("mutation", ["drop", "reorder", "duplicate"]) -def test_gemini_tool_state_fails_closed_when_public_calls_change(mutation: str) -> None: +def test_gemini_tool_state_warns_and_demotes_when_public_calls_change( + mutation: str, caplog: pytest.LogCaptureFixture +) -> None: client = CompletionClient(model="gemini/gemini-2.5-pro", api_key="account-a") try: with patch("litellm.completion", return_value=_gemini_response()): @@ -232,8 +236,10 @@ def test_gemini_tool_state_fails_closed_when_public_calls_change(mutation: str) replayed = next(message for message in prepared if message.get("tool_calls")) assert "provider_specific_fields" not in replayed assert all("provider_specific_fields" not in call for call in replayed["tool_calls"]) + assert replayed["content"] == "Inspect the value." assert GEMINI_SIGNATURE not in json.dumps(prepared) assert GEMINI_SIGNATURE_2 not in json.dumps(prepared) + assert "public tool-call sequence changed" in caplog.text finally: client.close() @@ -259,7 +265,59 @@ def test_non_gemini_tool_call_id_with_thought_substring_is_unchanged() -> None: client.close() -def test_cross_provider_replay_hides_opaque_state_but_keeps_reasoning_text() -> None: +def test_missing_thought_signature_is_normal_but_malformed_signature_fails( + caplog: pytest.LogCaptureFixture, +) -> None: + scope = replay_scope("gemini/gemini-2.5-pro", "chat", {}) + assert ( + capture_chat_state( + Message(role="assistant", content=None, tool_calls=[_tool_call()]), scope + ) + is None + ) + assert not caplog.text + + malformed = Message( + role="assistant", + content=None, + tool_calls=[_tool_call("call_1", {"thought_signature": 42})], + ) + with pytest.raises(ReasoningReplayError, match="thought_signature"): + capture_chat_state(malformed, scope) + + with pytest.raises(ReasoningReplayError, match="expected a mapping"): + capture_chat_state({"tool_calls": ["not-a-tool-call"]}, scope) + + +def test_unknown_provider_state_warns_but_unknown_envelope_state_fails( + caplog: pytest.LogCaptureFixture, +) -> None: + scope = replay_scope("gemini/gemini-2.5-pro", "chat", {}) + assert ( + capture_chat_state( + {"provider_specific_fields": {"future_reasoning_state": "opaque"}}, scope + ) + is None + ) + assert "retention may need updating" in caplog.text + + state = { + "version": 1, + "scope": scope, + "format": "litellm-chat", + "payload": { + "thinking_blocks": [{"type": "thinking", "signature": "opaque"}], + "future_reasoning_state": "opaque", + }, + } + carrier = ReplayCarryingMessage({"role": "assistant", "content": "answer"}, state) + with pytest.raises(ReasoningReplayError, match="malformed or unsupported"): + prepare_chat_messages([carrier], scope) + + +def test_cross_provider_replay_warns_hides_opaque_state_and_keeps_reasoning_text( + caplog: pytest.LogCaptureFixture, +) -> None: source = CompletionClient(model="gemini/gemini-2.5-pro", api_key="account-a") target = CompletionClient(model="anthropic/claude-sonnet-4", api_key="account-a") try: @@ -277,6 +335,7 @@ def test_cross_provider_replay_hides_opaque_state_but_keeps_reasoning_text() -> assert [call["id"] for call in assistant["tool_calls"]] == ["call_1", "call_2"] assert GEMINI_SIGNATURE not in json.dumps(replayed) assert GEMINI_SIGNATURE_2 not in json.dumps(replayed) + assert "is incompatible with" in caplog.text finally: source.close() target.close() @@ -412,6 +471,68 @@ def test_state_only_turn_drops_empty_carrier_across_api_styles() -> None: ) +def test_state_only_carrier_mutation_warns_and_preserves_public_content( + caplog: pytest.LogCaptureFixture, +) -> None: + chat_scope = "chat:openai:sha256:model" + chat_state = { + "version": 1, + "scope": chat_scope, + "format": "litellm-chat", + "payload": {"reasoning_items": [{"opaque": "chat-secret"}], "state_only": True}, + } + chat_carrier = ReplayCarryingMessage( + {"role": "assistant", "content": "middleware text"}, + chat_state, + "portable reasoning", + ) + assert prepare_chat_messages([chat_carrier], chat_scope) == [ + {"role": "assistant", "content": "portable reasoning\n\nmiddleware text"} + ] + + responses_scope = "responses:openai:sha256:model" + responses_state = { + "version": 1, + "scope": responses_scope, + "format": "openai-responses", + "payload": { + "items": [RESPONSES_REASONING], + "order": [{"type": "reasoning", "index": 0}], + "state_only": True, + }, + } + assert prepare_responses_batch( + [{"role": "assistant", "content": "middleware text"}], + responses_state, + responses_scope, + "portable reasoning", + ) == [{"role": "assistant", "content": "portable reasoning\n\nmiddleware text"}] + assert caplog.text.count("empty public carrier changed") == 2 + + +def test_legacy_state_warns_while_malformed_current_state_fails( + caplog: pytest.LogCaptureFixture, +) -> None: + legacy = ReplayCarryingMessage( + {"role": "assistant", "content": "answer"}, + {"reasoning_items": [{"opaque": "legacy"}]}, + "portable reasoning", + ) + assert prepare_chat_messages([legacy], None)[0]["content"] == ("portable reasoning\n\nanswer") + assert "unsupported or legacy version" in caplog.text + + scope = "chat:openai:sha256:model" + malformed_states = [ + {"version": 1, "scope": scope, "format": "litellm-chat", "payload": []}, + {"version": 1, "scope": scope, "format": "typo", "payload": {}}, + {"version": 1, "scope": scope, "format": "litellm-chat", "payload": {}, "extra": 1}, + ] + for state in malformed_states: + malformed = ReplayCarryingMessage({"role": "assistant", "content": "answer"}, state) + with pytest.raises(ReasoningReplayError, match="Malformed version-1"): + prepare_chat_messages([malformed], scope) + + def test_responses_demotion_keeps_native_output_content_valid() -> None: native_message = { "type": "message", @@ -440,9 +561,10 @@ def test_unverified_closed_provider_routes_have_no_opaque_replay_scope( assert replay_scope(model, api_style, {"api_key": "account-a"}) is None -def test_non_openai_responses_scope_cannot_capture_or_restore_opaque_items() -> None: +def test_non_openai_responses_scope_rejects_capture_and_restore() -> None: fabricated_scope = "responses:anthropic:sha256:untrusted" - assert capture_responses_state([RESPONSES_REASONING], fabricated_scope) is None + with pytest.raises(ReasoningReplayError, match="only supports.*OpenAI and Azure"): + capture_responses_state([RESPONSES_REASONING], fabricated_scope) state = { "version": 1, @@ -450,12 +572,8 @@ def test_non_openai_responses_scope_cannot_capture_or_restore_opaque_items() -> "format": "openai-responses", "payload": {"items": [RESPONSES_REASONING], "order": []}, } - assert prepare_responses_batch( - [RESPONSES_MESSAGE], state, fabricated_scope, "Check the evidence." - ) == [ - {"role": "assistant", "content": "Check the evidence."}, - RESPONSES_MESSAGE, - ] + with pytest.raises(ReasoningReplayError, match="only supported for OpenAI and Azure"): + prepare_responses_batch([RESPONSES_MESSAGE], state, fabricated_scope, "Check the evidence.") @pytest.mark.parametrize( diff --git a/tests/unifiedllm/test_reasoning_state_replay.py b/tests/unifiedllm/test_reasoning_state_replay.py index cc6187457..cef4f79b6 100644 --- a/tests/unifiedllm/test_reasoning_state_replay.py +++ b/tests/unifiedllm/test_reasoning_state_replay.py @@ -25,9 +25,11 @@ from nooa.runtime.middleware import LLMCallContext from nooa.unifiedllm import CompletionClient, LLMResponse, ResponsesClient, Tool from nooa.unifiedllm.replay_state import ( - prepare_chat_messages, + ReasoningReplayError, + capture_responses_state, prepare_responses_batch, replay_scope, + responses_reasoning_text, ) REASONING = { @@ -290,7 +292,7 @@ def test_reasoning_only_carrier_edit_drops_state_but_keeps_text() -> None: client.close() -def test_malformed_matching_responses_payload_is_not_forwarded(caplog) -> None: +def test_malformed_matching_responses_payload_raises() -> None: scope = replay_scope("openai/gpt-5.6", "responses", {}) state = { "version": 1, @@ -305,20 +307,26 @@ def test_malformed_matching_responses_payload_is_not_forwarded(caplog) -> None: }, } - 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 + with pytest.raises(ReasoningReplayError, match="Malformed OpenAI Responses"): + prepare_responses_batch( + [{"role": "assistant", "content": "public"}], + state, + scope, + "portable reasoning", + ) -def test_split_responses_replay_batch_keeps_public_call_but_drops_state() -> None: +def test_split_responses_replay_batch_warns_and_demotes_text( + caplog: pytest.LogCaptureFixture, +) -> None: """Middleware may edit public items, but cannot reattach state to new neighbors.""" client = ResponsesClient(model="openai/gpt-5.6", api_key="account-a") + reasoning = { + **REASONING, + "summary": [{"type": "summary_text", "text": "Check the evidence."}], + } try: - with patch("litellm.responses", return_value=_responses(REASONING, CALL, CALL_2)): + 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 @@ -327,13 +335,60 @@ def test_split_responses_replay_batch_keeps_public_call_but_drops_state() -> Non 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 [item.get("call_id") for item in replay if item.get("type") == "function_call"] == [ + "call_1" + ] + assert replay[0] == {"role": "assistant", "content": "Check the evidence."} + assert reasoning not in replay assert "provider-secret" not in repr(replay) + assert "middleware split its public carrier batch" in caplog.text finally: client.close() +def test_changed_responses_carrier_warns_and_demotes_text( + caplog: pytest.LogCaptureFixture, +) -> None: + scope = replay_scope("openai/gpt-5.6", "responses", {}) + reasoning = { + **REASONING, + "summary": [{"type": "summary_text", "text": "Check the evidence."}], + } + state = capture_responses_state([reasoning, CALL], scope) + changed_call = {**CALL, "call_id": "changed"} + + assert prepare_responses_batch([changed_call], state, scope, "Check the evidence.") == [ + {"role": "assistant", "content": "Check the evidence."}, + changed_call, + ] + assert "public carriers changed" in caplog.text + + +def test_responses_envelope_cannot_replay_a_non_reasoning_item_as_state() -> None: + scope = replay_scope("openai/gpt-5.6", "responses", {}) + state = { + "version": 1, + "scope": scope, + "format": "openai-responses", + "payload": { + "items": [CALL], + "order": [ + {"type": "reasoning", "index": 0}, + {"type": "function_call", "call_id": "call_1"}, + ], + }, + } + + with pytest.raises(ReasoningReplayError, match="Malformed Responses reasoning item"): + prepare_responses_batch([CALL], state, scope) + + with pytest.raises(ReasoningReplayError, match="unsupported output type.*future_public"): + capture_responses_state([REASONING, {"type": "future_public"}], scope) + + with pytest.raises(ReasoningReplayError, match="Malformed Responses reasoning summary"): + responses_reasoning_text([{**REASONING, "summary": ""}]) + + def test_reasoning_only_response_replays_without_empty_assistant_message() -> None: client = ResponsesClient(model="openai/gpt-5.6", api_key="account-a") try: @@ -604,122 +659,27 @@ def test_chat_state_is_captured_replayed_and_api_style_scoped() -> None: 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 +def test_unresolved_route_fails_when_provider_returns_opaque_state( + 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 first.llm_state["payload"]["carrier"] == {"content": "", "tool_calls": []} - 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"}]) + with pytest.raises(ReasoningReplayError, match="only supports.*OpenAI and Azure"): + client.call([{"role": "user", "content": "think"}]) - assert response.llm_state is None assert "include" not in call.call_args.kwargs + assert "could not resolve model 'unknown-route'" in caplog.text finally: client.close() -def test_direct_reasoning_items_cannot_bypass_envelope_gate() -> None: +def test_direct_reasoning_items_warn_and_cannot_bypass_envelope_gate( + caplog: pytest.LogCaptureFixture, +) -> None: client = ResponsesClient(model="openai/gpt-5.6", api_key="account-a") try: with patch("litellm.responses", return_value=_responses(MESSAGE)) as call: @@ -727,6 +687,7 @@ def test_direct_reasoning_items_cannot_bypass_envelope_gate() -> None: assert REASONING not in call.call_args.kwargs["input"] assert "provider-secret" not in repr(call.call_args.kwargs["input"]) + assert "untrusted reasoning item" in caplog.text finally: client.close() @@ -823,7 +784,7 @@ async def test_responses_reasoning_uses_per_call_override() -> None: @pytest.mark.parametrize("model", ["anthropic/claude-sonnet-4-5", "gemini/gemini-2.5-pro"]) -def test_non_openai_chat_provider_cannot_receive_openai_reasoning_state(model: str) -> None: +def test_current_envelope_with_wrong_provider_payload_fails(model: str) -> None: scope = replay_scope(model, "chat", {"api_key": "account-a"}) assert scope is not None client = CompletionClient(model=model, api_key="account-a") @@ -834,13 +795,14 @@ def test_non_openai_chat_provider_cannot_receive_openai_reasoning_state(model: s "payload": {"reasoning_items": [REASONING]}, } try: - with patch("litellm.completion", return_value=_chat_response()) as call: + with ( + patch("litellm.completion", return_value=_chat_response()) as call, + pytest.raises(ReasoningReplayError, match="malformed or unsupported"), + ): 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) + call.assert_not_called() finally: client.close() From f03a2a8fe6d33a9d54b53c2115e9f655d47c8c1c Mon Sep 17 00:00:00 2001 From: Paul Furgale Date: Thu, 10 Sep 2026 15:40:38 +0000 Subject: [PATCH 10/13] fix(llm): bind chat state to public carriers Signed-off-by: Paul Furgale --- src/nooa/strategies/pure_python.py | 4 +- src/nooa/unifiedllm/replay_state.py | 232 ++++++++++++------ tests/strategies/test_error_recovery_gl106.py | 9 +- .../test_cross_provider_reasoning.py | 154 +++++++++++- .../test_litellm_responses_bridge.py | 11 +- .../unifiedllm/test_reasoning_state_replay.py | 2 +- 6 files changed, 327 insertions(+), 85 deletions(-) diff --git a/src/nooa/strategies/pure_python.py b/src/nooa/strategies/pure_python.py index 55adc6247..3b8c8c4a3 100644 --- a/src/nooa/strategies/pure_python.py +++ b/src/nooa/strategies/pure_python.py @@ -308,7 +308,9 @@ async def execute(self, runtime: RuntimeServices, call: "CurrentCall") -> Any: generate_event_id: str | None = None try: code, generate_event_id = await self._generate_code(runtime, session) - except ReasoningReplayError: + except ReasoningReplayError as e: + turn_final = True + turn_exception = type(e).__name__ raise except _HTTPX_TIMEOUT_EXCEPTIONS as e: # Catch httpx timeout exceptions and preserve them diff --git a/src/nooa/unifiedllm/replay_state.py b/src/nooa/unifiedllm/replay_state.py index f087be0e7..aebc50b99 100644 --- a/src/nooa/unifiedllm/replay_state.py +++ b/src/nooa/unifiedllm/replay_state.py @@ -209,6 +209,57 @@ def _scope_provider(scope: str | None) -> str | None: return parts[1] +def _chat_public_carrier(message: Any, scope: str | None) -> dict[str, Any] | None: + """Project 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": public_tool_call_id(call, scope), + "name": name, + "arguments": arguments, + } + ) + + content = _field(message, "content") + if content is not None and not isinstance(content, str): + return None + # The canonical formatter uses None for an empty assistant tool-call turn + # and an empty string for an assistant turn with no public carrier. + content = (content or None) if calls else (content or "") + return {"content": content, "tool_calls": calls} + + +def _valid_chat_carrier(value: Any) -> bool: + if not isinstance(value, dict) or set(value) != {"content", "tool_calls"}: + return False + if value.get("content") is not None and not isinstance(value.get("content"), str): + return False + calls = value.get("tool_calls") + return isinstance(calls, list) and all( + isinstance(call, dict) + and set(call) == {"id", "name", "arguments"} + and all(isinstance(call.get(key), str) for key in ("id", "name", "arguments")) + for call in calls + ) + + def _sanitize_chat_payload(payload: dict[str, Any], scope: str | None) -> dict[str, Any]: """Validate the small provider-specific envelope owned by NOOA.""" provider = _scope_provider(scope) @@ -245,34 +296,37 @@ def _sanitize_chat_payload(payload: dict[str, Any], scope: str | None) -> dict[s for item in tool_state: fields = item.get("provider_specific_fields") if isinstance(item, dict) else None signature = fields.get("thought_signature") if isinstance(fields, dict) else None - calls.append( - {"provider_specific_fields": {"thought_signature": signature}} - if isinstance(signature, str) and signature - else None + inline_signature = ( + item.get("inline_thought_signature") if isinstance(item, dict) else None ) + call: dict[str, Any] = {} + if isinstance(signature, str) and signature: + call["provider_specific_fields"] = {"thought_signature": signature} + if isinstance(inline_signature, str) and inline_signature: + call["inline_thought_signature"] = inline_signature + if signature and inline_signature and signature != inline_signature: + raise ReasoningReplayError("Conflicting stored tool-call thought signatures.") + calls.append(call or None) if any(item is not None for item in calls): clean["tool_calls"] = calls - tool_call_ids = payload.get("tool_call_ids") - if ( - clean - and isinstance(tool_call_ids, list) - and tool_call_ids - and all(isinstance(call_id, str) and call_id for call_id in tool_call_ids) - and len(set(tool_call_ids)) == len(tool_call_ids) - ): - clean["tool_call_ids"] = tool_call_ids - - if clean and payload.get("state_only") is True: + has_state = bool(clean) + carrier = payload.get("carrier") + if has_state and _valid_chat_carrier(carrier): + clean["carrier"] = carrier + if has_state and payload.get("state_only") is True: clean["state_only"] = True if ( - not clean + not has_state or clean != payload or ( "tool_calls" in clean - and len(clean["tool_calls"]) != len(clean.get("tool_call_ids", ())) + and len(clean["tool_calls"]) != len(clean.get("carrier", {}).get("tool_calls", ())) + ) + or ( + clean.get("state_only") is True + and ("tool_calls" in clean or clean.get("carrier") != {"content": "", "tool_calls": []}) ) - or (clean.get("state_only") is True and ("tool_calls" in clean or "tool_call_ids" in clean)) ): raise ReasoningReplayError( f"Opaque chat reasoning state is malformed or unsupported for provider {provider!r}." @@ -280,7 +334,7 @@ def _sanitize_chat_payload(payload: dict[str, Any], scope: str | None) -> dict[s return clean -def _tool_call_state(tool_call: Any) -> dict[str, Any] | None: +def _tool_call_state(tool_call: Any, scope: str | None) -> dict[str, Any] | None: dumped = opaque_item(tool_call) if not isinstance(dumped, dict): raise ReasoningReplayError("Malformed provider tool call: expected a mapping.") @@ -297,17 +351,26 @@ def _tool_call_state(tool_call: Any) -> dict[str, Any] | None: ): raise ReasoningReplayError("Malformed tool-call thought_signature.") call_id = dumped.get("id") - inline_signature = None + inline_candidate = None if isinstance(call_id, str) and _INLINE_THOUGHT_SIGNATURE_SEPARATOR in call_id: - inline_signature = call_id.split(_INLINE_THOUGHT_SIGNATURE_SEPARATOR, 1)[1] - if not inline_signature: + inline_candidate = call_id.split(_INLINE_THOUGHT_SIGNATURE_SEPARATOR, 1)[1] + if not inline_candidate: raise ReasoningReplayError("Malformed inline tool-call thought signature.") - if signature and inline_signature and signature != inline_signature: + if signature and inline_candidate and signature != inline_candidate: raise ReasoningReplayError("Conflicting thought signatures on one provider tool call.") - signature = signature or inline_signature - if signature is None: + inline_signature = ( + inline_candidate + if _scope_provider(scope) == "gemini" or signature == inline_candidate + else None + ) + if signature is None and inline_signature is None: return None - return {"provider_specific_fields": {"thought_signature": signature}} + state: dict[str, Any] = {} + if signature is not None: + state["provider_specific_fields"] = {"thought_signature": signature} + if inline_signature is not None: + state["inline_thought_signature"] = inline_signature + return state def public_tool_call_id(value: Any, scope: str | None) -> str: @@ -315,7 +378,8 @@ def public_tool_call_id(value: Any, scope: str | None) -> str: call_id = _field(value, "id", "") if not isinstance(call_id, str) or not call_id: raise ReasoningReplayError("Malformed provider tool call: expected a non-empty id.") - if _scope_provider(scope) != "gemini": + state = _tool_call_state(value, scope) + if state is None or "inline_thought_signature" not in state: return call_id return call_id.split(_INLINE_THOUGHT_SIGNATURE_SEPARATOR, 1)[0] @@ -346,26 +410,23 @@ def capture_chat_state(message: Any, scope: str | None) -> dict | None: if raw_tool_calls_value is not None and not isinstance(raw_tool_calls_value, list): raise ReasoningReplayError("Malformed tool_calls in provider response.") raw_tool_calls = raw_tool_calls_value or [] - tool_state = [_tool_call_state(call) for call in raw_tool_calls] + tool_state = [_tool_call_state(call, scope) for call in raw_tool_calls] if any(item is not None for item in tool_state): payload["tool_calls"] = tool_state - if payload and raw_tool_calls: - tool_call_ids = [public_tool_call_id(call, scope) for call in raw_tool_calls] - if len(set(tool_call_ids)) != len(tool_call_ids): - raise ReasoningReplayError( - "Cannot retain opaque reasoning state for duplicate provider tool-call ids." - ) - payload["tool_call_ids"] = tool_call_ids - if not payload: return None - payload = _sanitize_chat_payload(payload, scope) + carrier = _chat_public_carrier(message, scope) + if carrier is None or len({call["id"] for call in carrier["tool_calls"]}) != len( + carrier["tool_calls"] + ): + raise ReasoningReplayError( + "Cannot retain opaque reasoning state for a malformed public assistant carrier." + ) + payload["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 + payload = _sanitize_chat_payload(payload, scope) return _envelope(scope, _CHAT_FORMAT, payload) @@ -467,26 +528,57 @@ def _strip_inline_signature(value: Any) -> Any: return value -def _strip_chat_state(message: dict[str, Any], *, strip_inline_signatures: bool) -> None: +def _strip_chat_state( + message: dict[str, Any], source_scope: str | None, target_scope: str | None +) -> None: removed = False + gemini_wire = "gemini" in { + _scope_provider(source_scope), + _scope_provider(target_scope), + } for key in ("reasoning_items", "thinking_blocks", "provider_specific_fields"): removed = key in message or removed message.pop(key, None) + content = message.get("content") + if isinstance(content, list): + public_blocks = [ + block + for block in content + if not ( + isinstance(block, dict) and block.get("type") in {"thinking", "redacted_thinking"} + ) + ] + removed = len(public_blocks) != len(content) or removed + message["content"] = public_blocks tool_calls = message.get("tool_calls") if isinstance(tool_calls, list): for call in tool_calls: if not isinstance(call, dict): continue - if strip_inline_signatures: - call["id"] = _strip_inline_signature(call.get("id")) + call_id = call.get("id") + fields = call.get("provider_specific_fields") + explicit_signature = ( + fields.get("thought_signature") if isinstance(fields, dict) else None + ) + inline_candidate = ( + call_id.split(_INLINE_THOUGHT_SIGNATURE_SEPARATOR, 1)[1] + if isinstance(call_id, str) and _INLINE_THOUGHT_SIGNATURE_SEPARATOR in call_id + else None + ) + if gemini_wire or (inline_candidate and inline_candidate == explicit_signature): + public_id = _strip_inline_signature(call_id) + removed = public_id != call_id or removed + call["id"] = public_id removed = "provider_specific_fields" in call or removed call.pop("provider_specific_fields", None) function = call.get("function") if isinstance(function, dict): removed = "provider_specific_fields" in function or removed function.pop("provider_specific_fields", None) - if strip_inline_signatures and "tool_call_id" in message: - message["tool_call_id"] = _strip_inline_signature(message["tool_call_id"]) + if gemini_wire and "tool_call_id" in message: + public_id = _strip_inline_signature(message["tool_call_id"]) + removed = public_id != message["tool_call_id"] or removed + message["tool_call_id"] = public_id if removed: logger.warning( "Removed untrusted provider reasoning fields from a public chat message; " @@ -495,33 +587,13 @@ def _strip_chat_state(message: dict[str, Any], *, strip_inline_signatures: bool) def _restore_chat_state(message: dict[str, Any], payload: dict[str, Any]) -> bool: - if payload.get("state_only") is True and message != {"role": "assistant", "content": ""}: + if _chat_public_carrier(message, None) != payload.get("carrier"): logger.warning( - "Opaque reasoning state was not replayed because its empty public carrier " + "Opaque reasoning state was not replayed because its public assistant carrier " "changed; portable reasoning text will be replayed instead." ) return False - expected_call_ids = payload.get("tool_call_ids") - if expected_call_ids is not None: - tool_calls = message.get("tool_calls") - if ( - not isinstance(tool_calls, list) - or [call.get("id") if isinstance(call, dict) else None for call in tool_calls] - != expected_call_ids - ): - logger.warning( - "Opaque reasoning state was not replayed because the public tool-call " - "sequence changed; portable reasoning text will be replayed instead." - ) - return False - elif message.get("tool_calls"): - logger.warning( - "Opaque reasoning state was not replayed because tool calls were added to its " - "public carrier; portable reasoning text will be replayed instead." - ) - return False - for key in ("reasoning_items", "thinking_blocks", "provider_specific_fields"): if key in payload: message[key] = payload[key] @@ -532,9 +604,12 @@ def _restore_chat_state(message: dict[str, Any], payload: dict[str, Any]) -> boo # The payload validator and identity check above guarantee equal lists. for call, state in zip(cast(list[Any], tool_calls), cast(list[Any], tool_state), strict=True): if state is not None: - cast(dict[str, Any], call)["provider_specific_fields"] = cast(dict[str, Any], state)[ - "provider_specific_fields" - ] + call = cast(dict[str, Any], call) + state = cast(dict[str, Any], state) + if "provider_specific_fields" in state: + call["provider_specific_fields"] = state["provider_specific_fields"] + if inline_signature := state.get("inline_thought_signature"): + call["id"] = f"{call['id']}{_INLINE_THOUGHT_SIGNATURE_SEPARATOR}{inline_signature}" return True @@ -618,6 +693,7 @@ def responses_reasoning_text(output: list[Any]) -> str | None: 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]] = [] + private_call_ids: dict[str, str] = {} for original in messages: state = carried_state(original) reasoning = carried_reasoning(original) @@ -630,14 +706,24 @@ def prepare_chat_messages(messages: list[dict[str, Any]], scope: str | None) -> ) _strip_chat_state( message, - strip_inline_signatures=( - isinstance(source_scope, str) and _scope_provider(source_scope) == "gemini" - ), + source_scope if isinstance(source_scope, str) else None, + scope, ) payload = _matching_payload(state, scope, _CHAT_FORMAT) if payload is not None: payload = _sanitize_chat_payload(payload, scope) restored = payload is not None and _restore_chat_state(message, payload) + if restored: + public_calls = cast(dict[str, Any], payload)["carrier"]["tool_calls"] + private_calls = message.get("tool_calls") or [] + for public_call, private_call in zip(public_calls, private_calls, strict=True): + public_id = public_call["id"] + private_id = private_call["id"] + if public_id != private_id: + private_call_ids[public_id] = private_id + tool_call_id = message.get("tool_call_id") + if isinstance(tool_call_id, str) and tool_call_id in private_call_ids: + message["tool_call_id"] = private_call_ids.pop(tool_call_id) if not restored: demote_reasoning_text(message, reasoning) if ( diff --git a/tests/strategies/test_error_recovery_gl106.py b/tests/strategies/test_error_recovery_gl106.py index 805fa28ee..23d705e1c 100644 --- a/tests/strategies/test_error_recovery_gl106.py +++ b/tests/strategies/test_error_recovery_gl106.py @@ -129,9 +129,16 @@ async def compute(self, value: int) -> int: ... failing_llm = ReasoningReplayFailingLLM() + agent = TestAgent(llm=failing_llm) + after_turns = [] + agent.event_manager.on("AfterTurn", after_turns.append) with pytest.raises(ReasoningReplayError, match="malformed retained reasoning state"): - await TestAgent(llm=failing_llm).compute(1) + await agent.compute(1) assert failing_llm.attempts == 1 + if strategy_name == "pure_python": + assert after_turns[-1].is_final is True + assert after_turns[-1].success is False + assert after_turns[-1].exception_type == "ReasoningReplayError" # --------------------------------------------------------------------------- diff --git a/tests/unifiedllm/test_cross_provider_reasoning.py b/tests/unifiedllm/test_cross_provider_reasoning.py index dd856657e..9357ffa90 100644 --- a/tests/unifiedllm/test_cross_provider_reasoning.py +++ b/tests/unifiedllm/test_cross_provider_reasoning.py @@ -190,7 +190,10 @@ def test_gemini_signatures_round_trip_without_becoming_public_call_ids() -> None assert [call.id for call in first.tool_calls] == ["call_1", "call_2"] assert first.llm_state is not None assert "discard-me" not in json.dumps(first.llm_state) - assert first.llm_state["payload"]["tool_call_ids"] == ["call_1", "call_2"] + assert [call["id"] for call in first.llm_state["payload"]["carrier"]["tool_calls"]] == [ + "call_1", + "call_2", + ] assistant = next( message for message in completion.call_args_list[1].kwargs["messages"] @@ -203,7 +206,10 @@ def test_gemini_signatures_round_trip_without_becoming_public_call_ids() -> None assistant["provider_specific_fields"]["thought_signatures"] is first.llm_state["payload"]["provider_specific_fields"]["thought_signatures"] ) - assert [call["id"] for call in assistant["tool_calls"]] == ["call_1", "call_2"] + assert [call["id"] for call in assistant["tool_calls"]] == [ + f"call_1__thought__{GEMINI_SIGNATURE}", + f"call_2__thought__{GEMINI_SIGNATURE_2}", + ] assert [ call["provider_specific_fields"]["thought_signature"] for call in assistant["tool_calls"] @@ -213,7 +219,72 @@ def test_gemini_signatures_round_trip_without_becoming_public_call_ids() -> None target.close() -@pytest.mark.parametrize("mutation", ["drop", "reorder", "duplicate"]) +def test_gateway_routed_gemini_keeps_inline_signatures_private() -> None: + client = CompletionClient( + model="openai/gcp/google/gemini-3.1-pro-preview", + api_base="https://inference-api.example/v1", + ) + try: + with patch( + "litellm.completion", side_effect=[_gemini_response(), _gemini_response()] + ) as completion: + first = client.call([{"role": "user", "content": "run"}], tools=[TOOL]) + client.call(_render(first), tools=[TOOL]) + + assert [call.id for call in first.tool_calls] == ["call_1", "call_2"] + assert first.llm_state is not None + tool_state = first.llm_state["payload"]["tool_calls"] + assert [call["inline_thought_signature"] for call in tool_state] == [ + GEMINI_SIGNATURE, + GEMINI_SIGNATURE_2, + ] + assistant = next( + message + for message in completion.call_args_list[1].kwargs["messages"] + if message.get("role") == "assistant" + ) + assert [call["id"] for call in assistant["tool_calls"]] == [ + f"call_1__thought__{GEMINI_SIGNATURE}", + f"call_2__thought__{GEMINI_SIGNATURE_2}", + ] + tool_results = [ + message + for message in completion.call_args_list[1].kwargs["messages"] + if message.get("role") == "tool" + ] + assert [message["tool_call_id"] for message in tool_results] == [ + f"call_1__thought__{GEMINI_SIGNATURE}", + f"call_2__thought__{GEMINI_SIGNATURE_2}", + ] + finally: + client.close() + + +def test_incompatible_gemini_state_keeps_tool_result_ids_public() -> None: + source = CompletionClient( + model="openai/gcp/google/gemini-3.1-pro-preview", + api_base="https://inference-api.example/v1", + ) + try: + with patch("litellm.completion", return_value=_gemini_response()): + first = source.call([{"role": "user", "content": "run"}], tools=[TOOL]) + + prepared = prepare_chat_messages( + _render(first), replay_scope("openai/gpt-5.6", "chat", {}) + ) + assistant = next(message for message in prepared if message.get("tool_calls")) + tool_results = [message for message in prepared if message.get("role") == "tool"] + assert [call["id"] for call in assistant["tool_calls"]] == ["call_1", "call_2"] + assert [message["tool_call_id"] for message in tool_results] == [ + "call_1", + "call_2", + ] + assert "__thought__" not in json.dumps(prepared) + finally: + source.close() + + +@pytest.mark.parametrize("mutation", ["drop", "reorder", "duplicate", "text", "name", "arguments"]) def test_gemini_tool_state_warns_and_demotes_when_public_calls_change( mutation: str, caplog: pytest.LogCaptureFixture ) -> None: @@ -229,21 +300,82 @@ def test_gemini_tool_state_warns_and_demotes_when_public_calls_change( assistant["tool_calls"].pop() elif mutation == "reorder": assistant["tool_calls"].reverse() - else: + elif mutation == "duplicate": assistant["tool_calls"][1]["id"] = assistant["tool_calls"][0]["id"] + elif mutation == "text": + assistant["content"] = "replacement" + elif mutation == "name": + assistant["tool_calls"][0]["function"]["name"] = "replacement" + else: + assistant["tool_calls"][0]["function"]["arguments"] = '{"code":"changed"}' prepared = prepare_chat_messages(rendered, first.llm_state["scope"]) replayed = next(message for message in prepared if message.get("tool_calls")) assert "provider_specific_fields" not in replayed assert all("provider_specific_fields" not in call for call in replayed["tool_calls"]) - assert replayed["content"] == "Inspect the value." + assert replayed["content"].startswith("Inspect the value.") + if mutation == "text": + assert replayed["content"].endswith("replacement") assert GEMINI_SIGNATURE not in json.dumps(prepared) assert GEMINI_SIGNATURE_2 not in json.dumps(prepared) - assert "public tool-call sequence changed" in caplog.text + assert "public assistant carrier changed" in caplog.text finally: client.close() +def test_public_thinking_content_blocks_are_stripped(caplog: pytest.LogCaptureFixture) -> None: + messages = [ + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "private", "signature": "secret-a"}, + {"type": "text", "text": "public"}, + {"type": "redacted_thinking", "data": "secret-b"}, + ], + } + ] + + prepared = prepare_chat_messages(messages, None) + assert prepared == [{"role": "assistant", "content": [{"type": "text", "text": "public"}]}] + assert "secret-a" not in repr(prepared) + assert "secret-b" not in repr(prepared) + assert "Removed untrusted provider reasoning fields" in caplog.text + + +def test_public_inline_signature_is_stripped_when_private_field_confirms_it() -> None: + raw_id = f"call_1__thought__{GEMINI_SIGNATURE}" + messages = [ + { + "role": "assistant", + "content": None, + "tool_calls": [_tool_call(raw_id, {"thought_signature": GEMINI_SIGNATURE})], + } + ] + + prepared = prepare_chat_messages(messages, None) + assert prepared[0]["tool_calls"][0]["id"] == "call_1" + assert "provider_specific_fields" not in prepared[0]["tool_calls"][0] + + +def test_direct_gemini_inline_signatures_cannot_bypass_the_envelope() -> None: + raw_id = f"call_1__thought__{GEMINI_SIGNATURE}" + messages = [ + { + "role": "assistant", + "content": None, + "tool_calls": [_tool_call(raw_id)], + }, + {"role": "tool", "tool_call_id": raw_id, "content": "complete"}, + ] + + prepared = prepare_chat_messages( + messages, replay_scope("gemini/gemini-2.5-pro", "chat", {}) + ) + assert prepared[0]["tool_calls"][0]["id"] == "call_1" + assert prepared[1]["tool_call_id"] == "call_1" + assert GEMINI_SIGNATURE not in json.dumps(prepared) + + def test_non_gemini_tool_call_id_with_thought_substring_is_unchanged() -> None: call_id = "call_business__thought__phase" response = _chat_response( @@ -254,6 +386,7 @@ def test_non_gemini_tool_call_id_with_thought_substring_is_unchanged() -> None: with patch("litellm.completion", return_value=response): first = client.call([{"role": "user", "content": "run"}], tools=[TOOL]) assert first.tool_calls[0].id == call_id + assert first.llm_state is None rendered = _render(first) assistant = next(message for message in rendered if message.get("tool_calls")) @@ -479,7 +612,11 @@ def test_state_only_carrier_mutation_warns_and_preserves_public_content( "version": 1, "scope": chat_scope, "format": "litellm-chat", - "payload": {"reasoning_items": [{"opaque": "chat-secret"}], "state_only": True}, + "payload": { + "reasoning_items": [{"opaque": "chat-secret"}], + "carrier": {"content": "", "tool_calls": []}, + "state_only": True, + }, } chat_carrier = ReplayCarryingMessage( {"role": "assistant", "content": "middleware text"}, @@ -507,7 +644,8 @@ def test_state_only_carrier_mutation_warns_and_preserves_public_content( responses_scope, "portable reasoning", ) == [{"role": "assistant", "content": "portable reasoning\n\nmiddleware text"}] - assert caplog.text.count("empty public carrier changed") == 2 + assert "public assistant carrier changed" in caplog.text + assert "empty public carrier changed" in caplog.text def test_legacy_state_warns_while_malformed_current_state_fails( diff --git a/tests/unifiedllm/test_litellm_responses_bridge.py b/tests/unifiedllm/test_litellm_responses_bridge.py index 8c1cf44d2..3437f03f2 100644 --- a/tests/unifiedllm/test_litellm_responses_bridge.py +++ b/tests/unifiedllm/test_litellm_responses_bridge.py @@ -99,7 +99,16 @@ async def test_default_reasoning_tool_call_uses_responses_bridge(model: str) -> assert result.llm_state["format"] == "litellm-chat" assert result.llm_state["payload"] == { "reasoning_items": [REASONING_ITEM], - "tool_call_ids": ["fc_123"], + "carrier": { + "content": None, + "tool_calls": [ + { + "id": "fc_123", + "name": "execute_python", + "arguments": '{"code":"print(1)"}', + } + ], + }, } finally: await client.aclose() diff --git a/tests/unifiedllm/test_reasoning_state_replay.py b/tests/unifiedllm/test_reasoning_state_replay.py index cef4f79b6..95768c965 100644 --- a/tests/unifiedllm/test_reasoning_state_replay.py +++ b/tests/unifiedllm/test_reasoning_state_replay.py @@ -379,7 +379,7 @@ def test_responses_envelope_cannot_replay_a_non_reasoning_item_as_state() -> Non }, } - with pytest.raises(ReasoningReplayError, match="Malformed Responses reasoning item"): + with pytest.raises(ReasoningReplayError, match="Malformed OpenAI Responses"): prepare_responses_batch([CALL], state, scope) with pytest.raises(ReasoningReplayError, match="unsupported output type.*future_public"): From 4d34239039acc22cc0deb0ebaf81038d91ddf219 Mon Sep 17 00:00:00 2001 From: Paul Furgale Date: Thu, 10 Sep 2026 18:24:12 +0000 Subject: [PATCH 11/13] fix(llm): keep replay results consistent and simplify redaction Signed-off-by: Paul Furgale --- src/nooa/strategies/codeact.py | 1 + src/nooa/unifiedllm/http_logging.py | 28 +----------------- src/nooa/unifiedllm/replay_state.py | 22 ++++++++++---- tests/strategies/test_error_recovery_gl106.py | 3 +- .../test_cross_provider_reasoning.py | 29 +++++++++++++------ tests/unifiedllm/test_http_logging.py | 13 +++++++-- 6 files changed, 51 insertions(+), 45 deletions(-) diff --git a/src/nooa/strategies/codeact.py b/src/nooa/strategies/codeact.py index afe20c6e9..02516094f 100644 --- a/src/nooa/strategies/codeact.py +++ b/src/nooa/strategies/codeact.py @@ -922,6 +922,7 @@ async def _run_generation( **self._build_sampling_kwargs(), ) except ReasoningReplayError: + turn_state.is_final = True raise except BlockSyntaxError as e: self._handle_block_syntax_error(e, session, runtime) diff --git a/src/nooa/unifiedllm/http_logging.py b/src/nooa/unifiedllm/http_logging.py index fa27ab430..3b309d9c2 100644 --- a/src/nooa/unifiedllm/http_logging.py +++ b/src/nooa/unifiedllm/http_logging.py @@ -5,34 +5,8 @@ from datetime import UTC, datetime from pathlib import Path -from nooa._llm_state import LLM_STATE_KEY from nooa.tracing._secret_scrubber import REDACTED, _is_sensitive_key, scrub_value -_OPAQUE_STATE_KEYS = { - LLM_STATE_KEY, - "encrypted_content", - "signature", - "thought_signature", - "thought_signatures", - "thoughtSignature", - "thoughtSignatures", -} - - -def _redact_opaque_state(value): - """Remove provider-owned reasoning state from debug logs.""" - if isinstance(value, dict): - redacted_thinking = value.get("type") == "redacted_thinking" - return { - key: REDACTED - if key in _OPAQUE_STATE_KEYS or (redacted_thinking and key == "data") - else _redact_opaque_state(item) - for key, item in value.items() - } - if isinstance(value, list): - return [_redact_opaque_state(item) for item in value] - return value - def enable_http_request_logging( output_dir: str | Path = ".", @@ -133,7 +107,7 @@ def _redact_body(body): # in an OAuth authorization-code exchange. if "code" in scrubbed: scrubbed["code"] = REDACTED - return _redact_opaque_state(scrubbed) + return scrubbed def _write_jsonl_entry(entry: dict): """Append a JSON entry to the JSONL error file.""" diff --git a/src/nooa/unifiedllm/replay_state.py b/src/nooa/unifiedllm/replay_state.py index aebc50b99..477288713 100644 --- a/src/nooa/unifiedllm/replay_state.py +++ b/src/nooa/unifiedllm/replay_state.py @@ -530,7 +530,8 @@ def _strip_inline_signature(value: Any) -> Any: def _strip_chat_state( message: dict[str, Any], source_scope: str | None, target_scope: str | None -) -> None: +) -> dict[str, str]: + public_call_ids: dict[str, str] = {} removed = False gemini_wire = "gemini" in { _scope_provider(source_scope), @@ -568,6 +569,8 @@ def _strip_chat_state( if gemini_wire or (inline_candidate and inline_candidate == explicit_signature): public_id = _strip_inline_signature(call_id) removed = public_id != call_id or removed + if isinstance(call_id, str) and public_id != call_id: + public_call_ids[call_id] = public_id call["id"] = public_id removed = "provider_specific_fields" in call or removed call.pop("provider_specific_fields", None) @@ -584,6 +587,7 @@ def _strip_chat_state( "Removed untrusted provider reasoning fields from a public chat message; " "replay opaque state through a persisted LLMResponse instead." ) + return public_call_ids def _restore_chat_state(message: dict[str, Any], payload: dict[str, Any]) -> bool: @@ -693,6 +697,7 @@ def responses_reasoning_text(output: list[Any]) -> str | None: 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]] = [] + public_call_ids: dict[str, str] = {} private_call_ids: dict[str, str] = {} for original in messages: state = carried_state(original) @@ -704,10 +709,12 @@ def prepare_chat_messages(messages: list[dict[str, Any]], scope: str | None) -> if isinstance(state, dict) and state.get("version") == _STATE_VERSION else None ) - _strip_chat_state( - message, - source_scope if isinstance(source_scope, str) else None, - scope, + public_call_ids.update( + _strip_chat_state( + message, + source_scope if isinstance(source_scope, str) else None, + scope, + ) ) payload = _matching_payload(state, scope, _CHAT_FORMAT) if payload is not None: @@ -722,6 +729,11 @@ def prepare_chat_messages(messages: list[dict[str, Any]], scope: str | None) -> if public_id != private_id: private_call_ids[public_id] = private_id tool_call_id = message.get("tool_call_id") + # Structural signature evidence on a raw assistant call also applies to + # its matching result, even when neither route identifies Gemini. + if isinstance(tool_call_id, str) and tool_call_id in public_call_ids: + tool_call_id = public_call_ids.pop(tool_call_id) + message["tool_call_id"] = tool_call_id if isinstance(tool_call_id, str) and tool_call_id in private_call_ids: message["tool_call_id"] = private_call_ids.pop(tool_call_id) if not restored: diff --git a/tests/strategies/test_error_recovery_gl106.py b/tests/strategies/test_error_recovery_gl106.py index 23d705e1c..bcd1aac6f 100644 --- a/tests/strategies/test_error_recovery_gl106.py +++ b/tests/strategies/test_error_recovery_gl106.py @@ -135,7 +135,8 @@ async def compute(self, value: int) -> int: with pytest.raises(ReasoningReplayError, match="malformed retained reasoning state"): await agent.compute(1) assert failing_llm.attempts == 1 - if strategy_name == "pure_python": + if strategy_name in {"codeact", "pure_python"}: + assert len(after_turns) == 1 assert after_turns[-1].is_final is True assert after_turns[-1].success is False assert after_turns[-1].exception_type == "ReasoningReplayError" diff --git a/tests/unifiedllm/test_cross_provider_reasoning.py b/tests/unifiedllm/test_cross_provider_reasoning.py index 9357ffa90..b12842517 100644 --- a/tests/unifiedllm/test_cross_provider_reasoning.py +++ b/tests/unifiedllm/test_cross_provider_reasoning.py @@ -269,9 +269,7 @@ def test_incompatible_gemini_state_keeps_tool_result_ids_public() -> None: with patch("litellm.completion", return_value=_gemini_response()): first = source.call([{"role": "user", "content": "run"}], tools=[TOOL]) - prepared = prepare_chat_messages( - _render(first), replay_scope("openai/gpt-5.6", "chat", {}) - ) + prepared = prepare_chat_messages(_render(first), replay_scope("openai/gpt-5.6", "chat", {})) assistant = next(message for message in prepared if message.get("tool_calls")) tool_results = [message for message in prepared if message.get("role") == "tool"] assert [call["id"] for call in assistant["tool_calls"]] == ["call_1", "call_2"] @@ -342,19 +340,28 @@ def test_public_thinking_content_blocks_are_stripped(caplog: pytest.LogCaptureFi assert "Removed untrusted provider reasoning fields" in caplog.text -def test_public_inline_signature_is_stripped_when_private_field_confirms_it() -> None: +@pytest.mark.parametrize("target_model", [None, "openai/gpt-4o"]) +def test_public_inline_signature_is_stripped_when_private_field_confirms_it( + target_model: str | None, +) -> None: raw_id = f"call_1__thought__{GEMINI_SIGNATURE}" messages = [ { "role": "assistant", "content": None, "tool_calls": [_tool_call(raw_id, {"thought_signature": GEMINI_SIGNATURE})], - } + }, + {"role": "tool", "tool_call_id": raw_id, "content": "complete"}, ] - prepared = prepare_chat_messages(messages, None) + scope = replay_scope(target_model, "chat", {}) if target_model else None + prepared = prepare_chat_messages(messages, scope) assert prepared[0]["tool_calls"][0]["id"] == "call_1" assert "provider_specific_fields" not in prepared[0]["tool_calls"][0] + assert prepared[1]["tool_call_id"] == "call_1" + assert GEMINI_SIGNATURE not in json.dumps(prepared) + assert messages[0]["tool_calls"][0]["id"] == raw_id + assert messages[1]["tool_call_id"] == raw_id def test_direct_gemini_inline_signatures_cannot_bypass_the_envelope() -> None: @@ -368,9 +375,7 @@ def test_direct_gemini_inline_signatures_cannot_bypass_the_envelope() -> None: {"role": "tool", "tool_call_id": raw_id, "content": "complete"}, ] - prepared = prepare_chat_messages( - messages, replay_scope("gemini/gemini-2.5-pro", "chat", {}) - ) + prepared = prepare_chat_messages(messages, replay_scope("gemini/gemini-2.5-pro", "chat", {})) assert prepared[0]["tool_calls"][0]["id"] == "call_1" assert prepared[1]["tool_call_id"] == "call_1" assert GEMINI_SIGNATURE not in json.dumps(prepared) @@ -394,6 +399,12 @@ def test_non_gemini_tool_call_id_with_thought_substring_is_unchanged() -> None: prepared = prepare_chat_messages(rendered, None) assistant = next(message for message in prepared if message.get("tool_calls")) assert assistant["tool_calls"][0]["id"] == call_id + paired = prepare_chat_messages( + [assistant, {"role": "tool", "tool_call_id": call_id, "content": "complete"}], + None, + ) + assert paired[0]["tool_calls"][0]["id"] == call_id + assert paired[1]["tool_call_id"] == call_id finally: client.close() diff --git a/tests/unifiedllm/test_http_logging.py b/tests/unifiedllm/test_http_logging.py index 26e274f29..73be8be14 100644 --- a/tests/unifiedllm/test_http_logging.py +++ b/tests/unifiedllm/test_http_logging.py @@ -9,10 +9,12 @@ import httpx import pytest -from nooa.unifiedllm.http_logging import _redact_opaque_state, enable_http_request_logging +from nooa.unifiedllm.http_logging import enable_http_request_logging -def test_opaque_reasoning_state_is_redacted_from_http_debug_payloads() -> None: +def test_opaque_reasoning_state_is_redacted_from_http_debug_payloads( + tmp_path, secret_header_server +) -> None: payload = { "input": [ {"encrypted_content": "provider-secret"}, @@ -23,7 +25,12 @@ def test_opaque_reasoning_state_is_redacted_from_http_debug_payloads() -> None: ] } - redacted = _redact_opaque_state(payload) + 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]"}, From ff9b2c0a4fd344a61be556b002fcaf7e2ddb48df Mon Sep 17 00:00:00 2001 From: Paul Furgale Date: Thu, 10 Sep 2026 18:57:58 +0000 Subject: [PATCH 12/13] fix(llm): distinguish portable summaries from opaque state Signed-off-by: Paul Furgale --- src/nooa/unifiedllm/replay_state.py | 13 +++++- .../test_cross_provider_reasoning.py | 40 ++++++++++++++++++- 2 files changed, 49 insertions(+), 4 deletions(-) diff --git a/src/nooa/unifiedllm/replay_state.py b/src/nooa/unifiedllm/replay_state.py index 477288713..17b144060 100644 --- a/src/nooa/unifiedllm/replay_state.py +++ b/src/nooa/unifiedllm/replay_state.py @@ -354,7 +354,7 @@ def _tool_call_state(tool_call: Any, scope: str | None) -> dict[str, Any] | None inline_candidate = None if isinstance(call_id, str) and _INLINE_THOUGHT_SIGNATURE_SEPARATOR in call_id: inline_candidate = call_id.split(_INLINE_THOUGHT_SIGNATURE_SEPARATOR, 1)[1] - if not inline_candidate: + if not inline_candidate and (_scope_provider(scope) == "gemini" or signature): raise ReasoningReplayError("Malformed inline tool-call thought signature.") if signature and inline_candidate and signature != inline_candidate: raise ReasoningReplayError("Conflicting thought signatures on one provider tool call.") @@ -622,7 +622,11 @@ def capture_responses_state(output: list[Any], scope: str | None) -> dict | None # may expose a similarly shaped API through a gateway, but that is not # evidence that their opaque state is wire-compatible. if _scope_provider(scope) not in {"openai", "azure"}: - if any(response_item_type(item) == "reasoning" for item in output): + if any( + response_item_type(item) == "reasoning" + and _field(item, "encrypted_content") is not None + for item in output + ): raise ReasoningReplayError( "The provider returned Responses reasoning state, but NOOA only supports " "opaque Responses replay for OpenAI and Azure routes." @@ -637,6 +641,11 @@ def capture_responses_state(output: list[Any], scope: str | None) -> dict | None for item in output: item_type = response_item_type(item) if item_type == "reasoning": + # Summaries are valid without encrypted content (for example when + # a gateway does not support include). They remain portable text, + # not malformed opaque state. + if _field(item, "encrypted_content") is None: + continue order.append({"type": "reasoning", "index": len(items)}) items.append(opaque_item(item)) elif item_type == "function_call": diff --git a/tests/unifiedllm/test_cross_provider_reasoning.py b/tests/unifiedllm/test_cross_provider_reasoning.py index b12842517..a2d6887e2 100644 --- a/tests/unifiedllm/test_cross_provider_reasoning.py +++ b/tests/unifiedllm/test_cross_provider_reasoning.py @@ -8,6 +8,7 @@ from unittest.mock import AsyncMock, patch import pytest +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import Choices, Message, ModelResponse from nooa._llm_state import ReplayCarryingMessage @@ -381,8 +382,8 @@ def test_direct_gemini_inline_signatures_cannot_bypass_the_envelope() -> None: assert GEMINI_SIGNATURE not in json.dumps(prepared) -def test_non_gemini_tool_call_id_with_thought_substring_is_unchanged() -> None: - call_id = "call_business__thought__phase" +@pytest.mark.parametrize("call_id", ["call_business__thought__phase", "call_business__thought__"]) +def test_non_gemini_tool_call_id_with_thought_substring_is_unchanged(call_id: str) -> None: response = _chat_response( Message(role="assistant", content=None, tool_calls=[_tool_call(call_id)]) ) @@ -536,6 +537,41 @@ def _responses_output(*items: dict) -> SimpleNamespace: } +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("has_answer", [False, True]) +async def test_summary_without_encrypted_content_is_portable_text(is_async, has_answer) -> None: + summary = { + key: value for key, value in RESPONSES_REASONING.items() if key != "encrypted_content" + } + raw = ResponsesAPIResponse( + id="resp", + created_at=0, + model="gpt-5.6", + status="completed", + output=[summary, RESPONSES_MESSAGE] if has_answer else [summary], + ) + async with ResponsesClient( + model="openai/gpt-5.6", api_key="test", api_base="https://gateway.example/v1" + ) as client: + target = "litellm.aresponses" if is_async else "litellm.responses" + with patch(target, return_value=raw) as request: + messages = [{"role": "user", "content": "think"}] + first = await client.acall(messages) if is_async else client.call(messages) + assert request.call_args.kwargs.get("include") is None + assert first.reasoning == "Check the evidence." + assert first.llm_state is None + restored = LLMResponse.model_validate_json(first.model_dump_json()) + with patch(target, return_value=raw) as replay: + rendered = _render(restored, responses=True) + if is_async: + await client.acall(rendered) + else: + client.call(rendered) + expected = "Check the evidence." + ("\n\nAnswer." if has_answer else "") + assert replay.call_args.kwargs["input"] == [{"role": "assistant", "content": expected}] + + def test_responses_summary_stays_exact_on_match_and_demotes_on_model_change() -> None: source = ResponsesClient(model="openai/gpt-5.6", api_key="account-a") target = ResponsesClient(model="openai/gpt-5.7", api_key="account-a") From b7945795ffdddff9d069de002a1823f4917f2f8a Mon Sep 17 00:00:00 2001 From: Paul Furgale Date: Fri, 11 Sep 2026 06:30:21 +0000 Subject: [PATCH 13/13] fix(tracing): preserve readable reasoning in OTLP exports Signed-off-by: Paul Furgale --- src/nooa/tracing/_litellm_patch.py | 20 ++++++ tests/tracing/test_reasoning_otlp.py | 94 ++++++++++++++++++++++++++++ 2 files changed, 114 insertions(+) create mode 100644 tests/tracing/test_reasoning_otlp.py diff --git a/src/nooa/tracing/_litellm_patch.py b/src/nooa/tracing/_litellm_patch.py index 07487a7b6..4d0d92f56 100644 --- a/src/nooa/tracing/_litellm_patch.py +++ b/src/nooa/tracing/_litellm_patch.py @@ -9,6 +9,7 @@ defined in OpenInference semantic conventions). 3. Missing reasoning_content capture for reasoning models (DeepSeek, o1, Nemotron, etc.) 4. Extract tags from content for models that embed reasoning (Nemotron, QwQ) +5. Preserve Responses reasoning summaries when journal mode strips message attributes Bug: https://github.com/Arize-ai/openinference/issues (to be filed) Affected version: openinference-instrumentation-litellm v0.1.28+ @@ -251,6 +252,22 @@ def _patched_get_attributes_from_message_param( yield (MessageAttributes.MESSAGE_TOOL_CALL_ID, tool_call_id) +def _patched_get_attributes_from_response_output(result: Any) -> dict[str, Any]: + """Keep readable Responses summaries on the same OTLP field as Chat reasoning.""" + from openinference.instrumentation.litellm._responses_attributes import ( + _get_attributes_from_response_output, + ) + + from nooa.unifiedllm.replay_state import responses_reasoning_text + + attributes = _get_attributes_from_response_output(result) + if reasoning := responses_reasoning_text(result.output): + # Message attributes may be stripped in journal mode. This field remains + # on the span and contains only visible text, never encrypted state. + attributes["llm.reasoning_content"] = reasoning + return attributes + + def apply_litellm_patch() -> None: """Apply monkey patches to fix litellm instrumentation bugs. @@ -258,6 +275,7 @@ def apply_litellm_patch() -> None: 1. _set_output_message_value - fixes null/empty content handling, adds reasoning_content, and stamps llm.cost.* (not emitted by the instrumentor) 2. _get_attributes_from_message_param - adds missing tool_call.id capture + 3. _get_attributes_from_response_output - keeps Responses summaries on OTLP It also enables ``litellm.return_response_headers`` so the gateway's ``x-litellm-response-cost`` headers are retained on the response for cost @@ -278,6 +296,8 @@ def apply_litellm_patch() -> None: # Patch 1: Fix null/empty content handling litellm._set_output_message_value = _patched_set_output_message_value + # This imported private helper is the Responses instrumentor's patch point. + litellm._get_attributes_from_response_output = _patched_get_attributes_from_response_output # pyright: ignore[reportPrivateImportUsage] # Patch 2: Fix missing tool_call.id # Store the original function so our patch can call it diff --git a/tests/tracing/test_reasoning_otlp.py b/tests/tracing/test_reasoning_otlp.py new file mode 100644 index 000000000..75a26ca41 --- /dev/null +++ b/tests/tracing/test_reasoning_otlp.py @@ -0,0 +1,94 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Readable reasoning survives scrubbing and the journal-mode OTLP strip.""" + +import json +from contextlib import nullcontext +from types import SimpleNamespace + +import pytest +from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.utils import ModelResponse +from openinference.instrumentation import litellm as instrumentation +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor + +from nooa.tracing._litellm_patch import apply_litellm_patch +from nooa.tracing._otlp_http_exporter import OtlpJsonHttpExporter +from nooa.tracing._secret_scrubber import SecretScrubSpanProcessor + +REASONING = "Check the evidence before answering." +OPAQUE = "opaque-provider-state-must-not-leave" + + +@pytest.mark.parametrize("strip_messages", [False, True]) +@pytest.mark.parametrize("shape", ["reasoning_content", "reasoning", "think_tags", "responses"]) +def test_plain_reasoning_reaches_serialized_otlp_without_opaque_state( + shape, strip_messages, monkeypatch +): + monkeypatch.delenv("NEMO_TRACE_KEEP_LLM_VALUES", raising=False) + captured = [] + + def send(request, timeout): + captured.append(json.loads(request.data)) + return nullcontext(SimpleNamespace(status=200)) + + monkeypatch.setattr("urllib.request.urlopen", send) + exporter = OtlpJsonHttpExporter(strip_llm_messages=strip_messages) + provider = TracerProvider() + provider.add_span_processor(SecretScrubSpanProcessor(SimpleSpanProcessor(exporter))) + apply_litellm_patch() + + if shape == "responses": + response = ResponsesAPIResponse.model_validate( + { + "id": "resp_test", + "created_at": 1, + "model": "gpt-5.6", + "object": "response", + "status": "completed", + "output": [ + { + "type": "reasoning", + "id": "rs_test", + "encrypted_content": OPAQUE, + "summary": [{"type": "summary_text", "text": REASONING}], + } + ], + } + ) + else: + message = {"role": "assistant", "content": "answer"} + if shape == "think_tags": + message["content"] = f"{REASONING}answer" + else: + message[shape] = REASONING + response = ModelResponse( + model="test", + choices=[{"index": 0, "finish_reason": "stop", "message": message}], + ) + + try: + with provider.get_tracer("reasoning-test").start_as_current_span("llm") as span: + span.set_attribute("openinference.span.kind", "LLM") + # Exercise the same redaction boundary for all three opaque forms. + span.set_attribute( + "input.value", + json.dumps( + [ + {"encrypted_content": OPAQUE}, + {"type": "thinking", "thinking": REASONING, "signature": OPAQUE}, + {"thought_signature": OPAQUE}, + ] + ), + ) + instrumentation._finalize_span(span, response) + finally: + provider.shutdown() + + assert len(captured) == 1 + wire = json.dumps(captured) + assert OPAQUE not in wire + exported_span = captured[0]["resourceSpans"][0]["scopeSpans"][0]["spans"][0] + attributes = {item["key"]: item["value"] for item in exported_span["attributes"]} + assert attributes["llm.reasoning_content"]["stringValue"] == REASONING