diff --git a/scripts/run_benchmark.py b/scripts/run_benchmark.py index 9e1cf47..ae07e08 100644 --- a/scripts/run_benchmark.py +++ b/scripts/run_benchmark.py @@ -175,6 +175,7 @@ def _make_mock_provider() -> MagicMock: content=_MOCK_ACTION_PLAN, model="mock-model", usage={"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150}, ) + provider.acomplete.return_value = provider.complete.return_value return provider diff --git a/src/rle/agents/base_role.py b/src/rle/agents/base_role.py index bb13dd3..551f8fe 100644 --- a/src/rle/agents/base_role.py +++ b/src/rle/agents/base_role.py @@ -2,14 +2,17 @@ from __future__ import annotations +import asyncio import json import logging +import time from abc import abstractmethod from typing import Any, ClassVar, Tuple from felix_agent_sdk import LLMAgent, LLMResult, LLMTask from felix_agent_sdk.communication import Spoke from felix_agent_sdk.core import HelixGeometry +from felix_agent_sdk.events.types import EventType from felix_agent_sdk.providers.base import BaseProvider from felix_agent_sdk.providers.types import ChatMessage, CompletionResult, MessageRole from felix_agent_sdk.tokens.budget import TokenBudget @@ -227,14 +230,8 @@ def _get_spoke_context(self) -> list[dict[str, Any]]: }) return context - def _call_provider( - self, - system_prompt: str, - user_prompt: str, - temperature: float, - max_tokens: int, - ) -> CompletionResult: - """Override to pass extra provider kwargs and no-think prefix.""" + def _build_messages(self, system_prompt: str, user_prompt: str) -> list[ChatMessage]: + """Build the message list, including the no-think prefill when enabled.""" messages = [ ChatMessage(role=MessageRole.SYSTEM, content=system_prompt), ChatMessage(role=MessageRole.USER, content=user_prompt), @@ -246,6 +243,22 @@ def _call_provider( messages.append( ChatMessage(role=MessageRole.ASSISTANT, content=""), ) + return messages + + def _record_completion(self, result: CompletionResult) -> CompletionResult: + self._last_raw_output = result.content + self._last_usage = result.usage if hasattr(result, "usage") else None + return result + + def _call_provider( + self, + system_prompt: str, + user_prompt: str, + temperature: float, + max_tokens: int, + ) -> CompletionResult: + """Override to pass extra provider kwargs and no-think prefix.""" + messages = self._build_messages(system_prompt, user_prompt) def _do_complete() -> CompletionResult: return self.provider.complete( @@ -261,9 +274,40 @@ def _do_complete() -> CompletionResult: else: result = _do_complete() - self._last_raw_output = result.content - self._last_usage = result.usage if hasattr(result, "usage") else None - return result + return self._record_completion(result) + + async def _acall_provider( + self, + system_prompt: str, + user_prompt: str, + temperature: float, + max_tokens: int, + ) -> CompletionResult: + """Async mirror of ``_call_provider`` using the provider's native + ``acomplete()`` (felix 0.3.0). Providers without async support fall + back to running the sync path in a worker thread. + """ + messages = self._build_messages(system_prompt, user_prompt) + + async def _do_complete() -> CompletionResult: + return await self.provider.acomplete( + messages, + temperature=temperature, + max_tokens=max_tokens, + **self._provider_kwargs, + ) + + try: + if self._weave_op is not None: + result: CompletionResult = await self._weave_op(_do_complete)() + else: + result = await _do_complete() + except NotImplementedError: + return await asyncio.to_thread( + self._call_provider, system_prompt, user_prompt, temperature, max_tokens, + ) + + return self._record_completion(result) # ------------------------------------------------------------------ # Abstract methods — subclasses must implement @@ -603,20 +647,118 @@ def deliberate( self._last_action_plan = plan return plan - def _retry_with_correction(self, bad_output: str, error: str) -> LLMResult: - """Retry with a short correction prompt asking the LLM to fix its JSON.""" - system_prompt = ( - "You previously produced invalid JSON. Fix the error and return ONLY " - "valid JSON matching the required schema. No explanation, no markdown." + async def adeliberate( + self, + state: GameState, + current_time: float, + context_history: list[dict[str, Any]] | None = None, + ) -> ActionPlan: + """Async mirror of ``deliberate()`` — the LLM call runs on the event + loop via the provider's native ``acomplete()`` (felix 0.3.0), so a + timeout cancellation actually aborts the request instead of leaving + an orphaned worker thread. + """ + if not hasattr(self, "_state") or self._state.value == "waiting": + self.spawn(current_time) + self.update_position(current_time) + + spoke_context = self._get_spoke_context() + effective_context = spoke_context if spoke_context else (context_history or []) + task = self.build_task(state, effective_context) + result = await self.aprocess_task(task) + + try: + plan = self.parse_action_plan(result, state.colony.tick) + except ActionPlanParseError as first_error: + logger.warning( + "%s: parse failed, retrying with correction prompt: %s", + self.ROLE_NAME, first_error.reason, + ) + try: + retry_result = await self._aretry_with_correction( + result.content, first_error.reason, + ) + plan = self.parse_action_plan(retry_result, state.colony.tick) + except (ActionPlanParseError, Exception) as retry_error: + raise ActionPlanParseError( + result.content, f"Retry also failed: {retry_error}", + ) from retry_error + + self._last_action_plan = plan + return plan + + async def aprocess_task(self, task: LLMTask) -> LLMResult: + """Async mirror of ``LLMAgent.process_task()`` using ``_acall_provider``. + + Replicates the SDK's sync pipeline (events, adaptive temperature, + confidence, accounting) — the SDK exposes async only at the provider + layer, not on LLMAgent. + """ + start = time.perf_counter() + + self.emit_event( + EventType.TASK_STARTED, + {"task_id": task.task_id, "agent_type": self.agent_type}, + source=f"agent:{self.agent_id}", ) - user_prompt = ( + + temperature = self.get_adaptive_temperature() + system_prompt, user_prompt = self.create_position_aware_prompt(task) + + completion = await self._acall_provider( + system_prompt, user_prompt, temperature, self.max_tokens, + ) + + confidence = self.calculate_confidence(completion.content) + self.record_confidence(confidence) + + elapsed = time.perf_counter() - start + self.total_tokens_used += completion.total_tokens + self.total_processing_time += elapsed + + result = LLMResult( + agent_id=self.agent_id, + task_id=task.task_id, + content=completion.content, + position_info=self.get_position_info(), + completion_result=completion, + processing_time=elapsed, + confidence=confidence, + temperature_used=temperature, + token_budget_used=completion.total_tokens, + ) + self.processing_results.append(result) + + self.emit_event( + EventType.TASK_COMPLETED, + { + "task_id": task.task_id, + "agent_type": self.agent_type, + "confidence": round(confidence, 4), + "temperature": round(temperature, 4), + "tokens": completion.total_tokens, + "processing_time": round(elapsed, 4), + "phase": result.position_info.get("phase", ""), + }, + source=f"agent:{self.agent_id}", + ) + + return result + + _CORRECTION_SYSTEM_PROMPT = ( + "You previously produced invalid JSON. Fix the error and return ONLY " + "valid JSON matching the required schema. No explanation, no markdown." + ) + + @staticmethod + def _correction_user_prompt(bad_output: str, error: str) -> str: + return ( f"Your previous output had this error: {error}\n\n" f"Original output (fix this):\n{bad_output[:2000]}\n\n" "Return ONLY the corrected JSON object." ) - completion = self._call_provider( - system_prompt, user_prompt, temperature=0.1, max_tokens=self.max_tokens, - ) + + def _make_retry_result(self, completion: CompletionResult) -> LLMResult: self.total_tokens_used += completion.total_tokens return LLMResult( agent_id=self.agent_id, @@ -629,3 +771,21 @@ def _retry_with_correction(self, bad_output: str, error: str) -> LLMResult: temperature_used=0.1, token_budget_used=completion.total_tokens, ) + + def _retry_with_correction(self, bad_output: str, error: str) -> LLMResult: + """Retry with a short correction prompt asking the LLM to fix its JSON.""" + completion = self._call_provider( + self._CORRECTION_SYSTEM_PROMPT, + self._correction_user_prompt(bad_output, error), + temperature=0.1, max_tokens=self.max_tokens, + ) + return self._make_retry_result(completion) + + async def _aretry_with_correction(self, bad_output: str, error: str) -> LLMResult: + """Async mirror of ``_retry_with_correction``.""" + completion = await self._acall_provider( + self._CORRECTION_SYSTEM_PROMPT, + self._correction_user_prompt(bad_output, error), + temperature=0.1, max_tokens=self.max_tokens, + ) + return self._make_retry_result(completion) diff --git a/src/rle/orchestration/game_loop.py b/src/rle/orchestration/game_loop.py index 723f3cc..41e352d 100644 --- a/src/rle/orchestration/game_loop.py +++ b/src/rle/orchestration/game_loop.py @@ -176,18 +176,17 @@ async def _deliberate_agent_with_timeout( self, agent: RimWorldRoleAgent, state: object, current_time: float, tick_num: int, ) -> tuple[RimWorldRoleAgent, ActionPlan | None]: - """Run one agent's deliberation in a worker thread with a hard timeout. + """Run one agent's deliberation with a hard timeout. - On timeout: emits a deliberation_timeout ERROR event, records a - provider_error entry in _deliberation_log, and returns (agent, None) - so the tick can continue. The hung thread keeps running until the - provider's own timeout cleans it up (Python threads can't be killed). + Deliberation is natively async (felix 0.3.0 ``acomplete()``), so a + timeout cancels the in-flight provider request instead of leaving an + orphaned worker thread. On timeout: emits a deliberation_timeout + ERROR event, records it in _deliberation_log, and returns + (agent, None) so the tick can continue. """ try: return await asyncio.wait_for( - asyncio.to_thread( - self._deliberate_agent, agent, state, current_time, tick_num, - ), + self._deliberate_agent(agent, state, current_time, tick_num), timeout=self._role_timeout_s, ) except asyncio.TimeoutError: @@ -212,7 +211,7 @@ async def _deliberate_agent_with_timeout( async def _deliberate_parallel( self, state: object, current_time: float, tick_num: int, ) -> list[tuple[RimWorldRoleAgent, ActionPlan | None]]: - """Run role agents concurrently in worker threads with per-task timeout.""" + """Run role agents concurrently on the event loop with per-task timeout.""" return list(await asyncio.gather(*[ self._deliberate_agent_with_timeout(a, state, current_time, tick_num) for a in self._role_agents @@ -230,17 +229,17 @@ async def _deliberate_sequential( results.append((agent_result, plan)) return results - def _deliberate_agent( + async def _deliberate_agent( self, agent: RimWorldRoleAgent, state: object, current_time: float, tick_num: int, ) -> tuple[RimWorldRoleAgent, ActionPlan | None]: - """Run one agent's deliberation. Thread-safe for parallel execution. + """Run one agent's deliberation. Agents read inter-agent context from their CentralPost spoke internally. """ t0 = _time.monotonic() try: - plan = agent.deliberate(state, current_time) # type: ignore[arg-type] + plan = await agent.adeliberate(state, current_time) # type: ignore[arg-type] except ActionPlanParseError as e: latency_ms = round((_time.monotonic() - t0) * 1000, 1) logger.warning( diff --git a/src/rle/providers/claude_code.py b/src/rle/providers/claude_code.py index c5123a8..8456a2c 100644 --- a/src/rle/providers/claude_code.py +++ b/src/rle/providers/claude_code.py @@ -8,6 +8,7 @@ from __future__ import annotations +import asyncio import json import logging import os @@ -84,15 +85,10 @@ def _split_messages(self, messages: Sequence[ChatMessage]) -> tuple[str, str]: ) return "\n\n".join(system_parts), "\n\n".join(user_parts) - def complete( - self, - messages: Sequence[ChatMessage], - *, - temperature: float | None = None, - max_tokens: int | None = None, - stop_sequences: list[str] | None = None, - **kwargs: Any, - ) -> CompletionResult: + def _build_invocation( + self, messages: Sequence[ChatMessage], + ) -> tuple[list[str], str, dict[str, str]]: + """Build (cmd, stdin_payload, env) for a ``claude -p`` call.""" cli = self._resolve_cli() system_prompt, user_prompt = self._split_messages(messages) @@ -109,6 +105,18 @@ def complete( cmd.extend(["--system-prompt", system_prompt]) env = {k: v for k, v in os.environ.items() if k not in _STRIPPED_ENV_VARS} + return cmd, user_prompt, env + + def complete( + self, + messages: Sequence[ChatMessage], + *, + temperature: float | None = None, + max_tokens: int | None = None, + stop_sequences: list[str] | None = None, + **kwargs: Any, + ) -> CompletionResult: + cmd, user_prompt, env = self._build_invocation(messages) try: proc = subprocess.run( @@ -127,18 +135,68 @@ def complete( provider=self.provider_name, ) from e - if proc.returncode != 0: - detail = (proc.stderr or proc.stdout or "").strip()[-2000:] + return self._parse_cli_output(proc.returncode, proc.stdout, proc.stderr) + + async def acomplete( + self, + messages: Sequence[ChatMessage], + *, + temperature: float | None = None, + max_tokens: int | None = None, + stop_sequences: list[str] | None = None, + **kwargs: Any, + ) -> CompletionResult: + """Async version of complete(); the CLI subprocess runs without + blocking the event loop, and cancellation kills the subprocess.""" + cmd, user_prompt, env = self._build_invocation(messages) + + proc = await asyncio.create_subprocess_exec( + *cmd, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=self._workdir, + env=env, + ) + try: + stdout_b, stderr_b = await asyncio.wait_for( + proc.communicate(user_prompt.encode("utf-8")), + timeout=self.config.timeout, + ) + except asyncio.TimeoutError as e: + proc.kill() + await proc.wait() + raise ProviderError( + f"claude -p timed out after {self.config.timeout}s", + provider=self.provider_name, + ) from e + except asyncio.CancelledError: + proc.kill() + await proc.wait() + raise + + return self._parse_cli_output( + proc.returncode or 0, + stdout_b.decode("utf-8", errors="replace"), + stderr_b.decode("utf-8", errors="replace"), + ) + + def _parse_cli_output( + self, returncode: int, stdout: str, stderr: str, + ) -> CompletionResult: + """Translate raw ``claude -p`` output into a CompletionResult.""" + if returncode != 0: + detail = (stderr or stdout or "").strip()[-2000:] raise ProviderError( - f"claude -p exited with code {proc.returncode}: {detail}", + f"claude -p exited with code {returncode}: {detail}", provider=self.provider_name, ) try: - data: dict[str, Any] = json.loads(proc.stdout) + data: dict[str, Any] = json.loads(stdout) except json.JSONDecodeError as e: raise ProviderError( - f"claude -p returned non-JSON output: {proc.stdout[:500]!r}", + f"claude -p returned non-JSON output: {stdout[:500]!r}", provider=self.provider_name, ) from e diff --git a/tests/conftest.py b/tests/conftest.py index 6954d33..ed76706 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -296,4 +296,5 @@ def mock_provider() -> MagicMock: model="mock-model", usage={"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150}, ) + provider.acomplete.return_value = provider.complete.return_value return provider diff --git a/tests/integration/test_game_loop.py b/tests/integration/test_game_loop.py index dc95b82..3a76054 100644 --- a/tests/integration/test_game_loop.py +++ b/tests/integration/test_game_loop.py @@ -1,4 +1,4 @@ -"""Integration tests for the full RLE game loop.""" +"""Integration tests for the full RLE game loop.""" from __future__ import annotations @@ -152,6 +152,7 @@ def _make_mock_provider() -> MagicMock: model="mock", usage={"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150}, ) + provider.acomplete.return_value = provider.complete.return_value return provider @@ -229,7 +230,7 @@ async def test_agent_called_each_tick(self) -> None: loop = RLEGameLoop(config, client, [agent]) await loop.run(max_ticks=5) - assert provider.complete.call_count == 5 + assert provider.acomplete.call_count == 5 class TestMacroTime: @@ -307,7 +308,7 @@ async def test_all_agents_deliberate(self) -> None: result = await loop.run_tick() # Provider called once per agent (1 MapAnalyst + 6 role agents) - assert provider.complete.call_count == 7 + assert provider.acomplete.call_count == 7 # Result is merged plan from orchestrator assert result.plan.role == "orchestrator" @@ -325,7 +326,7 @@ async def test_multi_agent_10_ticks(self) -> None: assert len(results) == 3 # 7 agents * 3 ticks = 21 provider calls - assert provider.complete.call_count == 21 + assert provider.acomplete.call_count == 21 async def test_deliberation_log_property_returns_per_tick_records(self) -> None: provider = _make_mock_provider() @@ -390,21 +391,20 @@ async def test_role_timeout_emits_event_and_other_agents_continue( ) -> None: """A7: a hung deliberation must time out, emit ERROR with error_type=deliberation_timeout, and leave the rest of the tick - operational. The hung thread keeps running in the background — we - check only the orchestrator-visible behaviour.""" + operational. The in-flight provider request is cancelled by the + timeout — we check only the orchestrator-visible behaviour.""" # Slow provider: blocks for 5s on every call. role_timeout_s=0.05 # forces the timeout long before the call completes. slow_provider = MagicMock(spec=BaseProvider) - def _slow_complete(*_args: object, **_kwargs: object) -> CompletionResult: - import time as _t - _t.sleep(5.0) + async def _slow_acomplete(*_args: object, **_kwargs: object) -> CompletionResult: + await asyncio.sleep(5.0) return CompletionResult( content=ACTION_PLAN_JSON, model="mock", usage={"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, ) - slow_provider.complete.side_effect = _slow_complete + slow_provider.acomplete.side_effect = _slow_acomplete slow_provider.provider_name = "openai" agents = _make_all_agents(slow_provider) @@ -460,10 +460,10 @@ async def test_provider_call_event_carries_raw_output( for event in provider_calls: assert "raw_output" in event.data assert "raw_output_truncated" in event.data - # Mock provider returns the canned ACTION_PLAN_JSON — verify it's there + # Mock provider returns the canned ACTION_PLAN_JSON — verify it's there assert event.data["raw_output"] is not None assert "actions" in event.data["raw_output"] - # Mock output is well under 4KB → truncation flag stays False + # Mock output is well under 4KB → truncation flag stays False assert event.data["raw_output_truncated"] is False @@ -550,7 +550,7 @@ async def test_evaluator_stops_loop(self) -> None: ) results = await loop.run(max_ticks=100) - # population=3, initial=10 → survival=0.3 < 0.5 → defeat on first tick + # population=3, initial=10 → survival=0.3 < 0.5 → defeat on first tick assert len(results) == 1 assert loop.evaluation_result is not None assert loop.evaluation_result.outcome == "defeat" @@ -591,7 +591,7 @@ async def test_parallel_all_agents_produce_plans(self) -> None: loop = RLEGameLoop(config, client, agents, parallel=True) result = await loop.run_tick() - assert provider.complete.call_count == 7 # 1 MapAnalyst + 6 role agents + assert provider.acomplete.call_count == 7 # 1 MapAnalyst + 6 role agents assert result.plan.role == "orchestrator" assert loop._parse_successes == 7 assert loop._parse_failures == 0 @@ -629,7 +629,7 @@ async def test_sequential_mode_still_works(self) -> None: loop = RLEGameLoop(config, client, agents, parallel=False) result = await loop.run_tick() - assert provider.complete.call_count == 7 # 1 MapAnalyst + 6 role agents + assert provider.acomplete.call_count == 7 # 1 MapAnalyst + 6 role agents assert result.plan.role == "orchestrator" assert loop._parse_successes == 7 @@ -646,7 +646,7 @@ async def test_sequential_3_ticks(self) -> None: results = await loop.run(max_ticks=3) assert len(results) == 3 - assert provider.complete.call_count == 21 # 7 agents * 3 ticks + assert provider.acomplete.call_count == 21 # 7 agents * 3 ticks # ------------------------------------------------------------------ @@ -757,7 +757,7 @@ async def test_scheduled_incidents_fire_at_tick(self) -> None: for _ in range(3): await loop.run_tick() - # Test passes if all 3 ticks completed without raising — end-to-end + # Test passes if all 3 ticks completed without raising — end-to-end # wiring verification. See test_fire_scheduled_incidents_calls_client # below for assertion on actual call arguments. @@ -772,7 +772,7 @@ async def test_no_triggered_incidents_no_calls(self) -> None: client._client = httpx.AsyncClient( transport=_make_transport(), base_url="http://test", ) - # Empty list — no incidents should fire + # Empty list — no incidents should fire loop = RLEGameLoop(config, client, [agent], triggered_incidents=[]) result = await loop.run_tick() @@ -800,18 +800,18 @@ async def test_fire_scheduled_incidents_calls_client(self) -> None: ], ) - # Tick 4 — should NOT fire + # Tick 4 — should NOT fire await loop._fire_scheduled_incidents(4) mock_client.trigger_incident.assert_not_awaited() - # Tick 5 — SHOULD fire Plague + # Tick 5 — SHOULD fire Plague await loop._fire_scheduled_incidents(5) mock_client.trigger_incident.assert_awaited_once_with( "Plague", map_id=0, custom_letter_label="Plague outbreak", ) - # Tick 6 — should NOT fire again + # Tick 6 — should NOT fire again mock_client.trigger_incident.reset_mock() await loop._fire_scheduled_incidents(6) mock_client.trigger_incident.assert_not_awaited() @@ -836,6 +836,6 @@ async def test_fire_scheduled_incidents_swallows_exceptions(self) -> None: ], ) - # Should NOT raise — the game loop catches exceptions and logs + # Should NOT raise — the game loop catches exceptions and logs await loop._fire_scheduled_incidents(3) mock_client.trigger_incident.assert_awaited_once() diff --git a/tests/integration/test_scenario_run.py b/tests/integration/test_scenario_run.py index e32bb92..7f61b56 100644 --- a/tests/integration/test_scenario_run.py +++ b/tests/integration/test_scenario_run.py @@ -141,6 +141,7 @@ def _make_provider() -> MagicMock: content=ACTION_PLAN_JSON, model="mock", usage={"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150}, ) + provider.acomplete.return_value = provider.complete.return_value return provider diff --git a/tests/unit/test_base_role.py b/tests/unit/test_base_role.py index 1e3e543..6a43b8d 100644 --- a/tests/unit/test_base_role.py +++ b/tests/unit/test_base_role.py @@ -402,3 +402,55 @@ def test_full_pipeline( assert agent._last_action_plan is plan # Provider was called exactly once mock_provider.complete.assert_called_once() + + +# ------------------------------------------------------------------ +# adeliberate (async path via provider.acomplete, felix 0.3.0) +# ------------------------------------------------------------------ + + +class TestADeliberate: + async def test_full_pipeline_uses_acomplete( + self, mock_provider: MagicMock, helix: HelixGeometry, + sample_game_state: GameState, + ) -> None: + mock_provider.acomplete.return_value = mock_provider.complete.return_value + agent = _DummyRoleAgent("d-01", mock_provider, helix, spawn_time=0.0) + plan = await agent.adeliberate(sample_game_state, current_time=0.2) + assert isinstance(plan, ActionPlan) + assert plan.role == "dummy" + assert plan.tick == 720000 + assert agent._last_action_plan is plan + mock_provider.acomplete.assert_called_once() + mock_provider.complete.assert_not_called() + + async def test_acall_falls_back_to_sync_when_async_unsupported( + self, mock_provider: MagicMock, helix: HelixGeometry, + ) -> None: + mock_provider.acomplete.side_effect = NotImplementedError("no async") + agent = _DummyRoleAgent("d-01", mock_provider, helix, spawn_time=0.0) + result = await agent._acall_provider("sys", "user", 0.5, 100) + assert result is mock_provider.complete.return_value + mock_provider.complete.assert_called_once() + + async def test_acall_records_last_output_and_usage( + self, mock_provider: MagicMock, helix: HelixGeometry, + ) -> None: + completion = mock_provider.complete.return_value + mock_provider.acomplete.return_value = completion + agent = _DummyRoleAgent("d-01", mock_provider, helix, spawn_time=0.0) + await agent._acall_provider("sys", "user", 0.5, 100) + assert agent._last_raw_output == completion.content + assert agent._last_usage == completion.usage + + async def test_acall_applies_no_think_prefill( + self, mock_provider: MagicMock, helix: HelixGeometry, + ) -> None: + mock_provider.provider_name = "openai" + mock_provider.acomplete.return_value = mock_provider.complete.return_value + agent = _DummyRoleAgent("d-01", mock_provider, helix, spawn_time=0.0) + agent.set_no_think(True) + await agent._acall_provider("sys", "user", 0.5, 100) + messages = mock_provider.acomplete.call_args[0][0] + assert messages[-1].role == MessageRole.ASSISTANT + assert messages[-1].content == "" diff --git a/tests/unit/test_claude_code_provider.py b/tests/unit/test_claude_code_provider.py index 93a1f97..6911e6c 100644 --- a/tests/unit/test_claude_code_provider.py +++ b/tests/unit/test_claude_code_provider.py @@ -2,10 +2,11 @@ from __future__ import annotations +import asyncio import json import subprocess from typing import Any -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from felix_agent_sdk.providers.errors import ProviderError @@ -148,6 +149,82 @@ def test_assistant_messages_ignored(self) -> None: assert mock_run.call_args[1]["input"] == "Do the thing." +def _mock_async_proc( + stdout: str, returncode: int = 0, stderr: str = "", +) -> MagicMock: + proc = MagicMock() + proc.returncode = returncode + proc.communicate = AsyncMock( + return_value=(stdout.encode("utf-8"), stderr.encode("utf-8")), + ) + proc.kill = MagicMock() + proc.wait = AsyncMock() + return proc + + +class TestAcomplete: + async def test_parses_result_and_usage(self) -> None: + provider = ClaudeCodeProvider() + proc = _mock_async_proc(_cli_envelope(result="hello", input_tokens=50, output_tokens=7)) + with patch( + "rle.providers.claude_code.shutil.which", return_value="claude", + ), patch( + "rle.providers.claude_code.asyncio.create_subprocess_exec", + new=AsyncMock(return_value=proc), + ): + result = await provider.acomplete(MESSAGES) + assert result.content == "hello" + assert result.usage["total_tokens"] == 57 + + async def test_nonzero_exit_raises(self) -> None: + provider = ClaudeCodeProvider() + proc = _mock_async_proc("", returncode=1, stderr="boom") + with patch( + "rle.providers.claude_code.shutil.which", return_value="claude", + ), patch( + "rle.providers.claude_code.asyncio.create_subprocess_exec", + new=AsyncMock(return_value=proc), + ): + with pytest.raises(ProviderError, match="exited with code 1"): + await provider.acomplete(MESSAGES) + + async def test_timeout_kills_subprocess(self) -> None: + provider = ClaudeCodeProvider() + proc = _mock_async_proc("") + proc.communicate = AsyncMock(side_effect=asyncio.TimeoutError) + with patch( + "rle.providers.claude_code.shutil.which", return_value="claude", + ), patch( + "rle.providers.claude_code.asyncio.create_subprocess_exec", + new=AsyncMock(return_value=proc), + ): + with pytest.raises(ProviderError, match="timed out"): + await provider.acomplete(MESSAGES) + proc.kill.assert_called_once() + + async def test_cancellation_kills_subprocess(self) -> None: + provider = ClaudeCodeProvider() + proc = _mock_async_proc("") + + async def _hang(*args: Any, **kwargs: Any) -> tuple[bytes, bytes]: + await asyncio.sleep(30) + return b"", b"" + + proc.communicate = AsyncMock(side_effect=_hang) + with patch( + "rle.providers.claude_code.shutil.which", return_value="claude", + ), patch( + "rle.providers.claude_code.asyncio.create_subprocess_exec", + new=AsyncMock(return_value=proc), + ): + task = asyncio.ensure_future(provider.acomplete(MESSAGES)) + await asyncio.sleep(0) # let the task start + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + proc.kill.assert_called_once() + + class TestStreamAndTokens: def test_stream_yields_content_then_final(self) -> None: provider = ClaudeCodeProvider()