diff --git a/tests/test_services.py b/tests/test_services.py new file mode 100644 index 00000000..75501fd7 --- /dev/null +++ b/tests/test_services.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +import json +import socket +import threading +import time +from pathlib import Path + +from villani_code.services import ServiceManager, is_likely_service_command, probe_service_readiness +from villani_code.tools import execute_tool + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def test_service_classification_heuristics() -> None: + assert is_likely_service_command("python app.py") + assert is_likely_service_command("python -m flask run --port 5000") + assert is_likely_service_command("npm run dev") + assert is_likely_service_command("docker compose up") + assert not is_likely_service_command("python -m pytest -q") + assert not is_likely_service_command("python -c 'print(1)'") + assert not is_likely_service_command("echo ok") + + +def test_readiness_probe_tcp_success() -> None: + port = _free_port() + ready = threading.Event() + + def _listener() -> None: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.bind(("127.0.0.1", port)) + sock.listen(1) + ready.set() + conn, _addr = sock.accept() + conn.close() + + thread = threading.Thread(target=_listener, daemon=True) + thread.start() + assert ready.wait(timeout=2) + + result = probe_service_readiness(port=port, base_url=None, timeout_sec=2.0, interval_sec=0.05) + assert result["ready"] is True + assert result["method"] == "tcp" + + +def test_service_integration_start_hit_stop(tmp_path: Path) -> None: + manager = ServiceManager(tmp_path) + port = _free_port() + result = manager.start_service(f"python -m http.server {port}", cwd=tmp_path, maybe_port_hint=port, readiness_timeout_sec=6.0) + assert result["service_id"] + assert result["pid"] + assert result["readiness"]["ready"] is True + + import httpx + + response = httpx.get(f"http://127.0.0.1:{port}", timeout=2.0) + assert response.status_code == 200 + + stopped = manager.stop_service(result["service_id"]) + assert stopped["stopped"] is True + + +def test_persistent_service_command_does_not_block_shell_path(tmp_path: Path) -> None: + manager = ServiceManager(tmp_path) + port = _free_port() + started = time.monotonic() + result = execute_tool( + "Bash", + {"command": f"python -m http.server {port}", "cwd": ".", "timeout_sec": 20}, + tmp_path, + service_manager=manager, + ) + elapsed = time.monotonic() - started + assert elapsed < 4.0 + payload = json.loads(str(result["content"])) + assert payload["mode"] == "service" + assert payload["service"]["readiness"]["ready"] is True + + service_id = payload["service"]["service_id"] + manager.stop_service(service_id) + + +def test_short_lived_command_uses_normal_execution_path(tmp_path: Path) -> None: + manager = ServiceManager(tmp_path) + result = execute_tool( + "Bash", + {"command": "python -c \"print('ok')\"", "cwd": ".", "timeout_sec": 10}, + tmp_path, + service_manager=manager, + ) + payload = json.loads(str(result["content"])) + assert payload["mode"] == "command" + assert payload["exit_code"] == 0 + assert "ok" in payload["stdout"] diff --git a/villani_code/services.py b/villani_code/services.py new file mode 100644 index 00000000..a4a6bc3a --- /dev/null +++ b/villani_code/services.py @@ -0,0 +1,300 @@ +from __future__ import annotations + +import os +import re +import signal +import socket +import subprocess +import threading +import time +import uuid +from contextlib import closing +from dataclasses import dataclass, field +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +import httpx + + +_SERVICE_HINTS = ( + "flask run", + "uvicorn", + "gunicorn", + "streamlit run", + "python -m http.server", + "docker compose up", + "docker-compose up", +) + +_SERVICE_EXACT_PATTERNS = ( + re.compile(r"^python\s+[^\s]+\.py$", re.IGNORECASE), + re.compile(r"^python\s+-m\s+flask\s+run\b", re.IGNORECASE), + re.compile(r"^flask\s+run\b", re.IGNORECASE), + re.compile(r"^(npm|pnpm|yarn)\s+(start|run\s+dev|dev)\b", re.IGNORECASE), +) + +_PORT_PATTERN = re.compile(r"(?:--port|-p)\s+([0-9]{2,5})") +_BIND_PATTERN = re.compile(r"(?:--host|--bind)\s+([0-9a-zA-Z\.:_-]+)") + + +@dataclass +class ServiceRecord: + service_id: str + command: str + cwd: str + started_at: str + stdout_log_path: str + stderr_log_path: str + pid: int | None + process: subprocess.Popen[str] | None = field(repr=False, default=None) + base_url: str | None = None + port: int | None = None + status: str = "starting" + readiness: dict[str, Any] = field(default_factory=dict) + + +def is_likely_service_command(command: str) -> bool: + compact = " ".join(str(command).strip().split()) + if not compact: + return False + lowered = compact.lower() + if any(token in lowered for token in ("&&", ";", "|", "pytest", "mypy", "ruff", "python -c", "echo ")): + return False + if any(hint in lowered for hint in _SERVICE_HINTS): + return True + return any(pattern.search(compact) for pattern in _SERVICE_EXACT_PATTERNS) + + +def infer_service_endpoint(command: str, env: dict[str, str] | None = None, maybe_port_hint: int | None = None) -> tuple[int | None, str | None]: + command = str(command) + env = env or {} + port = maybe_port_hint + port_match = _PORT_PATTERN.search(command) + if port_match: + port = int(port_match.group(1)) + if port is None and env.get("PORT", "").isdigit(): + port = int(env["PORT"]) + if port is None and "http.server" in command: + tail = command.split("http.server", 1)[1].strip().split() + if tail and tail[0].isdigit(): + port = int(tail[0]) + host = "127.0.0.1" + bind_match = _BIND_PATTERN.search(command) + if bind_match: + parsed_host = bind_match.group(1) + if parsed_host and parsed_host not in {"0.0.0.0", "::"}: + host = parsed_host + if port is None: + return None, None + return port, f"http://{host}:{port}" + + +def probe_service_readiness( + *, + port: int | None, + base_url: str | None, + timeout_sec: float = 8.0, + interval_sec: float = 0.2, +) -> dict[str, Any]: + started = time.monotonic() + attempts: list[dict[str, Any]] = [] + while time.monotonic() - started <= timeout_sec: + attempt: dict[str, Any] = {"at": time.monotonic(), "tcp_ok": None, "http_ok": None} + if port is not None: + try: + with closing(socket.create_connection(("127.0.0.1", port), timeout=0.5)): + attempt["tcp_ok"] = True + except OSError as exc: + attempt["tcp_ok"] = False + attempt["tcp_error"] = str(exc) + if base_url: + try: + response = httpx.get(base_url, timeout=0.75) + attempt["http_ok"] = 200 <= response.status_code < 500 + attempt["http_status"] = int(response.status_code) + except Exception as exc: # noqa: BLE001 + attempt["http_ok"] = False + attempt["http_error"] = str(exc) + + attempts.append(attempt) + tcp_ready = attempt.get("tcp_ok") in {True, None} + http_ready = attempt.get("http_ok") in {True, None} + if tcp_ready and http_ready: + return { + "ready": True, + "method": "tcp+http" if port and base_url else ("tcp" if port else "http"), + "attempts": attempts, + "elapsed_sec": round(time.monotonic() - started, 3), + } + time.sleep(interval_sec) + + return { + "ready": False, + "method": "tcp+http" if port and base_url else ("tcp" if port else "http"), + "attempts": attempts, + "elapsed_sec": round(time.monotonic() - started, 3), + "failure": "service not ready before timeout", + } + + +class ServiceManager: + def __init__(self, repo: Path): + self.repo = repo + self._services: dict[str, ServiceRecord] = {} + self._lock = threading.Lock() + + def list_service_ids(self) -> list[str]: + with self._lock: + return list(self._services.keys()) + + def list_services(self) -> list[dict[str, Any]]: + with self._lock: + return [self._to_dict(service) for service in self._services.values()] + + def get_service_status(self, service_id: str) -> dict[str, Any]: + with self._lock: + service = self._services.get(service_id) + if service is None: + return {"service_id": service_id, "status": "missing", "is_error": True} + self._refresh_status(service) + return self._to_dict(service) + + def start_service( + self, + command: str, + cwd: Path, + env: dict[str, str] | None = None, + maybe_port_hint: int | None = None, + readiness_timeout_sec: float = 8.0, + event_callback: Any | None = None, + ) -> dict[str, Any]: + service_id = f"svc-{uuid.uuid4().hex[:10]}" + logs_dir = self.repo / ".villani_code" / "services" + logs_dir.mkdir(parents=True, exist_ok=True) + stdout_path = logs_dir / f"{service_id}.out.log" + stderr_path = logs_dir / f"{service_id}.err.log" + started_at = datetime.now(timezone.utc).isoformat() + run_env = os.environ.copy() + if env: + run_env.update(env) + + port, base_url = infer_service_endpoint(command, run_env, maybe_port_hint) + if callable(event_callback): + event_callback({"type": "service_launch_requested", "command": command, "cwd": str(cwd), "port": port, "base_url": base_url}) + + with stdout_path.open("w", encoding="utf-8") as stdout_handle, stderr_path.open("w", encoding="utf-8") as stderr_handle: + kwargs: dict[str, Any] = { + "shell": True, + "cwd": str(cwd), + "env": run_env, + "stdout": stdout_handle, + "stderr": stderr_handle, + "text": True, + } + if os.name == "nt": + kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP # type: ignore[attr-defined] + else: + kwargs["start_new_session"] = True + process = subprocess.Popen(command, **kwargs) + + record = ServiceRecord( + service_id=service_id, + command=command, + cwd=str(cwd), + started_at=started_at, + stdout_log_path=str(stdout_path), + stderr_log_path=str(stderr_path), + pid=process.pid, + process=process, + base_url=base_url, + port=port, + ) + with self._lock: + self._services[service_id] = record + + if callable(event_callback): + event_callback({"type": "service_process_spawned", "service_id": service_id, "pid": process.pid}) + + readiness = probe_service_readiness(port=port, base_url=base_url, timeout_sec=readiness_timeout_sec) + record.readiness = readiness + if process.poll() is not None: + record.status = "exited" + else: + record.status = "running" if readiness.get("ready") else "running_not_ready" + + if callable(event_callback): + for attempt in readiness.get("attempts", []): + event_callback({"type": "service_readiness_probe_attempt", "service_id": service_id, **attempt}) + event_callback( + { + "type": "service_readiness_succeeded" if readiness.get("ready") else "service_readiness_failed", + "service_id": service_id, + "readiness": readiness, + } + ) + + payload = self._to_dict(record) + payload["is_error"] = not bool(readiness.get("ready")) + if payload["is_error"]: + payload["failure"] = readiness.get("failure", "service readiness failed") + return payload + + def stop_service(self, service_id: str, event_callback: Any | None = None) -> dict[str, Any]: + with self._lock: + service = self._services.get(service_id) + if service is None: + return {"service_id": service_id, "status": "missing", "stopped": False, "is_error": True} + + process = service.process + if process is None or process.poll() is not None: + service.status = "exited" + return {**self._to_dict(service), "stopped": False, "already_exited": True} + + try: + if os.name == "nt": + process.terminate() + else: + os.killpg(os.getpgid(process.pid), signal.SIGTERM) + process.wait(timeout=3) + stopped = True + except Exception: # noqa: BLE001 + stopped = False + try: + process.kill() + process.wait(timeout=2) + stopped = True + except Exception: # noqa: BLE001 + stopped = False + service.status = "stopped" if stopped else "running" + if callable(event_callback): + event_callback({"type": "service_stopped", "service_id": service_id, "stopped": stopped}) + return {**self._to_dict(service), "stopped": stopped, "is_error": not stopped} + + def cleanup_services_started_after(self, baseline: set[str], event_callback: Any | None = None) -> None: + for service_id in self.list_service_ids(): + if service_id in baseline: + continue + self.stop_service(service_id, event_callback=event_callback) + + def _refresh_status(self, service: ServiceRecord) -> None: + if service.process is None: + return + if service.process.poll() is not None and service.status not in {"stopped", "exited"}: + service.status = "exited" + + def _to_dict(self, service: ServiceRecord) -> dict[str, Any]: + self._refresh_status(service) + return { + "service_id": service.service_id, + "status": service.status, + "command": service.command, + "cwd": service.cwd, + "started_at": service.started_at, + "pid": service.pid, + "base_url": service.base_url, + "port": service.port, + "stdout_log_path": service.stdout_log_path, + "stderr_log_path": service.stderr_log_path, + "readiness": service.readiness, + } diff --git a/villani_code/state.py b/villani_code/state.py index 1a2c80d4..5e35da73 100644 --- a/villani_code/state.py +++ b/villani_code/state.py @@ -35,6 +35,7 @@ from villani_code.runtime_safety import ensure_runtime_dependencies_not_shadowed from villani_code.retrieval import Retriever from villani_code.skills import discover_skills +from villani_code.services import ServiceManager from villani_code.streaming import StreamCoalescer, assemble_anthropic_stream from villani_code.tools import tool_specs from villani_code.transcripts import save_transcript @@ -538,6 +539,7 @@ def __init__( self._mission_state: MissionState | None = None self._event_recorder: RuntimeEventRecorder | None = None self._current_turn_index: int | None = None + self._service_manager = ServiceManager(self.repo) if self.small_model: self._init_small_model_support() @@ -964,6 +966,7 @@ def run( ], } ) + mission_service_baseline = set(self._service_manager.list_service_ids()) previous_attributed = set() def _attributed_changed_files() -> list[str]: @@ -1026,6 +1029,7 @@ def _finish_bounded( self._save_session_snapshot(messages) mission_status = "completed" if completed else ("interrupted" if reason in {"max_seconds", "max_turns", "max_tool_calls"} else "failed") self._update_mission_state(status=mission_status, changed_files=all_changes, compact_summary=summarize_mission_state(self._mission_state) if self._mission_state else "") + self._service_manager.cleanup_services_started_after(mission_service_baseline, event_callback=self.event_callback) if self._event_recorder is not None: self._event_recorder.write_digest() if self._debug_recorder is not None: @@ -1210,6 +1214,7 @@ def _budget_reason( if post: response.setdefault("content", []).append({"type": "text", "text": post}) self._save_session_snapshot(messages) + self._service_manager.cleanup_services_started_after(mission_service_baseline, event_callback=self.event_callback) if self._event_recorder is not None: self._event_recorder.write_digest() if self._debug_recorder is not None: @@ -1338,6 +1343,7 @@ def _budget_reason( transcript_path = self._save_transcript_and_link(transcript) self._save_session_snapshot(messages) self._update_mission_state(status="completed", compact_summary=summarize_mission_state(self._mission_state) if self._mission_state else "") + self._service_manager.cleanup_services_started_after(mission_service_baseline, event_callback=self.event_callback) if self._event_recorder is not None: self._event_recorder.write_digest() if self._debug_recorder is not None: @@ -1666,14 +1672,18 @@ def _build_tool_result_event_payload( if isinstance(decoded, dict): payload["command"] = decoded.get("command") payload["exit_code"] = decoded.get("exit_code") + payload["mode"] = decoded.get("mode") payload["stdout"] = decoded.get("stdout") payload["stderr"] = decoded.get("stderr") + payload["service"] = decoded.get("service") payload["result_payload"] = { **base_result_payload, "command": decoded.get("command"), + "mode": decoded.get("mode"), "exit_code": decoded.get("exit_code"), "stdout": decoded.get("stdout"), "stderr": decoded.get("stderr"), + "service": decoded.get("service"), } elif tool_name == "Read": text_content = str(content) diff --git a/villani_code/state_tooling.py b/villani_code/state_tooling.py index c84c64ca..8144efa8 100644 --- a/villani_code/state_tooling.py +++ b/villani_code/state_tooling.py @@ -665,6 +665,8 @@ def _debug_callback_with_turn(event_type: str, payload: dict[str, Any]) -> None: unsafe=runner.unsafe, debug_callback=_debug_callback_with_turn, tool_call_id=stable_tool_use_id, + service_manager=getattr(runner, "_service_manager", None), + event_callback=runner.event_callback, ) runner.event_callback( { diff --git a/villani_code/tools.py b/villani_code/tools.py index 423807ed..c2b6b210 100644 --- a/villani_code/tools.py +++ b/villani_code/tools.py @@ -6,6 +6,8 @@ import subprocess from pathlib import Path from typing import Any + +from villani_code.services import ServiceManager, is_likely_service_command from urllib.parse import urlparse import httpx @@ -143,6 +145,8 @@ def execute_tool( unsafe: bool = False, debug_callback: Any | None = None, tool_call_id: str = "", + service_manager: ServiceManager | None = None, + event_callback: Any | None = None, ) -> dict[str, Any]: model = TOOL_MODELS.get(name) if not model: @@ -164,7 +168,7 @@ def execute_tool( if name == "Search": return _ok(_run_search(parsed, repo)) if name == "Bash": - return _ok(_run_bash(parsed, repo, unsafe=unsafe, debug_callback=debug_callback, tool_call_id=tool_call_id)) + return _ok(_run_bash(parsed, repo, unsafe=unsafe, debug_callback=debug_callback, tool_call_id=tool_call_id, service_manager=service_manager, event_callback=event_callback)) if name == "Write": return _ok(_run_write(parsed, repo, debug_callback=debug_callback, tool_call_id=tool_call_id)) if name == "Patch": @@ -235,13 +239,45 @@ def _run_search(data: SearchInput, repo: Path) -> str: return proc.stdout -def _run_bash(data: BashInput, repo: Path, unsafe: bool, debug_callback: Any | None = None, tool_call_id: str = "") -> str: +def _run_bash( + data: BashInput, + repo: Path, + unsafe: bool, + debug_callback: Any | None = None, + tool_call_id: str = "", + service_manager: ServiceManager | None = None, + event_callback: Any | None = None, +) -> str: lowered = data.command.lower() if not unsafe: for bad in DENYLIST: if bad in lowered: raise ValueError(f"Refusing command: {bad.strip()}") cwd = _safe_path(repo, data.cwd) + if service_manager is not None and is_likely_service_command(data.command): + if callable(debug_callback): + debug_callback("service_launch_requested", {"command": data.command, "cwd": data.cwd, "tool_call_id": tool_call_id}) + service_result = service_manager.start_service( + data.command, + cwd=cwd, + env=None, + readiness_timeout_sec=max(1.0, min(float(data.timeout_sec), 20.0)), + event_callback=event_callback, + ) + if callable(debug_callback): + debug_callback("service_process_spawned", {"tool_call_id": tool_call_id, **service_result}) + return json.dumps( + { + "command": data.command, + "mode": "service", + "exit_code": None, + "stdout": "", + "stderr": "", + "service": service_result, + }, + indent=2, + ) + if callable(debug_callback): debug_callback("command_started", {"command": data.command, "cwd": data.cwd, "tool_call_id": tool_call_id}) proc = subprocess.run(data.command, shell=True, cwd=str(cwd), capture_output=True, text=True, timeout=data.timeout_sec) @@ -258,7 +294,7 @@ def _run_bash(data: BashInput, repo: Path, unsafe: bool, debug_callback: Any | N "tool_call_id": tool_call_id, }, ) - return json.dumps({"command": data.command, "exit_code": proc.returncode, "stdout": proc.stdout, "stderr": proc.stderr}, indent=2) + return json.dumps({"command": data.command, "mode": "command", "exit_code": proc.returncode, "stdout": proc.stdout, "stderr": proc.stderr}, indent=2) def _run_write(data: WriteInput, repo: Path, debug_callback: Any | None = None, tool_call_id: str = "") -> str: