refactor(tests): re-architect tests + fix CI failures (#5946)

* refactor: re-architect tests to mirror the codebase

* Update tests.yml

* fix: add missing tool_error imports after registry refactor

* fix(tests): replace patch.dict with monkeypatch to prevent env var leaks under xdist

patch.dict(os.environ) can leak TERMINAL_ENV across xdist workers,
causing test_code_execution tests to hit the Modal remote path.

* fix(tests): fix update_check and telegram xdist failures

- test_update_check: replace patch("hermes_cli.banner.os.getenv") with
  monkeypatch.setenv("HERMES_HOME") — banner.py no longer imports os
  directly, it uses get_hermes_home() from hermes_constants.

- test_telegram_conflict/approval_buttons: provide real exception classes
  for telegram.error mock (NetworkError, TimedOut, BadRequest) so the
  except clause in connect() doesn't fail with "catching classes that do
  not inherit from BaseException" when xdist pollutes sys.modules.

* fix(tests): accept unavailable_models kwarg in _prompt_model_selection mock
This commit is contained in:
Siddharth Balyan
2026-04-07 17:19:07 -07:00
committed by GitHub
parent 99ff375f7a
commit f3006ebef9
110 changed files with 153 additions and 150 deletions

View File

View File

@@ -0,0 +1,268 @@
"""Tests for #1630 — gateway infinite 400 failure loop prevention.
Verifies that:
1. Generic 400 errors with large sessions are treated as context-length errors
and trigger compression instead of aborting.
2. The gateway does not persist messages when the agent fails early, preventing
the session from growing on each failure.
3. Context-overflow failures produce helpful error messages suggesting /compact.
"""
import pytest
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
# ---------------------------------------------------------------------------
# Test 1: Agent heuristic — generic 400 with large session → compression
# ---------------------------------------------------------------------------
class TestGeneric400Heuristic:
"""The agent should treat a generic 400 with a large session as a
probable context-length error and trigger compression, not abort."""
def _make_agent(self):
"""Create a minimal AIAgent for testing error handling."""
with (
patch("run_agent.get_tool_definitions", return_value=[]),
patch("run_agent.check_toolset_requirements", return_value={}),
patch("run_agent.OpenAI"),
):
from run_agent import AIAgent
a = AIAgent(
api_key="test-key-12345",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
)
a.client = MagicMock()
a._cached_system_prompt = "You are helpful."
a._use_prompt_caching = False
a.tool_delay = 0
a.compression_enabled = False
return a
def test_generic_400_with_small_session_is_client_error(self):
"""A generic 400 with a small session should still be treated
as a non-retryable client error (not context overflow)."""
error_msg = "error"
status_code = 400
approx_tokens = 1000 # Small session
api_messages = [{"role": "user", "content": "hi"}]
# Simulate the phrase matching
is_context_length_error = any(phrase in error_msg for phrase in [
'context length', 'context size', 'maximum context',
'token limit', 'too many tokens', 'reduce the length',
'exceeds the limit', 'context window',
'request entity too large',
'prompt is too long',
])
assert not is_context_length_error
# The heuristic should NOT trigger for small sessions
ctx_len = 200000
is_large_session = approx_tokens > ctx_len * 0.4 or len(api_messages) > 80
is_generic_error = len(error_msg.strip()) < 30
assert not is_large_session # Small session → heuristic doesn't fire
def test_generic_400_with_large_token_count_triggers_heuristic(self):
"""A generic 400 with high token count should be treated as
probable context overflow."""
error_msg = "error"
status_code = 400
ctx_len = 200000
approx_tokens = 100000 # > 40% of 200k
api_messages = [{"role": "user", "content": "hi"}] * 20
is_context_length_error = any(phrase in error_msg for phrase in [
'context length', 'context size', 'maximum context',
])
assert not is_context_length_error
# Heuristic check
is_large_session = approx_tokens > ctx_len * 0.4 or len(api_messages) > 80
is_generic_error = len(error_msg.strip()) < 30
assert is_large_session
assert is_generic_error
# Both conditions true → should be treated as context overflow
def test_generic_400_with_many_messages_triggers_heuristic(self):
"""A generic 400 with >80 messages should trigger the heuristic
even if estimated tokens are low."""
error_msg = "error"
status_code = 400
ctx_len = 200000
approx_tokens = 5000 # Low token estimate
api_messages = [{"role": "user", "content": "x"}] * 100 # > 80 messages
is_large_session = approx_tokens > ctx_len * 0.4 or len(api_messages) > 80
is_generic_error = len(error_msg.strip()) < 30
assert is_large_session
assert is_generic_error
def test_specific_error_message_bypasses_heuristic(self):
"""A 400 with a specific, long error message should NOT trigger
the heuristic even with a large session."""
error_msg = "invalid model: anthropic/claude-nonexistent-model is not available"
status_code = 400
ctx_len = 200000
approx_tokens = 100000
is_generic_error = len(error_msg.strip()) < 30
assert not is_generic_error # Long specific message → heuristic doesn't fire
def test_descriptive_context_error_caught_by_phrases(self):
"""Descriptive context-length errors should still be caught by
the existing phrase matching (not the heuristic)."""
error_msg = "prompt is too long: 250000 tokens > 200000 maximum"
is_context_length_error = any(phrase in error_msg for phrase in [
'context length', 'context size', 'maximum context',
'token limit', 'too many tokens', 'reduce the length',
'exceeds the limit', 'context window',
'request entity too large',
'prompt is too long',
])
assert is_context_length_error
# ---------------------------------------------------------------------------
# Test 2: Gateway skips persistence on failed agent results
# ---------------------------------------------------------------------------
class TestGatewaySkipsPersistenceOnFailure:
"""When the agent returns failed=True with no final_response,
the gateway should NOT persist messages to the transcript."""
def test_agent_failed_early_detected(self):
"""The agent_failed_early flag is True when failed=True and
no final_response."""
agent_result = {
"failed": True,
"final_response": None,
"messages": [],
"error": "Non-retryable client error",
}
agent_failed_early = (
agent_result.get("failed")
and not agent_result.get("final_response")
)
assert agent_failed_early
def test_agent_with_response_not_failed_early(self):
"""When the agent has a final_response, it's not a failed-early
scenario even if failed=True."""
agent_result = {
"failed": True,
"final_response": "Here is a partial response",
"messages": [],
}
agent_failed_early = (
agent_result.get("failed")
and not agent_result.get("final_response")
)
assert not agent_failed_early
def test_successful_agent_not_failed_early(self):
"""A successful agent result should not trigger skip."""
agent_result = {
"final_response": "Hello!",
"messages": [{"role": "assistant", "content": "Hello!"}],
}
agent_failed_early = (
agent_result.get("failed")
and not agent_result.get("final_response")
)
assert not agent_failed_early
# ---------------------------------------------------------------------------
# Test 3: Context-overflow error messages
# ---------------------------------------------------------------------------
class TestContextOverflowErrorMessages:
"""The gateway should produce helpful error messages when the failure
looks like a context overflow."""
def test_detects_context_keywords(self):
"""Error messages containing context-related keywords should be
identified as context failures."""
keywords = [
"context length exceeded",
"too many tokens in the prompt",
"request entity too large",
"payload too large for model",
"context window exceeded",
]
for error_str in keywords:
_is_ctx_fail = any(p in error_str.lower() for p in (
"context", "token", "too large", "too long",
"exceed", "payload",
))
assert _is_ctx_fail, f"Should detect: {error_str}"
def test_detects_generic_400_with_large_history(self):
"""A generic 400 error code in the string with a large history
should be flagged as context failure."""
error_str = "error code: 400 - {'type': 'error', 'message': 'Error'}"
history_len = 100 # Large session
_is_ctx_fail = any(p in error_str.lower() for p in (
"context", "token", "too large", "too long",
"exceed", "payload",
)) or (
"400" in error_str.lower()
and history_len > 50
)
assert _is_ctx_fail
def test_unrelated_error_not_flagged(self):
"""Unrelated errors should not be flagged as context failures."""
error_str = "invalid api key: authentication failed"
history_len = 10
_is_ctx_fail = any(p in error_str.lower() for p in (
"context", "token", "too large", "too long",
"exceed", "payload",
)) or (
"400" in error_str.lower()
and history_len > 50
)
assert not _is_ctx_fail
# ---------------------------------------------------------------------------
# Test 4: Agent skips persistence for large failed sessions
# ---------------------------------------------------------------------------
class TestAgentSkipsPersistenceForLargeFailedSessions:
"""When a 400 error occurs and the session is large, the agent
should skip persisting to prevent the growth loop."""
def test_large_session_400_skips_persistence(self):
"""Status 400 + high token count should skip persistence."""
status_code = 400
approx_tokens = 60000 # > 50000 threshold
api_messages = [{"role": "user", "content": "x"}] * 10
should_skip = status_code == 400 and (approx_tokens > 50000 or len(api_messages) > 80)
assert should_skip
def test_small_session_400_persists_normally(self):
"""Status 400 + small session should still persist."""
status_code = 400
approx_tokens = 5000 # < 50000
api_messages = [{"role": "user", "content": "x"}] * 10 # < 80
should_skip = status_code == 400 and (approx_tokens > 50000 or len(api_messages) > 80)
assert not should_skip
def test_non_400_error_persists_normally(self):
"""Non-400 errors should always persist normally."""
status_code = 401 # Auth error
approx_tokens = 100000 # Large session, but not a 400
api_messages = [{"role": "user", "content": "x"}] * 100
should_skip = status_code == 400 and (approx_tokens > 50000 or len(api_messages) > 80)
assert not should_skip

View File

@@ -0,0 +1,475 @@
"""Tests for payload/context-length → compression retry logic in AIAgent.
Verifies that:
- HTTP 413 errors trigger history compression and retry
- HTTP 400 context-length errors trigger compression (not generic 4xx abort)
- Preflight compression proactively compresses oversized sessions before API calls
"""
import pytest
#pytestmark = pytest.mark.skip(reason="Hangs in non-interactive environments")
import uuid
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
from agent.context_compressor import SUMMARY_PREFIX
from run_agent import AIAgent
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _make_tool_defs(*names: str) -> list:
return [
{
"type": "function",
"function": {
"name": n,
"description": f"{n} tool",
"parameters": {"type": "object", "properties": {}},
},
}
for n in names
]
def _mock_response(content="Hello", finish_reason="stop", tool_calls=None, usage=None):
msg = SimpleNamespace(
content=content,
tool_calls=tool_calls,
reasoning_content=None,
reasoning=None,
)
choice = SimpleNamespace(message=msg, finish_reason=finish_reason)
resp = SimpleNamespace(choices=[choice], model="test/model")
resp.usage = SimpleNamespace(**usage) if usage else None
return resp
def _make_413_error(*, use_status_code=True, message="Request entity too large"):
"""Create an exception that mimics a 413 HTTP error."""
err = Exception(message)
if use_status_code:
err.status_code = 413
return err
@pytest.fixture()
def agent():
with (
patch("run_agent.get_tool_definitions", return_value=_make_tool_defs("web_search")),
patch("run_agent.check_toolset_requirements", return_value={}),
patch("run_agent.OpenAI"),
):
a = AIAgent(
api_key="test-key-1234567890",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
)
a.client = MagicMock()
a._cached_system_prompt = "You are helpful."
a._use_prompt_caching = False
a.tool_delay = 0
a.compression_enabled = False
a.save_trajectories = False
return a
# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------
class TestHTTP413Compression:
"""413 errors should trigger compression, not abort as generic 4xx."""
def test_413_triggers_compression(self, agent):
"""A 413 error should call _compress_context and retry, not abort."""
# First call raises 413; second call succeeds after compression.
err_413 = _make_413_error()
ok_resp = _mock_response(content="Success after compression", finish_reason="stop")
agent.client.chat.completions.create.side_effect = [err_413, ok_resp]
# Prefill so there are multiple messages for compression to reduce
prefill = [
{"role": "user", "content": "previous question"},
{"role": "assistant", "content": "previous answer"},
]
with (
patch.object(agent, "_compress_context") as mock_compress,
patch.object(agent, "_persist_session"),
patch.object(agent, "_save_trajectory"),
patch.object(agent, "_cleanup_task_resources"),
):
# Compression reduces 3 messages down to 1
mock_compress.return_value = (
[{"role": "user", "content": "hello"}],
"compressed prompt",
)
result = agent.run_conversation("hello", conversation_history=prefill)
mock_compress.assert_called_once()
assert result["completed"] is True
assert result["final_response"] == "Success after compression"
def test_413_not_treated_as_generic_4xx(self, agent):
"""413 must NOT hit the generic 4xx abort path; it should attempt compression."""
err_413 = _make_413_error()
ok_resp = _mock_response(content="Recovered", finish_reason="stop")
agent.client.chat.completions.create.side_effect = [err_413, ok_resp]
prefill = [
{"role": "user", "content": "previous question"},
{"role": "assistant", "content": "previous answer"},
]
with (
patch.object(agent, "_compress_context") as mock_compress,
patch.object(agent, "_persist_session"),
patch.object(agent, "_save_trajectory"),
patch.object(agent, "_cleanup_task_resources"),
):
mock_compress.return_value = (
[{"role": "user", "content": "hello"}],
"compressed",
)
result = agent.run_conversation("hello", conversation_history=prefill)
# If 413 were treated as generic 4xx, result would have "failed": True
assert result.get("failed") is not True
assert result["completed"] is True
def test_413_error_message_detection(self, agent):
"""413 detected via error message string (no status_code attr)."""
err = _make_413_error(use_status_code=False, message="error code: 413")
ok_resp = _mock_response(content="OK", finish_reason="stop")
agent.client.chat.completions.create.side_effect = [err, ok_resp]
prefill = [
{"role": "user", "content": "previous question"},
{"role": "assistant", "content": "previous answer"},
]
with (
patch.object(agent, "_compress_context") as mock_compress,
patch.object(agent, "_persist_session"),
patch.object(agent, "_save_trajectory"),
patch.object(agent, "_cleanup_task_resources"),
):
mock_compress.return_value = (
[{"role": "user", "content": "hello"}],
"compressed",
)
result = agent.run_conversation("hello", conversation_history=prefill)
mock_compress.assert_called_once()
assert result["completed"] is True
def test_400_context_length_triggers_compression(self, agent):
"""A 400 with 'maximum context length' should trigger compression, not abort as generic 4xx.
OpenRouter returns HTTP 400 (not 413) for context-length errors. Before
the fix, this was caught by the generic 4xx handler which aborted
immediately — now it correctly triggers compression+retry.
"""
err_400 = Exception(
"Error code: 400 - {'error': {'message': "
"\"This endpoint's maximum context length is 204800 tokens. "
"However, you requested about 270460 tokens.\", 'code': 400}}"
)
err_400.status_code = 400
ok_resp = _mock_response(content="Recovered after compression", finish_reason="stop")
agent.client.chat.completions.create.side_effect = [err_400, ok_resp]
prefill = [
{"role": "user", "content": "previous question"},
{"role": "assistant", "content": "previous answer"},
]
with (
patch.object(agent, "_compress_context") as mock_compress,
patch.object(agent, "_persist_session"),
patch.object(agent, "_save_trajectory"),
patch.object(agent, "_cleanup_task_resources"),
):
mock_compress.return_value = (
[{"role": "user", "content": "hello"}],
"compressed prompt",
)
result = agent.run_conversation("hello", conversation_history=prefill)
mock_compress.assert_called_once()
# Must NOT have "failed": True (which would mean the generic 4xx handler caught it)
assert result.get("failed") is not True
assert result["completed"] is True
assert result["final_response"] == "Recovered after compression"
def test_400_reduce_length_triggers_compression(self, agent):
"""A 400 with 'reduce the length' should trigger compression."""
err_400 = Exception(
"Error code: 400 - Please reduce the length of the messages"
)
err_400.status_code = 400
ok_resp = _mock_response(content="OK", finish_reason="stop")
agent.client.chat.completions.create.side_effect = [err_400, ok_resp]
prefill = [
{"role": "user", "content": "previous question"},
{"role": "assistant", "content": "previous answer"},
]
with (
patch.object(agent, "_compress_context") as mock_compress,
patch.object(agent, "_persist_session"),
patch.object(agent, "_save_trajectory"),
patch.object(agent, "_cleanup_task_resources"),
):
mock_compress.return_value = (
[{"role": "user", "content": "hello"}],
"compressed",
)
result = agent.run_conversation("hello", conversation_history=prefill)
mock_compress.assert_called_once()
assert result["completed"] is True
def test_context_length_retry_rebuilds_request_after_compression(self, agent):
"""Retry must send the compressed transcript, not the stale oversized payload."""
err_400 = Exception(
"Error code: 400 - {'error': {'message': "
"\"This endpoint's maximum context length is 128000 tokens. "
"Please reduce the length of the messages.\"}}"
)
err_400.status_code = 400
ok_resp = _mock_response(content="Recovered after real compression", finish_reason="stop")
request_payloads = []
def _side_effect(**kwargs):
request_payloads.append(kwargs)
if len(request_payloads) == 1:
raise err_400
return ok_resp
agent.client.chat.completions.create.side_effect = _side_effect
prefill = [
{"role": "user", "content": "previous question"},
{"role": "assistant", "content": "previous answer"},
]
with (
patch.object(agent, "_compress_context") as mock_compress,
patch.object(agent, "_persist_session"),
patch.object(agent, "_save_trajectory"),
patch.object(agent, "_cleanup_task_resources"),
):
mock_compress.return_value = (
[{"role": "user", "content": "compressed summary"}],
"compressed prompt",
)
result = agent.run_conversation("hello", conversation_history=prefill)
assert result["completed"] is True
assert len(request_payloads) == 2
assert len(request_payloads[1]["messages"]) < len(request_payloads[0]["messages"])
assert request_payloads[1]["messages"][0] == {
"role": "system",
"content": "compressed prompt",
}
assert request_payloads[1]["messages"][1] == {
"role": "user",
"content": "compressed summary",
}
def test_413_cannot_compress_further(self, agent):
"""When compression can't reduce messages, return partial result."""
err_413 = _make_413_error()
agent.client.chat.completions.create.side_effect = [err_413]
with (
patch.object(agent, "_compress_context") as mock_compress,
patch.object(agent, "_persist_session"),
patch.object(agent, "_save_trajectory"),
patch.object(agent, "_cleanup_task_resources"),
):
# Compression returns same number of messages → can't compress further
mock_compress.return_value = (
[{"role": "user", "content": "hello"}],
"same prompt",
)
result = agent.run_conversation("hello")
assert result["completed"] is False
assert result.get("partial") is True
assert "413" in result["error"]
class TestPreflightCompression:
"""Preflight compression should compress history before the first API call."""
def test_preflight_compresses_oversized_history(self, agent):
"""When loaded history exceeds the model's context threshold, compress before API call."""
agent.compression_enabled = True
# Set a small context so the history is "oversized", but large enough
# that the compressed result (2 short messages) fits in a single pass.
agent.context_compressor.context_length = 2000
agent.context_compressor.threshold_tokens = 200
# Build a history that will be large enough to trigger preflight
# (each message ~50 chars ≈ 13 tokens, 40 messages ≈ 520 tokens > 200 threshold)
big_history = []
for i in range(20):
big_history.append({"role": "user", "content": f"Message number {i} with some extra text padding"})
big_history.append({"role": "assistant", "content": f"Response number {i} with extra padding here"})
ok_resp = _mock_response(content="After preflight", finish_reason="stop")
agent.client.chat.completions.create.side_effect = [ok_resp]
with (
patch.object(agent, "_compress_context") as mock_compress,
patch.object(agent, "_persist_session"),
patch.object(agent, "_save_trajectory"),
patch.object(agent, "_cleanup_task_resources"),
):
# Simulate compression reducing messages to a small set that fits
mock_compress.return_value = (
[
{"role": "user", "content": f"{SUMMARY_PREFIX}\nPrevious conversation"},
{"role": "user", "content": "hello"},
],
"new system prompt",
)
result = agent.run_conversation("hello", conversation_history=big_history)
# Preflight compression should have been called BEFORE the API call
mock_compress.assert_called_once()
assert result["completed"] is True
assert result["final_response"] == "After preflight"
def test_no_preflight_when_under_threshold(self, agent):
"""When history fits within context, no preflight compression needed."""
agent.compression_enabled = True
# Large context — history easily fits
agent.context_compressor.context_length = 1000000
agent.context_compressor.threshold_tokens = 850000
small_history = [
{"role": "user", "content": "hi"},
{"role": "assistant", "content": "hello"},
]
ok_resp = _mock_response(content="No compression needed", finish_reason="stop")
agent.client.chat.completions.create.side_effect = [ok_resp]
with (
patch.object(agent, "_compress_context") as mock_compress,
patch.object(agent, "_persist_session"),
patch.object(agent, "_save_trajectory"),
patch.object(agent, "_cleanup_task_resources"),
):
result = agent.run_conversation("hello", conversation_history=small_history)
mock_compress.assert_not_called()
assert result["completed"] is True
def test_no_preflight_when_compression_disabled(self, agent):
"""Preflight should not run when compression is disabled."""
agent.compression_enabled = False
agent.context_compressor.context_length = 100
agent.context_compressor.threshold_tokens = 85
big_history = [
{"role": "user", "content": "x" * 1000},
{"role": "assistant", "content": "y" * 1000},
] * 10
ok_resp = _mock_response(content="OK", finish_reason="stop")
agent.client.chat.completions.create.side_effect = [ok_resp]
with (
patch.object(agent, "_compress_context") as mock_compress,
patch.object(agent, "_persist_session"),
patch.object(agent, "_save_trajectory"),
patch.object(agent, "_cleanup_task_resources"),
):
result = agent.run_conversation("hello", conversation_history=big_history)
mock_compress.assert_not_called()
class TestToolResultPreflightCompression:
"""Compression should trigger when tool results push context past the threshold."""
def test_large_tool_results_trigger_compression(self, agent):
"""When tool results push estimated tokens past threshold, compress before next call."""
agent.compression_enabled = True
agent.context_compressor.context_length = 200_000
agent.context_compressor.threshold_tokens = 130_000 # below the 135k reported usage
agent.context_compressor.last_prompt_tokens = 130_000
agent.context_compressor.last_completion_tokens = 5_000
tc = SimpleNamespace(
id="tc1", type="function",
function=SimpleNamespace(name="web_search", arguments='{"query":"test"}'),
)
tool_resp = _mock_response(
content=None, finish_reason="stop", tool_calls=[tc],
usage={"prompt_tokens": 130_000, "completion_tokens": 5_000, "total_tokens": 135_000},
)
ok_resp = _mock_response(
content="Done after compression", finish_reason="stop",
usage={"prompt_tokens": 50_000, "completion_tokens": 100, "total_tokens": 50_100},
)
agent.client.chat.completions.create.side_effect = [tool_resp, ok_resp]
large_result = "x" * 100_000
with (
patch("run_agent.handle_function_call", return_value=large_result),
patch.object(agent, "_compress_context") as mock_compress,
patch.object(agent, "_persist_session"),
patch.object(agent, "_save_trajectory"),
patch.object(agent, "_cleanup_task_resources"),
):
mock_compress.return_value = (
[{"role": "user", "content": "hello"}], "compressed prompt",
)
result = agent.run_conversation("hello")
mock_compress.assert_called_once()
assert result["completed"] is True
def test_anthropic_prompt_too_long_safety_net(self, agent):
"""Anthropic 'prompt is too long' error triggers compression as safety net."""
err_400 = Exception(
"Error code: 400 - {'type': 'error', 'error': {'type': 'invalid_request_error', "
"'message': 'prompt is too long: 233153 tokens > 200000 maximum'}}"
)
err_400.status_code = 400
ok_resp = _mock_response(content="Recovered", finish_reason="stop")
agent.client.chat.completions.create.side_effect = [err_400, ok_resp]
prefill = [
{"role": "user", "content": "previous"},
{"role": "assistant", "content": "answer"},
]
with (
patch.object(agent, "_compress_context") as mock_compress,
patch.object(agent, "_persist_session"),
patch.object(agent, "_save_trajectory"),
patch.object(agent, "_cleanup_task_resources"),
):
mock_compress.return_value = (
[{"role": "user", "content": "hello"}], "compressed",
)
result = agent.run_conversation("hello", conversation_history=prefill)
mock_compress.assert_called_once()
assert result["completed"] is True

View File

@@ -0,0 +1,294 @@
"""Tests for issue #860 — SQLite session transcript deduplication.
Verifies that:
1. _flush_messages_to_session_db uses _last_flushed_db_idx to avoid re-writing
2. Multiple _persist_session calls don't duplicate messages
3. append_to_transcript(skip_db=True) skips SQLite but writes JSONL
4. The gateway doesn't double-write messages the agent already persisted
"""
import json
import os
import sqlite3
import tempfile
from pathlib import Path
from unittest.mock import MagicMock, patch
import pytest
# ---------------------------------------------------------------------------
# Test: _flush_messages_to_session_db only writes new messages
# ---------------------------------------------------------------------------
class TestFlushDeduplication:
"""Verify _flush_messages_to_session_db tracks what it already wrote."""
def _make_agent(self, session_db):
"""Create a minimal AIAgent with a real session DB."""
with patch.dict(os.environ, {"OPENROUTER_API_KEY": "test-key"}):
from run_agent import AIAgent
agent = AIAgent(
model="test/model",
quiet_mode=True,
session_db=session_db,
session_id="test-session-860",
skip_context_files=True,
skip_memory=True,
)
return agent
def test_flush_writes_only_new_messages(self):
"""First flush writes all new messages, second flush writes none."""
from hermes_state import SessionDB
with tempfile.TemporaryDirectory() as tmpdir:
db_path = Path(tmpdir) / "test.db"
db = SessionDB(db_path=db_path)
agent = self._make_agent(db)
conversation_history = [
{"role": "user", "content": "old message"},
]
messages = list(conversation_history) + [
{"role": "user", "content": "new question"},
{"role": "assistant", "content": "new answer"},
]
# First flush — should write 2 new messages
agent._flush_messages_to_session_db(messages, conversation_history)
rows = db.get_messages(agent.session_id)
assert len(rows) == 2, f"Expected 2 messages, got {len(rows)}"
# Second flush with SAME messages — should write 0 new messages
agent._flush_messages_to_session_db(messages, conversation_history)
rows = db.get_messages(agent.session_id)
assert len(rows) == 2, f"Expected still 2 messages after second flush, got {len(rows)}"
def test_flush_writes_incrementally(self):
"""Messages added between flushes are written exactly once."""
from hermes_state import SessionDB
with tempfile.TemporaryDirectory() as tmpdir:
db_path = Path(tmpdir) / "test.db"
db = SessionDB(db_path=db_path)
agent = self._make_agent(db)
conversation_history = []
messages = [
{"role": "user", "content": "hello"},
]
# First flush — 1 message
agent._flush_messages_to_session_db(messages, conversation_history)
rows = db.get_messages(agent.session_id)
assert len(rows) == 1
# Add more messages
messages.append({"role": "assistant", "content": "hi there"})
messages.append({"role": "user", "content": "follow up"})
# Second flush — should write only 2 new messages
agent._flush_messages_to_session_db(messages, conversation_history)
rows = db.get_messages(agent.session_id)
assert len(rows) == 3, f"Expected 3 total messages, got {len(rows)}"
def test_persist_session_multiple_calls_no_duplication(self):
"""Multiple _persist_session calls don't duplicate DB entries."""
from hermes_state import SessionDB
with tempfile.TemporaryDirectory() as tmpdir:
db_path = Path(tmpdir) / "test.db"
db = SessionDB(db_path=db_path)
agent = self._make_agent(db)
# Stub out _save_session_log to avoid file I/O
agent._save_session_log = MagicMock()
conversation_history = [{"role": "user", "content": "old"}]
messages = list(conversation_history) + [
{"role": "user", "content": "q1"},
{"role": "assistant", "content": "a1"},
{"role": "user", "content": "q2"},
{"role": "assistant", "content": "a2"},
]
# Simulate multiple persist calls (like the agent's many exit paths)
for _ in range(5):
agent._persist_session(messages, conversation_history)
rows = db.get_messages(agent.session_id)
assert len(rows) == 4, f"Expected 4 messages, got {len(rows)} (duplication bug!)"
def test_flush_reset_after_compression(self):
"""After compression creates a new session, flush index resets."""
from hermes_state import SessionDB
with tempfile.TemporaryDirectory() as tmpdir:
db_path = Path(tmpdir) / "test.db"
db = SessionDB(db_path=db_path)
agent = self._make_agent(db)
# Write some messages
messages = [
{"role": "user", "content": "msg1"},
{"role": "assistant", "content": "reply1"},
]
agent._flush_messages_to_session_db(messages, [])
old_session = agent.session_id
assert agent._last_flushed_db_idx == 2
# Simulate what _compress_context does: new session, reset idx
agent.session_id = "compressed-session-new"
db.create_session(session_id=agent.session_id, source="test")
agent._last_flushed_db_idx = 0
# Now flush compressed messages to new session
compressed_messages = [
{"role": "user", "content": "summary of conversation"},
]
agent._flush_messages_to_session_db(compressed_messages, [])
new_rows = db.get_messages(agent.session_id)
assert len(new_rows) == 1
# Old session should still have its 2 messages
old_rows = db.get_messages(old_session)
assert len(old_rows) == 2
# ---------------------------------------------------------------------------
# Test: append_to_transcript skip_db parameter
# ---------------------------------------------------------------------------
class TestAppendToTranscriptSkipDb:
"""Verify skip_db=True writes JSONL but not SQLite."""
@pytest.fixture()
def store(self, tmp_path):
from gateway.config import GatewayConfig
from gateway.session import SessionStore
config = GatewayConfig()
with patch("gateway.session.SessionStore._ensure_loaded"):
s = SessionStore(sessions_dir=tmp_path, config=config)
s._db = None # no SQLite for these JSONL-focused tests
s._loaded = True
return s
def test_skip_db_writes_jsonl_only(self, store, tmp_path):
"""With skip_db=True, message appears in JSONL but not SQLite."""
session_id = "test-skip-db"
msg = {"role": "assistant", "content": "hello world"}
store.append_to_transcript(session_id, msg, skip_db=True)
# JSONL should have the message
jsonl_path = store.get_transcript_path(session_id)
assert jsonl_path.exists()
with open(jsonl_path) as f:
lines = f.readlines()
assert len(lines) == 1
parsed = json.loads(lines[0])
assert parsed["content"] == "hello world"
def test_skip_db_prevents_sqlite_write(self, tmp_path):
"""With skip_db=True and a real DB, message does NOT appear in SQLite."""
from gateway.config import GatewayConfig
from gateway.session import SessionStore
from hermes_state import SessionDB
db_path = tmp_path / "test_skip.db"
db = SessionDB(db_path=db_path)
config = GatewayConfig()
with patch("gateway.session.SessionStore._ensure_loaded"):
store = SessionStore(sessions_dir=tmp_path, config=config)
store._db = db
store._loaded = True
session_id = "test-skip-db-real"
db.create_session(session_id=session_id, source="test")
msg = {"role": "assistant", "content": "hello world"}
store.append_to_transcript(session_id, msg, skip_db=True)
# SQLite should NOT have the message
rows = db.get_messages(session_id)
assert len(rows) == 0, f"Expected 0 DB rows with skip_db=True, got {len(rows)}"
# But JSONL should have it
jsonl_path = store.get_transcript_path(session_id)
with open(jsonl_path) as f:
lines = f.readlines()
assert len(lines) == 1
def test_default_writes_both(self, tmp_path):
"""Without skip_db, message appears in both JSONL and SQLite."""
from gateway.config import GatewayConfig
from gateway.session import SessionStore
from hermes_state import SessionDB
db_path = tmp_path / "test_both.db"
db = SessionDB(db_path=db_path)
config = GatewayConfig()
with patch("gateway.session.SessionStore._ensure_loaded"):
store = SessionStore(sessions_dir=tmp_path, config=config)
store._db = db
store._loaded = True
session_id = "test-default-write"
db.create_session(session_id=session_id, source="test")
msg = {"role": "user", "content": "test message"}
store.append_to_transcript(session_id, msg)
# JSONL should have the message
jsonl_path = store.get_transcript_path(session_id)
with open(jsonl_path) as f:
lines = f.readlines()
assert len(lines) == 1
# SQLite should also have the message
rows = db.get_messages(session_id)
assert len(rows) == 1
# ---------------------------------------------------------------------------
# Test: _last_flushed_db_idx initialization
# ---------------------------------------------------------------------------
class TestFlushIdxInit:
"""Verify _last_flushed_db_idx is properly initialized."""
def test_init_zero(self):
"""Agent starts with _last_flushed_db_idx = 0."""
with patch.dict(os.environ, {"OPENROUTER_API_KEY": "test-key"}):
from run_agent import AIAgent
agent = AIAgent(
model="test/model",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
)
assert agent._last_flushed_db_idx == 0
def test_no_session_db_noop(self):
"""Without session_db, flush is a no-op and doesn't crash."""
with patch.dict(os.environ, {"OPENROUTER_API_KEY": "test-key"}):
from run_agent import AIAgent
agent = AIAgent(
model="test/model",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
)
messages = [{"role": "user", "content": "test"}]
agent._flush_messages_to_session_db(messages, [])
# Should not crash, idx should remain 0
assert agent._last_flushed_db_idx == 0

View File

@@ -0,0 +1,263 @@
"""Unit tests for AIAgent pre/post-LLM-call guardrails.
Covers three static methods on AIAgent (inspired by PR #1321 — @alireza78a):
- _sanitize_api_messages() — Phase 1: orphaned tool pair repair
- _cap_delegate_task_calls() — Phase 2a: subagent concurrency limit
- _deduplicate_tool_calls() — Phase 2b: identical call deduplication
"""
import types
from run_agent import AIAgent
from tools.delegate_tool import MAX_CONCURRENT_CHILDREN
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def make_tc(name: str, arguments: str = "{}") -> types.SimpleNamespace:
"""Create a minimal tool_call SimpleNamespace mirroring the OpenAI SDK object."""
tc = types.SimpleNamespace()
tc.function = types.SimpleNamespace(name=name, arguments=arguments)
return tc
def tool_result(call_id: str, content: str = "ok") -> dict:
return {"role": "tool", "tool_call_id": call_id, "content": content}
def assistant_dict_call(call_id: str, name: str = "terminal") -> dict:
"""Dict-style tool_call (as stored in message history)."""
return {"id": call_id, "function": {"name": name, "arguments": "{}"}}
# ---------------------------------------------------------------------------
# Phase 1 — _sanitize_api_messages
# ---------------------------------------------------------------------------
class TestSanitizeApiMessages:
def test_orphaned_result_removed(self):
msgs = [
{"role": "assistant", "tool_calls": [assistant_dict_call("c1")]},
tool_result("c1"),
tool_result("c_ORPHAN"),
]
out = AIAgent._sanitize_api_messages(msgs)
assert len(out) == 2
assert all(m.get("tool_call_id") != "c_ORPHAN" for m in out)
def test_orphaned_call_gets_stub_result(self):
msgs = [
{"role": "assistant", "tool_calls": [assistant_dict_call("c2")]},
]
out = AIAgent._sanitize_api_messages(msgs)
assert len(out) == 2
stub = out[1]
assert stub["role"] == "tool"
assert stub["tool_call_id"] == "c2"
assert stub["content"]
def test_clean_messages_pass_through(self):
msgs = [
{"role": "user", "content": "hello"},
{"role": "assistant", "tool_calls": [assistant_dict_call("c3")]},
tool_result("c3"),
{"role": "assistant", "content": "done"},
]
out = AIAgent._sanitize_api_messages(msgs)
assert out == msgs
def test_mixed_orphaned_result_and_orphaned_call(self):
msgs = [
{"role": "assistant", "tool_calls": [
assistant_dict_call("c4"),
assistant_dict_call("c5"),
]},
tool_result("c4"),
tool_result("c_DANGLING"),
]
out = AIAgent._sanitize_api_messages(msgs)
ids = [m.get("tool_call_id") for m in out if m.get("role") == "tool"]
assert "c_DANGLING" not in ids
assert "c4" in ids
assert "c5" in ids
def test_empty_list_is_safe(self):
assert AIAgent._sanitize_api_messages([]) == []
def test_no_tool_messages(self):
msgs = [
{"role": "user", "content": "hi"},
{"role": "assistant", "content": "hello"},
]
out = AIAgent._sanitize_api_messages(msgs)
assert out == msgs
def test_sdk_object_tool_calls(self):
tc_obj = types.SimpleNamespace(id="c6", function=types.SimpleNamespace(
name="terminal", arguments="{}"
))
msgs = [
{"role": "assistant", "tool_calls": [tc_obj]},
]
out = AIAgent._sanitize_api_messages(msgs)
assert len(out) == 2
assert out[1]["tool_call_id"] == "c6"
# ---------------------------------------------------------------------------
# Phase 2a — _cap_delegate_task_calls
# ---------------------------------------------------------------------------
class TestCapDelegateTaskCalls:
def test_excess_delegates_truncated(self):
tcs = [make_tc("delegate_task") for _ in range(MAX_CONCURRENT_CHILDREN + 2)]
out = AIAgent._cap_delegate_task_calls(tcs)
delegate_count = sum(1 for tc in out if tc.function.name == "delegate_task")
assert delegate_count == MAX_CONCURRENT_CHILDREN
def test_non_delegate_calls_preserved(self):
tcs = (
[make_tc("delegate_task") for _ in range(MAX_CONCURRENT_CHILDREN + 1)]
+ [make_tc("terminal"), make_tc("web_search")]
)
out = AIAgent._cap_delegate_task_calls(tcs)
names = [tc.function.name for tc in out]
assert "terminal" in names
assert "web_search" in names
def test_at_limit_passes_through(self):
tcs = [make_tc("delegate_task") for _ in range(MAX_CONCURRENT_CHILDREN)]
out = AIAgent._cap_delegate_task_calls(tcs)
assert out is tcs
def test_below_limit_passes_through(self):
tcs = [make_tc("delegate_task") for _ in range(MAX_CONCURRENT_CHILDREN - 1)]
out = AIAgent._cap_delegate_task_calls(tcs)
assert out is tcs
def test_no_delegate_calls_unchanged(self):
tcs = [make_tc("terminal"), make_tc("web_search")]
out = AIAgent._cap_delegate_task_calls(tcs)
assert out is tcs
def test_empty_list_safe(self):
assert AIAgent._cap_delegate_task_calls([]) == []
def test_original_list_not_mutated(self):
tcs = [make_tc("delegate_task") for _ in range(MAX_CONCURRENT_CHILDREN + 2)]
original_len = len(tcs)
AIAgent._cap_delegate_task_calls(tcs)
assert len(tcs) == original_len
def test_interleaved_order_preserved(self):
delegates = [make_tc("delegate_task", f'{{"task":"{i}"}}')
for i in range(MAX_CONCURRENT_CHILDREN + 1)]
t1 = make_tc("terminal", '{"cmd":"ls"}')
w1 = make_tc("web_search", '{"q":"x"}')
tcs = [delegates[0], t1, delegates[1], w1] + delegates[2:]
out = AIAgent._cap_delegate_task_calls(tcs)
expected = [delegates[0], t1, delegates[1], w1] + delegates[2:MAX_CONCURRENT_CHILDREN]
assert len(out) == len(expected)
for i, (actual, exp) in enumerate(zip(out, expected)):
assert actual is exp, f"mismatch at index {i}"
# ---------------------------------------------------------------------------
# Phase 2b — _deduplicate_tool_calls
# ---------------------------------------------------------------------------
class TestDeduplicateToolCalls:
def test_duplicate_pair_deduplicated(self):
tcs = [
make_tc("web_search", '{"query":"foo"}'),
make_tc("web_search", '{"query":"foo"}'),
]
out = AIAgent._deduplicate_tool_calls(tcs)
assert len(out) == 1
def test_multiple_duplicates(self):
tcs = [
make_tc("web_search", '{"q":"a"}'),
make_tc("web_search", '{"q":"a"}'),
make_tc("terminal", '{"cmd":"ls"}'),
make_tc("terminal", '{"cmd":"ls"}'),
make_tc("terminal", '{"cmd":"pwd"}'),
]
out = AIAgent._deduplicate_tool_calls(tcs)
assert len(out) == 3
def test_same_tool_different_args_kept(self):
tcs = [
make_tc("terminal", '{"cmd":"ls"}'),
make_tc("terminal", '{"cmd":"pwd"}'),
]
out = AIAgent._deduplicate_tool_calls(tcs)
assert out is tcs
def test_different_tools_same_args_kept(self):
tcs = [
make_tc("tool_a", '{"x":1}'),
make_tc("tool_b", '{"x":1}'),
]
out = AIAgent._deduplicate_tool_calls(tcs)
assert out is tcs
def test_clean_list_unchanged(self):
tcs = [
make_tc("web_search", '{"q":"x"}'),
make_tc("terminal", '{"cmd":"ls"}'),
]
out = AIAgent._deduplicate_tool_calls(tcs)
assert out is tcs
def test_empty_list_safe(self):
assert AIAgent._deduplicate_tool_calls([]) == []
def test_first_occurrence_kept(self):
tc1 = make_tc("terminal", '{"cmd":"ls"}')
tc2 = make_tc("terminal", '{"cmd":"ls"}')
out = AIAgent._deduplicate_tool_calls([tc1, tc2])
assert len(out) == 1
assert out[0] is tc1
def test_original_list_not_mutated(self):
tcs = [
make_tc("web_search", '{"q":"dup"}'),
make_tc("web_search", '{"q":"dup"}'),
]
original_len = len(tcs)
AIAgent._deduplicate_tool_calls(tcs)
assert len(tcs) == original_len
# ---------------------------------------------------------------------------
# _get_tool_call_id_static
# ---------------------------------------------------------------------------
class TestGetToolCallIdStatic:
def test_dict_with_valid_id(self):
assert AIAgent._get_tool_call_id_static({"id": "call_123"}) == "call_123"
def test_dict_with_none_id(self):
assert AIAgent._get_tool_call_id_static({"id": None}) == ""
def test_dict_without_id_key(self):
assert AIAgent._get_tool_call_id_static({"function": {}}) == ""
def test_object_with_valid_id(self):
tc = types.SimpleNamespace(id="call_456")
assert AIAgent._get_tool_call_id_static(tc) == "call_456"
def test_object_with_none_id(self):
tc = types.SimpleNamespace(id=None)
assert AIAgent._get_tool_call_id_static(tc) == ""
def test_object_without_id_attr(self):
tc = types.SimpleNamespace()
assert AIAgent._get_tool_call_id_static(tc) == ""

View File

@@ -0,0 +1,505 @@
"""
Tests for environments/agent_loop.py — HermesAgentLoop.
Tests the multi-turn agent engine using mocked servers, without needing
real API keys or running servers.
"""
import asyncio
import json
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional
from unittest.mock import MagicMock
import pytest
# Ensure repo root is importable
sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent))
try:
from environments.agent_loop import (
AgentResult,
HermesAgentLoop,
ToolError,
_extract_reasoning_from_message,
resize_tool_pool,
)
except ImportError:
pytest.skip("atroposlib not installed", allow_module_level=True)
# ─── Mock server infrastructure ─────────────────────────────────────────
@dataclass
class MockFunction:
name: str
arguments: str
@dataclass
class MockToolCall:
id: str
function: MockFunction
type: str = "function"
@dataclass
class MockMessage:
content: Optional[str]
role: str = "assistant"
tool_calls: Optional[List[MockToolCall]] = None
reasoning_content: Optional[str] = None
reasoning: Optional[str] = None
reasoning_details: Optional[list] = None
@dataclass
class MockChoice:
message: MockMessage
finish_reason: str = "stop"
index: int = 0
@dataclass
class MockChatCompletion:
choices: List[MockChoice]
id: str = "chatcmpl-mock"
model: str = "mock-model"
class MockServer:
"""
Mock server that returns pre-configured responses in sequence.
Mimics the chat_completion() interface.
"""
def __init__(self, responses: List[MockChatCompletion]):
self.responses = responses
self.call_count = 0
self.call_history: List[Dict[str, Any]] = []
async def chat_completion(self, **kwargs) -> MockChatCompletion:
self.call_history.append(kwargs)
if self.call_count >= len(self.responses):
# Return a simple text response if we run out
return MockChatCompletion(
choices=[MockChoice(message=MockMessage(content="Done."))]
)
resp = self.responses[self.call_count]
self.call_count += 1
return resp
def make_text_response(content: str) -> MockChatCompletion:
"""Create a simple text-only response (no tool calls)."""
return MockChatCompletion(
choices=[MockChoice(message=MockMessage(content=content))]
)
def make_tool_response(
tool_name: str,
arguments: dict,
content: str = "",
tool_call_id: str = "call_001",
) -> MockChatCompletion:
"""Create a response with a single tool call."""
return MockChatCompletion(
choices=[
MockChoice(
message=MockMessage(
content=content,
tool_calls=[
MockToolCall(
id=tool_call_id,
function=MockFunction(
name=tool_name,
arguments=json.dumps(arguments),
),
)
],
),
finish_reason="tool_calls",
)
]
)
# ─── Tests ───────────────────────────────────────────────────────────────
class TestAgentResult:
def test_defaults(self):
result = AgentResult(messages=[])
assert result.messages == []
assert result.managed_state is None
assert result.turns_used == 0
assert result.finished_naturally is False
assert result.reasoning_per_turn == []
assert result.tool_errors == []
class TestExtractReasoning:
def test_reasoning_content_field(self):
msg = MockMessage(content="hello", reasoning_content="I think...")
assert _extract_reasoning_from_message(msg) == "I think..."
def test_reasoning_field(self):
msg = MockMessage(content="hello", reasoning="Let me consider...")
assert _extract_reasoning_from_message(msg) == "Let me consider..."
def test_reasoning_details(self):
detail = MagicMock()
detail.text = "Detail reasoning"
msg = MockMessage(content="hello", reasoning_details=[detail])
assert _extract_reasoning_from_message(msg) == "Detail reasoning"
def test_reasoning_details_dict_format(self):
msg = MockMessage(
content="hello",
reasoning_details=[{"text": "Dict reasoning"}],
)
assert _extract_reasoning_from_message(msg) == "Dict reasoning"
def test_no_reasoning(self):
msg = MockMessage(content="hello")
assert _extract_reasoning_from_message(msg) is None
def test_reasoning_content_takes_priority(self):
msg = MockMessage(
content="hello",
reasoning_content="First",
reasoning="Second",
)
assert _extract_reasoning_from_message(msg) == "First"
class TestHermesAgentLoop:
"""Test the agent loop with mock servers."""
@pytest.fixture
def basic_tools(self):
"""Minimal tool schema for testing."""
return [
{
"type": "function",
"function": {
"name": "terminal",
"description": "Run a command",
"parameters": {
"type": "object",
"properties": {
"command": {
"type": "string",
"description": "Command to run",
}
},
"required": ["command"],
},
},
},
{
"type": "function",
"function": {
"name": "read_file",
"description": "Read a file",
"parameters": {
"type": "object",
"properties": {
"path": {"type": "string"},
},
"required": ["path"],
},
},
},
]
@pytest.fixture
def valid_names(self):
return {"terminal", "read_file", "todo"}
@pytest.mark.asyncio
async def test_simple_text_response(self, basic_tools, valid_names):
"""Model responds with text only, no tool calls."""
server = MockServer([make_text_response("Hello! How can I help?")])
agent = HermesAgentLoop(
server=server,
tool_schemas=basic_tools,
valid_tool_names=valid_names,
max_turns=10,
)
messages = [{"role": "user", "content": "Hi"}]
result = await agent.run(messages)
assert result.finished_naturally is True
assert result.turns_used == 1
assert len(result.messages) >= 2 # user + assistant
assert result.messages[-1]["role"] == "assistant"
assert result.messages[-1]["content"] == "Hello! How can I help?"
@pytest.mark.asyncio
async def test_tool_call_then_text(self, basic_tools, valid_names):
"""Model calls a tool, then responds with text."""
server = MockServer([
make_tool_response("todo", {"todos": [{"id": "1", "content": "test", "status": "pending"}]}),
make_text_response("I created a todo for you."),
])
agent = HermesAgentLoop(
server=server,
tool_schemas=basic_tools,
valid_tool_names=valid_names,
max_turns=10,
)
messages = [{"role": "user", "content": "Create a todo"}]
result = await agent.run(messages)
assert result.finished_naturally is True
assert result.turns_used == 2
# Should have: user, assistant (tool_call), tool (result), assistant (text)
roles = [m["role"] for m in result.messages]
assert roles == ["user", "assistant", "tool", "assistant"]
@pytest.mark.asyncio
async def test_max_turns_reached(self, basic_tools, valid_names):
"""Model keeps calling tools until max_turns is hit."""
# Create responses that always call a tool
responses = [
make_tool_response("todo", {"todos": [{"id": str(i), "content": f"task {i}", "status": "pending"}]}, tool_call_id=f"call_{i}")
for i in range(10)
]
server = MockServer(responses)
agent = HermesAgentLoop(
server=server,
tool_schemas=basic_tools,
valid_tool_names=valid_names,
max_turns=3,
)
messages = [{"role": "user", "content": "Keep going"}]
result = await agent.run(messages)
assert result.finished_naturally is False
assert result.turns_used == 3
@pytest.mark.asyncio
async def test_unknown_tool_name(self, basic_tools, valid_names):
"""Model calls a tool not in valid_tool_names."""
server = MockServer([
make_tool_response("nonexistent_tool", {"arg": "val"}),
make_text_response("OK, that didn't work."),
])
agent = HermesAgentLoop(
server=server,
tool_schemas=basic_tools,
valid_tool_names=valid_names,
max_turns=10,
)
messages = [{"role": "user", "content": "Call something weird"}]
result = await agent.run(messages)
# Should record a tool error
assert len(result.tool_errors) >= 1
assert result.tool_errors[0].tool_name == "nonexistent_tool"
@pytest.mark.asyncio
async def test_empty_response(self, basic_tools, valid_names):
"""Server returns empty response."""
server = MockServer([MockChatCompletion(choices=[])])
agent = HermesAgentLoop(
server=server,
tool_schemas=basic_tools,
valid_tool_names=valid_names,
max_turns=10,
)
messages = [{"role": "user", "content": "Hi"}]
result = await agent.run(messages)
assert result.finished_naturally is False
assert result.turns_used == 1
@pytest.mark.asyncio
async def test_api_error_handling(self, basic_tools, valid_names):
"""Server raises an exception."""
class FailingServer:
async def chat_completion(self, **kwargs):
raise ConnectionError("Server unreachable")
agent = HermesAgentLoop(
server=FailingServer(),
tool_schemas=basic_tools,
valid_tool_names=valid_names,
max_turns=10,
)
messages = [{"role": "user", "content": "Hi"}]
result = await agent.run(messages)
assert result.finished_naturally is False
assert result.turns_used == 1
@pytest.mark.asyncio
async def test_tools_passed_to_server(self, basic_tools, valid_names):
"""Verify tools are passed in the chat_completion kwargs."""
server = MockServer([make_text_response("OK")])
agent = HermesAgentLoop(
server=server,
tool_schemas=basic_tools,
valid_tool_names=valid_names,
max_turns=10,
)
messages = [{"role": "user", "content": "Hi"}]
await agent.run(messages)
assert len(server.call_history) == 1
assert "tools" in server.call_history[0]
assert server.call_history[0]["tools"] == basic_tools
@pytest.mark.asyncio
async def test_extra_body_forwarded(self, basic_tools, valid_names):
"""extra_body should be forwarded to server."""
extra = {"provider": {"ignore": ["DeepInfra"]}}
server = MockServer([make_text_response("OK")])
agent = HermesAgentLoop(
server=server,
tool_schemas=basic_tools,
valid_tool_names=valid_names,
max_turns=10,
extra_body=extra,
)
messages = [{"role": "user", "content": "Hi"}]
await agent.run(messages)
assert server.call_history[0].get("extra_body") == extra
@pytest.mark.asyncio
async def test_managed_state_returned(self, basic_tools, valid_names):
"""If server has get_state(), result should include managed_state."""
server = MockServer([make_text_response("OK")])
server.get_state = lambda: {"nodes": [{"test": True}]}
agent = HermesAgentLoop(
server=server,
tool_schemas=basic_tools,
valid_tool_names=valid_names,
max_turns=10,
)
messages = [{"role": "user", "content": "Hi"}]
result = await agent.run(messages)
assert result.managed_state is not None
assert "nodes" in result.managed_state
@pytest.mark.asyncio
async def test_no_managed_state_without_get_state(self, basic_tools, valid_names):
"""Regular server without get_state() should return None managed_state."""
server = MockServer([make_text_response("OK")])
agent = HermesAgentLoop(
server=server,
tool_schemas=basic_tools,
valid_tool_names=valid_names,
max_turns=10,
)
messages = [{"role": "user", "content": "Hi"}]
result = await agent.run(messages)
assert result.managed_state is None
@pytest.mark.asyncio
async def test_memory_tool_blocked(self, basic_tools):
"""Memory tool should return error in RL environments."""
valid = {"terminal", "read_file", "todo", "memory"}
server = MockServer([
make_tool_response("memory", {"action": "add", "target": "user", "content": "test"}),
make_text_response("Done"),
])
agent = HermesAgentLoop(
server=server,
tool_schemas=basic_tools,
valid_tool_names=valid,
max_turns=10,
)
messages = [{"role": "user", "content": "Remember this"}]
result = await agent.run(messages)
# Find the tool response
tool_msgs = [m for m in result.messages if m["role"] == "tool"]
assert len(tool_msgs) >= 1
tool_result = json.loads(tool_msgs[0]["content"])
assert "error" in tool_result
assert "not available" in tool_result["error"].lower()
@pytest.mark.asyncio
async def test_session_search_blocked(self, basic_tools):
"""session_search should return error in RL environments."""
valid = {"terminal", "read_file", "todo", "session_search"}
server = MockServer([
make_tool_response("session_search", {"query": "test"}),
make_text_response("Done"),
])
agent = HermesAgentLoop(
server=server,
tool_schemas=basic_tools,
valid_tool_names=valid,
max_turns=10,
)
messages = [{"role": "user", "content": "Search sessions"}]
result = await agent.run(messages)
tool_msgs = [m for m in result.messages if m["role"] == "tool"]
assert len(tool_msgs) >= 1
tool_result = json.loads(tool_msgs[0]["content"])
assert "error" in tool_result
@pytest.mark.asyncio
async def test_reasoning_content_preserved(self, basic_tools, valid_names):
"""Reasoning content should be extracted and preserved."""
resp = MockChatCompletion(
choices=[
MockChoice(
message=MockMessage(
content="The answer is 42.",
reasoning_content="Let me think about this step by step...",
)
)
]
)
server = MockServer([resp])
agent = HermesAgentLoop(
server=server,
tool_schemas=basic_tools,
valid_tool_names=valid_names,
max_turns=10,
)
messages = [{"role": "user", "content": "What is the meaning of life?"}]
result = await agent.run(messages)
assert len(result.reasoning_per_turn) == 1
assert result.reasoning_per_turn[0] == "Let me think about this step by step..."
class TestResizeToolPool:
def test_resize_works(self):
"""resize_tool_pool should not raise."""
resize_tool_pool(16) # Small pool for testing
resize_tool_pool(128) # Restore default
def test_resize_shuts_down_previous_executor(self, monkeypatch):
"""Replacing the global tool executor should shut down the old pool."""
import environments.agent_loop as agent_loop_module
old_executor = MagicMock()
new_executor = MagicMock()
monkeypatch.setattr(agent_loop_module, "_tool_executor", old_executor)
monkeypatch.setattr(
agent_loop_module.concurrent.futures,
"ThreadPoolExecutor",
MagicMock(return_value=new_executor),
)
resize_tool_pool(16)
old_executor.shutdown.assert_called_once_with(wait=False)
assert agent_loop_module._tool_executor is new_executor

View File

@@ -0,0 +1,552 @@
"""Integration tests for HermesAgentLoop tool calling.
Tests the full agent loop with real LLM calls via OpenRouter.
Uses stepfun/step-3.5-flash:free by default (zero cost), falls back
to anthropic/claude-sonnet-4 if the free model is unavailable.
These tests verify:
1. Single tool call: model calls a tool, gets result, responds
2. Multi-tool call: model calls multiple tools in one turn
3. Multi-turn: model calls tools across multiple turns
4. Unknown tool rejection: model calling a non-existent tool gets an error
5. Max turns: loop stops when max_turns is reached
6. No tools: model responds without calling any tools
7. Tool error handling: tool execution errors are captured
Run:
pytest tests/test_agent_loop_tool_calling.py -v
pytest tests/test_agent_loop_tool_calling.py -v -k "single" # run one test
"""
import asyncio
import json
import os
import sys
from pathlib import Path
from typing import Any, Dict, List, Set
from unittest.mock import patch
import pytest
# pytestmark removed — tests skip gracefully via OPENROUTER_API_KEY check on line 59
# Ensure repo root is importable
_repo_root = Path(__file__).resolve().parent.parent.parent
if str(_repo_root) not in sys.path:
sys.path.insert(0, str(_repo_root))
try:
from environments.agent_loop import AgentResult, HermesAgentLoop
from atroposlib.envs.server_handling.openai_server import OpenAIServer # noqa: F401
except ImportError:
pytest.skip("atroposlib not installed", allow_module_level=True)
# =========================================================================
# Test infrastructure
# =========================================================================
# Models to try, in order of preference (free first)
_MODELS = [
"stepfun/step-3.5-flash:free",
"google/gemini-2.0-flash-001",
"anthropic/claude-sonnet-4",
]
def _get_api_key():
key = os.getenv("OPENROUTER_API_KEY", "")
if not key:
pytest.skip("OPENROUTER_API_KEY not set")
return key
def _make_server(model: str = None):
"""Create an OpenAI server for testing."""
from atroposlib.envs.server_handling.openai_server import OpenAIServer
from atroposlib.envs.server_handling.server_manager import APIServerConfig
config = APIServerConfig(
base_url="https://openrouter.ai/api/v1",
model_name=model or _MODELS[0],
server_type="openai",
api_key=_get_api_key(),
health_check=False,
)
return OpenAIServer(config)
async def _try_models(test_fn):
"""Try running a test with each model until one works."""
last_error = None
for model in _MODELS:
try:
server = _make_server(model)
return await test_fn(server, model)
except Exception as e:
last_error = e
if "rate" in str(e).lower() or "limit" in str(e).lower():
continue # Rate limited, try next model
raise # Real error
pytest.skip(f"All models failed. Last error: {last_error}")
# =========================================================================
# Fake tools for testing
# =========================================================================
# Simple calculator tool
CALC_TOOL = {
"type": "function",
"function": {
"name": "calculate",
"description": "Calculate a math expression. Returns the numeric result.",
"parameters": {
"type": "object",
"properties": {
"expression": {
"type": "string",
"description": "Math expression to evaluate, e.g. '2 + 3'"
}
},
"required": ["expression"],
},
},
}
# Weather lookup tool
WEATHER_TOOL = {
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the current weather for a city. Returns temperature and conditions.",
"parameters": {
"type": "object",
"properties": {
"city": {
"type": "string",
"description": "City name, e.g. 'Tokyo'"
}
},
"required": ["city"],
},
},
}
# Lookup tool (always succeeds)
LOOKUP_TOOL = {
"type": "function",
"function": {
"name": "lookup",
"description": "Look up a fact. Returns a short answer string.",
"parameters": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "What to look up"
}
},
"required": ["query"],
},
},
}
# Error tool (always fails)
ERROR_TOOL = {
"type": "function",
"function": {
"name": "failing_tool",
"description": "A tool that always fails with an error.",
"parameters": {
"type": "object",
"properties": {
"input": {"type": "string"}
},
"required": ["input"],
},
},
}
def _fake_tool_handler(tool_name: str, args: Dict[str, Any], **kwargs) -> str:
"""Handle fake tool calls for testing."""
if tool_name == "calculate":
expr = args.get("expression", "0")
try:
# Safe eval for simple math
result = eval(expr, {"__builtins__": {}}, {})
return json.dumps({"result": result})
except Exception as e:
return json.dumps({"error": str(e)})
elif tool_name == "get_weather":
city = args.get("city", "Unknown")
# Return canned weather
return json.dumps({
"city": city,
"temperature": 22,
"conditions": "sunny",
"humidity": 45,
})
elif tool_name == "lookup":
query = args.get("query", "")
return json.dumps({"answer": f"The answer to '{query}' is 42."})
elif tool_name == "failing_tool":
raise RuntimeError("This tool always fails!")
return json.dumps({"error": f"Unknown tool: {tool_name}"})
# =========================================================================
# Tests
# =========================================================================
@pytest.mark.asyncio
async def test_single_tool_call():
"""Model should call a single tool, get the result, and respond."""
async def _run(server, model):
agent = HermesAgentLoop(
server=server,
tool_schemas=[WEATHER_TOOL],
valid_tool_names={"get_weather"},
max_turns=5,
temperature=0.0,
max_tokens=500,
)
messages = [
{"role": "user", "content": "What's the weather in Tokyo? Use the get_weather tool."},
]
with patch("environments.agent_loop.handle_function_call", side_effect=_fake_tool_handler):
result = await agent.run(messages)
assert isinstance(result, AgentResult)
assert result.turns_used >= 2, f"Expected at least 2 turns (tool call + response), got {result.turns_used}"
# Verify a tool call happened
tool_calls_found = False
for msg in result.messages:
if msg.get("role") == "assistant" and msg.get("tool_calls"):
for tc in msg["tool_calls"]:
if tc["function"]["name"] == "get_weather":
tool_calls_found = True
args = json.loads(tc["function"]["arguments"])
assert "city" in args
assert tool_calls_found, "Model should have called get_weather"
# Verify tool result is in conversation
tool_results = [m for m in result.messages if m.get("role") == "tool"]
assert len(tool_results) >= 1, "Should have at least one tool result"
# Verify the final response references the weather
final_msg = result.messages[-1]
assert final_msg["role"] == "assistant"
assert final_msg["content"], "Final response should have content"
return result
await _try_models(_run)
@pytest.mark.asyncio
async def test_multi_tool_single_turn():
"""Model should call multiple tools in a single turn."""
async def _run(server, model):
agent = HermesAgentLoop(
server=server,
tool_schemas=[WEATHER_TOOL, CALC_TOOL],
valid_tool_names={"get_weather", "calculate"},
max_turns=5,
temperature=0.0,
max_tokens=500,
)
messages = [
{"role": "user", "content": (
"I need two things at once: "
"1) What's the weather in Paris? Use get_weather. "
"2) What is 15 * 7? Use calculate. "
"Call BOTH tools in a single response."
)},
]
with patch("environments.agent_loop.handle_function_call", side_effect=_fake_tool_handler):
result = await agent.run(messages)
# Count distinct tools called
tools_called = set()
for msg in result.messages:
if msg.get("role") == "assistant" and msg.get("tool_calls"):
for tc in msg["tool_calls"]:
tools_called.add(tc["function"]["name"])
# At minimum, both tools should have been called (maybe in different turns)
assert "get_weather" in tools_called, f"get_weather not called. Called: {tools_called}"
assert "calculate" in tools_called, f"calculate not called. Called: {tools_called}"
return result
await _try_models(_run)
@pytest.mark.asyncio
async def test_multi_turn_conversation():
"""Agent should handle multiple turns of tool calls."""
async def _run(server, model):
agent = HermesAgentLoop(
server=server,
tool_schemas=[LOOKUP_TOOL, CALC_TOOL],
valid_tool_names={"lookup", "calculate"},
max_turns=10,
temperature=0.0,
max_tokens=500,
)
messages = [
{"role": "user", "content": (
"First, use the lookup tool to look up 'meaning of life'. "
"Then use calculate to compute 6 * 7. "
"Do these in separate tool calls, one at a time."
)},
]
with patch("environments.agent_loop.handle_function_call", side_effect=_fake_tool_handler):
result = await agent.run(messages)
# Should have used both tools
tools_called = set()
for msg in result.messages:
if msg.get("role") == "assistant" and msg.get("tool_calls"):
for tc in msg["tool_calls"]:
tools_called.add(tc["function"]["name"])
assert "lookup" in tools_called, f"lookup not called. Called: {tools_called}"
assert "calculate" in tools_called, f"calculate not called. Called: {tools_called}"
# Should finish naturally
assert result.finished_naturally, "Should finish naturally after answering"
return result
await _try_models(_run)
@pytest.mark.asyncio
async def test_unknown_tool_rejected():
"""If the model calls a tool not in valid_tool_names, it gets an error."""
async def _run(server, model):
# Only allow "calculate" but give schema for both
agent = HermesAgentLoop(
server=server,
tool_schemas=[CALC_TOOL, WEATHER_TOOL],
valid_tool_names={"calculate"}, # weather NOT allowed
max_turns=5,
temperature=0.0,
max_tokens=500,
)
messages = [
{"role": "user", "content": "What's the weather in London? Use get_weather."},
]
with patch("environments.agent_loop.handle_function_call", side_effect=_fake_tool_handler):
result = await agent.run(messages)
# Check if get_weather was called and rejected
if result.tool_errors:
weather_errors = [e for e in result.tool_errors if e.tool_name == "get_weather"]
assert len(weather_errors) > 0, "get_weather should have been rejected"
assert "Unknown tool" in weather_errors[0].error
return result
await _try_models(_run)
@pytest.mark.asyncio
async def test_max_turns_limit():
"""Agent should stop after max_turns even if model keeps calling tools."""
async def _run(server, model):
agent = HermesAgentLoop(
server=server,
tool_schemas=[LOOKUP_TOOL],
valid_tool_names={"lookup"},
max_turns=2, # Very low limit
temperature=0.0,
max_tokens=500,
)
messages = [
{"role": "user", "content": (
"Keep looking up facts. Look up 'fact 1', then 'fact 2', "
"then 'fact 3', then 'fact 4'. Do them one at a time."
)},
]
with patch("environments.agent_loop.handle_function_call", side_effect=_fake_tool_handler):
result = await agent.run(messages)
assert result.turns_used <= 2, f"Should stop at max_turns=2, used {result.turns_used}"
assert not result.finished_naturally, "Should NOT finish naturally (hit max_turns)"
return result
await _try_models(_run)
@pytest.mark.asyncio
async def test_no_tools_direct_response():
"""When no tools are useful, model should respond directly."""
async def _run(server, model):
agent = HermesAgentLoop(
server=server,
tool_schemas=[WEATHER_TOOL],
valid_tool_names={"get_weather"},
max_turns=5,
temperature=0.0,
max_tokens=200,
)
messages = [
{"role": "user", "content": "What is 2 + 2? Just answer directly, no tools needed."},
]
with patch("environments.agent_loop.handle_function_call", side_effect=_fake_tool_handler):
result = await agent.run(messages)
assert result.finished_naturally, "Should finish naturally with a direct response"
assert result.turns_used == 1, f"Should take exactly 1 turn for a direct answer, took {result.turns_used}"
final = result.messages[-1]
assert final["role"] == "assistant"
assert final["content"], "Should have text content"
assert "4" in final["content"], "Should contain the answer '4'"
return result
await _try_models(_run)
@pytest.mark.asyncio
async def test_tool_error_handling():
"""Tool execution errors should be captured and reported to the model."""
async def _run(server, model):
agent = HermesAgentLoop(
server=server,
tool_schemas=[ERROR_TOOL],
valid_tool_names={"failing_tool"},
max_turns=5,
temperature=0.0,
max_tokens=500,
)
messages = [
{"role": "user", "content": "Please call the failing_tool with input 'test'."},
]
with patch("environments.agent_loop.handle_function_call", side_effect=_fake_tool_handler):
result = await agent.run(messages)
# The tool error should be recorded
assert len(result.tool_errors) >= 1, "Should have at least one tool error"
assert "RuntimeError" in result.tool_errors[0].error or "always fails" in result.tool_errors[0].error
# The error should be in the conversation as a tool result
tool_results = [m for m in result.messages if m.get("role") == "tool"]
assert len(tool_results) >= 1
error_result = json.loads(tool_results[0]["content"])
assert "error" in error_result
return result
await _try_models(_run)
@pytest.mark.asyncio
async def test_agent_result_structure():
"""Verify the AgentResult has all expected fields populated."""
async def _run(server, model):
agent = HermesAgentLoop(
server=server,
tool_schemas=[CALC_TOOL],
valid_tool_names={"calculate"},
max_turns=5,
temperature=0.0,
max_tokens=300,
)
messages = [
{"role": "user", "content": "What is 3 + 4? Use the calculate tool."},
]
with patch("environments.agent_loop.handle_function_call", side_effect=_fake_tool_handler):
result = await agent.run(messages)
# Structural checks
assert isinstance(result, AgentResult)
assert isinstance(result.messages, list)
assert len(result.messages) >= 3, "Should have user + assistant(tool) + tool_result + assistant(final)"
assert isinstance(result.turns_used, int)
assert result.turns_used > 0
assert isinstance(result.finished_naturally, bool)
assert isinstance(result.tool_errors, list)
assert isinstance(result.reasoning_per_turn, list)
# Messages should follow OpenAI format
for msg in result.messages:
assert "role" in msg, f"Message missing 'role': {msg}"
assert msg["role"] in ("system", "user", "assistant", "tool"), f"Invalid role: {msg['role']}"
return result
await _try_models(_run)
@pytest.mark.asyncio
async def test_conversation_history_preserved():
"""The full conversation history should be in result.messages."""
async def _run(server, model):
agent = HermesAgentLoop(
server=server,
tool_schemas=[WEATHER_TOOL],
valid_tool_names={"get_weather"},
max_turns=5,
temperature=0.0,
max_tokens=500,
)
messages = [
{"role": "system", "content": "You are a helpful weather assistant."},
{"role": "user", "content": "What's the weather in Berlin? Use get_weather."},
]
with patch("environments.agent_loop.handle_function_call", side_effect=_fake_tool_handler):
result = await agent.run(messages)
# System message should be preserved
assert result.messages[0]["role"] == "system"
assert "weather assistant" in result.messages[0]["content"]
# User message should be preserved
assert result.messages[1]["role"] == "user"
assert "Berlin" in result.messages[1]["content"]
# Should have assistant + tool + assistant sequence
roles = [m["role"] for m in result.messages]
assert "tool" in roles, "Should have tool results in conversation"
return result
await _try_models(_run)

View File

@@ -0,0 +1,359 @@
"""Integration tests for HermesAgentLoop with a local vLLM server.
Tests the full Phase 2 flow: ManagedServer + tool calling with a real
vLLM backend, producing actual token IDs and logprobs for RL training.
Requires a running vLLM server. Start one from the atropos directory:
python -m example_trainer.vllm_api_server \
--model Qwen/Qwen3-4B-Thinking-2507 \
--port 9001 \
--gpu-memory-utilization 0.8 \
--max-model-len=32000
Tests are automatically skipped if the server is not reachable.
Run:
pytest tests/test_agent_loop_vllm.py -v
pytest tests/test_agent_loop_vllm.py -v -k "single"
"""
import asyncio
import json
import os
import sys
from pathlib import Path
from typing import Any, Dict
from unittest.mock import patch
import pytest
import requests
# Ensure repo root is importable
_repo_root = Path(__file__).resolve().parent.parent.parent
if str(_repo_root) not in sys.path:
sys.path.insert(0, str(_repo_root))
try:
from environments.agent_loop import AgentResult, HermesAgentLoop
except ImportError:
pytest.skip("atroposlib not installed", allow_module_level=True)
# =========================================================================
# Configuration
# =========================================================================
VLLM_HOST = "localhost"
VLLM_PORT = 9001
VLLM_BASE_URL = f"http://{VLLM_HOST}:{VLLM_PORT}"
VLLM_MODEL = "Qwen/Qwen3-4B-Thinking-2507"
def _vllm_is_running() -> bool:
"""Check if the vLLM server is reachable."""
try:
r = requests.get(f"{VLLM_BASE_URL}/health", timeout=3)
return r.status_code == 200
except Exception:
return False
# Skip all tests in this module if vLLM is not running
pytestmark = pytest.mark.skipif(
not _vllm_is_running(),
reason=(
f"vLLM server not reachable at {VLLM_BASE_URL}. "
"Start it with: python -m example_trainer.vllm_api_server "
f"--model {VLLM_MODEL} --port {VLLM_PORT} "
"--gpu-memory-utilization 0.8 --max-model-len=32000"
),
)
# =========================================================================
# Server setup
# =========================================================================
def _make_server_manager():
"""Create a ServerManager pointing to the local vLLM server."""
from atroposlib.envs.server_handling.server_manager import (
ServerManager,
APIServerConfig,
)
config = APIServerConfig(
base_url=VLLM_BASE_URL,
model_name=VLLM_MODEL,
server_type="vllm",
health_check=False,
)
sm = ServerManager([config], tool_parser="hermes")
sm.servers[0].server_healthy = True
return sm
def _get_tokenizer():
"""Load the tokenizer for the model."""
from transformers import AutoTokenizer
return AutoTokenizer.from_pretrained(VLLM_MODEL)
# =========================================================================
# Fake tools
# =========================================================================
WEATHER_TOOL = {
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the current weather for a city. Returns temperature and conditions.",
"parameters": {
"type": "object",
"properties": {
"city": {
"type": "string",
"description": "City name, e.g. 'Tokyo'",
}
},
"required": ["city"],
},
},
}
CALC_TOOL = {
"type": "function",
"function": {
"name": "calculate",
"description": "Calculate a math expression. Returns the numeric result.",
"parameters": {
"type": "object",
"properties": {
"expression": {
"type": "string",
"description": "Math expression, e.g. '2 + 3'",
}
},
"required": ["expression"],
},
},
}
def _fake_tool_handler(tool_name: str, args: Dict[str, Any], **kwargs) -> str:
"""Handle fake tool calls for testing."""
if tool_name == "get_weather":
city = args.get("city", "Unknown")
return json.dumps({
"city": city,
"temperature": 22,
"conditions": "sunny",
"humidity": 45,
})
elif tool_name == "calculate":
expr = args.get("expression", "0")
try:
result = eval(expr, {"__builtins__": {}}, {})
return json.dumps({"result": result})
except Exception as e:
return json.dumps({"error": str(e)})
return json.dumps({"error": f"Unknown tool: {tool_name}"})
# =========================================================================
# Tests
# =========================================================================
@pytest.mark.asyncio
async def test_vllm_single_tool_call():
"""vLLM model calls a tool, gets result, responds — full Phase 2 flow."""
sm = _make_server_manager()
tokenizer = _get_tokenizer()
async with sm.managed_server(tokenizer=tokenizer) as managed:
agent = HermesAgentLoop(
server=managed,
tool_schemas=[WEATHER_TOOL],
valid_tool_names={"get_weather"},
max_turns=5,
temperature=0.6,
max_tokens=1000,
)
messages = [
{"role": "user", "content": "What's the weather in Tokyo? Use the get_weather tool."},
]
with patch("environments.agent_loop.handle_function_call", side_effect=_fake_tool_handler):
result = await agent.run(messages)
assert isinstance(result, AgentResult)
assert result.turns_used >= 2, f"Expected at least 2 turns, got {result.turns_used}"
# Verify tool call happened
tool_calls_found = False
for msg in result.messages:
if msg.get("role") == "assistant" and msg.get("tool_calls"):
for tc in msg["tool_calls"]:
if tc["function"]["name"] == "get_weather":
tool_calls_found = True
args = json.loads(tc["function"]["arguments"])
assert "city" in args
assert tool_calls_found, "Model should have called get_weather"
# Verify tool results in conversation
tool_results = [m for m in result.messages if m.get("role") == "tool"]
assert len(tool_results) >= 1
@pytest.mark.asyncio
async def test_vllm_multi_tool_calls():
"""vLLM model calls multiple tools across turns."""
sm = _make_server_manager()
tokenizer = _get_tokenizer()
async with sm.managed_server(tokenizer=tokenizer) as managed:
agent = HermesAgentLoop(
server=managed,
tool_schemas=[WEATHER_TOOL, CALC_TOOL],
valid_tool_names={"get_weather", "calculate"},
max_turns=10,
temperature=0.6,
max_tokens=1000,
)
messages = [
{"role": "user", "content": (
"I need two things: "
"1) What's the weather in Paris? Use get_weather. "
"2) What is 15 * 7? Use calculate."
)},
]
with patch("environments.agent_loop.handle_function_call", side_effect=_fake_tool_handler):
result = await agent.run(messages)
# Both tools should be called
tools_called = set()
for msg in result.messages:
if msg.get("role") == "assistant" and msg.get("tool_calls"):
for tc in msg["tool_calls"]:
tools_called.add(tc["function"]["name"])
assert "get_weather" in tools_called, f"get_weather not called. Called: {tools_called}"
assert "calculate" in tools_called, f"calculate not called. Called: {tools_called}"
@pytest.mark.asyncio
async def test_vllm_managed_server_produces_nodes():
"""ManagedServer should produce SequenceNodes with tokens and logprobs."""
sm = _make_server_manager()
tokenizer = _get_tokenizer()
async with sm.managed_server(tokenizer=tokenizer) as managed:
agent = HermesAgentLoop(
server=managed,
tool_schemas=[WEATHER_TOOL],
valid_tool_names={"get_weather"},
max_turns=5,
temperature=0.6,
max_tokens=1000,
)
messages = [
{"role": "user", "content": "What's the weather in Berlin? Use get_weather."},
]
with patch("environments.agent_loop.handle_function_call", side_effect=_fake_tool_handler):
result = await agent.run(messages)
# Get the managed state — should have SequenceNodes
state = managed.get_state()
assert state is not None, "ManagedServer should return state"
nodes = state.get("nodes", [])
assert len(nodes) >= 1, f"Should have at least 1 node, got {len(nodes)}"
node = nodes[0]
assert hasattr(node, "tokens"), "Node should have tokens"
assert hasattr(node, "logprobs"), "Node should have logprobs"
assert len(node.tokens) > 0, "Tokens should not be empty"
assert len(node.logprobs) > 0, "Logprobs should not be empty"
assert len(node.tokens) == len(node.logprobs), (
f"Tokens ({len(node.tokens)}) and logprobs ({len(node.logprobs)}) should have same length"
)
@pytest.mark.asyncio
async def test_vllm_no_tools_direct_response():
"""vLLM model should respond directly when no tools are needed."""
sm = _make_server_manager()
tokenizer = _get_tokenizer()
async with sm.managed_server(tokenizer=tokenizer) as managed:
agent = HermesAgentLoop(
server=managed,
tool_schemas=[WEATHER_TOOL],
valid_tool_names={"get_weather"},
max_turns=5,
temperature=0.6,
max_tokens=500,
)
messages = [
{"role": "user", "content": "What is 2 + 2? Answer directly, no tools."},
]
with patch("environments.agent_loop.handle_function_call", side_effect=_fake_tool_handler):
result = await agent.run(messages)
assert result.finished_naturally, "Should finish naturally"
assert result.turns_used == 1, f"Should take 1 turn, took {result.turns_used}"
final = result.messages[-1]
assert final["role"] == "assistant"
assert final["content"], "Should have content"
@pytest.mark.asyncio
async def test_vllm_thinking_content_extracted():
"""Qwen3-Thinking model should produce reasoning content."""
sm = _make_server_manager()
tokenizer = _get_tokenizer()
async with sm.managed_server(
tokenizer=tokenizer,
preserve_think_blocks=True,
) as managed:
agent = HermesAgentLoop(
server=managed,
tool_schemas=[CALC_TOOL],
valid_tool_names={"calculate"},
max_turns=5,
temperature=0.6,
max_tokens=1000,
)
messages = [
{"role": "user", "content": "What is 123 * 456? Use the calculate tool."},
]
with patch("environments.agent_loop.handle_function_call", side_effect=_fake_tool_handler):
result = await agent.run(messages)
# Qwen3-Thinking should generate <think> blocks
# Check if any content contains thinking markers
has_thinking = False
for msg in result.messages:
content = msg.get("content", "") or ""
if "<think>" in content or "</think>" in content:
has_thinking = True
break
# Also check reasoning_per_turn
has_reasoning = any(r for r in result.reasoning_per_turn if r)
# At least one of these should be true for a thinking model
assert has_thinking or has_reasoning, (
"Qwen3-Thinking should produce <think> blocks or reasoning content"
)

View File

@@ -0,0 +1,490 @@
"""Tests for Anthropic error handling in the agent retry loop.
Covers all error paths in run_agent.py's run_conversation() for api_mode=anthropic_messages:
- 429 rate limit → retried with backoff
- 529 overloaded → retried with backoff
- 400 bad request → non-retryable, immediate fail
- 401 unauthorized → credential refresh + retry
- 500 server error → retried with backoff
- "prompt is too long" → context length error triggers compression
"""
import asyncio
import sys
import types
from types import SimpleNamespace
from unittest.mock import MagicMock, AsyncMock
import pytest
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 gateway.run as gateway_run
import run_agent
from gateway.config import Platform
from gateway.session import SessionSource
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _patch_agent_bootstrap(monkeypatch):
monkeypatch.setattr(
run_agent,
"get_tool_definitions",
lambda **kwargs: [
{
"type": "function",
"function": {
"name": "terminal",
"description": "Run shell commands.",
"parameters": {"type": "object", "properties": {}},
},
}
],
)
monkeypatch.setattr(run_agent, "check_toolset_requirements", lambda: {})
def _anthropic_response(text: str):
"""Simulate an Anthropic messages.create() response object."""
return SimpleNamespace(
content=[SimpleNamespace(type="text", text=text)],
stop_reason="end_turn",
usage=SimpleNamespace(input_tokens=10, output_tokens=5),
model="claude-sonnet-4-6-20250514",
)
class _RateLimitError(Exception):
"""Simulates Anthropic 429 rate limit error."""
def __init__(self):
super().__init__("Error code: 429 - Rate limit exceeded. Please retry after 30s.")
self.status_code = 429
class _OverloadedError(Exception):
"""Simulates Anthropic 529 overloaded error."""
def __init__(self):
super().__init__("Error code: 529 - API is temporarily overloaded.")
self.status_code = 529
class _BadRequestError(Exception):
"""Simulates Anthropic 400 bad request error (non-retryable)."""
def __init__(self):
super().__init__("Error code: 400 - Invalid model specified.")
self.status_code = 400
class _UnauthorizedError(Exception):
"""Simulates Anthropic 401 unauthorized error."""
def __init__(self):
super().__init__("Error code: 401 - Unauthorized. Invalid API key.")
self.status_code = 401
class _ServerError(Exception):
"""Simulates Anthropic 500 internal server error."""
def __init__(self):
super().__init__("Error code: 500 - Internal server error.")
self.status_code = 500
class _PromptTooLongError(Exception):
"""Simulates Anthropic prompt-too-long error (triggers context compression)."""
def __init__(self):
super().__init__("prompt is too long: 250000 tokens > 200000 maximum")
self.status_code = 400
class _FakeAnthropicClient:
def close(self):
pass
def _fake_build_anthropic_client(key, base_url=None):
return _FakeAnthropicClient()
def _make_agent_cls(error_cls, recover_after=None):
"""Create an AIAgent subclass that raises error_cls on API calls.
If recover_after is set, the agent succeeds after that many failures.
"""
class _Agent(run_agent.AIAgent):
def __init__(self, *args, **kwargs):
kwargs.setdefault("skip_context_files", True)
kwargs.setdefault("skip_memory", True)
kwargs.setdefault("max_iterations", 4)
super().__init__(*args, **kwargs)
self._cleanup_task_resources = lambda task_id: None
self._persist_session = lambda messages, history=None: None
self._save_trajectory = lambda messages, user_message, completed: None
self._save_session_log = lambda messages: None
def run_conversation(self, user_message, conversation_history=None, task_id=None):
calls = {"n": 0}
def _fake_api_call(api_kwargs):
calls["n"] += 1
if recover_after is not None and calls["n"] > recover_after:
return _anthropic_response("Recovered")
raise error_cls()
self._interruptible_api_call = _fake_api_call
return super().run_conversation(
user_message, conversation_history=conversation_history, task_id=task_id
)
return _Agent
def _run_with_agent(monkeypatch, agent_cls):
"""Run _run_agent through the gateway with the given agent class."""
_patch_agent_bootstrap(monkeypatch)
monkeypatch.setattr(
"agent.anthropic_adapter.build_anthropic_client", _fake_build_anthropic_client
)
monkeypatch.setattr(run_agent, "AIAgent", agent_cls)
monkeypatch.setattr(
gateway_run,
"_resolve_runtime_agent_kwargs",
lambda: {
"provider": "anthropic",
"api_mode": "anthropic_messages",
"base_url": "https://api.anthropic.com",
"api_key": "sk-ant-api03-test-key",
},
)
monkeypatch.setenv("HERMES_TOOL_PROGRESS", "false")
runner = gateway_run.GatewayRunner.__new__(gateway_run.GatewayRunner)
runner.adapters = {}
runner._ephemeral_system_prompt = ""
runner._prefill_messages = []
runner._reasoning_config = None
runner._provider_routing = {}
runner._fallback_model = None
runner._running_agents = {}
runner.hooks = MagicMock()
runner.hooks.emit = AsyncMock()
runner.hooks.loaded_hooks = []
runner._session_db = None
source = SessionSource(
platform=Platform.LOCAL,
chat_id="cli",
chat_name="CLI",
chat_type="dm",
user_id="test-user-1",
)
return asyncio.run(
runner._run_agent(
message="hello",
context_prompt="",
history=[],
source=source,
session_id="test-session",
session_key="agent:main:local:dm",
)
)
# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------
def test_429_rate_limit_is_retried_and_recovers(monkeypatch):
"""429 should be retried with backoff. First call fails, second succeeds."""
agent_cls = _make_agent_cls(_RateLimitError, recover_after=1)
result = _run_with_agent(monkeypatch, agent_cls)
assert result["final_response"] == "Recovered"
def test_529_overloaded_is_retried_and_recovers(monkeypatch):
"""529 should be retried with backoff. First call fails, second succeeds."""
agent_cls = _make_agent_cls(_OverloadedError, recover_after=1)
result = _run_with_agent(monkeypatch, agent_cls)
assert result["final_response"] == "Recovered"
def test_429_exhausts_all_retries_before_raising(monkeypatch):
"""429 must retry max_retries times, then return a failed result.
The agent no longer re-raises after exhausting retries — it returns a
result dict with the error in final_response. This changed when the
fallback-provider feature was added (the agent tries a fallback before
giving up, and returns a result dict either way).
"""
agent_cls = _make_agent_cls(_RateLimitError) # always fails
result = _run_with_agent(monkeypatch, agent_cls)
resp = str(result.get("final_response", ""))
assert "429" in resp or "retries" in resp.lower()
def test_400_bad_request_is_non_retryable(monkeypatch):
"""400 should fail immediately with only 1 API call (regression guard)."""
agent_cls = _make_agent_cls(_BadRequestError)
result = _run_with_agent(monkeypatch, agent_cls)
assert result["api_calls"] == 1
assert "400" in str(result.get("final_response", ""))
def test_500_server_error_is_retried_and_recovers(monkeypatch):
"""500 should be retried with backoff. First call fails, second succeeds."""
agent_cls = _make_agent_cls(_ServerError, recover_after=1)
result = _run_with_agent(monkeypatch, agent_cls)
assert result["final_response"] == "Recovered"
def test_401_credential_refresh_recovers(monkeypatch):
"""401 should trigger credential refresh and retry once."""
_patch_agent_bootstrap(monkeypatch)
monkeypatch.setattr(
"agent.anthropic_adapter.build_anthropic_client", _fake_build_anthropic_client
)
monkeypatch.setenv("HERMES_TOOL_PROGRESS", "false")
refresh_count = {"n": 0}
class _Auth401ThenSuccessAgent(run_agent.AIAgent):
def __init__(self, *args, **kwargs):
kwargs.setdefault("skip_context_files", True)
kwargs.setdefault("skip_memory", True)
kwargs.setdefault("max_iterations", 4)
super().__init__(*args, **kwargs)
self._cleanup_task_resources = lambda task_id: None
self._persist_session = lambda messages, history=None: None
self._save_trajectory = lambda messages, user_message, completed: None
self._save_session_log = lambda messages: None
def _try_refresh_anthropic_client_credentials(self) -> bool:
refresh_count["n"] += 1
return True # Simulate successful credential refresh
def run_conversation(self, user_message, conversation_history=None, task_id=None):
calls = {"n": 0}
def _fake_api_call(api_kwargs):
calls["n"] += 1
if calls["n"] == 1:
raise _UnauthorizedError()
return _anthropic_response("Auth refreshed")
self._interruptible_api_call = _fake_api_call
# Also patch streaming path — run_conversation now prefers
# streaming for health checking even without stream consumers.
self._interruptible_streaming_api_call = lambda api_kwargs, **kw: _fake_api_call(api_kwargs)
return super().run_conversation(
user_message, conversation_history=conversation_history, task_id=task_id
)
monkeypatch.setattr(run_agent, "AIAgent", _Auth401ThenSuccessAgent)
monkeypatch.setattr(
gateway_run,
"_resolve_runtime_agent_kwargs",
lambda: {
"provider": "anthropic",
"api_mode": "anthropic_messages",
"base_url": "https://api.anthropic.com",
"api_key": "sk-ant-api03-test-key",
},
)
runner = gateway_run.GatewayRunner.__new__(gateway_run.GatewayRunner)
runner.adapters = {}
runner._ephemeral_system_prompt = ""
runner._prefill_messages = []
runner._reasoning_config = None
runner._provider_routing = {}
runner._fallback_model = None
runner._running_agents = {}
runner.hooks = MagicMock()
runner.hooks.emit = AsyncMock()
runner.hooks.loaded_hooks = []
runner._session_db = None
source = SessionSource(
platform=Platform.LOCAL, chat_id="cli", chat_name="CLI",
chat_type="dm", user_id="test-user-1",
)
result = asyncio.run(
runner._run_agent(
message="hello", context_prompt="", history=[],
source=source, session_id="session-401",
session_key="agent:main:local:dm",
)
)
assert result["final_response"] == "Auth refreshed"
assert refresh_count["n"] == 1
def test_401_refresh_fails_is_non_retryable(monkeypatch):
"""401 with failed credential refresh should be treated as non-retryable."""
_patch_agent_bootstrap(monkeypatch)
monkeypatch.setattr(
"agent.anthropic_adapter.build_anthropic_client", _fake_build_anthropic_client
)
monkeypatch.setenv("HERMES_TOOL_PROGRESS", "false")
class _Auth401AlwaysFailAgent(run_agent.AIAgent):
def __init__(self, *args, **kwargs):
kwargs.setdefault("skip_context_files", True)
kwargs.setdefault("skip_memory", True)
kwargs.setdefault("max_iterations", 4)
super().__init__(*args, **kwargs)
self._cleanup_task_resources = lambda task_id: None
self._persist_session = lambda messages, history=None: None
self._save_trajectory = lambda messages, user_message, completed: None
self._save_session_log = lambda messages: None
def _try_refresh_anthropic_client_credentials(self) -> bool:
return False # Simulate failed credential refresh
def run_conversation(self, user_message, conversation_history=None, task_id=None):
def _fake_api_call(api_kwargs):
raise _UnauthorizedError()
self._interruptible_api_call = _fake_api_call
return super().run_conversation(
user_message, conversation_history=conversation_history, task_id=task_id
)
monkeypatch.setattr(run_agent, "AIAgent", _Auth401AlwaysFailAgent)
monkeypatch.setattr(
gateway_run,
"_resolve_runtime_agent_kwargs",
lambda: {
"provider": "anthropic",
"api_mode": "anthropic_messages",
"base_url": "https://api.anthropic.com",
"api_key": "sk-ant-api03-test-key",
},
)
runner = gateway_run.GatewayRunner.__new__(gateway_run.GatewayRunner)
runner.adapters = {}
runner._ephemeral_system_prompt = ""
runner._prefill_messages = []
runner._reasoning_config = None
runner._provider_routing = {}
runner._fallback_model = None
runner._running_agents = {}
runner.hooks = MagicMock()
runner.hooks.emit = AsyncMock()
runner.hooks.loaded_hooks = []
runner._session_db = None
source = SessionSource(
platform=Platform.LOCAL, chat_id="cli", chat_name="CLI",
chat_type="dm", user_id="test-user-1",
)
result = asyncio.run(
runner._run_agent(
message="hello", context_prompt="", history=[],
source=source, session_id="session-401-fail",
session_key="agent:main:local:dm",
)
)
# 401 after failed refresh → non-retryable (falls through to is_client_error)
assert result["api_calls"] == 1
assert "401" in str(result.get("final_response", "")) or "unauthorized" in str(result.get("final_response", "")).lower()
def test_prompt_too_long_triggers_compression(monkeypatch):
"""Anthropic 'prompt is too long' error should trigger context compression, not immediate fail."""
_patch_agent_bootstrap(monkeypatch)
monkeypatch.setattr(
"agent.anthropic_adapter.build_anthropic_client", _fake_build_anthropic_client
)
monkeypatch.setenv("HERMES_TOOL_PROGRESS", "false")
class _PromptTooLongThenSuccessAgent(run_agent.AIAgent):
compress_called = 0
def __init__(self, *args, **kwargs):
kwargs.setdefault("skip_context_files", True)
kwargs.setdefault("skip_memory", True)
kwargs.setdefault("max_iterations", 4)
super().__init__(*args, **kwargs)
self._cleanup_task_resources = lambda task_id: None
self._persist_session = lambda messages, history=None: None
self._save_trajectory = lambda messages, user_message, completed: None
self._save_session_log = lambda messages: None
def _compress_context(self, messages, system_message, approx_tokens=0, task_id=None):
type(self).compress_called += 1
# Simulate compression by dropping oldest non-system message
if len(messages) > 2:
compressed = [messages[0]] + messages[2:]
else:
compressed = messages
return compressed, system_message
def run_conversation(self, user_message, conversation_history=None, task_id=None):
calls = {"n": 0}
def _fake_api_call(api_kwargs):
calls["n"] += 1
if calls["n"] == 1:
raise _PromptTooLongError()
return _anthropic_response("Compressed and recovered")
self._interruptible_api_call = _fake_api_call
return super().run_conversation(
user_message, conversation_history=conversation_history, task_id=task_id
)
_PromptTooLongThenSuccessAgent.compress_called = 0
monkeypatch.setattr(run_agent, "AIAgent", _PromptTooLongThenSuccessAgent)
monkeypatch.setattr(
gateway_run,
"_resolve_runtime_agent_kwargs",
lambda: {
"provider": "anthropic",
"api_mode": "anthropic_messages",
"base_url": "https://api.anthropic.com",
"api_key": "sk-ant-api03-test-key",
},
)
runner = gateway_run.GatewayRunner.__new__(gateway_run.GatewayRunner)
runner.adapters = {}
runner._ephemeral_system_prompt = ""
runner._prefill_messages = []
runner._reasoning_config = None
runner._provider_routing = {}
runner._fallback_model = None
runner._running_agents = {}
runner.hooks = MagicMock()
runner.hooks.emit = AsyncMock()
runner.hooks.loaded_hooks = []
runner._session_db = None
source = SessionSource(
platform=Platform.LOCAL, chat_id="cli", chat_name="CLI",
chat_type="dm", user_id="test-user-1",
)
result = asyncio.run(
runner._run_agent(
message="hello", context_prompt="", history=[],
source=source, session_id="session-prompt-long",
session_key="agent:main:local:dm",
)
)
assert result["final_response"] == "Compressed and recovered"
assert _PromptTooLongThenSuccessAgent.compress_called >= 1

View File

@@ -0,0 +1,162 @@
"""Tests for the AsyncHttpxClientWrapper.__del__ neuter fix.
The OpenAI SDK's ``AsyncHttpxClientWrapper.__del__`` schedules
``aclose()`` via ``asyncio.get_running_loop().create_task()``. When GC
fires during CLI idle time, prompt_toolkit's event loop picks up the task
and crashes with "Event loop is closed" because the underlying TCP
transport is bound to a dead worker loop.
The three-layer defence:
1. ``neuter_async_httpx_del()`` replaces ``__del__`` with a no-op.
2. A custom asyncio exception handler silences residual errors.
3. ``cleanup_stale_async_clients()`` evicts stale cache entries.
"""
import asyncio
import threading
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
# ---------------------------------------------------------------------------
# Layer 1: neuter_async_httpx_del
# ---------------------------------------------------------------------------
class TestNeuterAsyncHttpxDel:
"""Verify neuter_async_httpx_del replaces __del__ on the SDK class."""
def test_del_becomes_noop(self):
"""After neuter, __del__ should do nothing (no RuntimeError)."""
from agent.auxiliary_client import neuter_async_httpx_del
try:
from openai._base_client import AsyncHttpxClientWrapper
except ImportError:
pytest.skip("openai SDK not installed")
# Save original so we can restore
original_del = AsyncHttpxClientWrapper.__del__
try:
neuter_async_httpx_del()
# The patched __del__ should be a no-op lambda
assert AsyncHttpxClientWrapper.__del__ is not original_del
# Calling it should not raise, even without a running loop
wrapper = MagicMock(spec=AsyncHttpxClientWrapper)
AsyncHttpxClientWrapper.__del__(wrapper) # Should be silent
finally:
# Restore original to avoid leaking into other tests
AsyncHttpxClientWrapper.__del__ = original_del
def test_neuter_idempotent(self):
"""Calling neuter twice doesn't break anything."""
from agent.auxiliary_client import neuter_async_httpx_del
try:
from openai._base_client import AsyncHttpxClientWrapper
except ImportError:
pytest.skip("openai SDK not installed")
original_del = AsyncHttpxClientWrapper.__del__
try:
neuter_async_httpx_del()
first_del = AsyncHttpxClientWrapper.__del__
neuter_async_httpx_del()
second_del = AsyncHttpxClientWrapper.__del__
# Both calls should succeed; the class should have a no-op
assert first_del is not original_del
assert second_del is not original_del
finally:
AsyncHttpxClientWrapper.__del__ = original_del
def test_neuter_graceful_without_sdk(self):
"""neuter_async_httpx_del doesn't raise if the openai SDK isn't installed."""
from agent.auxiliary_client import neuter_async_httpx_del
with patch.dict("sys.modules", {"openai._base_client": None}):
# Should not raise
neuter_async_httpx_del()
# ---------------------------------------------------------------------------
# Layer 3: cleanup_stale_async_clients
# ---------------------------------------------------------------------------
class TestCleanupStaleAsyncClients:
"""Verify stale cache entries are evicted and force-closed."""
def test_removes_stale_entries(self):
"""Entries with a closed loop should be evicted."""
from agent.auxiliary_client import (
_client_cache,
_client_cache_lock,
cleanup_stale_async_clients,
)
# Create a loop, close it, make a cache entry
loop = asyncio.new_event_loop()
loop.close()
mock_client = MagicMock()
# Give it _client attribute for _force_close_async_httpx
mock_client._client = MagicMock()
mock_client._client.is_closed = False
key = ("test_stale", True, "", "", id(loop))
with _client_cache_lock:
_client_cache[key] = (mock_client, "test-model", loop)
try:
cleanup_stale_async_clients()
with _client_cache_lock:
assert key not in _client_cache, "Stale entry should be removed"
finally:
# Clean up in case test fails
with _client_cache_lock:
_client_cache.pop(key, None)
def test_keeps_live_entries(self):
"""Entries with an open loop should be preserved."""
from agent.auxiliary_client import (
_client_cache,
_client_cache_lock,
cleanup_stale_async_clients,
)
loop = asyncio.new_event_loop() # NOT closed
mock_client = MagicMock()
key = ("test_live", True, "", "", id(loop))
with _client_cache_lock:
_client_cache[key] = (mock_client, "test-model", loop)
try:
cleanup_stale_async_clients()
with _client_cache_lock:
assert key in _client_cache, "Live entry should be preserved"
finally:
loop.close()
with _client_cache_lock:
_client_cache.pop(key, None)
def test_keeps_entries_without_loop(self):
"""Sync entries (cached_loop=None) should be preserved."""
from agent.auxiliary_client import (
_client_cache,
_client_cache_lock,
cleanup_stale_async_clients,
)
mock_client = MagicMock()
key = ("test_sync", False, "", "", 0)
with _client_cache_lock:
_client_cache[key] = (mock_client, "test-model", None)
try:
cleanup_stale_async_clients()
with _client_cache_lock:
assert key in _client_cache, "Sync entry should be preserved"
finally:
with _client_cache_lock:
_client_cache.pop(key, None)

View File

@@ -0,0 +1,199 @@
"""Tests for context compression boundary alignment.
Verifies that _align_boundary_backward correctly handles tool result groups
so that parallel tool calls are never split during compression.
"""
import pytest
from unittest.mock import patch, MagicMock
from agent.context_compressor import ContextCompressor
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _tc(call_id: str) -> dict:
"""Create a minimal tool_call dict."""
return {"id": call_id, "type": "function", "function": {"name": "test", "arguments": "{}"}}
def _tool_result(call_id: str, content: str = "result") -> dict:
"""Create a tool result message."""
return {"role": "tool", "tool_call_id": call_id, "content": content}
def _assistant_with_tools(*call_ids: str) -> dict:
"""Create an assistant message with tool_calls."""
return {"role": "assistant", "tool_calls": [_tc(cid) for cid in call_ids], "content": None}
def _make_compressor(**kwargs) -> ContextCompressor:
defaults = dict(
model="test-model",
threshold_percent=0.75,
protect_first_n=3,
protect_last_n=4,
quiet_mode=True,
)
defaults.update(kwargs)
with patch("agent.context_compressor.get_model_context_length", return_value=8000):
return ContextCompressor(**defaults)
# ---------------------------------------------------------------------------
# _align_boundary_backward tests
# ---------------------------------------------------------------------------
class TestAlignBoundaryBackward:
"""Test that compress-end boundary never splits a tool_call/result group."""
def test_boundary_at_clean_position(self):
"""Boundary after a user message — no adjustment needed."""
comp = _make_compressor()
messages = [
{"role": "system", "content": "sys"},
{"role": "user", "content": "hello"},
{"role": "assistant", "content": "hi"},
{"role": "user", "content": "do something"},
_assistant_with_tools("tc_1"),
_tool_result("tc_1", "done"),
{"role": "user", "content": "thanks"}, # idx=6
{"role": "assistant", "content": "np"},
]
# Boundary at 7, messages[6] = user — no adjustment
assert comp._align_boundary_backward(messages, 7) == 7
def test_boundary_after_assistant_with_tools(self):
"""Original case: boundary right after assistant with tool_calls."""
comp = _make_compressor()
messages = [
{"role": "system", "content": "sys"},
{"role": "user", "content": "hello"},
{"role": "assistant", "content": "hi"},
_assistant_with_tools("tc_1", "tc_2"), # idx=3
_tool_result("tc_1"), # idx=4
_tool_result("tc_2"), # idx=5
{"role": "user", "content": "next"},
]
# Boundary at 4, messages[3] = assistant with tool_calls → pull back to 3
assert comp._align_boundary_backward(messages, 4) == 3
def test_boundary_in_middle_of_tool_results(self):
"""THE BUG: boundary falls between tool results of the same group."""
comp = _make_compressor()
messages = [
{"role": "system", "content": "sys"},
{"role": "user", "content": "hello"},
{"role": "assistant", "content": "hi"},
{"role": "user", "content": "do 5 things"},
_assistant_with_tools("tc_A", "tc_B", "tc_C", "tc_D", "tc_E"), # idx=4
_tool_result("tc_A", "result A"), # idx=5
_tool_result("tc_B", "result B"), # idx=6
_tool_result("tc_C", "result C"), # idx=7
_tool_result("tc_D", "result D"), # idx=8
_tool_result("tc_E", "result E"), # idx=9
{"role": "user", "content": "ok"},
{"role": "assistant", "content": "done"},
]
# Boundary at 8 — in middle of tool results. messages[7] = tool result.
# Must walk back to idx=4 (the parent assistant).
assert comp._align_boundary_backward(messages, 8) == 4
def test_boundary_at_last_tool_result(self):
"""Boundary right after last tool result — messages[idx-1] is tool."""
comp = _make_compressor()
messages = [
{"role": "system", "content": "sys"},
{"role": "user", "content": "hello"},
{"role": "assistant", "content": "hi"},
_assistant_with_tools("tc_1", "tc_2", "tc_3"), # idx=3
_tool_result("tc_1"), # idx=4
_tool_result("tc_2"), # idx=5
_tool_result("tc_3"), # idx=6
{"role": "user", "content": "next"},
]
# Boundary at 7 — messages[6] is last tool result.
# Walk back: [6]=tool, [5]=tool, [4]=tool, [3]=assistant with tools → idx=3
assert comp._align_boundary_backward(messages, 7) == 3
def test_boundary_with_consecutive_tool_groups(self):
"""Two consecutive tool groups — only walk back to the nearest parent."""
comp = _make_compressor()
messages = [
{"role": "system", "content": "sys"},
{"role": "user", "content": "hello"},
_assistant_with_tools("tc_1"), # idx=2
_tool_result("tc_1"), # idx=3
{"role": "user", "content": "more"},
_assistant_with_tools("tc_2", "tc_3"), # idx=5
_tool_result("tc_2"), # idx=6
_tool_result("tc_3"), # idx=7
{"role": "user", "content": "done"},
]
# Boundary at 7 — messages[6] = tool result for tc_2 group
# Walk back: [6]=tool, [5]=assistant with tools → idx=5
assert comp._align_boundary_backward(messages, 7) == 5
# ---------------------------------------------------------------------------
# End-to-end: compression must not lose tool results
# ---------------------------------------------------------------------------
class TestCompressionToolResultPreservation:
"""Verify that compress() never silently drops tool results."""
def test_parallel_tool_results_not_lost(self):
"""The exact scenario that triggered silent data loss before the fix."""
comp = _make_compressor(protect_first_n=3, protect_last_n=4)
messages = [
{"role": "system", "content": "You are helpful."}, # 0
{"role": "user", "content": "Hello"}, # 1
{"role": "assistant", "content": "Hi there!"}, # 2 (end of head)
{"role": "user", "content": "Read 7 files for me"}, # 3
_assistant_with_tools("tc_A", "tc_B", "tc_C", "tc_D", "tc_E", "tc_F", "tc_G"), # 4
_tool_result("tc_A", "content of file A"), # 5
_tool_result("tc_B", "content of file B"), # 6
_tool_result("tc_C", "content of file C"), # 7
_tool_result("tc_D", "content of file D"), # 8
_tool_result("tc_E", "content of file E"), # 9
_tool_result("tc_F", "content of file F"), # 10
_tool_result("tc_G", "CRITICAL DATA in file G"), # 11 ← compress_end=15-4=11
{"role": "user", "content": "Now summarize them"}, # 12
{"role": "assistant", "content": "Here is the summary..."}, # 13
{"role": "user", "content": "Thanks"}, # 14
]
# 15 messages. compress_end = 15 - 4 = 11 (before fix: splits tool group)
fake_summary = "[Summary of earlier conversation]"
with patch.object(comp, "_generate_summary", return_value=fake_summary):
result = comp.compress(messages, current_tokens=7000)
# After compression, no tool results should be orphaned/lost.
# All tool results in the result must have a matching assistant tool_call.
assistant_call_ids = set()
for msg in result:
if msg.get("role") == "assistant":
for tc in msg.get("tool_calls") or []:
cid = tc.get("id", "")
if cid:
assistant_call_ids.add(cid)
tool_result_ids = set()
for msg in result:
if msg.get("role") == "tool":
cid = msg.get("tool_call_id")
if cid:
tool_result_ids.add(cid)
# Every tool result must have a parent — no orphans
orphaned = tool_result_ids - assistant_call_ids
assert not orphaned, f"Orphaned tool results found (data loss!): {orphaned}"
# Every assistant tool_call must have a real result (not a stub)
for msg in result:
if msg.get("role") == "tool":
assert msg["content"] != "[Result from earlier conversation — see context summary above]", \
f"Stub result found for {msg.get('tool_call_id')} — real result was lost"

