fix(openai): address Codex review round 2 (PR #67)

Three P2 issues from the post-fix review pass:

1. **input_tokens double-counted on cache hit.** ``input_tokens`` is now
   the uncached portion (``prompt_tokens - cached_tokens``), matching
   the Responses/faux providers' accounting. Without this, cubebox cost
   aggregation counted the cached prefix twice (full-price input AND
   cache_read).

2. **on_payload couldn't disable stream_options.include_usage.** Some
   OpenAI-compatible backends reject ``stream_options`` entirely.
   Switch from unconditional override to ``setdefault``-style: only set
   ``include_usage=True`` when the caller hasn't already configured it.

3. **$ref resolution lost sibling metadata.** Pydantic emits
   ``$ref`` alongside ``description`` / ``default`` on
   Optional[Enum] fields (JSON Schema 2020-12 allows sibling keys).
   The previous code returned only the resolved $def, dropping the
   field-level description and default. Merge siblings onto the
   resolved schema; the $def's own values win on key collision.

Tests: 6 new unit tests in test_openai_extras_and_schema.py — two per
fix. openai.py coverage stays at 100%.

Author: xfgong

cubepi/providers/openai.py

@@ -99,7 +99,14 @@ async def _produce() -> None:
                     kwargs["extra_body"] = {**self._extra_body, **kwargs["extra_body"]}
 
                 # Request per-stream usage so we can populate AssistantMessage.usage.
-                kwargs.setdefault("stream_options", {})["include_usage"] = True
+                # Only set ``include_usage`` if the caller hasn't already
+                # configured it via on_payload — some OpenAI-compatible
+                # backends reject ``stream_options`` entirely, so callers
+                # need to be able to opt out by setting it to False (or
+                # removing the key).
+                so = kwargs.setdefault("stream_options", {})
+                if "include_usage" not in so:
+                    so["include_usage"] = True
 
                 response = await self._client.chat.completions.create(**kwargs)
 
@@ -139,15 +146,23 @@ async def _produce() -> None:
                     # trailing chunk with no choices and usage populated).
                     if getattr(chunk, "usage", None) is not None:
                         u = chunk.usage
-                        partial.usage = Usage(
-                            input_tokens=getattr(u, "prompt_tokens", 0) or 0,
-                            output_tokens=getattr(u, "completion_tokens", 0) or 0,
-                            cache_read_tokens=getattr(
+                        prompt_tokens = getattr(u, "prompt_tokens", 0) or 0
+                        cached_tokens = (
+                            getattr(
                                 getattr(u, "prompt_tokens_details", None),
                                 "cached_tokens",
                                 0,
                             )
-                            or 0,
+                            or 0
+                        )
+                        # ``input_tokens`` is the uncached prompt portion; the
+                        # cached prefix is reported separately. This matches
+                        # the Responses/faux providers' accounting so cost
+                        # aggregation across providers doesn't double-count.
+                        partial.usage = Usage(
+                            input_tokens=max(prompt_tokens - cached_tokens, 0),
+                            output_tokens=getattr(u, "completion_tokens", 0) or 0,
+                            cache_read_tokens=cached_tokens,
                         )
 
                     if response_id is None and getattr(chunk, "id", None):
@@ -468,13 +483,38 @@ def _normalise_tool_schema(
             if "$ref" in schema:
                 ref_name = schema["$ref"].split("/")[-1]
                 if defs and ref_name in defs:
-                    return N(
+                    resolved = N(
                         defs[ref_name],
                         defs=defs,
                         top=False,
                         strip_title=not in_any_of,
                         in_any_of=False,
                     )
+                    # JSON Schema 2020-12 allows sibling keys alongside
+                    # $ref; Pydantic emits ``description`` / ``default`` /
+                    # validators on Optional[Enum] fields, etc. Merge them
+                    # onto the resolved definition so field-level metadata
+                    # isn't silently dropped when we inline.
+                    siblings = {
+                        k: v
+                        for k, v in schema.items()
+                        if k != "$ref" and not (k == "title" and not in_any_of)
+                    }
+                    if siblings and isinstance(resolved, dict):
+                        merged: dict[str, Any] = dict(resolved)
+                        for k, v in siblings.items():
+                            merged.setdefault(
+                                k,
+                                N(
+                                    v,
+                                    defs=defs,
+                                    top=False,
+                                    strip_title=not in_any_of,
+                                    in_any_of=False,
+                                ),
+                            )
+                        return merged
+                    return resolved
                 # Unknown ref — leave as-is.
                 return schema
 

tests/providers/test_openai_extras_and_schema.py

@@ -259,3 +259,202 @@ def test_normalise_works_on_pydantic_generated_schema() -> None:
     # Property titles stripped.
     for prop in out["properties"].values():
         assert "title" not in prop
+
+
+# ---------------------------------------------------------------------------
+# Codex review #5 — input_tokens excludes cached_tokens (cache-hit accounting)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_usage_subtracts_cached_tokens_from_input() -> None:
+    """``input_tokens`` is the uncached portion; cached prefix reported
+    separately, matching Responses/faux accounting (avoids double-count
+    in cubebox cost aggregation)."""
+    usage_chunk = SimpleNamespace(
+        id="x",
+        choices=[],
+        usage=SimpleNamespace(
+            prompt_tokens=1000,
+            completion_tokens=50,
+            prompt_tokens_details=SimpleNamespace(cached_tokens=800),
+        ),
+    )
+    final = _make_chunk(finish_reason="stop")
+
+    with patch("openai.AsyncOpenAI") as mock_openai:
+        client = MagicMock()
+        mock_openai.return_value = client
+        client.chat = MagicMock()
+        client.chat.completions = MagicMock()
+
+        async def create(**_):
+            return _async_iter([usage_chunk, final])
+
+        client.chat.completions.create = create
+
+        provider = OpenAIProvider(api_key="x")
+        provider._client = client
+
+        ms = await provider.stream(
+            _model(),
+            [UserMessage(content=[TextContent(text="hi")])],
+        )
+        async for _ in ms:
+            pass
+        msg = await ms.result()
+
+    assert msg.usage.input_tokens == 200  # 1000 - 800
+    assert msg.usage.cache_read_tokens == 800
+    assert msg.usage.output_tokens == 50
+
+
+@pytest.mark.asyncio
+async def test_usage_no_cached_field_input_tokens_is_full_prompt() -> None:
+    """No prompt_tokens_details → cached_tokens=0, input_tokens=prompt_tokens."""
+    usage_chunk = SimpleNamespace(
+        id="x",
+        choices=[],
+        usage=SimpleNamespace(
+            prompt_tokens=200,
+            completion_tokens=10,
+            prompt_tokens_details=None,
+        ),
+    )
+    final = _make_chunk(finish_reason="stop")
+
+    with patch("openai.AsyncOpenAI") as mock_openai:
+        client = MagicMock()
+        mock_openai.return_value = client
+        client.chat = MagicMock()
+        client.chat.completions = MagicMock()
+
+        async def create(**_):
+            return _async_iter([usage_chunk, final])
+
+        client.chat.completions.create = create
+
+        provider = OpenAIProvider(api_key="x")
+        provider._client = client
+
+        ms = await provider.stream(
+            _model(),
+            [UserMessage(content=[TextContent(text="hi")])],
+        )
+        async for _ in ms:
+            pass
+        msg = await ms.result()
+
+    assert msg.usage.input_tokens == 200
+    assert msg.usage.cache_read_tokens == 0
+
+
+# ---------------------------------------------------------------------------
+# Codex review #6 — on_payload may opt out of stream_options.include_usage
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_on_payload_can_disable_include_usage() -> None:
+    """Caller-supplied ``stream_options.include_usage=False`` must survive
+    the default-inject step."""
+    captured: dict[str, Any] = {}
+
+    async def on_payload(payload, model):
+        payload["stream_options"] = {"include_usage": False}
+        return payload
+
+    with patch("openai.AsyncOpenAI") as mock_openai:
+        client = MagicMock()
+        mock_openai.return_value = client
+        client.chat = MagicMock()
+        client.chat.completions = MagicMock()
+
+        async def capture(**kwargs):
+            captured.update(kwargs)
+            return _async_iter([_make_chunk(finish_reason="stop")])
+
+        client.chat.completions.create = capture
+
+        provider = OpenAIProvider(api_key="x")
+        provider._client = client
+
+        ms = await provider.stream(
+            _model(),
+            [UserMessage(content=[TextContent(text="hi")])],
+            options=StreamOptions(on_payload=on_payload),
+        )
+        async for _ in ms:
+            pass
+        await ms.result()
+
+    assert captured["stream_options"] == {"include_usage": False}
+
+
+@pytest.mark.asyncio
+async def test_default_adds_include_usage_when_not_set() -> None:
+    """Without on_payload customization, default adds include_usage=True."""
+    captured: dict[str, Any] = {}
+    with patch("openai.AsyncOpenAI") as mock_openai:
+        client = MagicMock()
+        mock_openai.return_value = client
+        client.chat = MagicMock()
+        client.chat.completions = MagicMock()
+
+        async def capture(**kwargs):
+            captured.update(kwargs)
+            return _async_iter([_make_chunk(finish_reason="stop")])
+
+        client.chat.completions.create = capture
+
+        provider = OpenAIProvider(api_key="x")
+        provider._client = client
+
+        ms = await provider.stream(
+            _model(),
+            [UserMessage(content=[TextContent(text="hi")])],
+        )
+        async for _ in ms:
+            pass
+        await ms.result()
+
+    assert captured["stream_options"] == {"include_usage": True}
+
+
+# ---------------------------------------------------------------------------
+# Codex review #7 — preserve sibling metadata when $ref-resolving
+# ---------------------------------------------------------------------------
+
+
+def test_normalise_ref_resolution_preserves_sibling_description() -> None:
+    """Pydantic emits ``$ref`` alongside ``description`` / ``default`` on
+    Optional[Enum] fields; both must survive inlining."""
+    schema = {
+        "$defs": {"Color": {"title": "Color", "enum": ["r", "g"], "type": "string"}},
+        "properties": {
+            "c": {
+                "$ref": "#/$defs/Color",
+                "description": "primary colour",
+                "default": "r",
+            }
+        },
+    }
+    out = OpenAIProvider._normalise_tool_schema(schema)
+    resolved = out["properties"]["c"]
+    assert resolved["enum"] == ["r", "g"]
+    assert resolved["description"] == "primary colour"
+    assert resolved["default"] == "r"
+
+
+def test_normalise_ref_resolution_resolved_def_wins_on_key_collision() -> None:
+    """If both $def and sibling supply the same key, the $def value wins
+    (sibling is treated as additional metadata, not an override)."""
+    schema = {
+        "$defs": {"X": {"type": "string", "description": "from def"}},
+        "properties": {
+            "x": {"$ref": "#/$defs/X", "description": "from sibling"},
+        },
+    }
+    out = OpenAIProvider._normalise_tool_schema(schema)
+    # setdefault means the def's value sticks.
+    assert out["properties"]["x"]["description"] == "from def"