refactor(agent/transports): merge token-usage/compaction accounting into one handler, contextlib.suppress for pass-except blocks, lift extra_content dump

This commit is contained in:
Teknium
2026-09-02 18:54:23 -07:00
parent 18e73f891e
commit b6f60bf123
4 changed files with 51 additions and 73 deletions

View File

@@ -176,6 +176,19 @@ def _model_consumes_thought_signature(model: Any) -> bool:
return "gemini" in m or "gemma" in m
def _dump_extra_content(extra: Any) -> Any:
"""Plain-dict form of a pydantic ``extra_content``; older pydantic lacks ``warnings=``, so retry without it."""
if hasattr(extra, "model_dump"):
for dump_kwargs in ({"warnings": False}, {}):
try:
return extra.model_dump(**dump_kwargs)
except TypeError:
continue
except Exception:
break
return extra
def _swap_developer_role(sanitized: list, model_lower: str) -> list:
"""GPT-5/Codex models take a ``developer`` role instead of ``system``."""
if (
@@ -461,25 +474,13 @@ class ChatCompletionsTransport(ProviderTransport):
elif function_name in _alias_map:
function_name = _alias_map[function_name]
function_arguments = getattr(tc_function, "arguments", None)
tc_provider_data: dict[str, Any] = {}
extra = getattr(tc, "extra_content", None)
if extra is None and hasattr(tc, "model_extra"):
extra = (tc.model_extra if isinstance(tc.model_extra, dict) else {}).get("extra_content")
if extra is not None:
if hasattr(extra, "model_dump"):
for dump_kwargs in ({"warnings": False}, {}):
try:
extra = extra.model_dump(**dump_kwargs)
break
except TypeError:
continue # older pydantic: retry without ``warnings``
except Exception:
break
tc_provider_data["extra_content"] = extra
tool_calls.append(ToolCall(
id=getattr(tc, "id", None), name=function_name,
arguments=function_arguments if function_arguments is not None else "{}",
provider_data=tc_provider_data or None,
provider_data=None if extra is None else {"extra_content": _dump_extra_content(extra)},
))
usage = Usage.from_openai(response.usage) if hasattr(response, "usage") and response.usage else None

View File

@@ -9,6 +9,7 @@ Opt-in runtime gated behind ``model.openai_runtime == "codex_app_server"``.
from __future__ import annotations
import contextlib
import json
import os
import queue
@@ -117,20 +118,16 @@ class CodexAppServerClient:
if self._closed:
return
self._closed = True
try:
with contextlib.suppress(Exception):
if self._proc.stdin and not self._proc.stdin.closed:
self._proc.stdin.close()
except Exception:
pass
try:
self._proc.terminate()
self._proc.wait(timeout=timeout)
except subprocess.TimeoutExpired:
try:
with contextlib.suppress(Exception):
self._proc.kill()
self._proc.wait(timeout=1.0)
except Exception:
pass
def __enter__(self) -> "CodexAppServerClient":
return self
@@ -234,21 +231,17 @@ class CodexAppServerClient:
with self._pending_lock:
pending = self._pending.pop(msg["id"], None)
if pending is not None:
try:
with contextlib.suppress(queue.Full): # pragma: no cover - defensive
pending.put_nowait(msg)
except queue.Full: # pragma: no cover - defensive
pass
elif "method" in msg: # server-initiated request (has id) or notification
(self._server_requests if "id" in msg else self._notifications).put(msg)
def _read_stderr(self) -> None:
if self._proc.stderr is None:
return
try:
with contextlib.suppress(Exception): # pragma: no cover
for line in iter(self._proc.stderr.readline, b""):
self._append_stderr(line.decode("utf-8", "replace").rstrip())
except Exception: # pragma: no cover
pass
def parse_codex_version(output: str) -> Optional[tuple[int, int, int]]:
@@ -275,7 +268,7 @@ def check_codex_binary(
version = parse_codex_version(proc.stdout)
if version is None:
return False, f"could not parse codex version from: {proc.stdout!r}"
have, need = ".".join(map(str, version)), ".".join(map(str, min_version))
have = ".".join(map(str, version))
if version < min_version:
return False, f"codex {have} is older than required {need}. Run: npm i -g @openai/codex"
return False, f"codex {have} is older than required {'.'.join(map(str, min_version))}. Run: npm i -g @openai/codex"
return True, have

View File

@@ -9,6 +9,7 @@ reader threads feed queues that this adapter polls, like the chat_completions lo
from __future__ import annotations
import contextlib
import logging
import os
import threading
@@ -207,10 +208,8 @@ class CodexAppServerSession:
with self._active_turn_lock:
self._active_turn_id = None
if self._client is not None:
try:
with contextlib.suppress(Exception): # pragma: no cover - best-effort cleanup
self._client.close()
except Exception: # pragma: no cover - best-effort cleanup
pass
self._client = None
self._thread_id = None
@@ -310,8 +309,7 @@ class CodexAppServerSession:
self._on_event(note)
except Exception: # pragma: no cover - display callback
logger.debug("on_event callback raised", exc_info=True)
_apply_token_usage_notification(result, note)
_apply_compaction_notification(result, note)
_apply_accounting_notification(result, note)
self._track_pending_file_change(note)
projection = projector.project(note)
if projection.messages:
@@ -647,34 +645,29 @@ def _summarize_file_changes(raw_changes: list) -> str:
return f"{counts}: {preview}" if preview else counts
def _apply_token_usage_notification(result: TurnResult, note: dict) -> None:
"""Capture token usage (codex emits it as thread/tokenUsage/updated, not on turn/completed)."""
if not isinstance(note, dict) or note.get("method") != "thread/tokenUsage/updated":
return
token_usage = (note.get("params") or {}).get("tokenUsage") or {}
if not isinstance(token_usage, dict):
return
last, window = token_usage.get("last"), token_usage.get("modelContextWindow")
if isinstance(last, dict):
result.token_usage_last = dict(last)
if isinstance(window, int) and window > 0:
result.model_context_window = window
def _apply_compaction_notification(result: TurnResult, note: dict) -> None:
"""Capture compaction boundaries: a contextCompaction item (recent) or deprecated thread/compacted (older)."""
def _apply_accounting_notification(result: TurnResult, note: dict) -> None:
"""Capture token usage (thread/tokenUsage/updated, not turn/completed) and compaction
boundaries (a contextCompaction item on recent builds, deprecated thread/compacted on older)."""
if not isinstance(note, dict):
return
method = note.get("method") or ""
params = note.get("params") or {}
if not isinstance(params, dict):
return
item = params.get("item") if method in {"item/started", "item/completed"} else None
if method != "thread/compacted" and not (isinstance(item, dict) and item.get("type") == "contextCompaction"):
if method == "thread/tokenUsage/updated":
token_usage = params.get("tokenUsage") or {}
if isinstance(token_usage, dict):
last, window = token_usage.get("last"), token_usage.get("modelContextWindow")
if isinstance(last, dict):
result.token_usage_last = dict(last)
if isinstance(window, int) and window > 0:
result.model_context_window = window
return
result.compacted = True
result.thread_id = params.get("threadId") or result.thread_id
result.turn_id = params.get("turnId") or result.turn_id
item = params.get("item") if method in {"item/started", "item/completed"} else None
if method == "thread/compacted" or (isinstance(item, dict) and item.get("type") == "contextCompaction"):
result.compacted = True
result.thread_id = params.get("threadId") or result.thread_id
result.turn_id = params.get("turnId") or result.turn_id
# Hermes approval choice -> codex decision (app-server-protocol v2). "deny" and

View File

@@ -157,36 +157,27 @@ def get_vertex_credentials(credentials_path: Optional[str] = None) -> Tuple[Opti
try:
cached = _creds_cache.get(cache_key)
if cached is None:
loaded = _load_credentials(resolved_path, sa_raw)
if loaded is None:
cached = _load_credentials(resolved_path, sa_raw)
if cached is None:
return None, None
creds, project_id = loaded
_creds_cache[cache_key] = (creds, project_id)
_creds_cache[cache_key] = cached
# A rotation leaves the old signature's entry behind; keep at most
# one Credentials per file so stale identities can't be reused.
for k in [k for k in _creds_cache if k != cache_key and k[0] == cache_key[0]]:
_creds_cache.pop(k, None)
else:
creds, project_id = cached
creds, project_id = cached
if _needs_refresh(creds):
creds.refresh(google.auth.transport.requests.Request())
override_project = _resolve_project_override()
return creds.token, override_project or project_id
return creds.token, _resolve_project_override() or project_id
except Exception as e:
logger.error(f"Failed to resolve Vertex AI credentials: {e}")
_creds_cache.pop(cache_key, None)
# If ADC failed (e.g. expired refresh token), try the SA file before
# giving up — it may have been added after startup. Keyed on the
# resolved path being absent (this attempt was ADC), not on the cache key.
if not resolved_path:
sa_path = _resolve_credentials_path(credentials_path)
if sa_path:
logger.info("ADC failed, retrying with service account: %s", sa_path)
return get_vertex_credentials(sa_path)
# If ADC failed (e.g. expired refresh token), try the SA file before giving
# up — it may have been added after startup. Keyed on this attempt being ADC.
sa_path = None if resolved_path else _resolve_credentials_path(credentials_path)
if sa_path:
logger.info("ADC failed, retrying with service account: %s", sa_path)
return get_vertex_credentials(sa_path)
return None, None