View File

@@ -0,0 +1,202 @@
"""Tests for context compression persistence in the gateway.
Verifies that when context compression fires during run_conversation(),
the compressed messages are properly persisted to both SQLite (via the
agent) and JSONL (via the gateway).
Bug scenario (pre-fix):
1. Gateway loads 200-message history, passes to agent
2. Agent's run_conversation() compresses to ~30 messages mid-run
3. _compress_context() resets _last_flushed_db_idx = 0
4. On exit, _flush_messages_to_session_db() calculates:
flush_from = max(len(conversation_history=200), _last_flushed_db_idx=0) = 200
5. messages[200:] is empty (only ~30 messages after compression)
6. Nothing written to new session's SQLite — compressed context lost
7. Gateway's history_offset was still 200, producing empty new_messages
8. Fallback wrote only user/assistant pair — summary lost
"""
import os
import tempfile
from pathlib import Path
from unittest.mock import MagicMock, patch
import pytest
# ---------------------------------------------------------------------------
# Part 1: Agent-side — _flush_messages_to_session_db after compression
# ---------------------------------------------------------------------------
class TestFlushAfterCompression:
"""Verify that compressed messages are flushed to the new session's SQLite
even when conversation_history (from the original session) is longer than
the compressed messages list."""
def _make_agent(self, session_db):
with patch.dict(os.environ, {"OPENROUTER_API_KEY": "test-key"}):
from run_agent import AIAgent
agent = AIAgent(
model="test/model",
quiet_mode=True,
session_db=session_db,
session_id="original-session",
skip_context_files=True,
skip_memory=True,
)
return agent
def test_flush_after_compression_with_long_history(self):
"""The actual bug: conversation_history longer than compressed messages.
Before the fix, flush_from = max(len(conversation_history), 0) = 200,
but messages only has ~30 entries, so messages[200:] is empty.
After the fix, conversation_history is cleared to None after compression,
so flush_from = max(0, 0) = 0, and ALL compressed messages are written.
"""
from hermes_state import SessionDB
with tempfile.TemporaryDirectory() as tmpdir:
db_path = Path(tmpdir) / "test.db"
db = SessionDB(db_path=db_path)
agent = self._make_agent(db)
# Simulate the original long history (200 messages)
original_history = [
{"role": "user" if i % 2 == 0 else "assistant",
"content": f"message {i}"}
for i in range(200)
]
# First, flush original messages to the original session
agent._flush_messages_to_session_db(original_history, [])
original_rows = db.get_messages("original-session")
assert len(original_rows) == 200
# Now simulate compression: new session, reset idx, shorter messages
agent.session_id = "compressed-session"
db.create_session(session_id="compressed-session", source="test")
agent._last_flushed_db_idx = 0
# The compressed messages (summary + tail + new turn)
compressed_messages = [
{"role": "user", "content": "[CONTEXT COMPACTION] Summary of work..."},
{"role": "user", "content": "What should we do next?"},
{"role": "assistant", "content": "Let me check..."},
{"role": "user", "content": "new question"},
{"role": "assistant", "content": "new answer"},
]
# THE BUG: passing the original history as conversation_history
# causes flush_from = max(200, 0) = 200, skipping everything.
# After the fix, conversation_history should be None.
agent._flush_messages_to_session_db(compressed_messages, None)
new_rows = db.get_messages("compressed-session")
assert len(new_rows) == 5, (
f"Expected 5 compressed messages in new session, got {len(new_rows)}. "
f"Compression persistence bug: messages not written to SQLite."
)
def test_flush_with_stale_history_loses_messages(self):
"""Demonstrates the bug condition: stale conversation_history causes data loss."""
from hermes_state import SessionDB
with tempfile.TemporaryDirectory() as tmpdir:
db_path = Path(tmpdir) / "test.db"
db = SessionDB(db_path=db_path)
agent = self._make_agent(db)
# Simulate compression reset
agent.session_id = "new-session"
db.create_session(session_id="new-session", source="test")
agent._last_flushed_db_idx = 0
compressed = [
{"role": "user", "content": "summary"},
{"role": "assistant", "content": "continuing..."},
]
# Bug: passing a conversation_history longer than compressed messages
stale_history = [{"role": "user", "content": f"msg{i}"} for i in range(100)]
agent._flush_messages_to_session_db(compressed, stale_history)
rows = db.get_messages("new-session")
# With the stale history, flush_from = max(100, 0) = 100
# But compressed only has 2 entries → messages[100:] = empty
assert len(rows) == 0, (
"Expected 0 messages with stale conversation_history "
"(this test verifies the bug condition exists)"
)
# ---------------------------------------------------------------------------
# Part 2: Gateway-side — history_offset after session split
# ---------------------------------------------------------------------------
class TestGatewayHistoryOffsetAfterSplit:
"""Verify that when the agent creates a new session during compression,
the gateway uses history_offset=0 so all compressed messages are written
to the JSONL transcript."""
def test_history_offset_zero_on_session_split(self):
"""When agent.session_id differs from the original, history_offset must be 0."""
# This tests the logic in gateway/run.py run_sync():
# _session_was_split = agent.session_id != session_id
# _effective_history_offset = 0 if _session_was_split else len(agent_history)
original_session_id = "session-abc"
agent_session_id = "session-compressed-xyz" # Different = compression happened
agent_history_len = 200
# Simulate the gateway's offset calculation (post-fix)
_session_was_split = (agent_session_id != original_session_id)
_effective_history_offset = 0 if _session_was_split else agent_history_len
assert _session_was_split is True
assert _effective_history_offset == 0
def test_history_offset_preserved_without_split(self):
"""When no compression happened, history_offset is the original length."""
session_id = "session-abc"
agent_session_id = "session-abc" # Same = no compression
agent_history_len = 200
_session_was_split = (agent_session_id != session_id)
_effective_history_offset = 0 if _session_was_split else agent_history_len
assert _session_was_split is False
assert _effective_history_offset == 200
def test_new_messages_extraction_after_split(self):
"""After compression with offset=0, new_messages should be ALL agent messages."""
# Simulates the gateway's new_messages calculation
agent_messages = [
{"role": "user", "content": "[CONTEXT COMPACTION] Summary..."},
{"role": "user", "content": "recent question"},
{"role": "assistant", "content": "recent answer"},
{"role": "user", "content": "new question"},
{"role": "assistant", "content": "new answer"},
]
history_offset = 0 # After fix: 0 on session split
new_messages = agent_messages[history_offset:] if len(agent_messages) > history_offset else []
assert len(new_messages) == 5, (
f"Expected all 5 messages with offset=0, got {len(new_messages)}"
)
def test_new_messages_empty_with_stale_offset(self):
"""Demonstrates the bug: stale offset produces empty new_messages."""
agent_messages = [
{"role": "user", "content": "summary"},
{"role": "assistant", "content": "answer"},
]
# Bug: offset is the pre-compression history length
history_offset = 200
new_messages = agent_messages[history_offset:] if len(agent_messages) > history_offset else []
assert len(new_messages) == 0, (
"Expected 0 messages with stale offset=200 (demonstrates the bug)"
)

View File

@@ -0,0 +1,91 @@
"""Tests that _try_activate_fallback updates the context compressor."""
from unittest.mock import MagicMock, patch
from run_agent import AIAgent
from agent.context_compressor import ContextCompressor
def _make_agent_with_compressor() -> AIAgent:
"""Build a minimal AIAgent with a context_compressor, skipping __init__."""
agent = AIAgent.__new__(AIAgent)
# Primary model settings
agent.model = "primary-model"
agent.provider = "openrouter"
agent.base_url = "https://openrouter.ai/api/v1"
agent.api_key = "sk-primary"
agent.api_mode = "chat_completions"
agent.client = MagicMock()
agent.quiet_mode = True
# Fallback config
agent._fallback_activated = False
agent._fallback_model = {
"provider": "openai",
"model": "gpt-4o",
}
agent._fallback_chain = [agent._fallback_model]
agent._fallback_index = 0
# Context compressor with primary model values
compressor = ContextCompressor(
model="primary-model",
threshold_percent=0.50,
base_url="https://openrouter.ai/api/v1",
api_key="sk-primary",
provider="openrouter",
quiet_mode=True,
)
agent.context_compressor = compressor
return agent
@patch("agent.auxiliary_client.resolve_provider_client")
@patch("agent.model_metadata.get_model_context_length", return_value=128_000)
def test_compressor_updated_on_fallback(mock_ctx_len, mock_resolve):
"""After fallback activation, the compressor must reflect the fallback model."""
agent = _make_agent_with_compressor()
assert agent.context_compressor.model == "primary-model"
fb_client = MagicMock()
fb_client.base_url = "https://api.openai.com/v1"
fb_client.api_key = "sk-fallback"
mock_resolve.return_value = (fb_client, None)
agent._is_direct_openai_url = lambda url: "api.openai.com" in url
agent._emit_status = lambda msg: None
result = agent._try_activate_fallback()
assert result is True
assert agent._fallback_activated is True
c = agent.context_compressor
assert c.model == "gpt-4o"
assert c.base_url == "https://api.openai.com/v1"
assert c.api_key == "sk-fallback"
assert c.provider == "openai"
assert c.context_length == 128_000
assert c.threshold_tokens == int(128_000 * c.threshold_percent)
@patch("agent.auxiliary_client.resolve_provider_client")
@patch("agent.model_metadata.get_model_context_length", return_value=128_000)
def test_compressor_not_present_does_not_crash(mock_ctx_len, mock_resolve):
"""If the agent has no compressor, fallback should still succeed."""
agent = _make_agent_with_compressor()
agent.context_compressor = None
fb_client = MagicMock()
fb_client.base_url = "https://api.openai.com/v1"
fb_client.api_key = "sk-fallback"
mock_resolve.return_value = (fb_client, None)
agent._is_direct_openai_url = lambda url: "api.openai.com" in url
agent._emit_status = lambda msg: None
result = agent._try_activate_fallback()
assert result is True

View File

@@ -0,0 +1,248 @@
"""Tests for context pressure warnings (user-facing, not injected into messages).
Covers:
- Display formatting (CLI and gateway variants)
- Flag tracking and threshold logic on AIAgent
- Flag reset after compression
- status_callback invocation
"""
import json
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
from agent.display import format_context_pressure, format_context_pressure_gateway
from run_agent import AIAgent
# ---------------------------------------------------------------------------
# Display formatting tests
# ---------------------------------------------------------------------------
class TestFormatContextPressure:
"""CLI context pressure display (agent/display.py).
The bar shows progress toward the compaction threshold, not the
raw context window. 60% = 60% of the way to compaction.
"""
def test_80_percent_uses_warning_icon(self):
line = format_context_pressure(0.80, 100_000, 0.50)
assert "" in line
assert "80% to compaction" in line
def test_90_percent_uses_warning_icon(self):
line = format_context_pressure(0.90, 100_000, 0.50)
assert "" in line
assert "90% to compaction" in line
def test_bar_length_scales_with_progress(self):
line_80 = format_context_pressure(0.80, 100_000, 0.50)
line_95 = format_context_pressure(0.95, 100_000, 0.50)
assert line_95.count("") > line_80.count("")
def test_shows_threshold_tokens(self):
line = format_context_pressure(0.80, 100_000, 0.50)
assert "100k" in line
def test_small_threshold(self):
line = format_context_pressure(0.80, 500, 0.50)
assert "500" in line
def test_shows_threshold_percent(self):
line = format_context_pressure(0.80, 100_000, 0.50)
assert "50%" in line
def test_approaching_hint(self):
line = format_context_pressure(0.80, 100_000, 0.50)
assert "compaction approaching" in line
def test_no_compaction_when_disabled(self):
line = format_context_pressure(0.85, 100_000, 0.50, compression_enabled=False)
assert "no auto-compaction" in line
def test_returns_string(self):
result = format_context_pressure(0.65, 128_000, 0.50)
assert isinstance(result, str)
def test_over_100_percent_capped(self):
"""Progress > 1.0 should cap both bar and percentage text at 100%."""
line = format_context_pressure(1.05, 100_000, 0.50)
assert "" in line
assert line.count("") == 20
assert "100%" in line
assert "105%" not in line
class TestFormatContextPressureGateway:
"""Gateway (plain text) context pressure display."""
def test_80_percent_warning(self):
msg = format_context_pressure_gateway(0.80, 0.50)
assert "80% to compaction" in msg
assert "50%" in msg
def test_90_percent_warning(self):
msg = format_context_pressure_gateway(0.90, 0.50)
assert "90% to compaction" in msg
assert "approaching" in msg
def test_no_compaction_warning(self):
msg = format_context_pressure_gateway(0.85, 0.50, compression_enabled=False)
assert "disabled" in msg
def test_no_ansi_codes(self):
msg = format_context_pressure_gateway(0.80, 0.50)
assert "\033[" not in msg
def test_has_progress_bar(self):
msg = format_context_pressure_gateway(0.80, 0.50)
assert "" in msg
def test_over_100_percent_capped(self):
"""Progress > 1.0 should cap percentage text at 100%."""
msg = format_context_pressure_gateway(1.09, 0.50)
assert "100% to compaction" in msg
assert "109%" not in msg
assert msg.count("") == 20
# ---------------------------------------------------------------------------
# AIAgent context pressure flag tests
# ---------------------------------------------------------------------------
def _make_tool_defs(*names):
return [
{
"type": "function",
"function": {
"name": n,
"description": f"{n} tool",
"parameters": {"type": "object", "properties": {}},
},
}
for n in names
]
@pytest.fixture()
def agent():
"""Minimal AIAgent with mocked internals."""
with (
patch("run_agent.get_tool_definitions", return_value=_make_tool_defs("web_search")),
patch("run_agent.check_toolset_requirements", return_value={}),
patch("run_agent.OpenAI"),
):
a = AIAgent(
api_key="test-key-1234567890",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
)
a.client = MagicMock()
return a
class TestContextPressureFlags:
"""Context pressure warning flag tracking on AIAgent."""
def test_flag_initialized_false(self, agent):
assert agent._context_pressure_warned is False
def test_emit_calls_status_callback(self, agent):
"""status_callback should be invoked with event type and message."""
cb = MagicMock()
agent.status_callback = cb
compressor = MagicMock()
compressor.context_length = 200_000
compressor.threshold_tokens = 100_000 # 50%
agent._emit_context_pressure(0.85, compressor)
cb.assert_called_once()
args = cb.call_args[0]
assert args[0] == "context_pressure"
assert "85% to compaction" in args[1]
def test_emit_no_callback_no_crash(self, agent):
"""No status_callback set — should not crash."""
agent.status_callback = None
compressor = MagicMock()
compressor.context_length = 200_000
compressor.threshold_tokens = 100_000
# Should not raise
agent._emit_context_pressure(0.60, compressor)
def test_emit_prints_for_cli_platform(self, agent, capsys):
"""CLI platform should always print context pressure, even in quiet_mode."""
agent.quiet_mode = True
agent.platform = "cli"
agent.status_callback = None
compressor = MagicMock()
compressor.context_length = 200_000
compressor.threshold_tokens = 100_000
agent._emit_context_pressure(0.85, compressor)
captured = capsys.readouterr()
assert "" in captured.out
assert "to compaction" in captured.out
def test_emit_skips_print_for_gateway_platform(self, agent, capsys):
"""Gateway platforms get the callback, not CLI print."""
agent.platform = "telegram"
agent.status_callback = None
compressor = MagicMock()
compressor.context_length = 200_000
compressor.threshold_tokens = 100_000
agent._emit_context_pressure(0.85, compressor)
captured = capsys.readouterr()
assert "" not in captured.out
def test_flag_reset_on_compression(self, agent):
"""After _compress_context, context pressure flag should reset."""
agent._context_pressure_warned = True
agent.compression_enabled = True
agent.context_compressor = MagicMock()
agent.context_compressor.compress.return_value = [
{"role": "user", "content": "Summary of conversation so far."}
]
agent.context_compressor.context_length = 200_000
agent.context_compressor.threshold_tokens = 100_000
agent._todo_store = MagicMock()
agent._todo_store.format_for_injection.return_value = None
agent._build_system_prompt = MagicMock(return_value="system prompt")
agent._cached_system_prompt = "old system prompt"
agent._session_db = None
messages = [
{"role": "user", "content": "hello"},
{"role": "assistant", "content": "hi there"},
]
agent._compress_context(messages, "system prompt")
assert agent._context_pressure_warned is False
def test_emit_callback_error_handled(self, agent):
"""If status_callback raises, it should be caught gracefully."""
cb = MagicMock(side_effect=RuntimeError("callback boom"))
agent.status_callback = cb
compressor = MagicMock()
compressor.context_length = 200_000
compressor.threshold_tokens = 100_000
# Should not raise
agent._emit_context_pressure(0.85, compressor)

View File

