refactor(agent): extract run_conversation prologue into agent/turn_context.py
Phase 1 of the god-file decomposition plan. run_conversation's ~470-line once-per-turn setup block (stdio guarding, retry-counter resets, user-message sanitization, todo/nudge hydration, system-prompt restore-or-build, crash-resilience persistence, preflight compression, the pre_llm_call hook, and external-memory prefetch) is moved verbatim into build_turn_context(), which returns a TurnContext dataclass the loop unpacks. Behavior-neutral move-and-name refactor: the builder mutates `agent` exactly as the inline code did; only the locals the loop reads back are returned. - run_conversation: 4602 -> 4217 LOC (-385) - agent/conversation_loop.py: 4965 -> ~4580 LOC - new agent/turn_context.py: focused, dependency-injected, unit-tested in isolation Tests: tests/run_agent/ 1570 passed / 0 failed under per-file process isolation. Relocation follow-ups: 413_compression mocks now patch both module references; nudge/on_turn_start source-inspection guards point at the extracted module.
This commit is contained in:
@@ -0,0 +1,187 @@
|
||||
"""Unit tests for the extracted turn prologue (``agent/turn_context.py``).
|
||||
|
||||
These exercise ``build_turn_context`` against a lightweight fake agent to
|
||||
confirm the prologue produces the right ``TurnContext`` and applies the
|
||||
``agent`` side effects the loop relies on — without spinning up a real
|
||||
``AIAgent`` or hitting any provider.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import types
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from agent.turn_context import TurnContext, build_turn_context
|
||||
|
||||
|
||||
class _FakeTodoStore:
|
||||
def has_items(self):
|
||||
return True
|
||||
|
||||
def _hydrate(self, *_a, **_k):
|
||||
pass
|
||||
|
||||
|
||||
class _FakeGuardrails:
|
||||
def __init__(self):
|
||||
self.reset_called = False
|
||||
|
||||
def reset_for_turn(self):
|
||||
self.reset_called = True
|
||||
|
||||
|
||||
class _FakeAgent:
|
||||
"""Minimal stand-in covering only what the prologue touches."""
|
||||
|
||||
def __init__(self):
|
||||
self.session_id = "sess-1"
|
||||
self.model = "test/model"
|
||||
self.provider = "openrouter"
|
||||
self.base_url = "https://openrouter.ai/api/v1"
|
||||
self.api_key = "sk-x"
|
||||
self.api_mode = "chat_completions"
|
||||
self.platform = "cli"
|
||||
self.quiet_mode = True
|
||||
self.max_iterations = 90
|
||||
self.tools = []
|
||||
self.valid_tool_names = set()
|
||||
self.compression_enabled = False
|
||||
self.context_compressor = types.SimpleNamespace(
|
||||
protect_first_n=2, protect_last_n=2
|
||||
)
|
||||
self._cached_system_prompt = "SYSTEM"
|
||||
self._memory_store = None
|
||||
self._memory_manager = None
|
||||
self._memory_nudge_interval = 0
|
||||
self._turns_since_memory = 0
|
||||
self._user_turn_count = 0
|
||||
self._todo_store = _FakeTodoStore()
|
||||
self._tool_guardrails = _FakeGuardrails()
|
||||
self._compression_warning = None
|
||||
self._interrupt_requested = False
|
||||
self._memory_write_origin = "assistant_tool"
|
||||
self._stream_context_scrubber = None
|
||||
self._stream_think_scrubber = None
|
||||
# Attributes the prologue assigns; recorded for assertions.
|
||||
self._invalid_tool_retries = -1
|
||||
self._vision_supported = None
|
||||
self._persist_calls = 0
|
||||
|
||||
# --- methods the prologue calls ---
|
||||
def _ensure_db_session(self):
|
||||
pass
|
||||
|
||||
def _restore_primary_runtime(self):
|
||||
pass
|
||||
|
||||
def _cleanup_dead_connections(self):
|
||||
return False
|
||||
|
||||
def _emit_status(self, _msg):
|
||||
pass
|
||||
|
||||
def _replay_compression_warning(self):
|
||||
pass
|
||||
|
||||
def _hydrate_todo_store(self, *_a, **_k):
|
||||
pass
|
||||
|
||||
def _safe_print(self, *_a, **_k):
|
||||
pass
|
||||
|
||||
def _persist_session(self, *_a, **_k):
|
||||
self._persist_calls += 1
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _stub_runtime_main():
|
||||
"""``build_turn_context`` calls ``auxiliary_client.set_runtime_main`` as a
|
||||
production side effect (telling aux tools the live main provider/model).
|
||||
That writes a module-level global these unit tests don't care about and
|
||||
which would otherwise leak into sibling tests (e.g. provider-parity
|
||||
resolution) when the per-test process isolation plugin is disabled. Stub
|
||||
it out so the prologue tests stay hermetic.
|
||||
"""
|
||||
with patch("agent.auxiliary_client.set_runtime_main", lambda *a, **k: None):
|
||||
yield
|
||||
|
||||
|
||||
def _build(agent, **overrides):
|
||||
kwargs = dict(
|
||||
agent=agent,
|
||||
user_message="hello",
|
||||
system_message=None,
|
||||
conversation_history=None,
|
||||
task_id=None,
|
||||
stream_callback=None,
|
||||
persist_user_message=None,
|
||||
restore_or_build_system_prompt=lambda *a, **k: None,
|
||||
install_safe_stdio=lambda: None,
|
||||
sanitize_surrogates=lambda s: s,
|
||||
summarize_user_message_for_log=lambda s: s,
|
||||
set_session_context=lambda _sid: None,
|
||||
set_current_write_origin=lambda _o: None,
|
||||
ra=lambda: types.SimpleNamespace(_set_interrupt=lambda *a, **k: None),
|
||||
)
|
||||
kwargs.update(overrides)
|
||||
return build_turn_context(**kwargs)
|
||||
|
||||
|
||||
def test_returns_turn_context_with_user_message_appended():
|
||||
agent = _FakeAgent()
|
||||
ctx = _build(agent)
|
||||
assert isinstance(ctx, TurnContext)
|
||||
assert ctx.user_message == "hello"
|
||||
# The user turn was appended and indexed.
|
||||
assert ctx.messages[-1] == {"role": "user", "content": "hello"}
|
||||
assert ctx.current_turn_user_idx == len(ctx.messages) - 1
|
||||
assert ctx.active_system_prompt == "SYSTEM"
|
||||
|
||||
|
||||
def test_applies_agent_side_effects():
|
||||
agent = _FakeAgent()
|
||||
_build(agent)
|
||||
# Retry counters reset, guardrails reset, vision re-armed, turn counted.
|
||||
assert agent._invalid_tool_retries == 0
|
||||
assert agent._tool_guardrails.reset_called is True
|
||||
assert agent._vision_supported is True
|
||||
assert agent._user_turn_count == 1
|
||||
# Crash-resilience persistence fired once.
|
||||
assert agent._persist_calls == 1
|
||||
# task/turn ids assigned on the agent.
|
||||
assert agent._current_task_id
|
||||
assert agent._current_turn_id
|
||||
|
||||
|
||||
def test_task_id_passthrough():
|
||||
agent = _FakeAgent()
|
||||
ctx = _build(agent, task_id="fixed-task")
|
||||
assert ctx.effective_task_id == "fixed-task"
|
||||
assert agent._current_task_id == "fixed-task"
|
||||
|
||||
|
||||
def test_persist_user_message_becomes_original():
|
||||
agent = _FakeAgent()
|
||||
ctx = _build(agent, user_message="api-prefixed", persist_user_message="clean")
|
||||
# original_user_message tracks the clean persist override.
|
||||
assert ctx.original_user_message == "clean"
|
||||
# but the appended user turn carries the full (sanitized) message.
|
||||
assert ctx.messages[-1]["content"] == "api-prefixed"
|
||||
|
||||
|
||||
def test_memory_nudge_fires_at_interval():
|
||||
agent = _FakeAgent()
|
||||
agent._memory_nudge_interval = 1
|
||||
agent.valid_tool_names = {"memory"}
|
||||
agent._memory_store = object()
|
||||
ctx = _build(agent)
|
||||
assert ctx.should_review_memory is True
|
||||
assert agent._turns_since_memory == 0 # reset after firing
|
||||
|
||||
|
||||
def test_no_review_when_memory_disabled():
|
||||
agent = _FakeAgent()
|
||||
ctx = _build(agent)
|
||||
assert ctx.should_review_memory is False
|
||||
@@ -553,6 +553,7 @@ class TestPreflightCompression:
|
||||
agent.status_callback = lambda ev, msg: status_messages.append((ev, msg))
|
||||
|
||||
with (
|
||||
patch("agent.turn_context.estimate_request_tokens_rough", return_value=114_000),
|
||||
patch("agent.conversation_loop.estimate_request_tokens_rough", return_value=114_000),
|
||||
patch.object(agent, "_compress_context") as mock_compress,
|
||||
patch.object(agent, "_persist_session"),
|
||||
@@ -604,6 +605,7 @@ class TestPreflightCompression:
|
||||
return 125_000 if _rough_calls["n"] == 1 else 40_000
|
||||
|
||||
with (
|
||||
patch("agent.turn_context.estimate_request_tokens_rough", side_effect=_rough_estimate),
|
||||
patch("agent.conversation_loop.estimate_request_tokens_rough", side_effect=_rough_estimate),
|
||||
patch.object(agent, "_compress_context") as mock_compress,
|
||||
patch.object(agent, "_persist_session"),
|
||||
@@ -728,6 +730,7 @@ class TestPreflightCompression:
|
||||
agent.client.chat.completions.create.side_effect = [ok_resp]
|
||||
|
||||
with (
|
||||
patch("agent.turn_context.estimate_request_tokens_rough", return_value=144_669),
|
||||
patch("agent.conversation_loop.estimate_request_tokens_rough", return_value=144_669),
|
||||
# Compression no-ops (returns input unchanged) — mirrors an aux
|
||||
# summary-model timeout where the messages can't be reduced.
|
||||
@@ -760,6 +763,7 @@ class TestPreflightCompression:
|
||||
agent.client.chat.completions.create.side_effect = [ok_resp]
|
||||
|
||||
with (
|
||||
patch("agent.turn_context.estimate_request_tokens_rough", return_value=144_669),
|
||||
patch("agent.conversation_loop.estimate_request_tokens_rough", return_value=144_669),
|
||||
patch.object(agent, "_compress_context", side_effect=lambda msgs, *a, **k: (msgs, agent._cached_system_prompt)),
|
||||
patch.object(agent, "_persist_session"),
|
||||
|
||||
@@ -117,25 +117,29 @@ def test_assistant_only_history_does_not_advance_user_turn_count():
|
||||
|
||||
|
||||
def test_production_code_contains_hydration_block():
|
||||
"""Smoke test: confirm the hydration code is actually wired into
|
||||
run_conversation(). If someone deletes it, tests above still pass
|
||||
against the inline replica — this fails them awake.
|
||||
"""Smoke test: confirm the hydration code is actually wired into the
|
||||
turn path. If someone deletes it, tests above still pass against the
|
||||
inline replica — this fails them awake.
|
||||
|
||||
After the run_agent.py refactor the agent-loop body lives in
|
||||
``agent/conversation_loop.py`` and uses ``agent.X`` rather than
|
||||
``self.X``. Assert the block is present in the extracted module
|
||||
specifically — if it ever drifts back into run_agent.py or
|
||||
disappears entirely, this guard fails loudly.
|
||||
The agent-loop prologue now lives in ``agent/turn_context.py``
|
||||
(``build_turn_context``), with the loop body in
|
||||
``agent/conversation_loop.py``. Assert the block is present in the
|
||||
turn subsystem — if it disappears entirely, this guard fails loudly.
|
||||
Either module counts so the guard tolerates legitimate relocation
|
||||
within the turn subsystem.
|
||||
"""
|
||||
from pathlib import Path
|
||||
repo = Path(__file__).resolve().parents[2]
|
||||
cl_path = repo / "agent" / "conversation_loop.py"
|
||||
src_cl = cl_path.read_text(encoding="utf-8")
|
||||
turn_src = "".join(
|
||||
(repo / "agent" / name).read_text(encoding="utf-8")
|
||||
for name in ("conversation_loop.py", "turn_context.py")
|
||||
)
|
||||
# Anchor on the unique comment + the modulo line.
|
||||
assert "Hydrate per-session nudge counters from persisted history" in src_cl, (
|
||||
f"Hydration comment missing from {cl_path}"
|
||||
assert "Hydrate per-session nudge counters from persisted history" in turn_src, (
|
||||
"Hydration comment missing from the turn subsystem "
|
||||
"(conversation_loop.py / turn_context.py)"
|
||||
)
|
||||
assert (
|
||||
"agent._turns_since_memory = prior_user_turns % agent._memory_nudge_interval"
|
||||
in src_cl
|
||||
), f"Hydration modulo assignment missing from {cl_path}"
|
||||
in turn_src
|
||||
), "Hydration modulo assignment missing from the turn subsystem"
|
||||
|
||||
@@ -6393,18 +6393,16 @@ class TestMemoryNudgeCounterPersistence:
|
||||
assert a._iters_since_skill == 0
|
||||
|
||||
def test_counters_not_reset_in_preamble(self):
|
||||
"""The run_conversation preamble must not zero the nudge counters."""
|
||||
"""The turn preamble must not zero the nudge counters."""
|
||||
import inspect
|
||||
from agent.conversation_loop import run_conversation as _rc
|
||||
src = inspect.getsource(_rc)
|
||||
# The preamble resets many fields (retry counts, budget, etc.)
|
||||
# before the main loop. Find that reset block and verify our
|
||||
# counters aren't in it. The reset block ends at iteration_budget.
|
||||
# The extracted body uses ``agent.X`` (not ``self.X``). Anchor
|
||||
# exactly on ``agent.iteration_budget = IterationBudget`` so an
|
||||
# unrelated identifier ending in ``iteration_budget`` (e.g.
|
||||
# ``_iteration_budget`` or ``shared_iteration_budget``) can't
|
||||
# match the boundary.
|
||||
from agent.turn_context import build_turn_context as _btc
|
||||
src = inspect.getsource(_btc)
|
||||
# The preamble (now in build_turn_context) resets many fields (retry
|
||||
# counts, budget, etc.) before returning. Find that reset block and
|
||||
# verify our counters aren't in it. The reset block ends at
|
||||
# iteration_budget. Anchor exactly on
|
||||
# ``agent.iteration_budget = IterationBudget`` so an unrelated
|
||||
# identifier ending in ``iteration_budget`` can't match the boundary.
|
||||
preamble_end = src.index("agent.iteration_budget = IterationBudget")
|
||||
preamble = src[:preamble_end]
|
||||
assert "agent._turns_since_memory = 0" not in preamble
|
||||
@@ -6490,23 +6488,23 @@ class TestMemoryProviderTurnStart:
|
||||
"""
|
||||
|
||||
def test_on_turn_start_called_before_prefetch(self):
|
||||
"""Source-level check: on_turn_start appears before prefetch_all in run_conversation."""
|
||||
"""Source-level check: on_turn_start appears before prefetch_all in the prologue."""
|
||||
import inspect
|
||||
from agent.conversation_loop import run_conversation as _rc
|
||||
src = inspect.getsource(_rc)
|
||||
from agent.turn_context import build_turn_context as _btc
|
||||
src = inspect.getsource(_btc)
|
||||
# Find the actual method calls, not comments
|
||||
idx_turn_start = src.index(".on_turn_start(")
|
||||
idx_prefetch = src.index(".prefetch_all(")
|
||||
assert idx_turn_start < idx_prefetch, (
|
||||
"on_turn_start() must be called before prefetch_all() in run_conversation "
|
||||
"on_turn_start() must be called before prefetch_all() in the turn prologue "
|
||||
"so that memory providers have the correct turn count for cadence checks"
|
||||
)
|
||||
|
||||
def test_on_turn_start_uses_user_turn_count(self):
|
||||
"""Source-level check: on_turn_start receives the user_turn_count."""
|
||||
import inspect
|
||||
from agent.conversation_loop import run_conversation as _rc
|
||||
src = inspect.getsource(_rc)
|
||||
from agent.turn_context import build_turn_context as _btc
|
||||
src = inspect.getsource(_btc)
|
||||
# The extracted body uses ``agent.X`` rather than ``self.X``;
|
||||
# assert the extracted-form spelling directly.
|
||||
assert "on_turn_start(agent._user_turn_count" in src
|
||||
|
||||
Reference in New Issue
Block a user