import os import sys from datetime import datetime, timezone from typing import Any, cast import pytest PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) if PROJECT_ROOT not in sys.path: sys.path.insert(0, PROJECT_ROOT) WORKSPACE_ROOT = os.path.join(PROJECT_ROOT, "workspace") from SimpleLLMFunc.hooks.events import ( ReActEventType, ToolCallEndEvent, ToolCallErrorEvent, ToolCallStartEvent, ) from SimpleLLMFunc.hooks.stream import EventYield, ResponseYield import tools.code_tools as code_tools_module from tools.code_tools import cad_code_generator class _FakeEmitter: def __init__(self): self.events: list[tuple[str, dict]] = [] async def emit(self, event_name: str, data): self.events.append((event_name, data)) @pytest.mark.asyncio async def test_cad_code_generator_bridges_nested_specialist_events(monkeypatch): monkeypatch.chdir(WORKSPACE_ROOT) async def fake_specialist(**kwargs): assert "Current working directory:" in kwargs["message"] assert "Skill root: use the preferred skill root below." in kwargs["message"] assert "validation_command: uv run python part/model.py" in kwargs["message"] assert "references/docs/api/README.md" in kwargs["message"] yield EventYield( event=ToolCallStartEvent( event_type=ReActEventType.TOOL_CALL_START, timestamp=datetime.now(timezone.utc), trace_id="trace-1", func_name="cad_code_generator_specialist", iteration=0, tool_name="read_file", tool_call_id="nested-1", arguments={"file_path": "part/model.py"}, tool_call=cast( Any, { "id": "nested-1", "type": "function", "function": {"name": "read_file", "arguments": "{}"}, }, ), ) ) yield EventYield( event=ToolCallEndEvent( event_type=ReActEventType.TOOL_CALL_END, timestamp=datetime.now(timezone.utc), trace_id="trace-1", func_name="cad_code_generator_specialist", iteration=0, tool_name="read_file", tool_call_id="nested-1", arguments={"file_path": "part/model.py"}, result="print('old')", execution_time=0.05, success=True, ) ) yield EventYield( event=ToolCallErrorEvent( event_type=ReActEventType.TOOL_CALL_ERROR, timestamp=datetime.now(timezone.utc), trace_id="trace-1", func_name="cad_code_generator_specialist", iteration=0, tool_name="execute_command", tool_call_id="nested-2", arguments={"command": "python model.py"}, error=RuntimeError("boom"), error_message="boom", error_type="RuntimeError", execution_time=0.12, ) ) yield ResponseYield(response="STATUS: SUCCESS\nSUMMARY: ok", messages=[]) monkeypatch.setattr( code_tools_module, "cad_code_generator_specialist", fake_specialist, ) monkeypatch.setattr( code_tools_module, "_read_latest_code", lambda path: "print('ok')\n" ) emitter = _FakeEmitter() result = await cad_code_generator( task="Create a cube as a new file. This is a create-new-file task.", target_file_path="part/model.py", event_emitter=emitter, ) event_names = [event_name for event_name, _ in emitter.events] assert event_names == [ "subagent_status", "subagent_tool_start", "subagent_tool_end", "subagent_tool_error", "subagent_response", "subagent_status", ] first_status_payload = emitter.events[0][1] assert first_status_payload["subagent_label"] == "CAD Code Specialist" assert first_status_payload["validation_command"].startswith( "uv run python part/model.py" ) assert "ls part/*.stl" in first_status_payload["validation_command"] assert ( "(ls part/*.step || ls part/*.stp)" in first_status_payload["validation_command"] ) assert "STATUS: SUCCESS" in result assert "Latest code" in result response_payload = emitter.events[4][1] assert response_payload["delta_text"] == "STATUS: SUCCESS\nSUMMARY: ok" @pytest.mark.asyncio async def test_cad_code_generator_retries_when_first_attempt_does_not_produce_code( monkeypatch, ): monkeypatch.chdir(WORKSPACE_ROOT) calls = {"count": 0} async def fake_specialist(**kwargs): calls["count"] += 1 if calls["count"] == 1: assert kwargs["history"] == [] assert "Current working directory:" in kwargs["message"] assert ( "validation_command: uv run python part/model.py" in kwargs["message"] ) assert "references/docs/api/README.md" in kwargs["message"] yield ResponseYield( response=( "STATUS: SUCCESS\n" "SUMMARY: planned but no file write\n" "```python\n" "print('draft only')\n" "```" ), messages=[], ) return assert "must write the required Python script directly" in kwargs["message"] assert "part/model.py" in kwargs["message"] assert kwargs["history"][0]["role"] == "user" assert "target_file: part/model.py" in kwargs["history"][0]["content"] assert kwargs["history"][1]["role"] == "assistant" assert "planned but no file write" in kwargs["history"][1]["content"] yield ResponseYield( response="STATUS: SUCCESS\nCODE_WRITTEN: YES\nSUMMARY: wrote code", messages=[], ) monkeypatch.setattr( code_tools_module, "cad_code_generator_specialist", fake_specialist, ) code_reads = {"count": 0} def fake_read_latest_code(path): code_reads["count"] += 1 if code_reads["count"] == 1: return None return "print('ok')\n" monkeypatch.setattr(code_tools_module, "_read_latest_code", fake_read_latest_code) result = await cad_code_generator( task="Create a cube as a new file. This is a create-new-file task.", target_file_path="part/model.py", event_emitter=None, ) assert calls["count"] == 2 assert "CODE_WRITTEN: YES" in result assert "Latest code" in result