"""zcbot `/v1` 黑盒客户端。""" from __future__ import annotations import ipaddress import os import time from typing import Any from urllib.parse import urljoin, urlparse import httpx from .models import RunObservation class EvalClientError(RuntimeError): """zcbot API 调用失败。""" def require_safe_base_url(base_url: str, *, allow_remote: bool) -> str: normalized = base_url.rstrip("/") + "/" parsed = urlparse(normalized) if parsed.scheme not in {"http", "https"} or not parsed.hostname: raise EvalClientError(f"无效 base_url: {base_url!r}") if allow_remote: return normalized host = parsed.hostname.lower() safe = host in {"localhost", "127.0.0.1", "::1"} if not safe: try: safe = ipaddress.ip_address(host).is_loopback except ValueError: safe = False if not safe: raise EvalClientError( f"默认只允许本机 zcbot，当前是 {host!r}；确认测试环境后传 --allow-remote" ) return normalized def _response_text(payload: Any) -> str: if isinstance(payload, str): return payload if isinstance(payload, list): parts: list[str] = [] for item in payload: if isinstance(item, dict) and isinstance(item.get("text"), str): parts.append(item["text"]) elif isinstance(item, str): parts.append(item) return "\n".join(parts) return "" if payload is None else str(payload) class ZcbotClient: def __init__( self, *, base_url: str, token: str, allow_remote: bool = False, request_timeout_s: float = 30.0, ) -> None: self.base_url = require_safe_base_url(base_url, allow_remote=allow_remote) if not token.strip(): raise EvalClientError("缺少 token；请设置配置中 token_env 指向的环境变量") self._client = httpx.Client( base_url=self.base_url, headers={"Authorization": f"Bearer {token.strip()}"}, timeout=request_timeout_s, follow_redirects=False, # 评测 JWT 不得被 HTTP(S)_PROXY 透明转发；部署内网应直连测试实例。 trust_env=False, ) @classmethod def from_config( cls, raw: dict[str, Any], *, allow_remote: bool = False ) -> "ZcbotClient": token_env = str(raw.get("token_env", "ZCBOT_EVAL_TOKEN")) return cls( base_url=str(raw.get("base_url", "http://127.0.0.1:8765")), token=os.getenv(token_env, ""), allow_remote=allow_remote, request_timeout_s=float(raw.get("request_timeout_s", 30)), ) def close(self) -> None: self._client.close() def __enter__(self) -> "ZcbotClient": return self def __exit__(self, *_args: object) -> None: self.close() def _json(self, method: str, path: str, **kwargs: Any) -> dict[str, Any]: response = self._client.request(method, path, **kwargs) if response.status_code >= 400: raise EvalClientError( f"{method} {path} -> {response.status_code}: {response.text[:500]}" ) try: body = response.json() except ValueError as exc: raise EvalClientError(f"{method} {path} 返回了非 JSON") from exc if not isinstance(body, dict): raise EvalClientError(f"{method} {path} 返回 JSON 顶层不是 object") return body def create_task( self, *, name: str, working_dir: str, skill: str, model_profile: str, ) -> dict[str, Any]: return self._json( "POST", "/v1/tasks", json={ "name": name, "working_dir": working_dir, "description": "zcbot evaluation task", "skill": skill, "model_profile": model_profile, }, ) def run_prompt( self, *, task_id: str, prompt: str, working_dir: str, timeout_s: float, ) -> RunObservation: started = time.monotonic() body = self._json( "POST", f"/v1/tasks/{task_id}/messages", json={"content": prompt} ) events_url = str(body.get("events_url", "")) if not events_url: raise EvalClientError("POST messages 未返回 events_url") stream_error = self._consume_events(events_url, timeout_s=timeout_s) # AgentLoop 的 done 事件可能早于 BG worker 最终把 task.run_status 写回 idle； # 短暂轮询终态，避免把已成功任务误记为 running。 terminal_deadline = time.monotonic() + min(15.0, timeout_s) while True: task = self._json("GET", f"/v1/tasks/{task_id}") if task.get("run_status") not in {"running", "cancelling"}: break if time.monotonic() >= terminal_deadline: break time.sleep(0.1) messages = self._json("GET", f"/v1/tasks/{task_id}/messages") response = "" for item in reversed(messages.get("messages", [])): payload = item.get("payload", {}) if isinstance(item, dict) else {} if isinstance(payload, dict) and payload.get("role") == "assistant": response = _response_text(payload.get("content")) break run_error = str(task.get("run_error") or stream_error or "") return RunObservation( response=response, duration_s=time.monotonic() - started, cost_cny=float(task.get("cost_cny") or 0), run_status=str(task.get("run_status") or ""), run_error=run_error, task_id=task_id, working_dir=working_dir, model_profile=str(task.get("model_profile") or ""), artifact_loader=lambda relative: self.download_artifact( working_dir, relative ), ) def _consume_events(self, path: str, *, timeout_s: float) -> str: deadline = time.monotonic() + timeout_s error = "" timeout = httpx.Timeout( connect=min(30.0, timeout_s), read=max(30.0, timeout_s), write=30.0, pool=30.0, ) with self._client.stream("GET", path, timeout=timeout) as response: if response.status_code >= 400: raise EvalClientError( f"GET {path} -> {response.status_code}: {response.read()[:500]!r}" ) event = "" data = "" for line in response.iter_lines(): if time.monotonic() > deadline: raise EvalClientError(f"任务运行超过 {timeout_s:.0f}s") if line.startswith("event:"): event = line[6:].strip() elif line.startswith("data:"): data = line[5:].strip() elif not line: if event == "error": error = data[:1000] if event in {"done", "error"}: return error event, data = "", "" return error def download_artifact(self, working_dir: str, relative_path: str) -> bytes: relative = relative_path.replace("\\", "/").strip("/") if not relative or ".." in relative.split("/"): raise FileNotFoundError(relative_path) path = f"{working_dir.strip('/')}/{relative}" response = self._client.get("/v1/files/download", params={"path": path}) if response.status_code == 404: raise FileNotFoundError(relative_path) if response.status_code >= 400: raise EvalClientError( f"GET /v1/files/download {path!r} -> {response.status_code}: " f"{response.text[:500]}" ) return response.content