Teknium
2a95791992
run_agent.py: delete the `# noqa: F401` re-export block (agent.process_bootstrap
OpenAI/_SafeWriter/_get_proxy_*, model_tools get_tool_definitions/
handle_function_call/check_toolset_requirements, FailoverReason,
_qwen_portal_headers/_routermint_headers, session_persistence names,
estimate_request_tokens_rough, ContextCompressor + friends, jittered_backoff,
prompt_builder names, message_sanitization names, tool_dispatch_helpers
names) — 41 names run_agent never used itself — and the `_STREAM_DIAG_HEADERS`
back-compat class alias (no in-tree reader). run_agent now imports only what
it uses (get_toolset_for_tool, is_local_endpoint, coalesce/uniquify tool-call
ids, cleanup_vm/get_active_env from terminal_tool_lifecycle).
agent/*: `_ra().X` late-binds that only reached a re-export now import the
defining module directly (agent_runtime_helpers -> process_bootstrap.OpenAI,
model_tools.handle_function_call, session_persistence._safe_session_filename_component;
agent_init -> model_tools.get_tool_definitions/check_toolset_requirements,
_lazy_headers("agent.client_lifecycle", ...) for qwen/routermint;
system_prompt -> agent.prompt_builder / model_tools directly, dropping its
own _ra() shim and the `_r` parameter threading). `_ra()` stays for
run_agent-resident names (logger, AIAgent, _hermes_home, _set_interrupt, ...).
toolsets.py: remove resolve_multiple_toolsets (shim-only, restored by
34abf954bd); tests/test_toolsets.py pins the same union behavior via
resolve_toolset over each name.
providers/__init__.py: drop the OMIT_TEMPERATURE re-export (no callers via the
package); ProviderProfile stays because __init__ uses it for annotations —
2 tests repointed to providers.base.
agent/iteration_budget.py: drop the "run_agent re-exports the class"
docstring pointer; 4 tests import IterationBudget from its home.
model_tools.py (arg_coercion names), agent/tool_executor.py, and
hermes_cli/cli_session_mixin.py repoints landed via a sibling commit on this
shared worktree.
Callers repointed: gateway/run.py, hermes_cli/cli_chat_turn_mixin.py,
hermes_cli/cli_tui_mixin.py, tui_gateway/session_workdir.py,
agent/transports/codex.py (one-line imports) + comment pointers in
tools/file_state.py, tools/schema_sanitizer.py, scripts/tool_search_livetest.py.
Tests: patch("run_agent.X") / monkeypatch.setattr(run_agent, "X") /
`from run_agent import X` -> defining module across 99 test files.
129 lines
5.1 KiB
Python
129 lines
5.1 KiB
Python
"""Tests for context token tracking in run_agent.py's usage extraction.
|
|
|
|
The context counter (status bar) must show the TOTAL prompt tokens including
|
|
Anthropic's cached portions. This is an integration test for the token
|
|
extraction in run_conversation(), not the ContextCompressor itself (which
|
|
is tested in tests/agent/test_context_compressor.py).
|
|
"""
|
|
|
|
import sys
|
|
import types
|
|
from types import SimpleNamespace
|
|
|
|
sys.modules.setdefault("fire", types.SimpleNamespace(Fire=lambda *a, **k: None))
|
|
sys.modules.setdefault("firecrawl", types.SimpleNamespace(Firecrawl=object))
|
|
sys.modules.setdefault("fal_client", types.SimpleNamespace())
|
|
|
|
import run_agent
|
|
|
|
|
|
def _patch_bootstrap(monkeypatch):
|
|
monkeypatch.setattr("model_tools.get_tool_definitions", lambda **kwargs: [{
|
|
"type": "function",
|
|
"function": {"name": "t", "description": "t", "parameters": {"type": "object", "properties": {}}},
|
|
}])
|
|
monkeypatch.setattr("model_tools.check_toolset_requirements", lambda: {})
|
|
|
|
|
|
class _FakeAnthropicClient:
|
|
def close(self):
|
|
pass
|
|
|
|
|
|
class _FakeOpenAIClient:
|
|
"""Fake OpenAI client returned by mocked resolve_provider_client."""
|
|
api_key = "fake-codex-key"
|
|
base_url = "https://api.openai.com/v1"
|
|
_default_headers = None
|
|
|
|
|
|
def _make_agent(monkeypatch, api_mode, provider, response_fn):
|
|
_patch_bootstrap(monkeypatch)
|
|
if api_mode == "anthropic_messages":
|
|
monkeypatch.setattr("agent.anthropic_adapter.build_anthropic_client", lambda k, b=None, **kwargs: _FakeAnthropicClient())
|
|
if provider == "openai-codex":
|
|
monkeypatch.setattr(
|
|
"agent.auxiliary_client.resolve_provider_client",
|
|
lambda *a, **kw: (_FakeOpenAIClient(), "test-model"),
|
|
)
|
|
|
|
class _A(run_agent.AIAgent):
|
|
def __init__(self, *a, **kw):
|
|
kw.update(skip_context_files=True, skip_memory=True, max_iterations=4)
|
|
super().__init__(*a, **kw)
|
|
self._cleanup_task_resources = self._persist_session = lambda *a, **k: None
|
|
self._save_trajectory = lambda *a, **k: None
|
|
|
|
def run_conversation(self, msg, conversation_history=None, task_id=None):
|
|
self._interruptible_api_call = lambda kw: response_fn()
|
|
self._disable_streaming = True
|
|
return super().run_conversation(msg, conversation_history=conversation_history, task_id=task_id)
|
|
|
|
return _A(model="test-model", api_key="test-key", base_url="http://localhost:1234/v1", provider=provider, api_mode=api_mode)
|
|
|
|
|
|
def _anthropic_resp(input_tok, output_tok, cache_read=0, cache_creation=0):
|
|
usage_fields = {"input_tokens": input_tok, "output_tokens": output_tok}
|
|
if cache_read:
|
|
usage_fields["cache_read_input_tokens"] = cache_read
|
|
if cache_creation:
|
|
usage_fields["cache_creation_input_tokens"] = cache_creation
|
|
return SimpleNamespace(
|
|
content=[SimpleNamespace(type="text", text="ok")],
|
|
stop_reason="end_turn",
|
|
usage=SimpleNamespace(**usage_fields),
|
|
model="claude-sonnet-4-6",
|
|
)
|
|
|
|
|
|
# -- Anthropic: cached tokens must be included --
|
|
|
|
def test_anthropic_cache_read_and_creation_added(monkeypatch):
|
|
agent = _make_agent(monkeypatch, "anthropic_messages", "anthropic",
|
|
lambda: _anthropic_resp(3, 10, cache_read=15000, cache_creation=2000))
|
|
agent.run_conversation("hi")
|
|
assert agent.context_compressor.last_prompt_tokens == 17003 # 3+15000+2000
|
|
assert agent.session_prompt_tokens == 17003
|
|
|
|
|
|
def test_anthropic_no_cache_fields(monkeypatch):
|
|
agent = _make_agent(monkeypatch, "anthropic_messages", "anthropic",
|
|
lambda: _anthropic_resp(500, 20))
|
|
agent.run_conversation("hi")
|
|
assert agent.context_compressor.last_prompt_tokens == 500
|
|
|
|
|
|
def test_anthropic_cache_read_only(monkeypatch):
|
|
agent = _make_agent(monkeypatch, "anthropic_messages", "anthropic",
|
|
lambda: _anthropic_resp(5, 15, cache_read=17666, cache_creation=15))
|
|
agent.run_conversation("hi")
|
|
assert agent.context_compressor.last_prompt_tokens == 17686 # 5+17666+15
|
|
|
|
|
|
# -- OpenAI: prompt_tokens already total --
|
|
|
|
def test_openai_prompt_tokens_unchanged(monkeypatch):
|
|
resp = lambda: SimpleNamespace(
|
|
choices=[SimpleNamespace(index=0, message=SimpleNamespace(
|
|
role="assistant", content="ok", tool_calls=None, reasoning_content=None,
|
|
), finish_reason="stop")],
|
|
usage=SimpleNamespace(prompt_tokens=5000, completion_tokens=100, total_tokens=5100),
|
|
model="gpt-4o",
|
|
)
|
|
agent = _make_agent(monkeypatch, "chat_completions", "openrouter", resp)
|
|
agent.run_conversation("hi")
|
|
assert agent.context_compressor.last_prompt_tokens == 5000
|
|
|
|
|
|
# -- Codex: no cache fields, getattr returns 0 --
|
|
|
|
def test_codex_no_cache_fields(monkeypatch):
|
|
resp = lambda: SimpleNamespace(
|
|
output=[SimpleNamespace(type="message", content=[SimpleNamespace(type="output_text", text="ok")])],
|
|
usage=SimpleNamespace(input_tokens=3000, output_tokens=50, total_tokens=3050),
|
|
status="completed", model="gpt-5-codex",
|
|
)
|
|
agent = _make_agent(monkeypatch, "codex_responses", "openai-codex", resp)
|
|
agent.run_conversation("hi")
|
|
assert agent.context_compressor.last_prompt_tokens == 3000
|