@@ -0,0 +1,127 @@
"""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(run_agent, "get_tool_definitions", lambda **kwargs: [{
"type": "function",
"function": {"name": "t", "description": "t", "parameters": {"type": "object", "properties": {}}},
}])
monkeypatch.setattr(run_agent, "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: _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 = self._save_session_log = lambda *a, **k: None
def run_conversation(self, msg, conversation_history=None, task_id=None):
self._interruptible_api_call = lambda kw: response_fn()
return super().run_conversation(msg, conversation_history=conversation_history, task_id=task_id)
return _A(model="test-model", api_key="test-key", 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

View File

@@ -0,0 +1,72 @@
import json
from types import SimpleNamespace
def _tool_call(name: str, arguments):
return SimpleNamespace(
id="call_1",
type="function",
function=SimpleNamespace(name=name, arguments=arguments),
)
def _response_with_tool_call(arguments):
assistant = SimpleNamespace(
content=None,
reasoning=None,
tool_calls=[_tool_call("read_file", arguments)],
)
choice = SimpleNamespace(message=assistant, finish_reason="tool_calls")
return SimpleNamespace(choices=[choice], usage=None)
class _FakeChatCompletions:
def __init__(self):
self.calls = 0
def create(self, **kwargs):
self.calls += 1
if self.calls == 1:
return _response_with_tool_call({"path": "README.md"})
return SimpleNamespace(
choices=[
SimpleNamespace(
message=SimpleNamespace(content="done", reasoning=None, tool_calls=[]),
finish_reason="stop",
)
],
usage=None,
)
class _FakeClient:
def __init__(self):
self.chat = SimpleNamespace(completions=_FakeChatCompletions())
def test_tool_call_validation_accepts_dict_arguments(monkeypatch):
from run_agent import AIAgent
monkeypatch.setattr("run_agent.OpenAI", lambda **kwargs: _FakeClient())
monkeypatch.setattr(
"run_agent.get_tool_definitions",
lambda *args, **kwargs: [{"function": {"name": "read_file"}}],
)
monkeypatch.setattr(
"run_agent.handle_function_call",
lambda name, args, task_id=None, **kwargs: json.dumps({"ok": True, "args": args}),
)
agent = AIAgent(
model="test-model",
api_key="test-key",
base_url="http://localhost:8080/v1",
platform="cli",
max_iterations=3,
quiet_mode=True,
skip_memory=True,
)
result = agent.run_conversation("read the file")
assert result["final_response"] == "done"

View File

@@ -0,0 +1,73 @@
"""Tests for KeyboardInterrupt handling in exit cleanup paths.
``except Exception`` does not catch ``KeyboardInterrupt`` (which inherits
from ``BaseException``). A second Ctrl+C during exit cleanup must not
abort remaining cleanup steps. These tests exercise the actual production
code paths — not a copy of the try/except pattern.
"""
import atexit
import weakref
from unittest.mock import MagicMock, patch, call
import pytest
class TestCronJobCleanup:
"""cron/scheduler.py — end_session + close in the finally block."""
def test_keyboard_interrupt_in_end_session_does_not_skip_close(self):
"""If end_session raises KeyboardInterrupt, close() must still run."""
mock_db = MagicMock()
mock_db.end_session.side_effect = KeyboardInterrupt
from cron import scheduler
job = {
"id": "test-job-1",
"name": "test cleanup",
"prompt": "hello",
"schedule": "0 9 * * *",
"model": "test/model",
}
with patch("hermes_state.SessionDB", return_value=mock_db), \
patch.object(scheduler, "_build_job_prompt", return_value="hello"), \
patch.object(scheduler, "_resolve_origin", return_value=None), \
patch.object(scheduler, "_resolve_delivery_target", return_value=None), \
patch("dotenv.load_dotenv", return_value=None), \
patch("run_agent.AIAgent") as MockAgent:
# Make the agent raise immediately so we hit the finally block
MockAgent.return_value.run_conversation.side_effect = RuntimeError("boom")
scheduler.run_job(job)
mock_db.end_session.assert_called_once()
mock_db.close.assert_called_once()
def test_keyboard_interrupt_in_close_does_not_propagate(self):
"""If close() raises KeyboardInterrupt, it must not escape run_job."""
mock_db = MagicMock()
mock_db.close.side_effect = KeyboardInterrupt
from cron import scheduler
job = {
"id": "test-job-2",
"name": "test close interrupt",
"prompt": "hello",
"schedule": "0 9 * * *",
"model": "test/model",
}
with patch("hermes_state.SessionDB", return_value=mock_db), \
patch.object(scheduler, "_build_job_prompt", return_value="hello"), \
patch.object(scheduler, "_resolve_origin", return_value=None), \
patch.object(scheduler, "_resolve_delivery_target", return_value=None), \
patch("dotenv.load_dotenv", return_value=None), \
patch("run_agent.AIAgent") as MockAgent:
MockAgent.return_value.run_conversation.side_effect = RuntimeError("boom")
# Must not raise
scheduler.run_job(job)
mock_db.end_session.assert_called_once()
mock_db.close.assert_called_once()

View File

@@ -0,0 +1,377 @@
"""Tests for the provider fallback model feature.
Verifies that AIAgent can switch to a configured fallback model/provider
when the primary fails after retries.
"""
import os
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
from run_agent import AIAgent
def _make_tool_defs(*names: str) -> list:
return [
{
"type": "function",
"function": {
"name": n,
"description": f"{n} tool",
"parameters": {"type": "object", "properties": {}},
},
}
for n in names
]
def _make_agent(fallback_model=None):
"""Create a minimal AIAgent with optional fallback config."""
with (
patch("run_agent.get_tool_definitions", return_value=_make_tool_defs("web_search")),
patch("run_agent.check_toolset_requirements", return_value={}),
patch("run_agent.OpenAI"),
):
agent = AIAgent(
api_key="test-key",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
fallback_model=fallback_model,
)
agent.client = MagicMock()
return agent
def _mock_resolve(base_url="https://openrouter.ai/api/v1", api_key="test-key"):
"""Helper to create a mock client for resolve_provider_client."""
mock_client = MagicMock()
mock_client.api_key = api_key
mock_client.base_url = base_url
return mock_client
# =============================================================================
# _try_activate_fallback()
# =============================================================================
class TestTryActivateFallback:
def test_returns_false_when_not_configured(self):
agent = _make_agent(fallback_model=None)
assert agent._try_activate_fallback() is False
assert agent._fallback_activated is False
def test_returns_false_for_empty_config(self):
agent = _make_agent(fallback_model={"provider": "", "model": ""})
assert agent._try_activate_fallback() is False
def test_returns_false_for_missing_provider(self):
agent = _make_agent(fallback_model={"model": "gpt-4.1"})
assert agent._try_activate_fallback() is False
def test_returns_false_for_missing_model(self):
agent = _make_agent(fallback_model={"provider": "openrouter"})
assert agent._try_activate_fallback() is False
def test_activates_openrouter_fallback(self):
agent = _make_agent(
fallback_model={"provider": "openrouter", "model": "anthropic/claude-sonnet-4"},
)
mock_client = _mock_resolve(
api_key="sk-or-fallback-key",
base_url="https://openrouter.ai/api/v1",
)
with patch(
"agent.auxiliary_client.resolve_provider_client",
return_value=(mock_client, "anthropic/claude-sonnet-4"),
):
result = agent._try_activate_fallback()
assert result is True
assert agent._fallback_activated is True
assert agent.model == "anthropic/claude-sonnet-4"
assert agent.provider == "openrouter"
assert agent.api_mode == "chat_completions"
assert agent.client is mock_client
def test_activates_zai_fallback(self):
agent = _make_agent(
fallback_model={"provider": "zai", "model": "glm-5"},
)
mock_client = _mock_resolve(
api_key="sk-zai-key",
base_url="https://open.z.ai/api/v1",
)
with patch(
"agent.auxiliary_client.resolve_provider_client",
return_value=(mock_client, "glm-5"),
):
result = agent._try_activate_fallback()
assert result is True
assert agent.model == "glm-5"
assert agent.provider == "zai"
assert agent.client is mock_client
def test_activates_kimi_fallback(self):
agent = _make_agent(
fallback_model={"provider": "kimi-coding", "model": "kimi-k2.5"},
)
mock_client = _mock_resolve(
api_key="sk-kimi-key",
base_url="https://api.moonshot.ai/v1",
)
with patch(
"agent.auxiliary_client.resolve_provider_client",
return_value=(mock_client, "kimi-k2.5"),
):
assert agent._try_activate_fallback() is True
assert agent.model == "kimi-k2.5"
assert agent.provider == "kimi-coding"
def test_activates_minimax_fallback(self):
agent = _make_agent(
fallback_model={"provider": "minimax", "model": "MiniMax-M2.7"},
)
mock_client = _mock_resolve(
api_key="sk-mm-key",
base_url="https://api.minimax.io/v1",
)
with patch(
"agent.auxiliary_client.resolve_provider_client",
return_value=(mock_client, "MiniMax-M2.7"),
):
assert agent._try_activate_fallback() is True
assert agent.model == "MiniMax-M2.7"
assert agent.provider == "minimax"
assert agent.client is mock_client
def test_only_fires_once(self):
agent = _make_agent(
fallback_model={"provider": "openrouter", "model": "anthropic/claude-sonnet-4"},
)
mock_client = _mock_resolve(
api_key="sk-or-key",
base_url="https://openrouter.ai/api/v1",
)
with patch(
"agent.auxiliary_client.resolve_provider_client",
return_value=(mock_client, "anthropic/claude-sonnet-4"),
):
assert agent._try_activate_fallback() is True
# Second attempt should return False
assert agent._try_activate_fallback() is False
def test_returns_false_when_no_api_key(self):
"""Fallback should fail gracefully when the API key env var is unset."""
agent = _make_agent(
fallback_model={"provider": "minimax", "model": "MiniMax-M2.7"},
)
with patch(
"agent.auxiliary_client.resolve_provider_client",
return_value=(None, None),
):
assert agent._try_activate_fallback() is False
assert agent._fallback_activated is False
def test_custom_base_url(self):
"""Custom base_url in config should override the provider default."""
agent = _make_agent(
fallback_model={
"provider": "custom",
"model": "my-model",
"base_url": "http://localhost:8080/v1",
"api_key_env": "MY_CUSTOM_KEY",
},
)
mock_client = _mock_resolve(
api_key="custom-secret",
base_url="http://localhost:8080/v1",
)
with patch(
"agent.auxiliary_client.resolve_provider_client",
return_value=(mock_client, "my-model"),
):
assert agent._try_activate_fallback() is True
assert agent.client is mock_client
assert agent.model == "my-model"
def test_prompt_caching_enabled_for_claude_on_openrouter(self):
agent = _make_agent(
fallback_model={"provider": "openrouter", "model": "anthropic/claude-sonnet-4"},
)
mock_client = _mock_resolve(
api_key="sk-or-key",
base_url="https://openrouter.ai/api/v1",
)
with patch(
"agent.auxiliary_client.resolve_provider_client",
return_value=(mock_client, "anthropic/claude-sonnet-4"),
):
agent._try_activate_fallback()
assert agent._use_prompt_caching is True
def test_prompt_caching_disabled_for_non_claude(self):
agent = _make_agent(
fallback_model={"provider": "openrouter", "model": "google/gemini-2.5-flash"},
)
mock_client = _mock_resolve(
api_key="sk-or-key",
base_url="https://openrouter.ai/api/v1",
)
with patch(
"agent.auxiliary_client.resolve_provider_client",
return_value=(mock_client, "google/gemini-2.5-flash"),
):
agent._try_activate_fallback()
assert agent._use_prompt_caching is False
def test_prompt_caching_disabled_for_non_openrouter(self):
agent = _make_agent(
fallback_model={"provider": "zai", "model": "glm-5"},
)
mock_client = _mock_resolve(
api_key="sk-zai-key",
base_url="https://open.z.ai/api/v1",
)
with patch(
"agent.auxiliary_client.resolve_provider_client",
return_value=(mock_client, "glm-5"),
):
agent._try_activate_fallback()
assert agent._use_prompt_caching is False
def test_zai_alt_env_var(self):
"""Z.AI should also check Z_AI_API_KEY as fallback env var."""
agent = _make_agent(
fallback_model={"provider": "zai", "model": "glm-5"},
)
mock_client = _mock_resolve(
api_key="sk-alt-key",
base_url="https://open.z.ai/api/v1",
)
with patch(
"agent.auxiliary_client.resolve_provider_client",
return_value=(mock_client, "glm-5"),
):
assert agent._try_activate_fallback() is True
assert agent.client is mock_client
def test_activates_codex_fallback(self):
"""OpenAI Codex fallback should use OAuth credentials and codex_responses mode."""
agent = _make_agent(
fallback_model={"provider": "openai-codex", "model": "gpt-5.3-codex"},
)
mock_client = _mock_resolve(
api_key="codex-oauth-token",
base_url="https://chatgpt.com/backend-api/codex",
)
with patch(
"agent.auxiliary_client.resolve_provider_client",
return_value=(mock_client, "gpt-5.3-codex"),
):
result = agent._try_activate_fallback()
assert result is True
assert agent.model == "gpt-5.3-codex"
assert agent.provider == "openai-codex"
assert agent.api_mode == "codex_responses"
assert agent.client is mock_client
def test_codex_fallback_fails_gracefully_without_credentials(self):
"""Codex fallback should return False if no OAuth credentials available."""
agent = _make_agent(
fallback_model={"provider": "openai-codex", "model": "gpt-5.3-codex"},
)
with patch(
"agent.auxiliary_client.resolve_provider_client",
return_value=(None, None),
):
assert agent._try_activate_fallback() is False
assert agent._fallback_activated is False
def test_activates_nous_fallback(self):
"""Nous Portal fallback should use OAuth credentials and chat_completions mode."""
agent = _make_agent(
fallback_model={"provider": "nous", "model": "nous-hermes-3"},
)
mock_client = _mock_resolve(
api_key="nous-agent-key-abc",
base_url="https://inference-api.nousresearch.com/v1",
)
with patch(
"agent.auxiliary_client.resolve_provider_client",
return_value=(mock_client, "nous-hermes-3"),
):
result = agent._try_activate_fallback()
assert result is True
assert agent.model == "nous-hermes-3"
assert agent.provider == "nous"
assert agent.api_mode == "chat_completions"
assert agent.client is mock_client
def test_nous_fallback_fails_gracefully_without_login(self):
"""Nous fallback should return False if not logged in."""
agent = _make_agent(
fallback_model={"provider": "nous", "model": "nous-hermes-3"},
)
with patch(
"agent.auxiliary_client.resolve_provider_client",
return_value=(None, None),
):
assert agent._try_activate_fallback() is False
assert agent._fallback_activated is False
# =============================================================================
# Fallback config init
# =============================================================================
class TestFallbackInit:
def test_fallback_stored_when_configured(self):
agent = _make_agent(
fallback_model={"provider": "openrouter", "model": "anthropic/claude-sonnet-4"},
)
assert agent._fallback_model is not None
assert agent._fallback_model["provider"] == "openrouter"
assert agent._fallback_activated is False
def test_fallback_none_when_not_configured(self):
agent = _make_agent(fallback_model=None)
assert agent._fallback_model is None
assert agent._fallback_activated is False
def test_fallback_none_for_non_dict(self):
agent = _make_agent(fallback_model="not-a-dict")
assert agent._fallback_model is None
# =============================================================================
# Provider credential resolution
# =============================================================================
class TestProviderCredentials:
"""Verify that each supported provider resolves via the centralized router."""
@pytest.mark.parametrize("provider,env_var,base_url_fragment", [
("openrouter", "OPENROUTER_API_KEY", "openrouter"),
("zai", "ZAI_API_KEY", "z.ai"),
("kimi-coding", "KIMI_API_KEY", "moonshot.ai"),
("minimax", "MINIMAX_API_KEY", "minimax.io"),
("minimax-cn", "MINIMAX_CN_API_KEY", "minimaxi.com"),
])
def test_provider_resolves(self, provider, env_var, base_url_fragment):
agent = _make_agent(
fallback_model={"provider": provider, "model": "test-model"},
)
mock_client = MagicMock()
mock_client.api_key = "test-api-key"
mock_client.base_url = f"https://{base_url_fragment}/v1"
with patch(
"agent.auxiliary_client.resolve_provider_client",
return_value=(mock_client, "test-model"),
):
result = agent._try_activate_fallback()
assert result is True, f"Failed to activate fallback for {provider}"
assert agent.client is mock_client
assert agent.model == "test-model"
assert agent.provider == provider

View File

@@ -0,0 +1,222 @@
"""Tests for flush_memories() working correctly across all provider modes.
Catches the bug where Codex mode called chat.completions.create on a
Responses-only client, which would fail silently or with a 404.
"""
import json
import os
import sys
import types
from types import SimpleNamespace
from unittest.mock import patch, MagicMock, call
import pytest
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
class _FakeOpenAI:
def __init__(self, **kwargs):
self.kwargs = kwargs
self.api_key = kwargs.get("api_key", "test")
self.base_url = kwargs.get("base_url", "http://test")
def close(self):
pass
def _make_agent(monkeypatch, api_mode="chat_completions", provider="openrouter"):
"""Build an AIAgent with mocked internals, ready for flush_memories testing."""
monkeypatch.setattr(run_agent, "get_tool_definitions", lambda **kw: [
{
"type": "function",
"function": {
"name": "memory",
"description": "Manage memories.",
"parameters": {
"type": "object",
"properties": {
"action": {"type": "string"},
"target": {"type": "string"},
"content": {"type": "string"},
},
},
},
},
])
monkeypatch.setattr(run_agent, "check_toolset_requirements", lambda: {})
monkeypatch.setattr(run_agent, "OpenAI", _FakeOpenAI)
agent = run_agent.AIAgent(
api_key="test-key",
base_url="https://test.example.com/v1",
provider=provider,
api_mode=api_mode,
max_iterations=4,
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
)
# Give it a valid memory store
agent._memory_store = MagicMock()
agent._memory_flush_min_turns = 1
agent._user_turn_count = 5
return agent
def _chat_response_with_memory_call():
"""Simulated chat completions response with a memory tool call."""
return SimpleNamespace(
choices=[SimpleNamespace(
message=SimpleNamespace(
content=None,
tool_calls=[SimpleNamespace(
function=SimpleNamespace(
name="memory",
arguments=json.dumps({
"action": "add",
"target": "notes",
"content": "User prefers dark mode.",
}),
),
)],
),
)],
usage=SimpleNamespace(prompt_tokens=100, completion_tokens=20, total_tokens=120),
)
class TestFlushMemoriesUsesAuxiliaryClient:
"""When an auxiliary client is available, flush_memories should use it
instead of self.client -- especially critical in Codex mode."""
def test_flush_uses_auxiliary_when_available(self, monkeypatch):
agent = _make_agent(monkeypatch, api_mode="codex_responses", provider="openai-codex")
mock_response = _chat_response_with_memory_call()
with patch("agent.auxiliary_client.call_llm", return_value=mock_response) as mock_call:
messages = [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi there"},
{"role": "user", "content": "Remember this"},
]
with patch("tools.memory_tool.memory_tool", return_value="Saved.") as mock_memory:
agent.flush_memories(messages)
mock_call.assert_called_once()
call_kwargs = mock_call.call_args
assert call_kwargs.kwargs.get("task") == "flush_memories"
def test_flush_uses_main_client_when_no_auxiliary(self, monkeypatch):
"""Non-Codex mode with no auxiliary falls back to self.client."""
agent = _make_agent(monkeypatch, api_mode="chat_completions", provider="openrouter")
agent.client = MagicMock()
agent.client.chat.completions.create.return_value = _chat_response_with_memory_call()
with patch("agent.auxiliary_client.call_llm", side_effect=RuntimeError("no provider")):
messages = [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi there"},
{"role": "user", "content": "Save this"},
]
with patch("tools.memory_tool.memory_tool", return_value="Saved."):
agent.flush_memories(messages)
agent.client.chat.completions.create.assert_called_once()
def test_flush_executes_memory_tool_calls(self, monkeypatch):
"""Verify that memory tool calls from the flush response actually get executed."""
agent = _make_agent(monkeypatch, api_mode="chat_completions", provider="openrouter")
mock_response = _chat_response_with_memory_call()
with patch("agent.auxiliary_client.call_llm", return_value=mock_response):
messages = [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi"},
{"role": "user", "content": "Note this"},
]
with patch("tools.memory_tool.memory_tool", return_value="Saved.") as mock_memory:
agent.flush_memories(messages)
mock_memory.assert_called_once()
call_kwargs = mock_memory.call_args
assert call_kwargs.kwargs["action"] == "add"
assert call_kwargs.kwargs["target"] == "notes"
assert "dark mode" in call_kwargs.kwargs["content"]
def test_flush_strips_artifacts_from_messages(self, monkeypatch):
"""After flush, the flush prompt and any response should be removed from messages."""
agent = _make_agent(monkeypatch, api_mode="chat_completions", provider="openrouter")
mock_response = _chat_response_with_memory_call()
with patch("agent.auxiliary_client.call_llm", return_value=mock_response):
messages = [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi"},
{"role": "user", "content": "Remember X"},
]
original_len = len(messages)
with patch("tools.memory_tool.memory_tool", return_value="Saved."):
agent.flush_memories(messages)
# Messages should not grow from the flush
assert len(messages) <= original_len
# No flush sentinel should remain
for msg in messages:
assert "_flush_sentinel" not in msg
class TestFlushMemoriesCodexFallback:
"""When no auxiliary client exists and we're in Codex mode, flush should
use the Codex Responses API path instead of chat.completions."""
def test_codex_mode_no_aux_uses_responses_api(self, monkeypatch):
agent = _make_agent(monkeypatch, api_mode="codex_responses", provider="openai-codex")
codex_response = SimpleNamespace(
output=[
SimpleNamespace(
type="function_call",
call_id="call_1",
name="memory",
arguments=json.dumps({
"action": "add",
"target": "notes",
"content": "Codex flush test",
}),
),
],
usage=SimpleNamespace(input_tokens=50, output_tokens=10, total_tokens=60),
status="completed",
model="gpt-5-codex",
)
with patch("agent.auxiliary_client.call_llm", side_effect=RuntimeError("no provider")), \
patch.object(agent, "_run_codex_stream", return_value=codex_response) as mock_stream, \
patch.object(agent, "_build_api_kwargs") as mock_build, \
patch("tools.memory_tool.memory_tool", return_value="Saved.") as mock_memory:
mock_build.return_value = {
"model": "gpt-5-codex",
"instructions": "test",
"input": [],
"tools": [],
"max_output_tokens": 4096,
}
messages = [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi"},
{"role": "user", "content": "Save this"},
]
agent.flush_memories(messages)
mock_stream.assert_called_once()
mock_memory.assert_called_once()
assert mock_memory.call_args.kwargs["content"] == "Codex flush test"

View File

@@ -0,0 +1,203 @@
#!/usr/bin/env python3
"""Interactive interrupt test that mimics the exact CLI flow.
Starts an agent in a thread with a mock delegate_task that takes a while,
then simulates the user typing a message via _interrupt_queue.
Logs every step to stderr (which isn't affected by redirect_stdout)
so we can see exactly where the interrupt gets lost.
"""
import contextlib
import io
import json
import logging
import queue
import sys
import threading
import time
import os
# Force stderr logging so redirect_stdout doesn't swallow it
logging.basicConfig(level=logging.DEBUG, stream=sys.stderr,
format="%(asctime)s [%(threadName)s] %(message)s")
log = logging.getLogger("interrupt_test")
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
from unittest.mock import MagicMock, patch
from run_agent import AIAgent, IterationBudget
from tools.interrupt import set_interrupt, is_interrupted
def make_slow_response(delay=2.0):
"""API response that takes a while."""
def create(**kwargs):
log.info(f" 🌐 Mock API call starting (will take {delay}s)...")
time.sleep(delay)
log.info(f" 🌐 Mock API call completed")
resp = MagicMock()
resp.choices = [MagicMock()]
resp.choices[0].message.content = "Done with the task"
resp.choices[0].message.tool_calls = None
resp.choices[0].message.refusal = None
resp.choices[0].finish_reason = "stop"
resp.usage.prompt_tokens = 100
resp.usage.completion_tokens = 10
resp.usage.total_tokens = 110
resp.usage.prompt_tokens_details = None
return resp
return create
def main() -> int:
set_interrupt(False)
# ─── Create parent agent ───
parent = AIAgent.__new__(AIAgent)
parent._interrupt_requested = False
parent._interrupt_message = None
parent._active_children = []
parent._active_children_lock = threading.Lock()
parent.quiet_mode = True
parent.model = "test/model"
parent.base_url = "http://localhost:1"
parent.api_key = "test"
parent.provider = "test"
parent.api_mode = "chat_completions"
parent.platform = "cli"
parent.enabled_toolsets = ["terminal", "file"]
parent.providers_allowed = None
parent.providers_ignored = None
parent.providers_order = None
parent.provider_sort = None
parent.max_tokens = None
parent.reasoning_config = None
parent.prefill_messages = None
parent._session_db = None
parent._delegate_depth = 0
parent._delegate_spinner = None
parent.tool_progress_callback = None
parent.iteration_budget = IterationBudget(max_total=100)
parent._client_kwargs = {"api_key": "test", "base_url": "http://localhost:1"}
# Monkey-patch parent.interrupt to log
_original_interrupt = AIAgent.interrupt
def logged_interrupt(self, message=None):
log.info(f"🔴 parent.interrupt() called with: {message!r}")
log.info(f" _active_children count: {len(self._active_children)}")
_original_interrupt(self, message)
log.info(f" After interrupt: _interrupt_requested={self._interrupt_requested}")
for i, child in enumerate(self._active_children):
log.info(f" Child {i}._interrupt_requested={child._interrupt_requested}")
parent.interrupt = lambda msg=None: logged_interrupt(parent, msg)
# ─── Simulate the exact CLI flow ───
interrupt_queue = queue.Queue()
child_running = threading.Event()
agent_result = [None]
def agent_thread_func():
"""Simulates the agent_thread in cli.py's chat() method."""
log.info("🟢 agent_thread starting")
with patch("run_agent.OpenAI") as MockOpenAI:
mock_client = MagicMock()
mock_client.chat.completions.create = make_slow_response(delay=3.0)
mock_client.close = MagicMock()
MockOpenAI.return_value = mock_client
from tools.delegate_tool import _run_single_child
# Signal that child is about to start
original_init = AIAgent.__init__
def patched_init(self_agent, *a, **kw):
log.info("🟡 Child AIAgent.__init__ called")
original_init(self_agent, *a, **kw)
child_running.set()
log.info(
f"🟡 Child started, parent._active_children = {len(parent._active_children)}"
)
with patch.object(AIAgent, "__init__", patched_init):
result = _run_single_child(
task_index=0,
goal="Do a slow thing",
context=None,
toolsets=["terminal"],
model="test/model",
max_iterations=3,
parent_agent=parent,
task_count=1,
override_provider="test",
override_base_url="http://localhost:1",
override_api_key="test",
override_api_mode="chat_completions",
)
agent_result[0] = result
log.info(f"🟢 agent_thread finished. Result status: {result.get('status')}")
# ─── Start agent thread (like chat() does) ───
agent_thread = threading.Thread(target=agent_thread_func, name="agent_thread", daemon=True)
agent_thread.start()
# ─── Wait for child to start ───
if not child_running.wait(timeout=10):
print("FAIL: Child never started", file=sys.stderr)
set_interrupt(False)
return 1
# Give child time to enter its main loop and start API call
time.sleep(1.0)
# ─── Simulate user typing a message (like handle_enter does) ───
log.info("📝 Simulating user typing 'Hey stop that'")
interrupt_queue.put("Hey stop that")
# ─── Simulate chat() polling loop (like the real chat() method) ───
log.info("📡 Starting interrupt queue polling (like chat())")
interrupt_msg = None
poll_count = 0
while agent_thread.is_alive():
try:
interrupt_msg = interrupt_queue.get(timeout=0.1)
if interrupt_msg:
log.info(f"📨 Got interrupt message from queue: {interrupt_msg!r}")
log.info(" Calling parent.interrupt()...")
parent.interrupt(interrupt_msg)
log.info(" parent.interrupt() returned. Breaking poll loop.")
break
except queue.Empty:
poll_count += 1
if poll_count % 20 == 0: # Log every 2s
log.info(f" Still polling ({poll_count} iterations)...")
# ─── Wait for agent to finish ───
log.info("⏳ Waiting for agent_thread to join...")
t0 = time.monotonic()
agent_thread.join(timeout=10)
elapsed = time.monotonic() - t0
log.info(f"✅ agent_thread joined after {elapsed:.2f}s")
# ─── Check results ───
result = agent_result[0]
if result:
log.info(f"Result status: {result['status']}")
log.info(f"Result duration: {result['duration_seconds']}s")
if result["status"] == "interrupted" and elapsed < 2.0:
print("✅ PASS: Interrupt worked correctly!", file=sys.stderr)
set_interrupt(False)
return 0
print(f"❌ FAIL: status={result['status']}, elapsed={elapsed:.2f}s", file=sys.stderr)
set_interrupt(False)
return 1
print("❌ FAIL: No result returned", file=sys.stderr)
set_interrupt(False)
return 1
if __name__ == "__main__":
sys.exit(main())

View File

@@ -0,0 +1,161 @@
"""Test interrupt propagation from parent to child agents.
Reproduces the CLI scenario: user sends a message while delegate_task is
running, main thread calls parent.interrupt(), child should stop.
"""
import json
import threading
import time
import unittest
from unittest.mock import MagicMock, patch, PropertyMock
from tools.interrupt import set_interrupt, is_interrupted, _interrupt_event
class TestInterruptPropagationToChild(unittest.TestCase):
"""Verify interrupt propagates from parent to child agent."""
def setUp(self):
set_interrupt(False)
def tearDown(self):
set_interrupt(False)
def test_parent_interrupt_sets_child_flag(self):
"""When parent.interrupt() is called, child._interrupt_requested should be set."""
from run_agent import AIAgent
parent = AIAgent.__new__(AIAgent)
parent._interrupt_requested = False
parent._interrupt_message = None
parent._active_children = []
parent._active_children_lock = threading.Lock()
parent.quiet_mode = True
child = AIAgent.__new__(AIAgent)
child._interrupt_requested = False
child._interrupt_message = None
child._active_children = []
child._active_children_lock = threading.Lock()
child.quiet_mode = True
parent._active_children.append(child)
parent.interrupt("new user message")
assert parent._interrupt_requested is True
assert child._interrupt_requested is True
assert child._interrupt_message == "new user message"
assert is_interrupted() is True
def test_child_clear_interrupt_at_start_clears_global(self):
"""child.clear_interrupt() at start of run_conversation clears the GLOBAL event.
This is the intended behavior at startup, but verify it doesn't
accidentally clear an interrupt intended for a running child.
"""
from run_agent import AIAgent
child = AIAgent.__new__(AIAgent)
child._interrupt_requested = True
child._interrupt_message = "msg"
child.quiet_mode = True
child._active_children = []
child._active_children_lock = threading.Lock()
# Global is set
set_interrupt(True)
assert is_interrupted() is True
# child.clear_interrupt() clears both
child.clear_interrupt()
assert child._interrupt_requested is False
assert is_interrupted() is False
def test_interrupt_during_child_api_call_detected(self):
"""Interrupt set during _interruptible_api_call is detected within 0.5s."""
from run_agent import AIAgent
child = AIAgent.__new__(AIAgent)
child._interrupt_requested = False
child._interrupt_message = None
child._active_children = []
child._active_children_lock = threading.Lock()
child.quiet_mode = True
child.api_mode = "chat_completions"
child.log_prefix = ""
child._client_kwargs = {"api_key": "test", "base_url": "http://localhost:1234"}
# Mock a slow API call
mock_client = MagicMock()
def slow_api_call(**kwargs):
time.sleep(5) # Would take 5s normally
return MagicMock()
mock_client.chat.completions.create = slow_api_call
mock_client.close = MagicMock()
child.client = mock_client
# Set interrupt after 0.2s from another thread
def set_interrupt_later():
time.sleep(0.2)
child.interrupt("stop!")
t = threading.Thread(target=set_interrupt_later, daemon=True)
t.start()
start = time.monotonic()
try:
child._interruptible_api_call({"model": "test", "messages": []})
self.fail("Should have raised InterruptedError")
except InterruptedError:
elapsed = time.monotonic() - start
# Should detect within ~0.5s (0.2s delay + 0.3s poll interval)
assert elapsed < 1.0, f"Took {elapsed:.2f}s to detect interrupt (expected < 1.0s)"
finally:
t.join(timeout=2)
set_interrupt(False)
def test_concurrent_interrupt_propagation(self):
"""Simulates exact CLI flow: parent runs delegate in thread, main thread interrupts."""
from run_agent import AIAgent
parent = AIAgent.__new__(AIAgent)
parent._interrupt_requested = False
parent._interrupt_message = None
parent._active_children = []
parent._active_children_lock = threading.Lock()
parent.quiet_mode = True
child = AIAgent.__new__(AIAgent)
child._interrupt_requested = False
child._interrupt_message = None
child._active_children = []
child._active_children_lock = threading.Lock()
child.quiet_mode = True
# Register child (simulating what _run_single_child does)
parent._active_children.append(child)
# Simulate child running (checking flag in a loop)
child_detected = threading.Event()
def simulate_child_loop():
while not child._interrupt_requested:
time.sleep(0.05)
child_detected.set()
child_thread = threading.Thread(target=simulate_child_loop, daemon=True)
child_thread.start()
# Small delay, then interrupt from "main thread"
time.sleep(0.1)
parent.interrupt("user typed something new")
# Child should detect within 200ms
detected = child_detected.wait(timeout=1.0)
assert detected, "Child never detected the interrupt!"
child_thread.join(timeout=1)
set_interrupt(False)
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,162 @@
"""Tests for _save_oversized_tool_result() — the large tool response handler.
When a tool returns more than _LARGE_RESULT_CHARS characters, the full content
is saved to a file and the model receives a preview + file path instead.
"""
import os
import re
import pytest
from run_agent import (
_save_oversized_tool_result,
_LARGE_RESULT_CHARS,
_LARGE_RESULT_PREVIEW_CHARS,
)
class TestSaveOversizedToolResult:
"""Unit tests for the large tool result handler."""
def test_small_result_returned_unchanged(self):
"""Results under the threshold pass through untouched."""
small = "x" * 1000
assert _save_oversized_tool_result("terminal", small) is small
def test_exactly_at_threshold_returned_unchanged(self):
"""Results exactly at the threshold pass through."""
exact = "y" * _LARGE_RESULT_CHARS
assert _save_oversized_tool_result("terminal", exact) is exact
def test_oversized_result_saved_to_file(self, tmp_path, monkeypatch):
"""Results over the threshold are written to a file."""
monkeypatch.setenv("HERMES_HOME", str(tmp_path / ".hermes"))
os.makedirs(tmp_path / ".hermes", exist_ok=True)
big = "A" * (_LARGE_RESULT_CHARS + 500)
result = _save_oversized_tool_result("terminal", big)
# Should contain the preview
assert result.startswith("A" * _LARGE_RESULT_PREVIEW_CHARS)
# Should mention the file path
assert "Full output saved to:" in result
# Should mention original size
assert f"{len(big):,}" in result
# Extract the file path and verify the file exists with full content
match = re.search(r"Full output saved to: (.+?)\n", result)
assert match, f"No file path found in result: {result[:300]}"
filepath = match.group(1)
assert os.path.isfile(filepath)
with open(filepath, "r", encoding="utf-8") as f:
saved = f.read()
assert saved == big
assert len(saved) == _LARGE_RESULT_CHARS + 500
def test_file_placed_in_cache_tool_responses(self, tmp_path, monkeypatch):
"""Saved file lives under HERMES_HOME/cache/tool_responses/."""
hermes_home = str(tmp_path / ".hermes")
monkeypatch.setenv("HERMES_HOME", hermes_home)
os.makedirs(hermes_home, exist_ok=True)
big = "B" * (_LARGE_RESULT_CHARS + 1)
result = _save_oversized_tool_result("web_search", big)
match = re.search(r"Full output saved to: (.+?)\n", result)
filepath = match.group(1)
expected_dir = os.path.join(hermes_home, "cache", "tool_responses")
assert filepath.startswith(expected_dir)
def test_filename_contains_tool_name(self, tmp_path, monkeypatch):
"""The saved filename includes a sanitized version of the tool name."""
monkeypatch.setenv("HERMES_HOME", str(tmp_path / ".hermes"))
os.makedirs(tmp_path / ".hermes", exist_ok=True)
big = "C" * (_LARGE_RESULT_CHARS + 1)
result = _save_oversized_tool_result("browser_navigate", big)
match = re.search(r"Full output saved to: (.+?)\n", result)
filename = os.path.basename(match.group(1))
assert filename.startswith("browser_navigate_")
assert filename.endswith(".txt")
def test_tool_name_sanitized(self, tmp_path, monkeypatch):
"""Special characters in tool names are replaced in the filename."""
monkeypatch.setenv("HERMES_HOME", str(tmp_path / ".hermes"))
os.makedirs(tmp_path / ".hermes", exist_ok=True)
big = "D" * (_LARGE_RESULT_CHARS + 1)
result = _save_oversized_tool_result("mcp:some/weird tool", big)
match = re.search(r"Full output saved to: (.+?)\n", result)
filename = os.path.basename(match.group(1))
# No slashes or colons in filename
assert "/" not in filename
assert ":" not in filename
def test_fallback_on_write_failure(self, tmp_path, monkeypatch):
"""When file write fails, falls back to destructive truncation."""
# Point HERMES_HOME to a path that will fail (file, not directory)
bad_path = str(tmp_path / "not_a_dir.txt")
with open(bad_path, "w") as f:
f.write("I'm a file, not a directory")
monkeypatch.setenv("HERMES_HOME", bad_path)
big = "E" * (_LARGE_RESULT_CHARS + 50_000)
result = _save_oversized_tool_result("terminal", big)
# Should still contain data (fallback truncation)
assert len(result) > 0
assert result.startswith("E" * 1000)
# Should mention the failure
assert "File save failed" in result
# Should be truncated to approximately _LARGE_RESULT_CHARS + error msg
assert len(result) < len(big)
def test_preview_length_capped(self, tmp_path, monkeypatch):
"""The inline preview is capped at _LARGE_RESULT_PREVIEW_CHARS."""
monkeypatch.setenv("HERMES_HOME", str(tmp_path / ".hermes"))
os.makedirs(tmp_path / ".hermes", exist_ok=True)
# Use distinct chars so we can measure the preview
big = "Z" * (_LARGE_RESULT_CHARS + 5000)
result = _save_oversized_tool_result("terminal", big)
# The preview section is the content before the "[Large tool response:" marker
marker_pos = result.index("[Large tool response:")
preview_section = result[:marker_pos].rstrip()
assert len(preview_section) == _LARGE_RESULT_PREVIEW_CHARS
def test_guidance_message_mentions_tools(self, tmp_path, monkeypatch):
"""The replacement message tells the model how to access the file."""
monkeypatch.setenv("HERMES_HOME", str(tmp_path / ".hermes"))
os.makedirs(tmp_path / ".hermes", exist_ok=True)
big = "F" * (_LARGE_RESULT_CHARS + 1)
result = _save_oversized_tool_result("terminal", big)
assert "read_file" in result
assert "search_files" in result
def test_empty_result_passes_through(self):
"""Empty strings are not oversized."""
assert _save_oversized_tool_result("terminal", "") == ""
def test_unicode_content_preserved(self, tmp_path, monkeypatch):
"""Unicode content is fully preserved in the saved file."""
monkeypatch.setenv("HERMES_HOME", str(tmp_path / ".hermes"))
os.makedirs(tmp_path / ".hermes", exist_ok=True)
# Mix of ASCII and multi-byte unicode to exceed threshold
unit = "Hello 世界! 🎉 " * 100 # ~1400 chars per repeat
big = unit * ((_LARGE_RESULT_CHARS // len(unit)) + 1)
assert len(big) > _LARGE_RESULT_CHARS
result = _save_oversized_tool_result("terminal", big)
match = re.search(r"Full output saved to: (.+?)\n", result)
filepath = match.group(1)
with open(filepath, "r", encoding="utf-8") as f:
saved = f.read()
assert saved == big

View File

@@ -0,0 +1,209 @@
"""Tests for Anthropic Sonnet long-context tier 429 handling.
When Claude Max users without "extra usage" hit the 1M context tier
on Sonnet, Anthropic returns HTTP 429 "Extra usage is required for long
context requests." This is NOT a transient rate limit — the agent should
reduce context_length to 200k and compress instead of retrying.
Only Sonnet is affected — Opus 1M is general access.
"""
import pytest
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
# ---------------------------------------------------------------------------
# Detection logic
# ---------------------------------------------------------------------------
class TestLongContextTierDetection:
"""Verify the detection heuristic matches the Anthropic error."""
@staticmethod
def _is_long_context_tier_error(status_code, error_msg, model="claude-sonnet-4.6"):
error_msg = error_msg.lower()
return (
status_code == 429
and "extra usage" in error_msg
and "long context" in error_msg
and "sonnet" in model.lower()
)
def test_matches_anthropic_error(self):
assert self._is_long_context_tier_error(
429,
"Extra usage is required for long context requests.",
)
def test_matches_lowercase(self):
assert self._is_long_context_tier_error(
429,
"extra usage is required for long context requests.",
)
def test_matches_openrouter_model_id(self):
assert self._is_long_context_tier_error(
429,
"Extra usage is required for long context requests.",
model="anthropic/claude-sonnet-4.6",
)
def test_matches_nous_model_id(self):
assert self._is_long_context_tier_error(
429,
"Extra usage is required for long context requests.",
model="claude-sonnet-4-6",
)
def test_rejects_opus(self):
"""Opus 1M is general access — should NOT trigger reduction."""
assert not self._is_long_context_tier_error(
429,
"Extra usage is required for long context requests.",
model="claude-opus-4.6",
)
def test_rejects_opus_openrouter(self):
assert not self._is_long_context_tier_error(
429,
"Extra usage is required for long context requests.",
model="anthropic/claude-opus-4.6",
)
def test_rejects_normal_429(self):
assert not self._is_long_context_tier_error(
429,
"Rate limit exceeded. Please retry after 30 seconds.",
)
def test_rejects_wrong_status(self):
assert not self._is_long_context_tier_error(
400,
"Extra usage is required for long context requests.",
)
def test_rejects_partial_match(self):
"""Both 'extra usage' AND 'long context' must be present."""
assert not self._is_long_context_tier_error(
429, "extra usage required"
)
assert not self._is_long_context_tier_error(
429, "long context requests not supported"
)
# ---------------------------------------------------------------------------
# Context reduction
# ---------------------------------------------------------------------------
class TestContextReduction:
"""When the long-context tier error fires, context_length should
drop to 200k and the reduced flag should be set correctly."""
def _make_compressor(self, context_length=1_000_000, threshold_percent=0.5):
c = SimpleNamespace(
context_length=context_length,
threshold_percent=threshold_percent,
threshold_tokens=int(context_length * threshold_percent),
_context_probed=False,
_context_probe_persistable=False,
)
return c
def test_reduces_1m_to_200k(self):
comp = self._make_compressor(1_000_000)
reduced_ctx = 200_000
if comp.context_length > reduced_ctx:
comp.context_length = reduced_ctx
comp.threshold_tokens = int(reduced_ctx * comp.threshold_percent)
comp._context_probed = True
comp._context_probe_persistable = False
assert comp.context_length == 200_000
assert comp.threshold_tokens == 100_000
assert comp._context_probed is True
# Must NOT persist — subscription tier, not model capability
assert comp._context_probe_persistable is False
def test_no_reduction_when_already_200k(self):
comp = self._make_compressor(200_000)
reduced_ctx = 200_000
original = comp.context_length
if comp.context_length > reduced_ctx:
comp.context_length = reduced_ctx
assert comp.context_length == original # unchanged
def test_no_reduction_when_below_200k(self):
comp = self._make_compressor(128_000)
reduced_ctx = 200_000
original = comp.context_length
if comp.context_length > reduced_ctx:
comp.context_length = reduced_ctx
assert comp.context_length == original # unchanged
# ---------------------------------------------------------------------------
# Integration: agent error handler path
# ---------------------------------------------------------------------------
class TestAgentErrorPath:
"""Verify the long-context 429 doesn't hit the generic rate-limit
or client-error handlers."""
def test_long_context_429_not_treated_as_rate_limit(self):
"""The error should be intercepted before the generic
is_rate_limited check fires a fallback switch."""
error_msg = "extra usage is required for long context requests."
status_code = 429
model = "claude-sonnet-4.6"
_is_long_context_tier_error = (
status_code == 429
and "extra usage" in error_msg
and "long context" in error_msg
and "sonnet" in model.lower()
)
assert _is_long_context_tier_error
def test_opus_429_falls_through_to_rate_limit(self):
"""Opus should NOT match — falls through to generic rate-limit."""
error_msg = "extra usage is required for long context requests."
status_code = 429
model = "claude-opus-4.6"
_is_long_context_tier_error = (
status_code == 429
and "extra usage" in error_msg
and "long context" in error_msg
and "sonnet" in model.lower()
)
assert not _is_long_context_tier_error
def test_normal_429_still_treated_as_rate_limit(self):
"""A normal 429 should NOT match the long-context check."""
error_msg = "rate limit exceeded"
status_code = 429
model = "claude-sonnet-4.6"
_is_long_context_tier_error = (
status_code == 429
and "extra usage" in error_msg
and "long context" in error_msg
and "sonnet" in model.lower()
)
assert not _is_long_context_tier_error
is_rate_limited = (
status_code == 429
or "rate limit" in error_msg
)
assert is_rate_limited

View File

@@ -0,0 +1,189 @@
import sys
import threading
import types
from types import SimpleNamespace
import httpx
import pytest
from openai import APIConnectionError
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
class FakeRequestClient:
def __init__(self, responder):
self._responder = responder
self._client = SimpleNamespace(is_closed=False)
self.chat = SimpleNamespace(
completions=SimpleNamespace(create=self._create)
)
self.responses = SimpleNamespace()
self.close_calls = 0
def _create(self, **kwargs):
return self._responder(**kwargs)
def close(self):
self.close_calls += 1
self._client.is_closed = True
class FakeSharedClient(FakeRequestClient):
pass
class OpenAIFactory:
def __init__(self, clients):
self._clients = list(clients)
self.calls = []
def __call__(self, **kwargs):
self.calls.append(dict(kwargs))
if not self._clients:
raise AssertionError("OpenAI factory exhausted")
return self._clients.pop(0)
def _build_agent(shared_client=None):
agent = run_agent.AIAgent.__new__(run_agent.AIAgent)
agent.api_mode = "chat_completions"
agent.provider = "openai-codex"
agent.base_url = "https://chatgpt.com/backend-api/codex"
agent.model = "gpt-5-codex"
agent.log_prefix = ""
agent.quiet_mode = True
agent._interrupt_requested = False
agent._interrupt_message = None
agent._client_lock = threading.RLock()
agent._client_kwargs = {"api_key": "***", "base_url": agent.base_url}
agent.client = shared_client or FakeSharedClient(lambda **kwargs: {"shared": True})
agent.stream_delta_callback = None
agent._stream_callback = None
agent.reasoning_callback = None
return agent
def _connection_error():
return APIConnectionError(
message="Connection error.",
request=httpx.Request("POST", "https://example.com/v1/chat/completions"),
)
def test_retry_after_api_connection_error_recreates_request_client(monkeypatch):
first_request = FakeRequestClient(lambda **kwargs: (_ for _ in ()).throw(_connection_error()))
second_request = FakeRequestClient(lambda **kwargs: {"ok": True})
factory = OpenAIFactory([first_request, second_request])
monkeypatch.setattr(run_agent, "OpenAI", factory)
agent = _build_agent()
with pytest.raises(APIConnectionError):
agent._interruptible_api_call({"model": agent.model, "messages": []})
result = agent._interruptible_api_call({"model": agent.model, "messages": []})
assert result == {"ok": True}
assert len(factory.calls) == 2
assert first_request.close_calls >= 1
assert second_request.close_calls >= 1
def test_closed_shared_client_is_recreated_before_request(monkeypatch):
stale_shared = FakeSharedClient(lambda **kwargs: (_ for _ in ()).throw(AssertionError("stale shared client used")))
stale_shared._client.is_closed = True
replacement_shared = FakeSharedClient(lambda **kwargs: {"replacement": True})
request_client = FakeRequestClient(lambda **kwargs: {"ok": "fresh-request-client"})
factory = OpenAIFactory([replacement_shared, request_client])
monkeypatch.setattr(run_agent, "OpenAI", factory)
agent = _build_agent(shared_client=stale_shared)
result = agent._interruptible_api_call({"model": agent.model, "messages": []})
assert result == {"ok": "fresh-request-client"}
assert agent.client is replacement_shared
assert stale_shared.close_calls >= 1
assert replacement_shared.close_calls == 0
assert len(factory.calls) == 2
def test_concurrent_requests_do_not_break_each_other_when_one_client_closes(monkeypatch):
first_started = threading.Event()
first_closed = threading.Event()
def first_responder(**kwargs):
first_started.set()
first_client.close()
first_closed.set()
raise _connection_error()
def second_responder(**kwargs):
assert first_started.wait(timeout=2)
assert first_closed.wait(timeout=2)
return {"ok": "second"}
first_client = FakeRequestClient(first_responder)
second_client = FakeRequestClient(second_responder)
factory = OpenAIFactory([first_client, second_client])
monkeypatch.setattr(run_agent, "OpenAI", factory)
agent = _build_agent()
results = {}
def run_call(name):
try:
results[name] = agent._interruptible_api_call({"model": agent.model, "messages": []})
except Exception as exc: # noqa: BLE001 - asserting exact type below
results[name] = exc
thread_one = threading.Thread(target=run_call, args=("first",), daemon=True)
thread_two = threading.Thread(target=run_call, args=("second",), daemon=True)
thread_one.start()
thread_two.start()
thread_one.join(timeout=5)
thread_two.join(timeout=5)
values = list(results.values())
assert sum(isinstance(value, APIConnectionError) for value in values) == 1
assert values.count({"ok": "second"}) == 1
assert len(factory.calls) == 2
def test_streaming_call_recreates_closed_shared_client_before_request(monkeypatch):
chunks = iter([
SimpleNamespace(
model="gpt-5-codex",
choices=[SimpleNamespace(delta=SimpleNamespace(content="Hello", tool_calls=None), finish_reason=None)],
),
SimpleNamespace(
model="gpt-5-codex",
choices=[SimpleNamespace(delta=SimpleNamespace(content=" world", tool_calls=None), finish_reason="stop")],
),
])
stale_shared = FakeSharedClient(lambda **kwargs: (_ for _ in ()).throw(AssertionError("stale shared client used")))
stale_shared._client.is_closed = True
replacement_shared = FakeSharedClient(lambda **kwargs: {"replacement": True})
request_client = FakeRequestClient(lambda **kwargs: chunks)
factory = OpenAIFactory([replacement_shared, request_client])
monkeypatch.setattr(run_agent, "OpenAI", factory)
agent = _build_agent(shared_client=stale_shared)
agent.stream_delta_callback = lambda _delta: None
# Force chat_completions mode so the streaming path uses
# chat.completions.create(stream=True) instead of Codex responses.stream()
agent.api_mode = "chat_completions"
response = agent._interruptible_streaming_api_call({"model": agent.model, "messages": []})
assert response.choices[0].message.content == "Hello world"
assert agent.client is replacement_shared
assert stale_shared.close_calls >= 1
assert request_client.close_calls >= 1
assert len(factory.calls) == 2

View File

@@ -0,0 +1,154 @@
"""Tests for percentage clamping at 100% across display paths.
PR #3480 capped context pressure percentage at 100% in agent/display.py
but missed the same unclamped pattern in 4 other files. When token counts
overshoot the context length (possible during streaming or before
compression fires), users see >100% in /stats, gateway status, and
memory tool output.
"""
import pytest
class TestContextCompressorUsagePercent:
"""agent/context_compressor.py — get_status() usage_percent"""
def test_usage_percent_capped_at_100(self):
"""Tokens exceeding context_length should still show max 100%."""
from agent.context_compressor import ContextCompressor
comp = ContextCompressor.__new__(ContextCompressor)
comp.last_prompt_tokens = 210_000 # exceeds context_length
comp.context_length = 200_000
comp.threshold_tokens = 160_000
comp.compression_count = 0
status = comp.get_status()
assert status["usage_percent"] <= 100
def test_usage_percent_normal(self):
"""Normal usage should show correct percentage."""
from agent.context_compressor import ContextCompressor
comp = ContextCompressor.__new__(ContextCompressor)
comp.last_prompt_tokens = 100_000
comp.context_length = 200_000
comp.threshold_tokens = 160_000
comp.compression_count = 0
status = comp.get_status()
assert status["usage_percent"] == 50.0
def test_usage_percent_zero_context_length(self):
"""Zero context_length should return 0, not crash."""
from agent.context_compressor import ContextCompressor
comp = ContextCompressor.__new__(ContextCompressor)
comp.last_prompt_tokens = 1000
comp.context_length = 0
comp.threshold_tokens = 0
comp.compression_count = 0
status = comp.get_status()
assert status["usage_percent"] == 0
class TestMemoryToolPercentClamp:
"""tools/memory_tool.py — _success_response and _render_block pct"""
def test_over_limit_clamped_at_100(self):
"""Percentage should be capped at 100 even if current > limit."""
# Simulate the calculation directly
current = 5500
limit = 5000
pct = min(100, int((current / limit) * 100)) if limit > 0 else 0
assert pct == 100
def test_normal_percentage(self):
current = 2500
limit = 5000
pct = min(100, int((current / limit) * 100)) if limit > 0 else 0
assert pct == 50
def test_zero_limit_returns_zero(self):
current = 100
limit = 0
pct = min(100, int((current / limit) * 100)) if limit > 0 else 0
assert pct == 0
class TestCLIStatsPercentClamp:
"""cli.py — /stats command percentage"""
def test_over_context_clamped_at_100(self):
"""Tokens exceeding context_length should show max 100%."""
last_prompt = 210_000
ctx_len = 200_000
pct = min(100, (last_prompt / ctx_len * 100)) if ctx_len else 0
assert pct == 100
def test_normal_context(self):
last_prompt = 100_000
ctx_len = 200_000
pct = min(100, (last_prompt / ctx_len * 100)) if ctx_len else 0
assert pct == 50.0
def test_zero_context_length(self):
last_prompt = 1000
ctx_len = 0
pct = min(100, (last_prompt / ctx_len * 100)) if ctx_len else 0
assert pct == 0
class TestGatewayStatsPercentClamp:
"""gateway/run.py — _format_usage_stats percentage"""
def test_over_context_clamped_at_100(self):
last_prompt_tokens = 210_000
context_length = 200_000
pct = min(100, last_prompt_tokens / context_length * 100) if context_length else 0
assert pct == 100
def test_normal_context(self):
last_prompt_tokens = 150_000
context_length = 200_000
pct = min(100, last_prompt_tokens / context_length * 100) if context_length else 0
assert pct == 75.0
class TestSourceLinesAreClamped:
"""Verify the actual source files have min(100, ...) applied."""
@staticmethod
def _read_file(rel_path: str) -> str:
import os
base = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
with open(os.path.join(base, rel_path)) as f:
return f.read()
def test_context_compressor_clamped(self):
src = self._read_file("agent/context_compressor.py")
assert "min(100," in src, (
"context_compressor.py usage_percent is not clamped with min(100, ...)"
)
def test_gateway_run_clamped(self):
src = self._read_file("gateway/run.py")
# Check that the stats handler has min(100, ...)
assert "min(100, ctx.last_prompt_tokens" in src, (
"gateway/run.py stats pct is not clamped with min(100, ...)"
)
def test_cli_clamped(self):
src = self._read_file("cli.py")
assert "min(100, (last_prompt" in src, (
"cli.py /stats pct is not clamped with min(100, ...)"
)
def test_memory_tool_clamped(self):
src = self._read_file("tools/memory_tool.py")
# Both _success_response and _render_block should have min(100, ...)
count = src.count("min(100, int((current / limit)")
assert count >= 2, (
f"memory_tool.py has only {count} clamped pct lines, expected >= 2"
)

View File

@@ -0,0 +1,424 @@
"""Tests for per-turn primary runtime restoration and transport recovery.
Verifies that:
1. Fallback is turn-scoped: a new turn restores the primary model/provider
2. The fallback chain index resets so all fallbacks are available again
3. Context compressor state is restored alongside the runtime
4. Transient transport errors get one recovery cycle before fallback
5. Recovery is skipped for aggregator providers (OpenRouter, Nous)
6. Non-transport errors don't trigger recovery
"""
import time
from types import SimpleNamespace
from unittest.mock import MagicMock, patch, PropertyMock
import pytest
from run_agent import AIAgent
def _make_tool_defs(*names: str) -> list:
return [
{
"type": "function",
"function": {
"name": n,
"description": f"{n} tool",
"parameters": {"type": "object", "properties": {}},
},
}
for n in names
]
def _make_agent(fallback_model=None, provider="custom", base_url="https://my-llm.example.com/v1"):
"""Create a minimal AIAgent with optional fallback config."""
with (
patch("run_agent.get_tool_definitions", return_value=_make_tool_defs("web_search")),
patch("run_agent.check_toolset_requirements", return_value={}),
patch("run_agent.OpenAI"),
):
agent = AIAgent(
api_key="test-key-12345678",
base_url=base_url,
provider=provider,
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
fallback_model=fallback_model,
)
agent.client = MagicMock()
return agent
def _mock_resolve(base_url="https://openrouter.ai/api/v1", api_key="fallback-key-1234"):
"""Helper to create a mock client for resolve_provider_client."""
mock_client = MagicMock()
mock_client.api_key = api_key
mock_client.base_url = base_url
return mock_client
# =============================================================================
# _primary_runtime snapshot
# =============================================================================
class TestPrimaryRuntimeSnapshot:
def test_snapshot_created_at_init(self):
agent = _make_agent()
assert hasattr(agent, "_primary_runtime")
rt = agent._primary_runtime
assert rt["model"] == agent.model
assert rt["provider"] == "custom"
assert rt["base_url"] == "https://my-llm.example.com/v1"
assert rt["api_mode"] == agent.api_mode
assert "client_kwargs" in rt
assert "compressor_context_length" in rt
def test_snapshot_includes_compressor_state(self):
agent = _make_agent()
rt = agent._primary_runtime
cc = agent.context_compressor
assert rt["compressor_model"] == cc.model
assert rt["compressor_provider"] == cc.provider
assert rt["compressor_context_length"] == cc.context_length
assert rt["compressor_threshold_tokens"] == cc.threshold_tokens
def test_snapshot_includes_anthropic_state_when_applicable(self):
"""Anthropic-mode agents should snapshot Anthropic-specific state."""
with (
patch("run_agent.get_tool_definitions", return_value=_make_tool_defs("web_search")),
patch("run_agent.check_toolset_requirements", return_value={}),
patch("run_agent.OpenAI"),
patch("agent.anthropic_adapter.build_anthropic_client", return_value=MagicMock()),
):
agent = AIAgent(
api_key="sk-ant-test-12345678",
base_url="https://api.anthropic.com",
provider="anthropic",
api_mode="anthropic_messages",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
)
rt = agent._primary_runtime
assert "anthropic_api_key" in rt
assert "anthropic_base_url" in rt
assert "is_anthropic_oauth" in rt
def test_snapshot_omits_anthropic_for_openai_mode(self):
agent = _make_agent(provider="custom")
rt = agent._primary_runtime
assert "anthropic_api_key" not in rt
# =============================================================================
# _restore_primary_runtime()
# =============================================================================
class TestRestorePrimaryRuntime:
def test_noop_when_not_fallback(self):
agent = _make_agent()
assert agent._fallback_activated is False
assert agent._restore_primary_runtime() is False
def test_restores_model_and_provider(self):
agent = _make_agent(
fallback_model={"provider": "openrouter", "model": "anthropic/claude-sonnet-4"},
)
original_model = agent.model
original_provider = agent.provider
# Simulate fallback activation
mock_client = _mock_resolve()
with patch("agent.auxiliary_client.resolve_provider_client", return_value=(mock_client, None)):
agent._try_activate_fallback()
assert agent._fallback_activated is True
assert agent.model == "anthropic/claude-sonnet-4"
assert agent.provider == "openrouter"
# Restore should bring back the primary
with patch("run_agent.OpenAI", return_value=MagicMock()):
result = agent._restore_primary_runtime()
assert result is True
assert agent._fallback_activated is False
assert agent.model == original_model
assert agent.provider == original_provider
def test_resets_fallback_index(self):
"""After restore, the full fallback chain should be available again."""
agent = _make_agent(
fallback_model=[
{"provider": "openrouter", "model": "model-a"},
{"provider": "anthropic", "model": "model-b"},
],
)
# Advance through the chain
mock_client = _mock_resolve()
with patch("agent.auxiliary_client.resolve_provider_client", return_value=(mock_client, None)):
agent._try_activate_fallback()
assert agent._fallback_index == 1 # consumed one entry
with patch("run_agent.OpenAI", return_value=MagicMock()):
agent._restore_primary_runtime()
assert agent._fallback_index == 0 # reset for next turn
def test_restores_compressor_state(self):
agent = _make_agent(
fallback_model={"provider": "openrouter", "model": "anthropic/claude-sonnet-4"},
)
original_ctx_len = agent.context_compressor.context_length
original_threshold = agent.context_compressor.threshold_tokens
# Simulate fallback modifying compressor
mock_client = _mock_resolve()
with patch("agent.auxiliary_client.resolve_provider_client", return_value=(mock_client, None)):
agent._try_activate_fallback()
# Manually simulate compressor being changed (as _try_activate_fallback does)
agent.context_compressor.context_length = 32000
agent.context_compressor.threshold_tokens = 25600
with patch("run_agent.OpenAI", return_value=MagicMock()):
agent._restore_primary_runtime()
assert agent.context_compressor.context_length == original_ctx_len
assert agent.context_compressor.threshold_tokens == original_threshold
def test_restores_prompt_caching_flag(self):
agent = _make_agent()
original_caching = agent._use_prompt_caching
# Simulate fallback changing the caching flag
agent._fallback_activated = True
agent._use_prompt_caching = not original_caching
with patch("run_agent.OpenAI", return_value=MagicMock()):
agent._restore_primary_runtime()
assert agent._use_prompt_caching == original_caching
def test_restore_survives_exception(self):
"""If client rebuild fails, the method returns False gracefully."""
agent = _make_agent()
agent._fallback_activated = True
with patch("run_agent.OpenAI", side_effect=Exception("connection refused")):
result = agent._restore_primary_runtime()
assert result is False
# =============================================================================
# _try_recover_primary_transport()
# =============================================================================
def _make_transport_error(error_type="ReadTimeout"):
"""Create an exception whose type().__name__ matches the given name."""
cls = type(error_type, (Exception,), {})
return cls("connection timed out")
class TestTryRecoverPrimaryTransport:
def test_recovers_on_read_timeout(self):
agent = _make_agent(provider="custom")
error = _make_transport_error("ReadTimeout")
with patch("run_agent.OpenAI", return_value=MagicMock()), \
patch("time.sleep"):
result = agent._try_recover_primary_transport(
error, retry_count=3, max_retries=3,
)
assert result is True
def test_recovers_on_connect_timeout(self):
agent = _make_agent(provider="custom")
error = _make_transport_error("ConnectTimeout")
with patch("run_agent.OpenAI", return_value=MagicMock()), \
patch("time.sleep"):
result = agent._try_recover_primary_transport(
error, retry_count=3, max_retries=3,
)
assert result is True
def test_recovers_on_pool_timeout(self):
agent = _make_agent(provider="zai")
error = _make_transport_error("PoolTimeout")
with patch("run_agent.OpenAI", return_value=MagicMock()), \
patch("time.sleep"):
result = agent._try_recover_primary_transport(
error, retry_count=3, max_retries=3,
)
assert result is True
def test_skipped_when_already_on_fallback(self):
agent = _make_agent(provider="custom")
agent._fallback_activated = True
error = _make_transport_error("ReadTimeout")
result = agent._try_recover_primary_transport(
error, retry_count=3, max_retries=3,
)
assert result is False
def test_skipped_for_non_transport_error(self):
"""Non-transport errors (ValueError, APIError, etc.) skip recovery."""
agent = _make_agent(provider="custom")
error = ValueError("invalid model")
result = agent._try_recover_primary_transport(
error, retry_count=3, max_retries=3,
)
assert result is False
def test_skipped_for_openrouter(self):
agent = _make_agent(provider="openrouter", base_url="https://openrouter.ai/api/v1")
error = _make_transport_error("ReadTimeout")
result = agent._try_recover_primary_transport(
error, retry_count=3, max_retries=3,
)
assert result is False
def test_skipped_for_nous_provider(self):
agent = _make_agent(provider="nous", base_url="https://inference.nous.nousresearch.com/v1")
error = _make_transport_error("ReadTimeout")
result = agent._try_recover_primary_transport(
error, retry_count=3, max_retries=3,
)
assert result is False
def test_allowed_for_anthropic_direct(self):
"""Direct Anthropic endpoint should get recovery."""
agent = _make_agent(provider="anthropic", base_url="https://api.anthropic.com")
# For non-anthropic_messages api_mode, it will use OpenAI client
error = _make_transport_error("ConnectError")
with patch("run_agent.OpenAI", return_value=MagicMock()), \
patch("time.sleep"):
result = agent._try_recover_primary_transport(
error, retry_count=3, max_retries=3,
)
assert result is True
def test_allowed_for_ollama(self):
agent = _make_agent(provider="ollama", base_url="http://localhost:11434/v1")
error = _make_transport_error("ConnectTimeout")
with patch("run_agent.OpenAI", return_value=MagicMock()), \
patch("time.sleep"):
result = agent._try_recover_primary_transport(
error, retry_count=3, max_retries=3,
)
assert result is True
def test_wait_time_scales_with_retry_count(self):
agent = _make_agent(provider="custom")
error = _make_transport_error("ReadTimeout")
with patch("run_agent.OpenAI", return_value=MagicMock()), \
patch("time.sleep") as mock_sleep:
agent._try_recover_primary_transport(
error, retry_count=3, max_retries=3,
)
# wait_time = min(3 + retry_count, 8) = min(6, 8) = 6
mock_sleep.assert_called_once_with(6)
def test_wait_time_capped_at_8(self):
agent = _make_agent(provider="custom")
error = _make_transport_error("ReadTimeout")
with patch("run_agent.OpenAI", return_value=MagicMock()), \
patch("time.sleep") as mock_sleep:
agent._try_recover_primary_transport(
error, retry_count=10, max_retries=3,
)
# wait_time = min(3 + 10, 8) = 8
mock_sleep.assert_called_once_with(8)
def test_closes_existing_client_before_rebuild(self):
agent = _make_agent(provider="custom")
old_client = agent.client
error = _make_transport_error("ReadTimeout")
with patch("run_agent.OpenAI", return_value=MagicMock()), \
patch("time.sleep"), \
patch.object(agent, "_close_openai_client") as mock_close:
agent._try_recover_primary_transport(
error, retry_count=3, max_retries=3,
)
mock_close.assert_called_once_with(
old_client, reason="primary_recovery", shared=True,
)
def test_survives_rebuild_failure(self):
"""If client rebuild fails, returns False gracefully."""
agent = _make_agent(provider="custom")
error = _make_transport_error("ReadTimeout")
with patch("run_agent.OpenAI", side_effect=Exception("socket error")), \
patch("time.sleep"):
result = agent._try_recover_primary_transport(
error, retry_count=3, max_retries=3,
)
assert result is False
# =============================================================================
# Integration: restore_primary_runtime called from run_conversation
# =============================================================================
class TestRestoreInRunConversation:
"""Verify the hook in run_conversation() calls _restore_primary_runtime."""
def test_restore_called_at_turn_start(self):
agent = _make_agent()
agent._fallback_activated = True
with patch.object(agent, "_restore_primary_runtime", return_value=True) as mock_restore, \
patch.object(agent, "run_conversation", wraps=None) as _:
# We can't easily run the full conversation, but we can verify
# the method exists and is callable
agent._restore_primary_runtime()
mock_restore.assert_called_once()
def test_full_cycle_fallback_then_restore(self):
"""Simulate: turn 1 activates fallback, turn 2 restores primary."""
agent = _make_agent(
fallback_model={"provider": "openrouter", "model": "anthropic/claude-sonnet-4"},
provider="custom",
)
# Turn 1: activate fallback
mock_client = _mock_resolve()
with patch("agent.auxiliary_client.resolve_provider_client", return_value=(mock_client, None)):
assert agent._try_activate_fallback() is True
assert agent._fallback_activated is True
assert agent.model == "anthropic/claude-sonnet-4"
assert agent.provider == "openrouter"
assert agent._fallback_index == 1
# Turn 2: restore primary
with patch("run_agent.OpenAI", return_value=MagicMock()):
assert agent._restore_primary_runtime() is True
assert agent._fallback_activated is False
assert agent._fallback_index == 0
assert agent.provider == "custom"
assert agent.base_url == "https://my-llm.example.com/v1"

View File

@@ -0,0 +1,156 @@
"""Tests for ordered provider fallback chain (salvage of PR #1761).
Extends the single-fallback tests in test_fallback_model.py to cover
the new list-based ``fallback_providers`` config format and chain
advancement through multiple providers.
"""
from unittest.mock import MagicMock, patch
from run_agent import AIAgent
def _make_agent(fallback_model=None):
"""Create a minimal AIAgent with optional fallback config."""
with (
patch("run_agent.get_tool_definitions", return_value=[]),
patch("run_agent.check_toolset_requirements", return_value={}),
patch("run_agent.OpenAI"),
):
agent = AIAgent(
api_key="test-key",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
fallback_model=fallback_model,
)
agent.client = MagicMock()
return agent
def _mock_client(base_url="https://openrouter.ai/api/v1", api_key="fb-key"):
mock = MagicMock()
mock.base_url = base_url
mock.api_key = api_key
return mock
# ── Chain initialisation ──────────────────────────────────────────────────
class TestFallbackChainInit:
def test_no_fallback(self):
agent = _make_agent(fallback_model=None)
assert agent._fallback_chain == []
assert agent._fallback_index == 0
assert agent._fallback_model is None
def test_single_dict_backwards_compat(self):
fb = {"provider": "openai", "model": "gpt-4o"}
agent = _make_agent(fallback_model=fb)
assert agent._fallback_chain == [fb]
assert agent._fallback_model == fb
def test_list_of_providers(self):
fbs = [
{"provider": "openai", "model": "gpt-4o"},
{"provider": "zai", "model": "glm-4.7"},
]
agent = _make_agent(fallback_model=fbs)
assert len(agent._fallback_chain) == 2
assert agent._fallback_model == fbs[0]
def test_invalid_entries_filtered(self):
fbs = [
{"provider": "openai", "model": "gpt-4o"},
{"provider": "", "model": "glm-4.7"},
{"provider": "zai"},
"not-a-dict",
]
agent = _make_agent(fallback_model=fbs)
assert len(agent._fallback_chain) == 1
assert agent._fallback_chain[0]["provider"] == "openai"
def test_empty_list(self):
agent = _make_agent(fallback_model=[])
assert agent._fallback_chain == []
assert agent._fallback_model is None
def test_invalid_dict_no_provider(self):
agent = _make_agent(fallback_model={"model": "gpt-4o"})
assert agent._fallback_chain == []
# ── Chain advancement ─────────────────────────────────────────────────────
class TestFallbackChainAdvancement:
def test_exhausted_returns_false(self):
agent = _make_agent(fallback_model=None)
assert agent._try_activate_fallback() is False
def test_advances_index(self):
fbs = [
{"provider": "openai", "model": "gpt-4o"},
{"provider": "zai", "model": "glm-4.7"},
]
agent = _make_agent(fallback_model=fbs)
with patch("agent.auxiliary_client.resolve_provider_client",
return_value=(_mock_client(), "gpt-4o")):
assert agent._try_activate_fallback() is True
assert agent._fallback_index == 1
assert agent.model == "gpt-4o"
assert agent._fallback_activated is True
def test_second_fallback_works(self):
fbs = [
{"provider": "openai", "model": "gpt-4o"},
{"provider": "zai", "model": "glm-4.7"},
]
agent = _make_agent(fallback_model=fbs)
with patch("agent.auxiliary_client.resolve_provider_client",
return_value=(_mock_client(), "resolved")):
assert agent._try_activate_fallback() is True
assert agent.model == "gpt-4o"
assert agent._try_activate_fallback() is True
assert agent.model == "glm-4.7"
assert agent._fallback_index == 2
def test_all_exhausted_returns_false(self):
fbs = [{"provider": "openai", "model": "gpt-4o"}]
agent = _make_agent(fallback_model=fbs)
with patch("agent.auxiliary_client.resolve_provider_client",
return_value=(_mock_client(), "gpt-4o")):
assert agent._try_activate_fallback() is True
assert agent._try_activate_fallback() is False
def test_skips_unconfigured_provider_to_next(self):
"""If resolve_provider_client returns None, skip to next in chain."""
fbs = [
{"provider": "broken", "model": "nope"},
{"provider": "openai", "model": "gpt-4o"},
]
agent = _make_agent(fallback_model=fbs)
with patch("agent.auxiliary_client.resolve_provider_client") as mock_rpc:
mock_rpc.side_effect = [
(None, None), # broken provider
(_mock_client(), "gpt-4o"), # fallback succeeds
]
assert agent._try_activate_fallback() is True
assert agent.model == "gpt-4o"
assert agent._fallback_index == 2
def test_skips_provider_that_raises_to_next(self):
"""If resolve_provider_client raises, skip to next in chain."""
fbs = [
{"provider": "broken", "model": "nope"},
{"provider": "openai", "model": "gpt-4o"},
]
agent = _make_agent(fallback_model=fbs)
with patch("agent.auxiliary_client.resolve_provider_client") as mock_rpc:
mock_rpc.side_effect = [
RuntimeError("auth failed"),
(_mock_client(), "gpt-4o"),
]
assert agent._try_activate_fallback() is True
assert agent.model == "gpt-4o"

View File

@@ -0,0 +1,850 @@
"""Provider parity tests: verify that AIAgent builds correct API kwargs
and handles responses properly for all supported providers.
Ensures changes to one provider path don't silently break another.
"""
import json
import os
import sys
import types
from types import SimpleNamespace
from unittest.mock import patch, MagicMock
import pytest
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())
from run_agent import AIAgent
# ── Helpers ──────────────────────────────────────────────────────────────────
def _tool_defs(*names):
return [
{
"type": "function",
"function": {
"name": n,
"description": f"{n} tool",
"parameters": {"type": "object", "properties": {}},
},
}
for n in names
]
class _FakeOpenAI:
def __init__(self, **kw):
self.api_key = kw.get("api_key", "test")
self.base_url = kw.get("base_url", "http://test")
def close(self):
pass
def _make_agent(monkeypatch, provider, api_mode="chat_completions", base_url="https://openrouter.ai/api/v1"):
monkeypatch.setattr("run_agent.get_tool_definitions", lambda **kw: _tool_defs("web_search", "terminal"))
monkeypatch.setattr("run_agent.check_toolset_requirements", lambda: {})
monkeypatch.setattr("run_agent.OpenAI", _FakeOpenAI)
return AIAgent(
api_key="test-key",
base_url=base_url,
provider=provider,
api_mode=api_mode,
max_iterations=4,
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
)
# ── _build_api_kwargs tests ─────────────────────────────────────────────────
class TestBuildApiKwargsOpenRouter:
def test_uses_chat_completions_format(self, monkeypatch):
agent = _make_agent(monkeypatch, "openrouter")
messages = [{"role": "user", "content": "hi"}]
kwargs = agent._build_api_kwargs(messages)
assert "messages" in kwargs
assert "model" in kwargs
assert kwargs["messages"][-1]["content"] == "hi"
def test_includes_reasoning_in_extra_body(self, monkeypatch):
agent = _make_agent(monkeypatch, "openrouter")
agent.model = "anthropic/claude-sonnet-4-20250514"
messages = [{"role": "user", "content": "hi"}]
kwargs = agent._build_api_kwargs(messages)
extra = kwargs.get("extra_body", {})
assert "reasoning" in extra
assert extra["reasoning"]["enabled"] is True
def test_includes_tools(self, monkeypatch):
agent = _make_agent(monkeypatch, "openrouter")
messages = [{"role": "user", "content": "hi"}]
kwargs = agent._build_api_kwargs(messages)
assert "tools" in kwargs
tool_names = [t["function"]["name"] for t in kwargs["tools"]]
assert "web_search" in tool_names
def test_no_responses_api_fields(self, monkeypatch):
agent = _make_agent(monkeypatch, "openrouter")
messages = [{"role": "user", "content": "hi"}]
kwargs = agent._build_api_kwargs(messages)
assert "input" not in kwargs
assert "instructions" not in kwargs
assert "store" not in kwargs
def test_strips_codex_only_tool_call_fields_from_chat_messages(self, monkeypatch):
agent = _make_agent(monkeypatch, "openrouter")
messages = [
{"role": "user", "content": "hi"},
{
"role": "assistant",
"content": "Checking now.",
"codex_reasoning_items": [
{"type": "reasoning", "id": "rs_1", "encrypted_content": "blob"},
],
"tool_calls": [
{
"id": "call_123",
"call_id": "call_123",
"response_item_id": "fc_123",
"type": "function",
"function": {"name": "terminal", "arguments": "{\"command\":\"pwd\"}"},
"extra_content": {"thought_signature": "opaque"},
}
],
},
{"role": "tool", "tool_call_id": "call_123", "content": "/tmp"},
]
kwargs = agent._build_api_kwargs(messages)
assistant_msg = kwargs["messages"][1]
tool_call = assistant_msg["tool_calls"][0]
assert "codex_reasoning_items" not in assistant_msg
assert tool_call["id"] == "call_123"
assert tool_call["function"]["name"] == "terminal"
assert tool_call["extra_content"] == {"thought_signature": "opaque"}
assert "call_id" not in tool_call
assert "response_item_id" not in tool_call
# Original stored history must remain unchanged for Responses replay mode.
assert messages[1]["tool_calls"][0]["call_id"] == "call_123"
assert messages[1]["tool_calls"][0]["response_item_id"] == "fc_123"
assert "codex_reasoning_items" in messages[1]
def test_should_sanitize_tool_calls_codex_vs_chat(self, monkeypatch):
"""Codex API should NOT sanitize, all other APIs should sanitize."""
# Codex mode should NOT need sanitization
codex_agent = _make_agent(monkeypatch, "openrouter")
codex_agent.api_mode = "codex_responses"
assert codex_agent._should_sanitize_tool_calls() is False
# Chat completions mode should need sanitization
chat_agent = _make_agent(monkeypatch, "openrouter")
chat_agent.api_mode = "chat_completions"
assert chat_agent._should_sanitize_tool_calls() is True
# Anthropic mode should need sanitization
anthropic_agent = _make_agent(monkeypatch, "openrouter")
anthropic_agent.api_mode = "anthropic_messages"
assert anthropic_agent._should_sanitize_tool_calls() is True
class TestDeveloperRoleSwap:
"""GPT-5 and Codex models should get 'developer' instead of 'system' role."""
@pytest.mark.parametrize("model", [
"openai/gpt-5",
"openai/gpt-5-turbo",
"openai/gpt-5.4",
"gpt-5-mini",
"openai/codex-mini",
"codex-mini-latest",
"openai/codex-pro",
])
def test_gpt5_codex_get_developer_role(self, monkeypatch, model):
agent = _make_agent(monkeypatch, "openrouter")
agent.model = model
messages = [
{"role": "system", "content": "You are helpful."},
{"role": "user", "content": "hi"},
]
kwargs = agent._build_api_kwargs(messages)
assert kwargs["messages"][0]["role"] == "developer"
assert kwargs["messages"][0]["content"] == "You are helpful."
assert kwargs["messages"][1]["role"] == "user"
@pytest.mark.parametrize("model", [
"anthropic/claude-opus-4.6",
"openai/gpt-4o",
"google/gemini-2.5-pro",
"deepseek/deepseek-chat",
"openai/o3-mini",
])
def test_non_matching_models_keep_system_role(self, monkeypatch, model):
agent = _make_agent(monkeypatch, "openrouter")
agent.model = model
messages = [
{"role": "system", "content": "You are helpful."},
{"role": "user", "content": "hi"},
]
kwargs = agent._build_api_kwargs(messages)
assert kwargs["messages"][0]["role"] == "system"
def test_no_system_message_no_crash(self, monkeypatch):
agent = _make_agent(monkeypatch, "openrouter")
agent.model = "openai/gpt-5"
messages = [{"role": "user", "content": "hi"}]
kwargs = agent._build_api_kwargs(messages)
assert kwargs["messages"][0]["role"] == "user"
def test_original_messages_not_mutated(self, monkeypatch):
agent = _make_agent(monkeypatch, "openrouter")
agent.model = "openai/gpt-5"
messages = [
{"role": "system", "content": "You are helpful."},
{"role": "user", "content": "hi"},
]
agent._build_api_kwargs(messages)
# Original messages must be untouched (internal representation stays "system")
assert messages[0]["role"] == "system"
def test_developer_role_via_nous_portal(self, monkeypatch):
agent = _make_agent(monkeypatch, "nous", base_url="https://inference-api.nousresearch.com/v1")
agent.model = "gpt-5"
messages = [
{"role": "system", "content": "You are helpful."},
{"role": "user", "content": "hi"},
]
kwargs = agent._build_api_kwargs(messages)
assert kwargs["messages"][0]["role"] == "developer"
class TestBuildApiKwargsAIGateway:
def test_uses_chat_completions_format(self, monkeypatch):
agent = _make_agent(monkeypatch, "ai-gateway", base_url="https://ai-gateway.vercel.sh/v1")
messages = [{"role": "user", "content": "hi"}]
kwargs = agent._build_api_kwargs(messages)
assert "messages" in kwargs
assert "model" in kwargs
assert kwargs["messages"][-1]["content"] == "hi"
def test_no_responses_api_fields(self, monkeypatch):
agent = _make_agent(monkeypatch, "ai-gateway", base_url="https://ai-gateway.vercel.sh/v1")
messages = [{"role": "user", "content": "hi"}]
kwargs = agent._build_api_kwargs(messages)
assert "input" not in kwargs
assert "instructions" not in kwargs
assert "store" not in kwargs
def test_includes_reasoning_in_extra_body(self, monkeypatch):
agent = _make_agent(monkeypatch, "ai-gateway", base_url="https://ai-gateway.vercel.sh/v1")
messages = [{"role": "user", "content": "hi"}]
kwargs = agent._build_api_kwargs(messages)
extra = kwargs.get("extra_body", {})
assert "reasoning" in extra
assert extra["reasoning"]["enabled"] is True
def test_includes_tools(self, monkeypatch):
agent = _make_agent(monkeypatch, "ai-gateway", base_url="https://ai-gateway.vercel.sh/v1")
messages = [{"role": "user", "content": "hi"}]
kwargs = agent._build_api_kwargs(messages)
assert "tools" in kwargs
tool_names = [t["function"]["name"] for t in kwargs["tools"]]
assert "web_search" in tool_names
class TestBuildApiKwargsNousPortal:
def test_includes_nous_product_tags(self, monkeypatch):
agent = _make_agent(monkeypatch, "nous", base_url="https://inference-api.nousresearch.com/v1")
messages = [{"role": "user", "content": "hi"}]
kwargs = agent._build_api_kwargs(messages)
extra = kwargs.get("extra_body", {})
assert extra.get("tags") == ["product=hermes-agent"]
def test_uses_chat_completions_format(self, monkeypatch):
agent = _make_agent(monkeypatch, "nous", base_url="https://inference-api.nousresearch.com/v1")
messages = [{"role": "user", "content": "hi"}]
kwargs = agent._build_api_kwargs(messages)
assert "messages" in kwargs
assert "input" not in kwargs
class TestBuildApiKwargsCustomEndpoint:
def test_uses_chat_completions_format(self, monkeypatch):
agent = _make_agent(monkeypatch, "custom", base_url="http://localhost:1234/v1")
messages = [{"role": "user", "content": "hi"}]
kwargs = agent._build_api_kwargs(messages)
assert "messages" in kwargs
assert "input" not in kwargs
def test_no_openrouter_extra_body(self, monkeypatch):
agent = _make_agent(monkeypatch, "custom", base_url="http://localhost:1234/v1")
messages = [{"role": "user", "content": "hi"}]
kwargs = agent._build_api_kwargs(messages)
extra = kwargs.get("extra_body", {})
assert "reasoning" not in extra
def test_fireworks_tool_call_payload_strips_codex_only_fields(self, monkeypatch):
agent = _make_agent(
monkeypatch,
"custom",
base_url="https://api.fireworks.ai/inference/v1",
)
messages = [
{"role": "user", "content": "hi"},
{
"role": "assistant",
"content": "Checking now.",
"codex_reasoning_items": [
{"type": "reasoning", "id": "rs_1", "encrypted_content": "blob"},
],
"tool_calls": [
{
"id": "call_fw_123",
"call_id": "call_fw_123",
"response_item_id": "fc_fw_123",
"type": "function",
"function": {
"name": "terminal",
"arguments": "{\"command\":\"pwd\"}",
},
}
],
},
{"role": "tool", "tool_call_id": "call_fw_123", "content": "/tmp"},
]
kwargs = agent._build_api_kwargs(messages)
assert kwargs["tools"][0]["function"]["name"] == "web_search"
assert "input" not in kwargs
assert kwargs.get("extra_body", {}) == {}
assistant_msg = kwargs["messages"][1]
tool_call = assistant_msg["tool_calls"][0]
assert "codex_reasoning_items" not in assistant_msg
assert tool_call["id"] == "call_fw_123"
assert tool_call["type"] == "function"
assert tool_call["function"]["name"] == "terminal"
assert "call_id" not in tool_call
assert "response_item_id" not in tool_call
class TestBuildApiKwargsCodex:
def test_uses_responses_api_format(self, monkeypatch):
agent = _make_agent(monkeypatch, "openai-codex", api_mode="codex_responses",
base_url="https://chatgpt.com/backend-api/codex")
messages = [{"role": "user", "content": "hi"}]
kwargs = agent._build_api_kwargs(messages)
assert "input" in kwargs
assert "instructions" in kwargs
assert "messages" not in kwargs
assert kwargs["store"] is False
def test_includes_reasoning_config(self, monkeypatch):
agent = _make_agent(monkeypatch, "openai-codex", api_mode="codex_responses",
base_url="https://chatgpt.com/backend-api/codex")
messages = [{"role": "user", "content": "hi"}]
kwargs = agent._build_api_kwargs(messages)
assert "reasoning" in kwargs
assert kwargs["reasoning"]["effort"] == "medium"
def test_includes_encrypted_content_in_include(self, monkeypatch):
agent = _make_agent(monkeypatch, "openai-codex", api_mode="codex_responses",
base_url="https://chatgpt.com/backend-api/codex")
messages = [{"role": "user", "content": "hi"}]
kwargs = agent._build_api_kwargs(messages)
assert "reasoning.encrypted_content" in kwargs.get("include", [])
def test_tools_converted_to_responses_format(self, monkeypatch):
agent = _make_agent(monkeypatch, "openai-codex", api_mode="codex_responses",
base_url="https://chatgpt.com/backend-api/codex")
messages = [{"role": "user", "content": "hi"}]
kwargs = agent._build_api_kwargs(messages)
tools = kwargs.get("tools", [])
assert len(tools) > 0
# Responses format has "name" at top level, not nested under "function"
assert "name" in tools[0]
assert "function" not in tools[0]
# ── Message conversion tests ────────────────────────────────────────────────
class TestChatMessagesToResponsesInput:
"""Verify _chat_messages_to_responses_input for Codex mode."""
def test_user_message_passes_through(self, monkeypatch):
agent = _make_agent(monkeypatch, "openai-codex", api_mode="codex_responses",
base_url="https://chatgpt.com/backend-api/codex")
messages = [{"role": "user", "content": "hello"}]
items = agent._chat_messages_to_responses_input(messages)
assert items == [{"role": "user", "content": "hello"}]
def test_system_messages_filtered(self, monkeypatch):
agent = _make_agent(monkeypatch, "openai-codex", api_mode="codex_responses",
base_url="https://chatgpt.com/backend-api/codex")
messages = [
{"role": "system", "content": "be helpful"},
{"role": "user", "content": "hello"},
]
items = agent._chat_messages_to_responses_input(messages)
assert len(items) == 1
assert items[0]["role"] == "user"
def test_assistant_tool_calls_become_function_call_items(self, monkeypatch):
agent = _make_agent(monkeypatch, "openai-codex", api_mode="codex_responses",
base_url="https://chatgpt.com/backend-api/codex")
messages = [{
"role": "assistant",
"content": "",
"tool_calls": [{
"id": "call_abc",
"call_id": "call_abc",
"function": {"name": "web_search", "arguments": '{"query": "test"}'},
}],
}]
items = agent._chat_messages_to_responses_input(messages)
fc_items = [i for i in items if i.get("type") == "function_call"]
assert len(fc_items) == 1
assert fc_items[0]["name"] == "web_search"
assert fc_items[0]["call_id"] == "call_abc"
def test_tool_results_become_function_call_output(self, monkeypatch):
agent = _make_agent(monkeypatch, "openai-codex", api_mode="codex_responses",
base_url="https://chatgpt.com/backend-api/codex")
messages = [{"role": "tool", "tool_call_id": "call_abc", "content": "result here"}]
items = agent._chat_messages_to_responses_input(messages)
assert items[0]["type"] == "function_call_output"
assert items[0]["call_id"] == "call_abc"
assert items[0]["output"] == "result here"
def test_encrypted_reasoning_replayed(self, monkeypatch):
"""Encrypted reasoning items from previous turns must be included in input."""
agent = _make_agent(monkeypatch, "openai-codex", api_mode="codex_responses",
base_url="https://chatgpt.com/backend-api/codex")
messages = [
{"role": "user", "content": "think about this"},
{
"role": "assistant",
"content": "I thought about it.",
"codex_reasoning_items": [
{"type": "reasoning", "id": "rs_abc", "encrypted_content": "gAAAA_test_blob"},
],
},
{"role": "user", "content": "continue"},
]
items = agent._chat_messages_to_responses_input(messages)
reasoning_items = [i for i in items if i.get("type") == "reasoning"]
assert len(reasoning_items) == 1
assert reasoning_items[0]["encrypted_content"] == "gAAAA_test_blob"
def test_no_reasoning_items_for_non_codex_messages(self, monkeypatch):
"""Messages without codex_reasoning_items should not inject anything."""
agent = _make_agent(monkeypatch, "openai-codex", api_mode="codex_responses",
base_url="https://chatgpt.com/backend-api/codex")
messages = [
{"role": "assistant", "content": "hi"},
{"role": "user", "content": "hello"},
]
items = agent._chat_messages_to_responses_input(messages)
reasoning_items = [i for i in items if i.get("type") == "reasoning"]
assert len(reasoning_items) == 0
# ── Response normalization tests ─────────────────────────────────────────────
class TestNormalizeCodexResponse:
"""Verify _normalize_codex_response extracts all fields correctly."""
def _make_codex_agent(self, monkeypatch):
return _make_agent(monkeypatch, "openai-codex", api_mode="codex_responses",
base_url="https://chatgpt.com/backend-api/codex")
def test_text_response(self, monkeypatch):
agent = self._make_codex_agent(monkeypatch)
response = SimpleNamespace(
output=[
SimpleNamespace(type="message", status="completed",
content=[SimpleNamespace(type="output_text", text="Hello!")],
phase="final_answer"),
],
status="completed",
)
msg, reason = agent._normalize_codex_response(response)
assert msg.content == "Hello!"
assert reason == "stop"
def test_reasoning_summary_extracted(self, monkeypatch):
agent = self._make_codex_agent(monkeypatch)
response = SimpleNamespace(
output=[
SimpleNamespace(type="reasoning",
encrypted_content="gAAAA_blob",
summary=[SimpleNamespace(type="summary_text", text="Thinking about math")],
id="rs_123", status=None),
SimpleNamespace(type="message", status="completed",
content=[SimpleNamespace(type="output_text", text="42")],
phase="final_answer"),
],
status="completed",
)
msg, reason = agent._normalize_codex_response(response)
assert msg.content == "42"
assert "math" in msg.reasoning
assert reason == "stop"
def test_encrypted_content_captured(self, monkeypatch):
agent = self._make_codex_agent(monkeypatch)
response = SimpleNamespace(
output=[
SimpleNamespace(type="reasoning",
encrypted_content="gAAAA_secret_blob_123",
summary=[SimpleNamespace(type="summary_text", text="Thinking")],
id="rs_456", status=None),
SimpleNamespace(type="message", status="completed",
content=[SimpleNamespace(type="output_text", text="done")],
phase="final_answer"),
],
status="completed",
)
msg, reason = agent._normalize_codex_response(response)
assert msg.codex_reasoning_items is not None
assert len(msg.codex_reasoning_items) == 1
assert msg.codex_reasoning_items[0]["encrypted_content"] == "gAAAA_secret_blob_123"
assert msg.codex_reasoning_items[0]["id"] == "rs_456"
def test_no_encrypted_content_when_missing(self, monkeypatch):
agent = self._make_codex_agent(monkeypatch)
response = SimpleNamespace(
output=[
SimpleNamespace(type="message", status="completed",
content=[SimpleNamespace(type="output_text", text="no reasoning")],
phase="final_answer"),
],
status="completed",
)
msg, reason = agent._normalize_codex_response(response)
assert msg.codex_reasoning_items is None
def test_tool_calls_extracted(self, monkeypatch):
agent = self._make_codex_agent(monkeypatch)
response = SimpleNamespace(
output=[
SimpleNamespace(type="function_call", status="completed",
call_id="call_xyz", name="web_search",
arguments='{"query":"test"}', id="fc_xyz"),
],
status="completed",
)
msg, reason = agent._normalize_codex_response(response)
assert reason == "tool_calls"
assert len(msg.tool_calls) == 1
assert msg.tool_calls[0].function.name == "web_search"
# ── Chat completions response handling (OpenRouter/Nous) ─────────────────────
class TestBuildAssistantMessage:
"""Verify _build_assistant_message works for all provider response formats."""
def test_openrouter_reasoning_fields(self, monkeypatch):
agent = _make_agent(monkeypatch, "openrouter")
msg = SimpleNamespace(
content="answer",
tool_calls=None,
reasoning="I thought about it",
reasoning_content=None,
reasoning_details=None,
)
result = agent._build_assistant_message(msg, "stop")
assert result["content"] == "answer"
assert result["reasoning"] == "I thought about it"
assert "codex_reasoning_items" not in result
def test_openrouter_reasoning_details_preserved_unmodified(self, monkeypatch):
"""reasoning_details must be passed back exactly as received for
multi-turn continuity (OpenRouter, Anthropic, OpenAI all need this)."""
agent = _make_agent(monkeypatch, "openrouter")
original_detail = {
"type": "thinking",
"thinking": "deep thoughts here",
"signature": "sig123_opaque_blob",
"encrypted_content": "some_provider_blob",
"extra_field": "should_not_be_dropped",
}
msg = SimpleNamespace(
content="answer",
tool_calls=None,
reasoning=None,
reasoning_content=None,
reasoning_details=[original_detail],
)
result = agent._build_assistant_message(msg, "stop")
stored = result["reasoning_details"][0]
# ALL fields must survive, not just type/text/signature
assert stored["signature"] == "sig123_opaque_blob"
assert stored["encrypted_content"] == "some_provider_blob"
assert stored["extra_field"] == "should_not_be_dropped"
assert stored["thinking"] == "deep thoughts here"
def test_codex_preserves_encrypted_reasoning(self, monkeypatch):
agent = _make_agent(monkeypatch, "openai-codex", api_mode="codex_responses",
base_url="https://chatgpt.com/backend-api/codex")
msg = SimpleNamespace(
content="result",
tool_calls=None,
reasoning="summary text",
reasoning_content=None,
reasoning_details=None,
codex_reasoning_items=[
{"type": "reasoning", "id": "rs_1", "encrypted_content": "gAAAA_blob"},
],
)
result = agent._build_assistant_message(msg, "stop")
assert result["codex_reasoning_items"] == [
{"type": "reasoning", "id": "rs_1", "encrypted_content": "gAAAA_blob"},
]
def test_plain_message_no_codex_items(self, monkeypatch):
agent = _make_agent(monkeypatch, "openrouter")
msg = SimpleNamespace(
content="simple",
tool_calls=None,
reasoning=None,
reasoning_content=None,
reasoning_details=None,
)
result = agent._build_assistant_message(msg, "stop")
assert "codex_reasoning_items" not in result
# ── Auxiliary client provider resolution ─────────────────────────────────────
class TestAuxiliaryClientProviderPriority:
"""Verify auxiliary client resolution doesn't break for any provider."""
def test_openrouter_always_wins(self, monkeypatch):
monkeypatch.setenv("OPENROUTER_API_KEY", "or-key")
from agent.auxiliary_client import get_text_auxiliary_client
with patch("agent.auxiliary_client.OpenAI") as mock:
client, model = get_text_auxiliary_client()
assert model == "google/gemini-3-flash-preview"
assert "openrouter" in str(mock.call_args.kwargs["base_url"]).lower()
def test_nous_when_no_openrouter(self, monkeypatch):
monkeypatch.delenv("OPENROUTER_API_KEY", raising=False)
from agent.auxiliary_client import get_text_auxiliary_client
with patch("agent.auxiliary_client._read_nous_auth", return_value={"access_token": "nous-tok"}), \
patch("agent.auxiliary_client.OpenAI") as mock:
client, model = get_text_auxiliary_client()
assert model == "google/gemini-3-flash-preview"
def test_custom_endpoint_when_no_nous(self, monkeypatch):
"""Custom endpoint is used when no OpenRouter/Nous keys are available.
Since the March 2026 config refactor, OPENAI_BASE_URL env var is no
longer consulted — base_url comes from config.yaml via
resolve_runtime_provider. Mock _resolve_custom_runtime directly.
"""
monkeypatch.delenv("OPENROUTER_API_KEY", raising=False)
monkeypatch.setenv("OPENAI_API_KEY", "local-key")
from agent.auxiliary_client import get_text_auxiliary_client
with patch("agent.auxiliary_client._read_nous_auth", return_value=None), \
patch("agent.auxiliary_client._resolve_custom_runtime",
return_value=("http://localhost:1234/v1", "local-key")), \
patch("agent.auxiliary_client.OpenAI") as mock:
client, model = get_text_auxiliary_client()
assert mock.call_args.kwargs["base_url"] == "http://localhost:1234/v1"
def test_codex_fallback_last_resort(self, monkeypatch):
monkeypatch.delenv("OPENROUTER_API_KEY", raising=False)
monkeypatch.delenv("OPENAI_BASE_URL", raising=False)
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
from agent.auxiliary_client import get_text_auxiliary_client, CodexAuxiliaryClient
with patch("agent.auxiliary_client._read_nous_auth", return_value=None), \
patch("agent.auxiliary_client._read_codex_access_token", return_value="codex-tok"), \
patch("agent.auxiliary_client.OpenAI"):
client, model = get_text_auxiliary_client()
assert model == "gpt-5.2-codex"
assert isinstance(client, CodexAuxiliaryClient)
# ── Provider routing tests ───────────────────────────────────────────────────
class TestProviderRouting:
"""Verify provider_routing config flows into extra_body.provider."""
def test_sort_throughput(self, monkeypatch):
agent = _make_agent(monkeypatch, "openrouter")
agent.provider_sort = "throughput"
kwargs = agent._build_api_kwargs([{"role": "user", "content": "hi"}])
assert kwargs["extra_body"]["provider"]["sort"] == "throughput"
def test_only_providers(self, monkeypatch):
agent = _make_agent(monkeypatch, "openrouter")
agent.providers_allowed = ["anthropic", "google"]
kwargs = agent._build_api_kwargs([{"role": "user", "content": "hi"}])
assert kwargs["extra_body"]["provider"]["only"] == ["anthropic", "google"]
def test_ignore_providers(self, monkeypatch):
agent = _make_agent(monkeypatch, "openrouter")
agent.providers_ignored = ["deepinfra"]
kwargs = agent._build_api_kwargs([{"role": "user", "content": "hi"}])
assert kwargs["extra_body"]["provider"]["ignore"] == ["deepinfra"]
def test_order_providers(self, monkeypatch):
agent = _make_agent(monkeypatch, "openrouter")
agent.providers_order = ["anthropic", "together"]
kwargs = agent._build_api_kwargs([{"role": "user", "content": "hi"}])
assert kwargs["extra_body"]["provider"]["order"] == ["anthropic", "together"]
def test_require_parameters(self, monkeypatch):
agent = _make_agent(monkeypatch, "openrouter")
agent.provider_require_parameters = True
kwargs = agent._build_api_kwargs([{"role": "user", "content": "hi"}])
assert kwargs["extra_body"]["provider"]["require_parameters"] is True
def test_data_collection_deny(self, monkeypatch):
agent = _make_agent(monkeypatch, "openrouter")
agent.provider_data_collection = "deny"
kwargs = agent._build_api_kwargs([{"role": "user", "content": "hi"}])
assert kwargs["extra_body"]["provider"]["data_collection"] == "deny"
def test_no_routing_when_unset(self, monkeypatch):
agent = _make_agent(monkeypatch, "openrouter")
kwargs = agent._build_api_kwargs([{"role": "user", "content": "hi"}])
assert "provider" not in kwargs.get("extra_body", {}).get("provider", {}) or \
kwargs.get("extra_body", {}).get("provider") is None or \
"only" not in kwargs.get("extra_body", {}).get("provider", {})
def test_combined_routing(self, monkeypatch):
agent = _make_agent(monkeypatch, "openrouter")
agent.provider_sort = "latency"
agent.providers_ignored = ["deepinfra"]
agent.provider_data_collection = "deny"
kwargs = agent._build_api_kwargs([{"role": "user", "content": "hi"}])
prov = kwargs["extra_body"]["provider"]
assert prov["sort"] == "latency"
assert prov["ignore"] == ["deepinfra"]
assert prov["data_collection"] == "deny"
def test_routing_not_injected_for_codex(self, monkeypatch):
"""Codex Responses API doesn't use extra_body.provider."""
agent = _make_agent(monkeypatch, "openai-codex", api_mode="codex_responses",
base_url="https://chatgpt.com/backend-api/codex")
agent.provider_sort = "throughput"
kwargs = agent._build_api_kwargs([{"role": "user", "content": "hi"}])
assert "extra_body" not in kwargs
assert "provider" not in kwargs or kwargs.get("provider") is None
# ── Codex reasoning items preflight tests ────────────────────────────────────
class TestCodexReasoningPreflight:
"""Verify reasoning items pass through preflight normalization."""
def test_reasoning_item_passes_through(self, monkeypatch):
agent = _make_agent(monkeypatch, "openai-codex", api_mode="codex_responses",
base_url="https://chatgpt.com/backend-api/codex")
raw_input = [
{"role": "user", "content": "hello"},
{"type": "reasoning", "encrypted_content": "abc123encrypted", "id": "r_001",
"summary": [{"type": "summary_text", "text": "Thinking about it"}]},
{"role": "assistant", "content": "hi there"},
]
normalized = agent._preflight_codex_input_items(raw_input)
reasoning_items = [i for i in normalized if i.get("type") == "reasoning"]
assert len(reasoning_items) == 1
assert reasoning_items[0]["encrypted_content"] == "abc123encrypted"
assert reasoning_items[0]["id"] == "r_001"
assert reasoning_items[0]["summary"] == [{"type": "summary_text", "text": "Thinking about it"}]
def test_reasoning_item_without_id(self, monkeypatch):
agent = _make_agent(monkeypatch, "openai-codex", api_mode="codex_responses",
base_url="https://chatgpt.com/backend-api/codex")
raw_input = [
{"type": "reasoning", "encrypted_content": "abc123"},
]
normalized = agent._preflight_codex_input_items(raw_input)
assert len(normalized) == 1
assert "id" not in normalized[0]
assert normalized[0]["summary"] == [] # default empty summary
def test_reasoning_item_empty_encrypted_skipped(self, monkeypatch):
agent = _make_agent(monkeypatch, "openai-codex", api_mode="codex_responses",
base_url="https://chatgpt.com/backend-api/codex")
raw_input = [
{"type": "reasoning", "encrypted_content": ""},
{"role": "user", "content": "hello"},
]
normalized = agent._preflight_codex_input_items(raw_input)
reasoning_items = [i for i in normalized if i.get("type") == "reasoning"]
assert len(reasoning_items) == 0
def test_reasoning_items_replayed_from_history(self, monkeypatch):
"""Reasoning items stored in codex_reasoning_items get replayed."""
agent = _make_agent(monkeypatch, "openai-codex", api_mode="codex_responses",
base_url="https://chatgpt.com/backend-api/codex")
messages = [
{"role": "user", "content": "hello"},
{
"role": "assistant",
"content": "hi",
"codex_reasoning_items": [
{"type": "reasoning", "encrypted_content": "enc123", "id": "r_1"},
],
},
{"role": "user", "content": "follow up"},
]
items = agent._chat_messages_to_responses_input(messages)
reasoning_items = [i for i in items if isinstance(i, dict) and i.get("type") == "reasoning"]
assert len(reasoning_items) == 1
assert reasoning_items[0]["encrypted_content"] == "enc123"
# ── Reasoning effort consistency tests ───────────────────────────────────────
class TestReasoningEffortDefaults:
"""Verify reasoning effort defaults to medium across all provider paths."""
def test_openrouter_default_medium(self, monkeypatch):
agent = _make_agent(monkeypatch, "openrouter")
agent.model = "anthropic/claude-sonnet-4-20250514"
kwargs = agent._build_api_kwargs([{"role": "user", "content": "hi"}])
reasoning = kwargs["extra_body"]["reasoning"]
assert reasoning["effort"] == "medium"
def test_codex_default_medium(self, monkeypatch):
agent = _make_agent(monkeypatch, "openai-codex", api_mode="codex_responses",
base_url="https://chatgpt.com/backend-api/codex")
kwargs = agent._build_api_kwargs([{"role": "user", "content": "hi"}])
assert kwargs["reasoning"]["effort"] == "medium"
def test_codex_reasoning_disabled(self, monkeypatch):
agent = _make_agent(monkeypatch, "openai-codex", api_mode="codex_responses",
base_url="https://chatgpt.com/backend-api/codex")
agent.reasoning_config = {"enabled": False}
kwargs = agent._build_api_kwargs([{"role": "user", "content": "hi"}])
assert "reasoning" not in kwargs
assert kwargs["include"] == []
def test_codex_reasoning_low(self, monkeypatch):
agent = _make_agent(monkeypatch, "openai-codex", api_mode="codex_responses",
base_url="https://chatgpt.com/backend-api/codex")
agent.reasoning_config = {"enabled": True, "effort": "low"}
kwargs = agent._build_api_kwargs([{"role": "user", "content": "hi"}])
assert kwargs["reasoning"]["effort"] == "low"
def test_openrouter_reasoning_config_override(self, monkeypatch):
agent = _make_agent(monkeypatch, "openrouter")
agent.model = "anthropic/claude-sonnet-4-20250514"
agent.reasoning_config = {"enabled": True, "effort": "medium"}
kwargs = agent._build_api_kwargs([{"role": "user", "content": "hi"}])
assert kwargs["extra_body"]["reasoning"]["effort"] == "medium"

View File

@@ -0,0 +1,186 @@
"""Test real interrupt propagation through delegate_task with actual AIAgent.
This uses a real AIAgent with mocked HTTP responses to test the complete
interrupt flow through _run_single_child → child.run_conversation().
"""
import json
import os
import threading
import time
import unittest
from unittest.mock import MagicMock, patch, PropertyMock
from tools.interrupt import set_interrupt, is_interrupted
def _make_slow_api_response(delay=5.0):
"""Create a mock that simulates a slow API response (like a real LLM call)."""
def slow_create(**kwargs):
# Simulate a slow API call
time.sleep(delay)
# Return a simple text response (no tool calls)
resp = MagicMock()
resp.choices = [MagicMock()]
resp.choices[0].message = MagicMock()
resp.choices[0].message.content = "Done"
resp.choices[0].message.tool_calls = None
resp.choices[0].message.refusal = None
resp.choices[0].finish_reason = "stop"
resp.usage = MagicMock()
resp.usage.prompt_tokens = 100
resp.usage.completion_tokens = 10
resp.usage.total_tokens = 110
resp.usage.prompt_tokens_details = None
return resp
return slow_create
class TestRealSubagentInterrupt(unittest.TestCase):
"""Test interrupt with real AIAgent child through delegate_tool."""
def setUp(self):
set_interrupt(False)
os.environ.setdefault("OPENAI_API_KEY", "test-key")
def tearDown(self):
set_interrupt(False)
def test_interrupt_child_during_api_call(self):
"""Real AIAgent child interrupted while making API call."""
from run_agent import AIAgent, IterationBudget
# Create a real parent agent (just enough to be a parent)
parent = AIAgent.__new__(AIAgent)
parent._interrupt_requested = False
parent._interrupt_message = None
parent._active_children = []
parent._active_children_lock = threading.Lock()
parent.quiet_mode = True
parent.model = "test/model"
parent.base_url = "http://localhost:1"
parent.api_key = "test"
parent.provider = "test"
parent.api_mode = "chat_completions"
parent.platform = "cli"
parent.enabled_toolsets = ["terminal", "file"]
parent.providers_allowed = None
parent.providers_ignored = None
parent.providers_order = None
parent.provider_sort = None
parent.max_tokens = None
parent.reasoning_config = None
parent.prefill_messages = None
parent._session_db = None
parent._delegate_depth = 0
parent._delegate_spinner = None
parent.tool_progress_callback = None
parent.iteration_budget = IterationBudget(max_total=100)
parent._client_kwargs = {"api_key": "test", "base_url": "http://localhost:1"}
from tools.delegate_tool import _run_single_child
child_started = threading.Event()
result_holder = [None]
error_holder = [None]
def run_delegate():
try:
# Patch the OpenAI client creation inside AIAgent.__init__
with patch('run_agent.OpenAI') as MockOpenAI:
mock_client = MagicMock()
# API call takes 5 seconds — should be interrupted before that
mock_client.chat.completions.create = _make_slow_api_response(delay=5.0)
mock_client.close = MagicMock()
MockOpenAI.return_value = mock_client
# Patch the instance method so it skips prompt assembly
with patch.object(AIAgent, '_build_system_prompt', return_value="You are a test agent"):
# Signal when child starts
original_run = AIAgent.run_conversation
def patched_run(self_agent, *args, **kwargs):
child_started.set()
return original_run(self_agent, *args, **kwargs)
with patch.object(AIAgent, 'run_conversation', patched_run):
# Build a real child agent (AIAgent is NOT patched here,
# only run_conversation and _build_system_prompt are)
child = AIAgent(
base_url="http://localhost:1",
api_key="test-key",
model="test/model",
provider="test",
api_mode="chat_completions",
max_iterations=5,
enabled_toolsets=["terminal"],
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
platform="cli",
)
child._delegate_depth = 1
parent._active_children.append(child)
result = _run_single_child(
task_index=0,
goal="Test task",
child=child,
parent_agent=parent,
)
result_holder[0] = result
except Exception as e:
import traceback
traceback.print_exc()
error_holder[0] = e
agent_thread = threading.Thread(target=run_delegate, daemon=True)
agent_thread.start()
# Wait for child to start run_conversation
started = child_started.wait(timeout=10)
if not started:
agent_thread.join(timeout=1)
if error_holder[0]:
raise error_holder[0]
self.fail("Child never started run_conversation")
# Give child time to enter main loop and start API call
time.sleep(0.5)
# Verify child is registered
print(f"Active children: {len(parent._active_children)}")
self.assertGreaterEqual(len(parent._active_children), 1,
"Child not registered in _active_children")
# Interrupt! (simulating what CLI does)
start = time.monotonic()
parent.interrupt("User typed a new message")
# Check propagation
child = parent._active_children[0] if parent._active_children else None
if child:
print(f"Child._interrupt_requested after parent.interrupt(): {child._interrupt_requested}")
self.assertTrue(child._interrupt_requested,
"Interrupt did not propagate to child!")
# Wait for delegate to finish (should be fast since interrupted)
agent_thread.join(timeout=5)
elapsed = time.monotonic() - start
if error_holder[0]:
raise error_holder[0]
result = result_holder[0]
self.assertIsNotNone(result, "Delegate returned no result")
print(f"Result status: {result['status']}, elapsed: {elapsed:.2f}s")
print(f"Full result: {result}")
# The child should have been interrupted, not completed the full 5s API call
self.assertLess(elapsed, 3.0,
f"Took {elapsed:.2f}s — interrupt was not detected quickly enough")
self.assertEqual(result["status"], "interrupted",
f"Expected 'interrupted', got '{result['status']}'")
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,54 @@
"""Verify that redirect_stdout in _run_single_child is process-wide.
This demonstrates that contextlib.redirect_stdout changes sys.stdout
for ALL threads, not just the current one. This means during subagent
execution, all output from other threads (including the CLI's process_thread)
is swallowed.
"""
import contextlib
import io
import sys
import threading
import time
import unittest
class TestRedirectStdoutIsProcessWide(unittest.TestCase):
def test_redirect_stdout_affects_other_threads(self):
"""contextlib.redirect_stdout changes sys.stdout for ALL threads."""
captured_from_other_thread = []
real_stdout = sys.stdout
other_thread_saw_devnull = threading.Event()
def other_thread_work():
"""Runs in a different thread, tries to use sys.stdout."""
time.sleep(0.2) # Let redirect_stdout take effect
# Check what sys.stdout is
if sys.stdout is not real_stdout:
other_thread_saw_devnull.set()
# Try to print — this should go to devnull
captured_from_other_thread.append(sys.stdout)
t = threading.Thread(target=other_thread_work, daemon=True)
t.start()
# redirect_stdout in main thread
devnull = io.StringIO()
with contextlib.redirect_stdout(devnull):
time.sleep(0.5) # Let the other thread check during redirect
t.join(timeout=2)
# The other thread should have seen devnull, NOT the real stdout
self.assertTrue(
other_thread_saw_devnull.is_set(),
"redirect_stdout was NOT process-wide — other thread still saw real stdout. "
"This test's premise is wrong."
)
print("Confirmed: redirect_stdout IS process-wide — affects all threads")
if __name__ == "__main__":
unittest.main()

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,90 @@
"""Tests for session_meta filtering — issue #4715.
Ensures that transcript-only session_meta messages never reach the
chat-completions API, via both the API-boundary guard in
_sanitize_api_messages() and the CLI session-restore paths.
"""
import logging
import types
from unittest.mock import MagicMock, patch
from run_agent import AIAgent
# ---------------------------------------------------------------------------
# Layer 1 — _sanitize_api_messages role-allowlist guard
# ---------------------------------------------------------------------------
class TestSanitizeApiMessagesRoleFilter:
def test_drops_session_meta_role(self):
msgs = [
{"role": "user", "content": "hello"},
{"role": "session_meta", "content": {"model": "gpt-4"}},
{"role": "assistant", "content": "hi"},
]
out = AIAgent._sanitize_api_messages(msgs)
assert len(out) == 2
assert all(m["role"] != "session_meta" for m in out)
def test_preserves_valid_roles(self):
msgs = [
{"role": "system", "content": "you are helpful"},
{"role": "user", "content": "hello"},
{"role": "assistant", "content": "hi"},
{"role": "tool", "tool_call_id": "c1", "content": "ok"},
]
# Need a matching assistant tool_call so the tool result isn't orphaned
msgs[2]["tool_calls"] = [{"id": "c1", "function": {"name": "t", "arguments": "{}"}}]
out = AIAgent._sanitize_api_messages(msgs)
roles = [m["role"] for m in out]
assert "system" in roles
assert "user" in roles
assert "assistant" in roles
assert "tool" in roles
def test_logs_warning_when_dropping(self, caplog):
msgs = [
{"role": "user", "content": "hello"},
{"role": "session_meta", "content": {"info": "test"}},
]
with caplog.at_level(logging.DEBUG, logger="run_agent"):
AIAgent._sanitize_api_messages(msgs)
assert any("invalid role" in r.message and "session_meta" in r.message for r in caplog.records)
def test_drops_multiple_invalid_roles(self):
msgs = [
{"role": "user", "content": "hello"},
{"role": "session_meta", "content": {}},
{"role": "transcript_note", "content": "note"},
{"role": "assistant", "content": "hi"},
]
out = AIAgent._sanitize_api_messages(msgs)
assert len(out) == 2
assert [m["role"] for m in out] == ["user", "assistant"]
# ---------------------------------------------------------------------------
# Layer 2 — CLI session-restore filters session_meta before loading
# ---------------------------------------------------------------------------
class TestCLISessionRestoreFiltering:
def test_restore_filters_session_meta(self):
"""Simulates the CLI restore path and verifies session_meta is removed."""
# Build a fake restored message list (as returned by get_messages_as_conversation)
fake_restored = [
{"role": "session_meta", "content": {"model": "gpt-4"}},
{"role": "user", "content": "hello"},
{"role": "assistant", "content": "hi there"},
{"role": "session_meta", "content": {"tools": []}},
]
# Apply the same filtering that the patched CLI code now does
filtered = [m for m in fake_restored if m.get("role") != "session_meta"]
assert len(filtered) == 2
assert all(m["role"] != "session_meta" for m in filtered)
assert filtered[0]["role"] == "user"
assert filtered[1]["role"] == "assistant"

View File

@@ -0,0 +1,121 @@
"""Tests for session reset completeness (fixes #2635).
/clear and /new must not carry stale state into the next session.
Two fields were added after reset_session_state() was written and were
therefore never cleared:
- ContextCompressor._previous_summary
- AIAgent._user_turn_count
"""
import sys
import types
from pathlib import Path
import pytest
# Ensure repo root is importable
sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent))
# Stub out optional heavy dependencies not installed in the test environment
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())
from run_agent import AIAgent
from agent.context_compressor import ContextCompressor
def _make_minimal_agent() -> AIAgent:
"""Return an AIAgent constructed with the absolute minimum args.
We pass dummy values that bypass network calls and filesystem access.
The object is never used to make API calls — only its attributes and
reset_session_state() are exercised.
"""
agent = AIAgent.__new__(AIAgent) # skip __init__ entirely
# Seed the exact attributes that reset_session_state() writes
agent.session_total_tokens = 0
agent.session_input_tokens = 0
agent.session_output_tokens = 0
agent.session_prompt_tokens = 0
agent.session_completion_tokens = 0
agent.session_cache_read_tokens = 0
agent.session_cache_write_tokens = 0
agent.session_reasoning_tokens = 0
agent.session_api_calls = 0
agent.session_estimated_cost_usd = 0.0
agent.session_cost_status = "unknown"
agent.session_cost_source = "none"
# The two fields under test
agent._user_turn_count = 0
agent.context_compressor = None # will be set per-test as needed
return agent
class TestResetSessionState:
"""reset_session_state() must clear ALL session-scoped state."""
def test_previous_summary_cleared_on_reset(self):
"""Compression summary from old session must not leak into new session."""
agent = _make_minimal_agent()
compressor = ContextCompressor.__new__(ContextCompressor)
compressor._previous_summary = "Old session summary about unrelated topic"
# Seed counter attributes that reset_session_state touches
compressor.last_prompt_tokens = 100
compressor.last_completion_tokens = 50
compressor.last_total_tokens = 150
compressor.compression_count = 3
compressor._context_probed = True
agent.context_compressor = compressor
agent.reset_session_state()
assert compressor._previous_summary is None, (
"_previous_summary must be None after reset; got: "
f"{compressor._previous_summary!r}"
)
def test_user_turn_count_cleared_on_reset(self):
"""Turn counter must reset to 0 on new session."""
agent = _make_minimal_agent()
agent._user_turn_count = 7 # simulates turns accumulated in previous session
agent.context_compressor = None
agent.reset_session_state()
assert agent._user_turn_count == 0, (
f"_user_turn_count must be 0 after reset; got: {agent._user_turn_count}"
)
def test_both_fields_cleared_together(self):
"""Both stale fields are cleared in a single reset_session_state() call."""
agent = _make_minimal_agent()
agent._user_turn_count = 3
compressor = ContextCompressor.__new__(ContextCompressor)
compressor._previous_summary = "Stale summary"
compressor.last_prompt_tokens = 0
compressor.last_completion_tokens = 0
compressor.last_total_tokens = 0
compressor.compression_count = 0
compressor._context_probed = False
agent.context_compressor = compressor
agent.reset_session_state()
assert agent._user_turn_count == 0
assert compressor._previous_summary is None
def test_reset_without_compressor_does_not_raise(self):
"""reset_session_state() must not raise when context_compressor is None."""
agent = _make_minimal_agent()
agent._user_turn_count = 2
agent.context_compressor = None
# Must not raise
agent.reset_session_state()
assert agent._user_turn_count == 0

