"""Full-server transport driver.""" from __future__ import annotations import asyncio import contextlib import json import threading import time from collections.abc import Callable from typing import Any import httpx from tests.e2e._harness_probes import cli_unavailable_reason from tests.harness_bench.driver import TurnResult, fill_snapshot_cost from tests.harness_bench.full_server import ( _DENY_REASON, _POLL_INTERVAL_S, _TOOL_NAME, SharedFullServer, ) from tests.harness_bench.profile import BenchProfile from tests.harness_bench.runtime_env import bench_creds_skip_reason, resolve_bench_env from tests.harness_bench.session_items import ( assistant_text, contains_user_text, function_calls, tool_output_states, ) _TOOL_PROMPT = f"List the files using the {_TOOL_NAME} tool, then tell me how many there are." # Interrupted turns persist this marker in session history. _CANCELLATION_MARKER = "interrupted" _LONG_PROMPT = ( "Write a very detailed 600-word essay about the history of computing, in full paragraphs." ) _STREAM_PROMPT = ( "Count from 1 to 30 in words, one number per line, and add a short note after each." ) _TERMINAL_EVENTS = frozenset({"response.completed", "response.failed", "response.cancelled"}) class FullServerDriver: """Drive turns through a live Omnigent server and runner.""" transport = "full-server" def __init__( self, profile: BenchProfile, *, databricks_profile: str | None, shared: SharedFullServer | None = None, ) -> None: self._profile = profile self._db_profile = databricks_profile self._shared = shared self._owns_shared = shared is None # Fixed-action policies must be baked into separate agent specs. self._policy_session_ids: dict[str, str] = {} self._session_id: str | None = None @property def _client(self) -> httpx.Client | None: return self._shared.client if self._shared is not None else None @staticmethod def unavailable(profile: BenchProfile, *, databricks_profile: str | None) -> str | None: """Return why this driver cannot run the profile, if applicable.""" # Native harnesses require host-daemon provisioning. if profile.transport == "native-tui": return ( f"{profile.harness!r} is a native-tui harness; the full-server transport " "registers via an agent bundle and cannot drive it (use --transport native-tui)" ) creds_skip = bench_creds_skip_reason(databricks_profile) if creds_skip is not None: return creds_skip if profile.cli_binary is not None: return cli_unavailable_reason(profile.cli_binary) return None def __enter__(self) -> FullServerDriver: if self._shared is None: self._shared = SharedFullServer(resolve_bench_env(self._db_profile)) self._shared.__enter__() agent_name = self._shared.register_agent(self._profile, policy_action=None) self._session_id = self._shared.create_session(agent_name) return self def __exit__(self, *exc: object) -> None: if self._owns_shared and self._shared is not None: self._shared.__exit__(*exc) # Blocking server operations run off the event loop. async def __aenter__(self) -> FullServerDriver: return await asyncio.to_thread(self.__enter__) async def __aexit__(self, *exc: object) -> None: await asyncio.to_thread(self.__exit__, *exc) async def run_basic_turn(self, marker: str) -> TurnResult: prompt = f"Reply with exactly the literal string {marker} and nothing else." return await asyncio.to_thread(self.run_turn, prompt) async def run_streaming_turn(self) -> TurnResult: return await asyncio.to_thread(self.streaming_probe_turn) async def run_tool_turn(self, *, deny: bool) -> TurnResult: return await asyncio.to_thread(lambda: self.tool_probe_turn(deny=deny)) async def run_policy_turn(self, *, action: str) -> TurnResult: return await asyncio.to_thread(lambda: self.policy_probe_turn(action=action)) async def run_interrupt_turn(self) -> TurnResult: return await asyncio.to_thread(self.interrupt_probe_turn) def _ensure_policy_session(self, action: str) -> str: """Return a session whose agent bakes in the requested policy action.""" assert self._shared is not None if action not in self._policy_session_ids: name = self._shared.register_agent(self._profile, policy_action=action) self._policy_session_ids[action] = self._shared.create_session(name) return self._policy_session_ids[action] def _poll_session( self, sid: str, result: TurnResult, *, timeout: float, scan_tools: bool = False, fill_cost: bool = False, stop_when: Callable[[], bool] | None = None, ) -> TurnResult: """Poll a session until it settles, fails, times out, or is stopped.""" assert self._client is not None deadline = time.monotonic() + timeout seen_running = False while time.monotonic() < deadline: if stop_when is not None and stop_when(): return result response = self._client.get(f"/v1/sessions/{sid}") response.raise_for_status() snapshot = response.json() status = snapshot.get("status") items = snapshot.get("items", []) if scan_tools: _scan_tool_items(items, result) if status in ("running", "waiting"): seen_running = True elif status == "failed": result.failed = True result.error = snapshot.get("last_task_error") or snapshot.get("error") return result elif status == "idle" and seen_running: result.completed = True result.text = assistant_text(items) if fill_cost: fill_snapshot_cost(result, snapshot) return result time.sleep(_POLL_INTERVAL_S) result.timed_out = True return result def tool_probe_turn(self, *, deny: bool, timeout: float = 180.0) -> TurnResult: """Drive a turn that calls the builtin tool; return a :class:`TurnResult`. On the full-server transport a tool call is real and server- dispatched. With *deny* the turn runs against a session whose agent bakes a ``tool_call`` deny policy, so the server blocks the call and the tool output carries the deny reason. Fills :attr:`TurnResult.tool_calls` (the builtin call) and :attr:`TurnResult.tool_call_denied` (whether the server blocked it), plus ``completed`` / ``failed`` / ``text``. """ assert self._client is not None sid = self._ensure_policy_session("deny") if deny else self._session_id assert sid is not None result = TurnResult() body = { "type": "message", "data": {"role": "user", "content": [{"type": "input_text", "text": _TOOL_PROMPT}]}, } self._client.post(f"/v1/sessions/{sid}/events", json=body).raise_for_status() return self._poll_session(sid, result, timeout=timeout, scan_tools=True) def policy_probe_turn(self, *, action: str, timeout: float = 90.0) -> TurnResult: """Drive a tool turn under a fixed tool_call policy *action*. ``"allow"``: the call proceeds (``tool_call_allowed`` from the non-blocked output). ``"ask"``: it parks on an elicitation; a background reader sets ``elicitation_requested`` off ``response.elicitation_request``. The ASK verdict is decided the moment that fires, so we resolve the elicitation (approval accept, to leave no dangling park) and return immediately rather than polling the turn to a terminal state. The timeout bounds the *worst* case (the model never calls the tool, so no elicitation fires): a bounded SKIP, not a 3-minute stall. """ assert self._client is not None sid = self._ensure_policy_session(action) result = TurnResult() elicitation_id: dict[str, str] = {} stop = threading.Event() def _watch() -> None: try: with self._client.stream( # type: ignore[union-attr] "GET", f"/v1/sessions/{sid}/stream", timeout=timeout ) as resp: for raw in resp.iter_lines(): if stop.is_set(): return line = raw.strip() if not line.startswith("data:"): continue try: frame = json.loads(line[len("data:") :].strip()) except (ValueError, TypeError): continue if frame.get("type") == "response.elicitation_request": result.elicitation_requested = True eid = frame.get("elicitation_id") if isinstance(eid, str): elicitation_id["id"] = eid else: # No parseable id: verdict is recorded, but we # can't resolve, so the turn parks to the deadline. result.error = "elicitation_request frame had no parseable id" return except httpx.HTTPError: # Best-effort watcher; an SSE read error must not fail the turn. pass watcher = None if action == "ask": watcher = threading.Thread(target=_watch, daemon=True) watcher.start() time.sleep(1.0) # register the subscription before the turn starts body = { "type": "message", "data": {"role": "user", "content": [{"type": "input_text", "text": _TOOL_PROMPT}]}, } self._client.post(f"/v1/sessions/{sid}/events", json=body).raise_for_status() def _elicitation_observed() -> bool: if action == "ask" and result.elicitation_requested: if "id" in elicitation_id: self._resolve_elicitation(sid, elicitation_id.pop("id")) return True return False self._poll_session( sid, result, timeout=timeout, scan_tools=True, stop_when=_elicitation_observed, ) stop.set() if watcher is not None: watcher.join(timeout=5.0) return result def _resolve_elicitation(self, sid: str, elicitation_id: str) -> None: """Accept an outstanding elicitation via an ``approval`` event so an ASK turn settles (best-effort; a raced resolve is harmless).""" assert self._client is not None # The server reads the id from inside `data` (SessionEventInput has no # top-level elicitation_id field), so it must be nested there or the # resolve is a silent no-op and the park dangles. with contextlib.suppress(httpx.HTTPError): self._client.post( f"/v1/sessions/{sid}/events", json={ "type": "approval", "data": {"elicitation_id": elicitation_id, "action": "accept"}, }, ) def streaming_probe_turn(self, *, timeout: float = 120.0) -> TurnResult: """Measure token-level streaming via the session SSE subscribe stream. The full-server stream (``GET /v1/sessions/{id}/stream``) is separate from the message POST, so a background thread subscribes and counts ``response.output_text.delta`` events while the main thread posts the turn. More than one delta means the harness streams incrementally. """ assert self._client is not None and self._session_id is not None sid = self._session_id result = TurnResult() done = threading.Event() def _read_stream() -> None: try: with self._client.stream( # type: ignore[union-attr] "GET", f"/v1/sessions/{sid}/stream", timeout=timeout ) as resp: for line in resp.iter_lines(): if not line.startswith("event:"): continue etype = line[len("event:") :].strip() if etype == "response.output_text.delta": result.text_delta_count += 1 elif etype in _TERMINAL_EVENTS: result.completed = etype == "response.completed" result.cancelled = etype == "response.cancelled" result.failed = etype == "response.failed" return except httpx.HTTPError as exc: result.error = repr(exc) finally: done.set() reader = threading.Thread(target=_read_stream, daemon=True) reader.start() time.sleep(1.0) # let the subscription register before the turn starts self._client.post( f"/v1/sessions/{sid}/events", json={ "type": "message", "data": { "role": "user", "content": [{"type": "input_text", "text": _STREAM_PROMPT}], }, }, ).raise_for_status() if not done.wait(timeout): result.timed_out = True return result def interrupt_probe_turn(self, *, timeout: float = 120.0) -> TurnResult: """Start a long turn, interrupt it mid-flight, and report the outcome. Posts an ``interrupt`` event once the turn is running (after a short hold so some text streams first), then waits for the server's cancellation marker. Sets :attr:`TurnResult.cancelled` when the marker appears — the honored-interrupt signal. """ assert self._client is not None and self._session_id is not None sid = self._session_id result = TurnResult() body = { "type": "message", "data": {"role": "user", "content": [{"type": "input_text", "text": _LONG_PROMPT}]}, } self._client.post(f"/v1/sessions/{sid}/events", json=body).raise_for_status() deadline = time.monotonic() + timeout interrupted = False while time.monotonic() < deadline: snap = self._client.get(f"/v1/sessions/{sid}").json() status = snap.get("status") items = snap.get("items", []) if status in ("running", "waiting") and not interrupted: # Let a little text stream so the interrupt lands mid-turn. time.sleep(1.5) self._client.post(f"/v1/sessions/{sid}/events", json={"type": "interrupt"}) interrupted = True if _has_cancellation_marker(items): result.cancelled = True result.text = assistant_text(items) break if status == "idle" and interrupted: # Settled after the interrupt; the marker lands just after. result.cancelled = _has_cancellation_marker(items) result.text = assistant_text(items) break time.sleep(_POLL_INTERVAL_S) else: result.timed_out = True return result def run_turn(self, prompt: str, *, timeout: float = 180.0) -> TurnResult: """Drive one basic turn through the full server, return a :class:`TurnResult`. Posts the user message and polls the session snapshot to a terminal state, filling text, completion, failure, timeout, and usage fields. """ assert self._client is not None and self._session_id is not None result = TurnResult() body: dict[str, Any] = { "type": "message", "data": {"role": "user", "content": [{"type": "input_text", "text": prompt}]}, } posted = self._client.post(f"/v1/sessions/{self._session_id}/events", json=body) if posted.status_code == 202 and posted.json().get("denied"): result.failed = True result.error = {"denied": True, "reason": posted.json().get("reason")} return result posted.raise_for_status() return self._poll_session( self._session_id, result, timeout=timeout, fill_cost=True, ) def _scan_tool_items(items: list[dict[str, Any]], result: TurnResult) -> None: """Populate tool_calls and tool_call_denied from session items.""" result.tool_calls = function_calls(items) result.tool_call_allowed, result.tool_call_denied = tool_output_states( items, deny_marker=_DENY_REASON ) def _has_cancellation_marker(items: list[dict[str, Any]]) -> bool: """Whether items include the synthetic 'interrupted' user message.""" return contains_user_text(items, _CANCELLATION_MARKER)