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:
0
tests/run_agent/__init__.py
Normal file
0
tests/run_agent/__init__.py
Normal file
268
tests/run_agent/test_1630_context_overflow_loop.py
Normal file
268
tests/run_agent/test_1630_context_overflow_loop.py
Normal 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
|
||||
475
tests/run_agent/test_413_compression.py
Normal file
475
tests/run_agent/test_413_compression.py
Normal 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
|
||||
294
tests/run_agent/test_860_dedup.py
Normal file
294
tests/run_agent/test_860_dedup.py
Normal 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
|
||||
263
tests/run_agent/test_agent_guardrails.py
Normal file
263
tests/run_agent/test_agent_guardrails.py
Normal 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) == ""
|
||||
505
tests/run_agent/test_agent_loop.py
Normal file
505
tests/run_agent/test_agent_loop.py
Normal 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
|
||||
552
tests/run_agent/test_agent_loop_tool_calling.py
Normal file
552
tests/run_agent/test_agent_loop_tool_calling.py
Normal 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)
|
||||
359
tests/run_agent/test_agent_loop_vllm.py
Normal file
359
tests/run_agent/test_agent_loop_vllm.py
Normal 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"
|
||||
)
|
||||
490
tests/run_agent/test_anthropic_error_handling.py
Normal file
490
tests/run_agent/test_anthropic_error_handling.py
Normal 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
|
||||
162
tests/run_agent/test_async_httpx_del_neuter.py
Normal file
162
tests/run_agent/test_async_httpx_del_neuter.py
Normal 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)
|
||||
199
tests/run_agent/test_compression_boundary.py
Normal file
199
tests/run_agent/test_compression_boundary.py
Normal 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"
|
||||
202
tests/run_agent/test_compression_persistence.py
Normal file
202
tests/run_agent/test_compression_persistence.py
Normal 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)"
|
||||
)
|
||||
91
tests/run_agent/test_compressor_fallback_update.py
Normal file
91
tests/run_agent/test_compressor_fallback_update.py
Normal 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
|
||||
248
tests/run_agent/test_context_pressure.py
Normal file
248
tests/run_agent/test_context_pressure.py
Normal 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)
|
||||
127
tests/run_agent/test_context_token_tracking.py
Normal file
127
tests/run_agent/test_context_token_tracking.py
Normal 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
|
||||
72
tests/run_agent/test_dict_tool_call_args.py
Normal file
72
tests/run_agent/test_dict_tool_call_args.py
Normal 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"
|
||||
73
tests/run_agent/test_exit_cleanup_interrupt.py
Normal file
73
tests/run_agent/test_exit_cleanup_interrupt.py
Normal 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()
|
||||
377
tests/run_agent/test_fallback_model.py
Normal file
377
tests/run_agent/test_fallback_model.py
Normal 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
|
||||
222
tests/run_agent/test_flush_memories_codex.py
Normal file
222
tests/run_agent/test_flush_memories_codex.py
Normal 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"
|
||||
203
tests/run_agent/test_interactive_interrupt.py
Normal file
203
tests/run_agent/test_interactive_interrupt.py
Normal 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())
|
||||
161
tests/run_agent/test_interrupt_propagation.py
Normal file
161
tests/run_agent/test_interrupt_propagation.py
Normal 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()
|
||||
162
tests/run_agent/test_large_tool_result.py
Normal file
162
tests/run_agent/test_large_tool_result.py
Normal 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
|
||||
209
tests/run_agent/test_long_context_tier_429.py
Normal file
209
tests/run_agent/test_long_context_tier_429.py
Normal 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
|
||||
189
tests/run_agent/test_openai_client_lifecycle.py
Normal file
189
tests/run_agent/test_openai_client_lifecycle.py
Normal 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
|
||||
154
tests/run_agent/test_percentage_clamp.py
Normal file
154
tests/run_agent/test_percentage_clamp.py
Normal 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"
|
||||
)
|
||||
424
tests/run_agent/test_primary_runtime_restore.py
Normal file
424
tests/run_agent/test_primary_runtime_restore.py
Normal 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"
|
||||
156
tests/run_agent/test_provider_fallback.py
Normal file
156
tests/run_agent/test_provider_fallback.py
Normal 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"
|
||||
850
tests/run_agent/test_provider_parity.py
Normal file
850
tests/run_agent/test_provider_parity.py
Normal 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"
|
||||
186
tests/run_agent/test_real_interrupt_subagent.py
Normal file
186
tests/run_agent/test_real_interrupt_subagent.py
Normal 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()
|
||||
54
tests/run_agent/test_redirect_stdout_issue.py
Normal file
54
tests/run_agent/test_redirect_stdout_issue.py
Normal 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()
|
||||
3474
tests/run_agent/test_run_agent.py
Normal file
3474
tests/run_agent/test_run_agent.py
Normal file
File diff suppressed because it is too large
Load Diff
1041
tests/run_agent/test_run_agent_codex_responses.py
Normal file
1041
tests/run_agent/test_run_agent_codex_responses.py
Normal file
File diff suppressed because it is too large
Load Diff
90
tests/run_agent/test_session_meta_filtering.py
Normal file
90
tests/run_agent/test_session_meta_filtering.py
Normal 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"
|
||||
121
tests/run_agent/test_session_reset_fix.py
Normal file
121
tests/run_agent/test_session_reset_fix.py
Normal 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
|
||||
816
tests/run_agent/test_streaming.py
Normal file
816
tests/run_agent/test_streaming.py
Normal 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)
|
||||
144
tests/run_agent/test_strict_api_validation.py
Normal file
144
tests/run_agent/test_strict_api_validation.py
Normal 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
|
||||
62
tests/run_agent/test_token_persistence_non_cli.py
Normal file
62
tests/run_agent/test_token_persistence_non_cli.py
Normal 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"
|
||||
262
tests/run_agent/test_tool_arg_coercion.py
Normal file
262
tests/run_agent/test_tool_arg_coercion.py
Normal 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)
|
||||
Reference in New Issue
Block a user