View File

@@ -0,0 +1,816 @@
"""Tests for streaming token delivery infrastructure.
Tests the unified streaming API call, delta callbacks, tool-call
suppression, provider fallback, and CLI streaming display.
"""
import json
import threading
import uuid
from types import SimpleNamespace
from unittest.mock import MagicMock, patch, PropertyMock
import pytest
# ── Helpers ──────────────────────────────────────────────────────────────
def _make_stream_chunk(
content=None, tool_calls=None, finish_reason=None,
model=None, reasoning_content=None, usage=None,
):
"""Build a mock streaming chunk matching OpenAI's ChatCompletionChunk shape."""
delta = SimpleNamespace(
content=content,
tool_calls=tool_calls,
reasoning_content=reasoning_content,
reasoning=None,
)
choice = SimpleNamespace(
index=0,
delta=delta,
finish_reason=finish_reason,
)
chunk = SimpleNamespace(
choices=[choice],
model=model,
usage=usage,
)
return chunk
def _make_tool_call_delta(index=0, tc_id=None, name=None, arguments=None, extra_content=None, model_extra=None):
"""Build a mock tool call delta."""
func = SimpleNamespace(name=name, arguments=arguments)
delta = SimpleNamespace(index=index, id=tc_id, function=func)
if extra_content is not None:
delta.extra_content = extra_content
if model_extra is not None:
delta.model_extra = model_extra
return delta
def _make_empty_chunk(model=None, usage=None):
"""Build a chunk with no choices (usage-only final chunk)."""
return SimpleNamespace(choices=[], model=model, usage=usage)
# ── Test: Streaming Accumulator ──────────────────────────────────────────
class TestStreamingAccumulator:
"""Verify that _interruptible_streaming_api_call accumulates content
and tool calls into a response matching the non-streaming shape."""
@patch("run_agent.AIAgent._create_request_openai_client")
@patch("run_agent.AIAgent._close_request_openai_client")
def test_text_only_response(self, mock_close, mock_create):
"""Text-only stream produces correct response shape."""
from run_agent import AIAgent
chunks = [
_make_stream_chunk(content="Hello"),
_make_stream_chunk(content=" world"),
_make_stream_chunk(content="!", finish_reason="stop", model="test-model"),
_make_empty_chunk(usage=SimpleNamespace(prompt_tokens=10, completion_tokens=3)),
]
mock_client = MagicMock()
mock_client.chat.completions.create.return_value = iter(chunks)
mock_create.return_value = mock_client
agent = AIAgent(
model="test/model",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
)
agent.api_mode = "chat_completions"
agent._interrupt_requested = False
response = agent._interruptible_streaming_api_call({})
assert response.choices[0].message.content == "Hello world!"
assert response.choices[0].message.tool_calls is None
assert response.choices[0].finish_reason == "stop"
assert response.usage is not None
assert response.usage.completion_tokens == 3
@patch("run_agent.AIAgent._create_request_openai_client")
@patch("run_agent.AIAgent._close_request_openai_client")
def test_tool_call_response(self, mock_close, mock_create):
"""Tool call stream accumulates ID, name, and arguments."""
from run_agent import AIAgent
chunks = [
_make_stream_chunk(tool_calls=[
_make_tool_call_delta(index=0, tc_id="call_123", name="terminal")
]),
_make_stream_chunk(tool_calls=[
_make_tool_call_delta(index=0, arguments='{"command":')
]),
_make_stream_chunk(tool_calls=[
_make_tool_call_delta(index=0, arguments=' "ls"}')
]),
_make_stream_chunk(finish_reason="tool_calls"),
]
mock_client = MagicMock()
mock_client.chat.completions.create.return_value = iter(chunks)
mock_create.return_value = mock_client
agent = AIAgent(
model="test/model",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
)
agent.api_mode = "chat_completions"
agent._interrupt_requested = False
response = agent._interruptible_streaming_api_call({})
tc = response.choices[0].message.tool_calls
assert tc is not None
assert len(tc) == 1
assert tc[0].id == "call_123"
assert tc[0].function.name == "terminal"
assert tc[0].function.arguments == '{"command": "ls"}'
@patch("run_agent.AIAgent._create_request_openai_client")
@patch("run_agent.AIAgent._close_request_openai_client")
def test_tool_call_extra_content_preserved(self, mock_close, mock_create):
"""Streamed tool calls preserve provider-specific extra_content metadata."""
from run_agent import AIAgent
chunks = [
_make_stream_chunk(tool_calls=[
_make_tool_call_delta(
index=0,
tc_id="call_gemini",
name="cronjob",
model_extra={
"extra_content": {
"google": {"thought_signature": "sig-123"}
}
},
)
]),
_make_stream_chunk(tool_calls=[
_make_tool_call_delta(index=0, arguments='{"task": "deep index on ."}')
]),
_make_stream_chunk(finish_reason="tool_calls"),
]
mock_client = MagicMock()
mock_client.chat.completions.create.return_value = iter(chunks)
mock_create.return_value = mock_client
agent = AIAgent(
model="test/model",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
)
agent.api_mode = "chat_completions"
agent._interrupt_requested = False
response = agent._interruptible_streaming_api_call({})
tc = response.choices[0].message.tool_calls
assert tc is not None
assert tc[0].extra_content == {
"google": {"thought_signature": "sig-123"}
}
@patch("run_agent.AIAgent._create_request_openai_client")
@patch("run_agent.AIAgent._close_request_openai_client")
def test_mixed_content_and_tool_calls(self, mock_close, mock_create):
"""Stream with both text and tool calls accumulates both."""
from run_agent import AIAgent
chunks = [
_make_stream_chunk(content="Let me check"),
_make_stream_chunk(tool_calls=[
_make_tool_call_delta(index=0, tc_id="call_456", name="web_search")
]),
_make_stream_chunk(tool_calls=[
_make_tool_call_delta(index=0, arguments='{"query": "test"}')
]),
_make_stream_chunk(finish_reason="tool_calls"),
]
mock_client = MagicMock()
mock_client.chat.completions.create.return_value = iter(chunks)
mock_create.return_value = mock_client
agent = AIAgent(
model="test/model",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
)
agent.api_mode = "chat_completions"
agent._interrupt_requested = False
response = agent._interruptible_streaming_api_call({})
assert response.choices[0].message.content == "Let me check"
assert len(response.choices[0].message.tool_calls) == 1
# ── Test: Streaming Callbacks ────────────────────────────────────────────
class TestStreamingCallbacks:
"""Verify that delta callbacks fire correctly."""
@patch("run_agent.AIAgent._create_request_openai_client")
@patch("run_agent.AIAgent._close_request_openai_client")
def test_deltas_fire_in_order(self, mock_close, mock_create):
"""Callbacks receive text deltas in order."""
from run_agent import AIAgent
chunks = [
_make_stream_chunk(content="a"),
_make_stream_chunk(content="b"),
_make_stream_chunk(content="c"),
_make_stream_chunk(finish_reason="stop"),
]
deltas = []
mock_client = MagicMock()
mock_client.chat.completions.create.return_value = iter(chunks)
mock_create.return_value = mock_client
agent = AIAgent(
model="test/model",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
stream_delta_callback=lambda t: deltas.append(t),
)
agent.api_mode = "chat_completions"
agent._interrupt_requested = False
agent._interruptible_streaming_api_call({})
assert deltas == ["a", "b", "c"]
@patch("run_agent.AIAgent._create_request_openai_client")
@patch("run_agent.AIAgent._close_request_openai_client")
def test_on_first_delta_fires_once(self, mock_close, mock_create):
"""on_first_delta callback fires exactly once."""
from run_agent import AIAgent
chunks = [
_make_stream_chunk(content="a"),
_make_stream_chunk(content="b"),
_make_stream_chunk(finish_reason="stop"),
]
first_delta_calls = []
mock_client = MagicMock()
mock_client.chat.completions.create.return_value = iter(chunks)
mock_create.return_value = mock_client
agent = AIAgent(
model="test/model",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
)
agent.api_mode = "chat_completions"
agent._interrupt_requested = False
agent._interruptible_streaming_api_call(
{}, on_first_delta=lambda: first_delta_calls.append(True)
)
assert len(first_delta_calls) == 1
@patch("run_agent.AIAgent._create_request_openai_client")
@patch("run_agent.AIAgent._close_request_openai_client")
def test_tool_only_does_not_fire_callback(self, mock_close, mock_create):
"""Tool-call-only stream does not fire the delta callback."""
from run_agent import AIAgent
chunks = [
_make_stream_chunk(tool_calls=[
_make_tool_call_delta(index=0, tc_id="call_789", name="terminal")
]),
_make_stream_chunk(tool_calls=[
_make_tool_call_delta(index=0, arguments='{"command": "ls"}')
]),
_make_stream_chunk(finish_reason="tool_calls"),
]
deltas = []
mock_client = MagicMock()
mock_client.chat.completions.create.return_value = iter(chunks)
mock_create.return_value = mock_client
agent = AIAgent(
model="test/model",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
stream_delta_callback=lambda t: deltas.append(t),
)
agent.api_mode = "chat_completions"
agent._interrupt_requested = False
agent._interruptible_streaming_api_call({})
assert deltas == []
@patch("run_agent.AIAgent._create_request_openai_client")
@patch("run_agent.AIAgent._close_request_openai_client")
def test_text_suppressed_when_tool_calls_present(self, mock_close, mock_create):
"""Text deltas are suppressed when tool calls are also in the stream."""
from run_agent import AIAgent
chunks = [
_make_stream_chunk(content="thinking..."),
_make_stream_chunk(tool_calls=[
_make_tool_call_delta(index=0, tc_id="call_abc", name="read_file")
]),
_make_stream_chunk(content=" more text"),
_make_stream_chunk(finish_reason="tool_calls"),
]
deltas = []
mock_client = MagicMock()
mock_client.chat.completions.create.return_value = iter(chunks)
mock_create.return_value = mock_client
agent = AIAgent(
model="test/model",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
stream_delta_callback=lambda t: deltas.append(t),
)
agent.api_mode = "chat_completions"
agent._interrupt_requested = False
response = agent._interruptible_streaming_api_call({})
# Text before tool call IS fired (we don't know yet it will have tools)
assert "thinking..." in deltas
# Text after tool call IS still routed to stream_delta_callback so that
# reasoning tag extraction can fire (PR #3566). Display-level suppression
# of non-reasoning text happens in the CLI's _stream_delta, not here.
assert " more text" in deltas
# Content is still accumulated in the response
assert response.choices[0].message.content == "thinking... more text"
# ── Test: Streaming Fallback ────────────────────────────────────────────
class TestStreamingFallback:
"""Verify fallback to non-streaming on ANY streaming error."""
@patch("run_agent.AIAgent._interruptible_api_call")
@patch("run_agent.AIAgent._create_request_openai_client")
@patch("run_agent.AIAgent._close_request_openai_client")
def test_stream_error_falls_back(self, mock_close, mock_create, mock_non_stream):
"""'not supported' error triggers fallback to non-streaming."""
from run_agent import AIAgent
mock_client = MagicMock()
mock_client.chat.completions.create.side_effect = Exception(
"Streaming is not supported for this model"
)
mock_create.return_value = mock_client
fallback_response = SimpleNamespace(
id="fallback",
model="test",
choices=[SimpleNamespace(
index=0,
message=SimpleNamespace(
role="assistant",
content="fallback response",
tool_calls=None,
reasoning_content=None,
),
finish_reason="stop",
)],
usage=None,
)
mock_non_stream.return_value = fallback_response
agent = AIAgent(
model="test/model",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
)
agent.api_mode = "chat_completions"
agent._interrupt_requested = False
response = agent._interruptible_streaming_api_call({})
assert response.choices[0].message.content == "fallback response"
mock_non_stream.assert_called_once()
@patch("run_agent.AIAgent._interruptible_api_call")
@patch("run_agent.AIAgent._create_request_openai_client")
@patch("run_agent.AIAgent._close_request_openai_client")
def test_any_stream_error_falls_back(self, mock_close, mock_create, mock_non_stream):
"""ANY streaming error triggers fallback — not just specific messages."""
from run_agent import AIAgent
mock_client = MagicMock()
mock_client.chat.completions.create.side_effect = Exception(
"Connection reset by peer"
)
mock_create.return_value = mock_client
fallback_response = SimpleNamespace(
id="fallback",
model="test",
choices=[SimpleNamespace(
index=0,
message=SimpleNamespace(
role="assistant",
content="fallback after connection error",
tool_calls=None,
reasoning_content=None,
),
finish_reason="stop",
)],
usage=None,
)
mock_non_stream.return_value = fallback_response
agent = AIAgent(
model="test/model",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
)
agent.api_mode = "chat_completions"
agent._interrupt_requested = False
response = agent._interruptible_streaming_api_call({})
assert response.choices[0].message.content == "fallback after connection error"
mock_non_stream.assert_called_once()
@patch("run_agent.AIAgent._interruptible_api_call")
@patch("run_agent.AIAgent._create_request_openai_client")
@patch("run_agent.AIAgent._close_request_openai_client")
def test_fallback_error_propagates(self, mock_close, mock_create, mock_non_stream):
"""When both streaming AND fallback fail, the fallback error propagates."""
from run_agent import AIAgent
mock_client = MagicMock()
mock_client.chat.completions.create.side_effect = Exception("stream broke")
mock_create.return_value = mock_client
mock_non_stream.side_effect = Exception("Rate limit exceeded")
agent = AIAgent(
model="test/model",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
)
agent.api_mode = "chat_completions"
agent._interrupt_requested = False
with pytest.raises(Exception, match="Rate limit exceeded"):
agent._interruptible_streaming_api_call({})
@patch("run_agent.AIAgent._interruptible_api_call")
@patch("run_agent.AIAgent._create_request_openai_client")
@patch("run_agent.AIAgent._close_request_openai_client")
def test_exhausted_transient_stream_error_falls_back(self, mock_close, mock_create, mock_non_stream):
"""Transient stream errors retry first, then fall back after retries are exhausted."""
from run_agent import AIAgent
import httpx
mock_client = MagicMock()
mock_client.chat.completions.create.side_effect = httpx.ConnectError("socket closed")
mock_create.return_value = mock_client
fallback_response = SimpleNamespace(
id="fallback",
model="test",
choices=[SimpleNamespace(
index=0,
message=SimpleNamespace(
role="assistant",
content="fallback after retries exhausted",
tool_calls=None,
reasoning_content=None,
),
finish_reason="stop",
)],
usage=None,
)
mock_non_stream.return_value = fallback_response
agent = AIAgent(
model="test/model",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
)
agent.api_mode = "chat_completions"
agent._interrupt_requested = False
response = agent._interruptible_streaming_api_call({})
assert response.choices[0].message.content == "fallback after retries exhausted"
assert mock_client.chat.completions.create.call_count == 3
mock_non_stream.assert_called_once()
assert mock_close.call_count >= 1
@patch("run_agent.AIAgent._interruptible_api_call")
@patch("run_agent.AIAgent._create_request_openai_client")
@patch("run_agent.AIAgent._close_request_openai_client")
def test_sse_connection_lost_retried_as_transient(self, mock_close, mock_create, mock_non_stream):
"""SSE 'Network connection lost' (APIError w/ no status_code) retries like httpx errors.
OpenRouter sends {"error":{"message":"Network connection lost."}} as an SSE
event when the upstream stream drops. The OpenAI SDK raises APIError from
this. It should be retried at the streaming level, same as httpx connection
errors, before falling back to non-streaming.
"""
from run_agent import AIAgent
import httpx
# Create an APIError that mimics what the OpenAI SDK raises from SSE error events.
# Key: no status_code attribute (unlike APIStatusError which has one).
from openai import APIError as OAIAPIError
sse_error = OAIAPIError(
message="Network connection lost.",
request=httpx.Request("POST", "https://openrouter.ai/api/v1/chat/completions"),
body={"message": "Network connection lost."},
)
mock_client = MagicMock()
mock_client.chat.completions.create.side_effect = sse_error
mock_create.return_value = mock_client
fallback_response = SimpleNamespace(
id="fallback",
model="test",
choices=[SimpleNamespace(
index=0,
message=SimpleNamespace(
role="assistant",
content="fallback after SSE retries",
tool_calls=None,
reasoning_content=None,
),
finish_reason="stop",
)],
usage=None,
)
mock_non_stream.return_value = fallback_response
agent = AIAgent(
model="test/model",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
)
agent.api_mode = "chat_completions"
agent._interrupt_requested = False
response = agent._interruptible_streaming_api_call({})
assert response.choices[0].message.content == "fallback after SSE retries"
# Should retry 3 times (default HERMES_STREAM_RETRIES=2 → 3 attempts)
# before falling back to non-streaming
assert mock_client.chat.completions.create.call_count == 3
mock_non_stream.assert_called_once()
# Connection cleanup should happen for each failed retry
assert mock_close.call_count >= 2
@patch("run_agent.AIAgent._interruptible_api_call")
@patch("run_agent.AIAgent._create_request_openai_client")
@patch("run_agent.AIAgent._close_request_openai_client")
def test_sse_non_connection_error_falls_back_immediately(self, mock_close, mock_create, mock_non_stream):
"""SSE errors that aren't connection-related still fall back immediately (no stream retry)."""
from run_agent import AIAgent
import httpx
from openai import APIError as OAIAPIError
sse_error = OAIAPIError(
message="Invalid model configuration.",
request=httpx.Request("POST", "https://openrouter.ai/api/v1/chat/completions"),
body={"message": "Invalid model configuration."},
)
mock_client = MagicMock()
mock_client.chat.completions.create.side_effect = sse_error
mock_create.return_value = mock_client
fallback_response = SimpleNamespace(
id="fallback",
model="test",
choices=[SimpleNamespace(
index=0,
message=SimpleNamespace(
role="assistant",
content="fallback no retry",
tool_calls=None,
reasoning_content=None,
),
finish_reason="stop",
)],
usage=None,
)
mock_non_stream.return_value = fallback_response
agent = AIAgent(
model="test/model",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
)
agent.api_mode = "chat_completions"
agent._interrupt_requested = False
response = agent._interruptible_streaming_api_call({})
assert response.choices[0].message.content == "fallback no retry"
# Should NOT retry — goes straight to non-streaming fallback
assert mock_client.chat.completions.create.call_count == 1
mock_non_stream.assert_called_once()
# ── Test: Reasoning Streaming ────────────────────────────────────────────
class TestReasoningStreaming:
"""Verify reasoning content is accumulated and callback fires."""
@patch("run_agent.AIAgent._create_request_openai_client")
@patch("run_agent.AIAgent._close_request_openai_client")
def test_reasoning_callback_fires(self, mock_close, mock_create):
"""Reasoning deltas fire the reasoning_callback."""
from run_agent import AIAgent
chunks = [
_make_stream_chunk(reasoning_content="Let me think"),
_make_stream_chunk(reasoning_content=" about this"),
_make_stream_chunk(content="The answer is 42"),
_make_stream_chunk(finish_reason="stop"),
]
reasoning_deltas = []
text_deltas = []
mock_client = MagicMock()
mock_client.chat.completions.create.return_value = iter(chunks)
mock_create.return_value = mock_client
agent = AIAgent(
model="test/model",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
stream_delta_callback=lambda t: text_deltas.append(t),
reasoning_callback=lambda t: reasoning_deltas.append(t),
)
agent.api_mode = "chat_completions"
agent._interrupt_requested = False
response = agent._interruptible_streaming_api_call({})
assert reasoning_deltas == ["Let me think", " about this"]
assert text_deltas == ["The answer is 42"]
assert response.choices[0].message.reasoning_content == "Let me think about this"
assert response.choices[0].message.content == "The answer is 42"
# ── Test: _has_stream_consumers ──────────────────────────────────────────
class TestHasStreamConsumers:
"""Verify _has_stream_consumers() detects registered callbacks."""
def test_no_consumers(self):
from run_agent import AIAgent
agent = AIAgent(
model="test/model",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
)
assert agent._has_stream_consumers() is False
def test_delta_callback_set(self):
from run_agent import AIAgent
agent = AIAgent(
model="test/model",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
stream_delta_callback=lambda t: None,
)
assert agent._has_stream_consumers() is True
def test_stream_callback_set(self):
from run_agent import AIAgent
agent = AIAgent(
model="test/model",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
)
agent._stream_callback = lambda t: None
assert agent._has_stream_consumers() is True
# ── Test: Codex stream fires callbacks ────────────────────────────────
class TestCodexStreamCallbacks:
"""Verify _run_codex_stream fires delta callbacks."""
def test_codex_text_delta_fires_callback(self):
from run_agent import AIAgent
deltas = []
agent = AIAgent(
model="test/model",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
stream_delta_callback=lambda t: deltas.append(t),
)
agent.api_mode = "codex_responses"
agent._interrupt_requested = False
# Mock the stream context manager
mock_event_text = SimpleNamespace(
type="response.output_text.delta",
delta="Hello from Codex!",
)
mock_event_done = SimpleNamespace(
type="response.completed",
delta="",
)
mock_stream = MagicMock()
mock_stream.__enter__ = MagicMock(return_value=mock_stream)
mock_stream.__exit__ = MagicMock(return_value=False)
mock_stream.__iter__ = MagicMock(return_value=iter([mock_event_text, mock_event_done]))
mock_stream.get_final_response.return_value = SimpleNamespace(
output=[SimpleNamespace(
type="message",
content=[SimpleNamespace(type="output_text", text="Hello from Codex!")],
)],
status="completed",
)
mock_client = MagicMock()
mock_client.responses.stream.return_value = mock_stream
response = agent._run_codex_stream({}, client=mock_client)
assert "Hello from Codex!" in deltas
def test_codex_remote_protocol_error_falls_back_to_create_stream(self):
from run_agent import AIAgent
import httpx
fallback_response = SimpleNamespace(
output=[SimpleNamespace(
type="message",
content=[SimpleNamespace(type="output_text", text="fallback from create stream")],
)],
status="completed",
)
mock_client = MagicMock()
mock_client.responses.stream.side_effect = httpx.RemoteProtocolError(
"peer closed connection without sending complete message body"
)
agent = AIAgent(
model="test/model",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
)
agent.api_mode = "codex_responses"
agent._interrupt_requested = False
with patch.object(agent, "_run_codex_create_stream_fallback", return_value=fallback_response) as mock_fallback:
response = agent._run_codex_stream({}, client=mock_client)
assert response is fallback_response
mock_fallback.assert_called_once_with({}, client=mock_client)

