Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions backend/druks/harnesses/codex.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import json
import logging
import os
import re
import shlex
from collections.abc import Collection, Iterator
from dataclasses import dataclass
Expand Down Expand Up @@ -43,6 +44,9 @@
_TOKEN_COUNT_MARKERS = ('"type":"token_count"', '"type": "token_count"')
_CHATGPT_HOST = "chatgpt.com"
_SUBSCRIPTION_TOKEN = "codex_subscription_token"
# OpenAI web search leaves this citation in model text. The pattern runs on
# serialized JSON. It stops at a quote, so a match stays in one string.
_CITATION_MARKER = re.compile(r' *\ue200cite(?:\ue202[^\ue201"]*)?\ue201')


@dataclass(frozen=True)
Expand Down Expand Up @@ -566,6 +570,9 @@ def parse(self, result: HarnessRunResult, *, artifact_dir: Path, run_id: str) ->
call_dir / "output.json",
name=self.name,
)
payload = json.loads(
_CITATION_MARKER.sub("", json.dumps(payload, ensure_ascii=False)),
)
(call_dir / "output.json").write_text(
json.dumps(payload, indent=2, sort_keys=True),
)
Expand Down
23 changes: 23 additions & 0 deletions backend/tests/test_cost_capture.py
Original file line number Diff line number Diff line change
Expand Up @@ -138,6 +138,29 @@ def test_failed_codex_call_writes_cost_before_raising(tmp_path: Path):
assert metadata["output_tokens"] == 1_000


def test_codex_parse_strips_citation_markers(tmp_path: Path):
call_dir = tmp_path / "call"
call_dir.mkdir()
citation = "\ue200cite\ue202turn905355view5\ue202turn430332view7\ue201"
output = {
"summary": f"Calls can require fresh execution. {citation}",
"sources": [f"Docs {citation} say so.", "\ue200cite\ue201"],
}
(call_dir / "output.json").write_text(json.dumps(output))
harness = CodexHarness(model="openai/gpt-5.5", fast_mode=False, effort=None)

payload = harness.parse(
HarnessRunResult(returncode=0, stdout=b"", stderr=b""),
artifact_dir=tmp_path,
run_id="call",
)

assert payload == {
"summary": "Calls can require fresh execution.",
"sources": ["Docs say so.", ""],
}


def test_failed_pi_call_writes_cost_before_raising(tmp_path: Path):
message = {
"type": "message_end",
Expand Down
Loading