View File

@@ -0,0 +1,144 @@
"""Test validation error prevention for strict APIs (Fireworks, etc.)"""
import sys
import types
from unittest.mock import patch, MagicMock
import pytest
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())
from run_agent import AIAgent
# ── Helpers ──────────────────────────────────────────────────────────────────
def _tool_defs(*names):
return [
{
"type": "function",
"function": {
"name": n,
"description": f"{n} tool",
"parameters": {"type": "object", "properties": {}},
},
}
for n in names
]
class _FakeOpenAI:
def __init__(self, **kw):
self.api_key = kw.get("api_key", "test")
self.base_url = kw.get("base_url", "http://test")
def close(self):
pass
def _make_agent(monkeypatch, provider, api_mode="chat_completions", base_url="https://openrouter.ai/api/v1"):
monkeypatch.setattr("run_agent.get_tool_definitions", lambda **kw: _tool_defs("web_search", "terminal"))
monkeypatch.setattr("run_agent.check_toolset_requirements", lambda: {})
monkeypatch.setattr("run_agent.OpenAI", _FakeOpenAI)
return AIAgent(
api_key="test",
base_url=base_url,
provider=provider,
api_mode=api_mode,
max_iterations=4,
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
)
class TestStrictApiValidation:
"""Verify tool_call field sanitization prevents 400 errors on strict APIs."""
def test_fireworks_compatible_messages_after_sanitization(self, monkeypatch):
"""Messages should be Fireworks-compatible after sanitization."""
agent = _make_agent(monkeypatch, "openrouter")
agent.api_mode = "chat_completions" # Fireworks uses chat completions
messages = [
{"role": "user", "content": "hi"},
{
"role": "assistant",
"content": "Checking now.",
"tool_calls": [
{
"id": "call_123",
"call_id": "call_123", # Codex-only field
"response_item_id": "fc_123", # Codex-only field
"type": "function",
"function": {"name": "terminal", "arguments": '{"command":"pwd"}'},
}
],
},
{"role": "tool", "tool_call_id": "call_123", "content": "/tmp"},
]
# After _build_api_kwargs, Codex fields should be stripped
kwargs = agent._build_api_kwargs(messages)
assistant_msg = kwargs["messages"][1]
tool_call = assistant_msg["tool_calls"][0]
# Fireworks rejects these fields
assert "call_id" not in tool_call
assert "response_item_id" not in tool_call
# Standard fields should remain
assert tool_call["id"] == "call_123"
assert tool_call["function"]["name"] == "terminal"
def test_codex_preserves_fields_for_replay(self, monkeypatch):
"""Codex mode should preserve fields for Responses API replay."""
agent = _make_agent(monkeypatch, "openrouter")
agent.api_mode = "codex_responses"
messages = [
{"role": "user", "content": "hi"},
{
"role": "assistant",
"content": "Checking now.",
"tool_calls": [
{
"id": "call_123",
"call_id": "call_123",
"response_item_id": "fc_123",
"type": "function",
"function": {"name": "terminal", "arguments": '{"command":"pwd"}'},
}
],
},
]
# In Codex mode, original messages should NOT be mutated
assert messages[1]["tool_calls"][0]["call_id"] == "call_123"
assert messages[1]["tool_calls"][0]["response_item_id"] == "fc_123"
def test_sanitize_method_with_fireworks_provider(self, monkeypatch):
"""Simulating Fireworks provider should trigger sanitization."""
agent = _make_agent(
monkeypatch,
"fireworks",
api_mode="chat_completions",
base_url="https://api.fireworks.ai/inference/v1"
)
# Should sanitize for Fireworks (chat_completions mode)
assert agent._should_sanitize_tool_calls() is True
def test_no_sanitize_for_codex_responses(self, monkeypatch):
"""Codex responses mode should NOT sanitize."""
agent = _make_agent(
monkeypatch,
"openai",
api_mode="codex_responses",
base_url="https://api.openai.com/v1"
)
# Should NOT sanitize for Codex
assert agent._should_sanitize_tool_calls() is False

View File

@@ -0,0 +1,62 @@
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
from run_agent import AIAgent
def _mock_response(*, usage: dict, content: str = "done"):
msg = SimpleNamespace(content=content, tool_calls=None)
choice = SimpleNamespace(message=msg, finish_reason="stop")
return SimpleNamespace(
choices=[choice],
model="test/model",
usage=SimpleNamespace(**usage),
)
def _make_agent(session_db, *, platform: str):
with (
patch("run_agent.get_tool_definitions", return_value=[]),
patch("run_agent.check_toolset_requirements", return_value={}),
patch("run_agent.OpenAI"),
):
agent = AIAgent(
api_key="test-key",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
session_db=session_db,
session_id=f"{platform}-session",
platform=platform,
)
agent.client = MagicMock()
agent.client.chat.completions.create.return_value = _mock_response(
usage={
"prompt_tokens": 11,
"completion_tokens": 7,
"total_tokens": 18,
}
)
return agent
def test_run_conversation_persists_tokens_for_telegram_sessions():
session_db = MagicMock()
agent = _make_agent(session_db, platform="telegram")
result = agent.run_conversation("hello")
assert result["final_response"] == "done"
session_db.update_token_counts.assert_called_once()
assert session_db.update_token_counts.call_args.args[0] == "telegram-session"
def test_run_conversation_persists_tokens_for_cron_sessions():
session_db = MagicMock()
agent = _make_agent(session_db, platform="cron")
result = agent.run_conversation("hello")
assert result["final_response"] == "done"
session_db.update_token_counts.assert_called_once()
assert session_db.update_token_counts.call_args.args[0] == "cron-session"

View File

@@ -0,0 +1,262 @@
"""Tests for tool argument type coercion.
When LLMs return tool call arguments, they frequently put numbers as strings
("42" instead of 42) and booleans as strings ("true" instead of true).
coerce_tool_args() fixes these type mismatches by comparing argument values
against the tool's JSON Schema before dispatch.
"""
import pytest
from unittest.mock import patch
from model_tools import (
coerce_tool_args,
_coerce_value,
_coerce_number,
_coerce_boolean,
)
# ── Low-level coercion helpers ────────────────────────────────────────────
class TestCoerceNumber:
"""Unit tests for _coerce_number."""
def test_integer_string(self):
assert _coerce_number("42") == 42
assert isinstance(_coerce_number("42"), int)
def test_negative_integer(self):
assert _coerce_number("-7") == -7
def test_zero(self):
assert _coerce_number("0") == 0
assert isinstance(_coerce_number("0"), int)
def test_float_string(self):
assert _coerce_number("3.14") == 3.14
assert isinstance(_coerce_number("3.14"), float)
def test_float_with_zero_fractional(self):
"""3.0 should become int(3) since there's no fractional part."""
assert _coerce_number("3.0") == 3
assert isinstance(_coerce_number("3.0"), int)
def test_integer_only_rejects_float(self):
"""When integer_only=True, "3.14" should stay as string."""
result = _coerce_number("3.14", integer_only=True)
assert result == "3.14"
assert isinstance(result, str)
def test_integer_only_accepts_whole(self):
assert _coerce_number("42", integer_only=True) == 42
def test_not_a_number(self):
assert _coerce_number("hello") == "hello"
def test_empty_string(self):
assert _coerce_number("") == ""
def test_large_number(self):
assert _coerce_number("1000000") == 1000000
def test_scientific_notation(self):
assert _coerce_number("1e5") == 100000
def test_inf_stays_string_for_integer_only(self):
"""Infinity should not be converted to int."""
result = _coerce_number("inf")
assert result == float("inf")
def test_negative_float(self):
assert _coerce_number("-2.5") == -2.5
class TestCoerceBoolean:
"""Unit tests for _coerce_boolean."""
def test_true_lowercase(self):
assert _coerce_boolean("true") is True
def test_false_lowercase(self):
assert _coerce_boolean("false") is False
def test_true_mixed_case(self):
assert _coerce_boolean("True") is True
def test_false_mixed_case(self):
assert _coerce_boolean("False") is False
def test_true_with_whitespace(self):
assert _coerce_boolean(" true ") is True
def test_not_a_boolean(self):
assert _coerce_boolean("yes") == "yes"
def test_one_zero_not_coerced(self):
"""'1' and '0' are not boolean values."""
assert _coerce_boolean("1") == "1"
assert _coerce_boolean("0") == "0"
def test_empty_string(self):
assert _coerce_boolean("") == ""
class TestCoerceValue:
"""Unit tests for _coerce_value."""
def test_integer_type(self):
assert _coerce_value("5", "integer") == 5
def test_number_type(self):
assert _coerce_value("3.14", "number") == 3.14
def test_boolean_type(self):
assert _coerce_value("true", "boolean") is True
def test_string_type_passthrough(self):
"""Strings expected as strings should not be coerced."""
assert _coerce_value("hello", "string") == "hello"
def test_unknown_type_passthrough(self):
assert _coerce_value("stuff", "object") == "stuff"
def test_union_type_prefers_first_match(self):
"""Union types try each in order."""
assert _coerce_value("42", ["integer", "string"]) == 42
def test_union_type_falls_through(self):
"""If no type matches, return original string."""
assert _coerce_value("hello", ["integer", "boolean"]) == "hello"
def test_union_with_string_preserves_original(self):
"""A non-numeric string in [number, string] should stay a string."""
assert _coerce_value("hello", ["number", "string"]) == "hello"
# ── Full coerce_tool_args with registry ───────────────────────────────────
class TestCoerceToolArgs:
"""Integration tests for coerce_tool_args using the tool registry."""
def _mock_schema(self, properties):
"""Build a minimal tool schema with the given properties."""
return {
"name": "test_tool",
"description": "test",
"parameters": {
"type": "object",
"properties": properties,
},
}
def test_coerces_integer_arg(self):
schema = self._mock_schema({"limit": {"type": "integer"}})
with patch("model_tools.registry.get_schema", return_value=schema):
args = {"limit": "10"}
result = coerce_tool_args("test_tool", args)
assert result["limit"] == 10
assert isinstance(result["limit"], int)
def test_coerces_boolean_arg(self):
schema = self._mock_schema({"merge": {"type": "boolean"}})
with patch("model_tools.registry.get_schema", return_value=schema):
args = {"merge": "true"}
result = coerce_tool_args("test_tool", args)
assert result["merge"] is True
def test_coerces_number_arg(self):
schema = self._mock_schema({"temperature": {"type": "number"}})
with patch("model_tools.registry.get_schema", return_value=schema):
args = {"temperature": "0.7"}
result = coerce_tool_args("test_tool", args)
assert result["temperature"] == 0.7
def test_leaves_string_args_alone(self):
schema = self._mock_schema({"path": {"type": "string"}})
with patch("model_tools.registry.get_schema", return_value=schema):
args = {"path": "/tmp/file.txt"}
result = coerce_tool_args("test_tool", args)
assert result["path"] == "/tmp/file.txt"
def test_leaves_already_correct_types(self):
schema = self._mock_schema({"limit": {"type": "integer"}})
with patch("model_tools.registry.get_schema", return_value=schema):
args = {"limit": 10}
result = coerce_tool_args("test_tool", args)
assert result["limit"] == 10
def test_unknown_tool_returns_args_unchanged(self):
with patch("model_tools.registry.get_schema", return_value=None):
args = {"limit": "10"}
result = coerce_tool_args("unknown_tool", args)
assert result["limit"] == "10"
def test_empty_args(self):
assert coerce_tool_args("test_tool", {}) == {}
def test_none_args(self):
assert coerce_tool_args("test_tool", None) is None
def test_preserves_non_string_values(self):
"""Lists, dicts, and other non-string values are never touched."""
schema = self._mock_schema({
"items": {"type": "array"},
"config": {"type": "object"},
})
with patch("model_tools.registry.get_schema", return_value=schema):
args = {"items": [1, 2, 3], "config": {"key": "val"}}
result = coerce_tool_args("test_tool", args)
assert result["items"] == [1, 2, 3]
assert result["config"] == {"key": "val"}
def test_extra_args_without_schema_left_alone(self):
"""Args not in the schema properties are not touched."""
schema = self._mock_schema({"limit": {"type": "integer"}})
with patch("model_tools.registry.get_schema", return_value=schema):
args = {"limit": "10", "extra": "42"}
result = coerce_tool_args("test_tool", args)
assert result["limit"] == 10
assert result["extra"] == "42" # no schema for extra, stays string
def test_mixed_coercion(self):
"""Multiple args coerced in the same call."""
schema = self._mock_schema({
"offset": {"type": "integer"},
"limit": {"type": "integer"},
"full": {"type": "boolean"},
"path": {"type": "string"},
})
with patch("model_tools.registry.get_schema", return_value=schema):
args = {
"offset": "1",
"limit": "500",
"full": "false",
"path": "readme.md",
}
result = coerce_tool_args("test_tool", args)
assert result["offset"] == 1
assert result["limit"] == 500
assert result["full"] is False
assert result["path"] == "readme.md"
def test_failed_coercion_preserves_original(self):
"""A non-parseable string stays as string even if schema says integer."""
schema = self._mock_schema({"limit": {"type": "integer"}})
with patch("model_tools.registry.get_schema", return_value=schema):
args = {"limit": "not_a_number"}
result = coerce_tool_args("test_tool", args)
assert result["limit"] == "not_a_number"
def test_real_read_file_schema(self):
"""Test against the actual read_file schema from the registry."""
# This uses the real registry — read_file should be registered
args = {"path": "foo.py", "offset": "10", "limit": "100"}
result = coerce_tool_args("read_file", args)
assert result["path"] == "foo.py"
assert result["offset"] == 10
assert isinstance(result["offset"], int)
assert result["limit"] == 100
assert isinstance(result["limit"], int)