wehub-resource-sync
403 lines
15 KiB
Python
403 lines
15 KiB
Python
"""SDK in-process transport driver and shared turn results."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import shutil
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from omnigent.runtime.harnesses.process_manager import HarnessProcessManager
|
|
from tests.e2e._harness_probes import cli_unavailable_reason
|
|
from tests.harness_bench.profile import BenchProfile
|
|
from tests.harness_bench.runtime_env import bench_creds_skip_reason, resolve_bench_env
|
|
|
|
|
|
class ProvisioningError(RuntimeError):
|
|
"""Expected environment failure that should skip one harness."""
|
|
|
|
|
|
POLICY_ALLOW = "POLICY_ACTION_ALLOW"
|
|
POLICY_DENY = "POLICY_ACTION_DENY"
|
|
|
|
# Denying another phase does not prove tool-call enforcement.
|
|
PHASE_TOOL_CALL = "PHASE_TOOL_CALL"
|
|
|
|
_CONV_ID = "conv_bench"
|
|
|
|
_STREAM_PROMPT = (
|
|
"Count from 1 to 30 in words, one number per line, and add a short note after each."
|
|
)
|
|
_LONG_PROMPT = (
|
|
"Write a very detailed 600-word essay about the history of computing, in full paragraphs."
|
|
)
|
|
_BENCH_TOOL_NAME = "bench_tool"
|
|
_BENCH_DENY_REASON = "bench-policy-deny"
|
|
_BENCH_TOOL_SPEC = [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": _BENCH_TOOL_NAME,
|
|
"description": "A bench probe tool. Call it when asked.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {"arg": {"type": "string"}},
|
|
"required": ["arg"],
|
|
},
|
|
},
|
|
}
|
|
]
|
|
|
|
# Infrastructure failures must not be reported as capability gaps.
|
|
_INFRA_ERROR_MARKERS: tuple[str, ...] = (
|
|
"403",
|
|
"401",
|
|
"Forbidden",
|
|
"Unauthorized",
|
|
"Invalid Token",
|
|
"invalid token",
|
|
"unexpected status",
|
|
"Connection",
|
|
"connection",
|
|
"Temporarily Unavailable",
|
|
"502",
|
|
"503",
|
|
"504",
|
|
"already processing",
|
|
"could not fetch a gateway token",
|
|
"provider auth command",
|
|
"empty token",
|
|
"Failed to resolve external API key auth",
|
|
"are logged in",
|
|
"AcpProcessExited",
|
|
"ACP subprocess",
|
|
"ACP session",
|
|
)
|
|
|
|
|
|
def _error_text(error: object) -> str:
|
|
if isinstance(error, dict):
|
|
return f"{error.get('message', '')} {error.get('code', '')}"
|
|
return str(error or "")
|
|
|
|
|
|
def infra_failure_reason(result: TurnResult) -> str | None:
|
|
"""Return a skip reason when failure reflects infrastructure, not capability."""
|
|
if not result.failed:
|
|
return None
|
|
text = _error_text(result.error)
|
|
if not any(marker in text for marker in _INFRA_ERROR_MARKERS):
|
|
return None
|
|
for code in ("403", "401"):
|
|
if code in text:
|
|
return (
|
|
f"auth rejected ({code} Invalid/Forbidden token); the harness "
|
|
"credential is stale or shadowed by an ambient env var. Refresh "
|
|
"the harness auth source (profile, API key, or token env var)"
|
|
)
|
|
if "already processing" in text:
|
|
return "session busy from a prior turn (sequencing, not a capability gap)"
|
|
if any(
|
|
marker in text
|
|
for marker in (
|
|
"could not fetch a gateway token",
|
|
"provider auth command",
|
|
"empty token",
|
|
"Failed to resolve external API key auth",
|
|
)
|
|
):
|
|
return (
|
|
"gateway/provider token could not be provisioned for this transport "
|
|
"(environment/auth gap, not a capability the harness lacks)"
|
|
)
|
|
if any(
|
|
marker in text
|
|
for marker in ("are logged in", "AcpProcessExited", "ACP subprocess", "ACP session")
|
|
):
|
|
return (
|
|
"vendor CLI not installed or not logged in (own-auth harness); "
|
|
"the agent process exited before a turn could run"
|
|
)
|
|
if "unexpected status" in text:
|
|
return "gateway returned an unexpected status (environment/auth issue)"
|
|
return "environment/connectivity error reaching the gateway"
|
|
|
|
|
|
@dataclass
|
|
class TurnResult:
|
|
"""Probe-observable state from one turn."""
|
|
|
|
events: list[dict[str, Any]] = field(default_factory=list)
|
|
text: str = ""
|
|
text_delta_count: int = 0
|
|
reasoning_delta_count: int = 0
|
|
tool_calls: list[dict[str, Any]] = field(default_factory=list)
|
|
policy_actions: list[tuple[str, str]] = field(default_factory=list)
|
|
tool_call_denied: bool = False
|
|
completed: bool = False
|
|
cancelled: bool = False
|
|
failed: bool = False
|
|
error: Any = None
|
|
timed_out: bool = False
|
|
total_tokens: int | None = None
|
|
total_cost_usd: float | None = None
|
|
elicitation_requested: bool = False
|
|
tool_call_allowed: bool = False
|
|
|
|
@property
|
|
def reached_terminal(self) -> bool:
|
|
return self.completed or self.cancelled or self.failed
|
|
|
|
@property
|
|
def event_types(self) -> list[str]:
|
|
return [e.get("type", "") for e in self.events]
|
|
|
|
|
|
def fill_snapshot_cost(result: TurnResult, snapshot: dict[str, Any]) -> None:
|
|
"""Copy observed usage and cost from a session snapshot."""
|
|
tokens = snapshot.get("last_total_tokens")
|
|
if isinstance(tokens, int):
|
|
result.total_tokens = tokens
|
|
cost = snapshot.get("total_cost_usd")
|
|
if isinstance(cost, (int, float)):
|
|
result.total_cost_usd = float(cost)
|
|
|
|
|
|
class SdkInprocDriver:
|
|
"""Drive turns through a harness wrap subprocess."""
|
|
|
|
transport = "sdk-inproc"
|
|
|
|
def __init__(self, profile: BenchProfile, *, databricks_profile: str | None) -> None:
|
|
self._profile = profile
|
|
self._databricks_profile = databricks_profile
|
|
self._pm: HarnessProcessManager | None = None
|
|
self._client: httpx.AsyncClient | None = None
|
|
self._tmp_parent: Path | None = None
|
|
|
|
@staticmethod
|
|
def unavailable(profile: BenchProfile, *, databricks_profile: str | None) -> str | None:
|
|
"""Return why this driver cannot run the profile, if applicable."""
|
|
if profile.transport != SdkInprocDriver.transport:
|
|
return (
|
|
f"transport {profile.transport!r} not supported by the "
|
|
f"{SdkInprocDriver.transport!r} driver"
|
|
)
|
|
creds_skip = bench_creds_skip_reason(databricks_profile)
|
|
if creds_skip is not None:
|
|
return creds_skip
|
|
if profile.cli_binary is not None:
|
|
reason = cli_unavailable_reason(profile.cli_binary)
|
|
if reason is not None:
|
|
return reason
|
|
return None
|
|
|
|
async def __aenter__(self) -> SdkInprocDriver:
|
|
self._tmp_parent = Path("/tmp") / f"omni-bench-{uuid.uuid4().hex[:8]}"
|
|
self._tmp_parent.mkdir(mode=0o700)
|
|
self._pm = HarnessProcessManager(tmp_parent=self._tmp_parent)
|
|
await self._pm.start()
|
|
p = self._profile
|
|
resolved = resolve_bench_env(self._databricks_profile)
|
|
wrap_env = {
|
|
f"{p.env_prefix}GATEWAY": "true",
|
|
f"{p.env_prefix}MODEL": p.model,
|
|
}
|
|
if resolved.db_profile:
|
|
wrap_env[f"{p.env_prefix}DATABRICKS_PROFILE"] = resolved.db_profile
|
|
self._client = await self._pm.get_client(_CONV_ID, p.harness, env=wrap_env)
|
|
return self
|
|
|
|
async def __aexit__(self, *exc: object) -> None:
|
|
if self._pm is not None:
|
|
await self._pm.shutdown()
|
|
if self._tmp_parent is not None:
|
|
shutil.rmtree(self._tmp_parent, ignore_errors=True)
|
|
|
|
async def run_turn(
|
|
self,
|
|
prompt: str,
|
|
*,
|
|
tools: list[dict[str, Any]] | None = None,
|
|
deny_phases: frozenset[str] = frozenset(),
|
|
policy_reason: str | None = None,
|
|
auto_tool_output: str | None = None,
|
|
interrupt_on_first_delta: bool = False,
|
|
timeout: float = 120.0,
|
|
) -> TurnResult:
|
|
"""Start one turn and drain its event stream."""
|
|
assert self._client is not None, "driver used outside its async context"
|
|
body: dict[str, Any] = {
|
|
"type": "message",
|
|
"role": "user",
|
|
"model": f"{self._profile.harness}-bench-agent",
|
|
"content": [{"type": "input_text", "text": prompt}],
|
|
}
|
|
if tools is not None:
|
|
body["tools"] = tools
|
|
|
|
result = TurnResult()
|
|
try:
|
|
await asyncio.wait_for(
|
|
self._drive(
|
|
body,
|
|
result,
|
|
deny_phases,
|
|
policy_reason,
|
|
auto_tool_output,
|
|
interrupt_on_first_delta,
|
|
),
|
|
timeout=timeout,
|
|
)
|
|
except (asyncio.TimeoutError, httpx.ReadTimeout):
|
|
result.timed_out = True
|
|
return result
|
|
|
|
async def run_basic_turn(self, marker: str) -> TurnResult:
|
|
return await self.run_turn(
|
|
f"Reply with exactly the literal string {marker} and nothing else."
|
|
)
|
|
|
|
async def run_streaming_turn(self) -> TurnResult:
|
|
return await self.run_turn(_STREAM_PROMPT)
|
|
|
|
async def run_tool_turn(self, *, deny: bool) -> TurnResult:
|
|
"""Provoke a tool call and optionally deny its policy evaluation."""
|
|
if deny:
|
|
return await self.run_turn(
|
|
f"Call the {_BENCH_TOOL_NAME} tool with arg='go'. It is required.",
|
|
tools=_BENCH_TOOL_SPEC,
|
|
deny_phases=frozenset({PHASE_TOOL_CALL}),
|
|
policy_reason=_BENCH_DENY_REASON,
|
|
timeout=150.0,
|
|
)
|
|
return await self.run_turn(
|
|
f"You must call the {_BENCH_TOOL_NAME} tool with arg='go', "
|
|
"then reply with the tool's output verbatim.",
|
|
tools=_BENCH_TOOL_SPEC,
|
|
auto_tool_output="bench-tool-ok",
|
|
timeout=150.0,
|
|
)
|
|
|
|
async def run_policy_turn(self, *, action: str) -> TurnResult:
|
|
"""Return unmeasured because wrap-direct cannot observe ALLOW or ASK."""
|
|
return TurnResult()
|
|
|
|
async def run_interrupt_turn(self) -> TurnResult:
|
|
return await self.run_turn(_LONG_PROMPT, interrupt_on_first_delta=True, timeout=120.0)
|
|
|
|
async def _drive(
|
|
self,
|
|
body: dict[str, Any],
|
|
result: TurnResult,
|
|
deny_phases: frozenset[str],
|
|
policy_reason: str | None,
|
|
auto_tool_output: str | None,
|
|
interrupt_on_first_delta: bool,
|
|
) -> None:
|
|
"""POST the turn and consume the SSE stream into *result* in place."""
|
|
client = self._client
|
|
assert client is not None
|
|
interrupted = False
|
|
async with client.stream("POST", f"/v1/sessions/{_CONV_ID}/events", json=body) as response:
|
|
response.raise_for_status()
|
|
buffer = ""
|
|
async for chunk in response.aiter_text():
|
|
buffer += chunk
|
|
while "\n\n" in buffer:
|
|
frame, _, buffer = buffer.partition("\n\n")
|
|
event = _decode_frame(frame)
|
|
if event is None:
|
|
continue
|
|
result.events.append(event)
|
|
etype = event.get("type", "")
|
|
|
|
if etype == "response.output_text.delta":
|
|
result.text += event.get("delta", "")
|
|
result.text_delta_count += 1
|
|
if interrupt_on_first_delta and not interrupted:
|
|
interrupted = True
|
|
await self._post({"type": "interrupt"})
|
|
elif etype in _REASONING_DELTA_TYPES:
|
|
result.reasoning_delta_count += 1
|
|
elif etype == "response.output_item.done":
|
|
# Action-required calls park until a tool result arrives.
|
|
item = event.get("item") or {}
|
|
if (
|
|
item.get("type") == "function_call"
|
|
and item.get("status") == "action_required"
|
|
):
|
|
call_id = item.get("call_id", "")
|
|
result.tool_calls.append(item)
|
|
if auto_tool_output is not None:
|
|
await self._post(
|
|
{
|
|
"type": "tool_result",
|
|
"call_id": call_id,
|
|
"output": auto_tool_output,
|
|
}
|
|
)
|
|
elif etype == "policy_evaluation.requested":
|
|
phase = str(event.get("phase", ""))
|
|
action = POLICY_DENY if phase in deny_phases else POLICY_ALLOW
|
|
verdict: dict[str, Any] = {
|
|
"type": "policy_verdict",
|
|
"evaluation_id": event["evaluation_id"],
|
|
"action": action,
|
|
}
|
|
if action == POLICY_DENY and policy_reason is not None:
|
|
verdict["reason"] = policy_reason
|
|
# A raced or rejected verdict was not delivered.
|
|
if await self._post(verdict):
|
|
result.policy_actions.append((phase, action))
|
|
if action == POLICY_DENY and phase == PHASE_TOOL_CALL:
|
|
result.tool_call_denied = True
|
|
elif etype == "response.completed":
|
|
result.completed = True
|
|
usage = (event.get("response") or {}).get("usage") or {}
|
|
tok = usage.get("total_tokens")
|
|
if isinstance(tok, int):
|
|
result.total_tokens = tok
|
|
cost = usage.get("cost_usd")
|
|
if isinstance(cost, (int, float)):
|
|
result.total_cost_usd = float(cost)
|
|
elif etype == "response.cancelled":
|
|
result.cancelled = True
|
|
elif etype == "response.failed":
|
|
result.failed = True
|
|
result.error = event.get("error") or event.get("response", {}).get("error")
|
|
|
|
async def _post(self, payload: dict[str, Any]) -> bool:
|
|
"""Post a downward event, tolerating races with turn completion."""
|
|
assert self._client is not None
|
|
try:
|
|
resp = await self._client.post(f"/v1/sessions/{_CONV_ID}/events", json=payload)
|
|
except httpx.HTTPError:
|
|
return False
|
|
return not resp.is_error
|
|
|
|
|
|
# Harness wraps use both reasoning-delta spellings.
|
|
_REASONING_DELTA_TYPES: frozenset[str] = frozenset(
|
|
{"response.reasoning.delta", "response.reasoning_summary_text.delta"}
|
|
)
|
|
|
|
|
|
def _decode_frame(frame: str) -> dict[str, Any] | None:
|
|
data_line = next(
|
|
(line for line in frame.splitlines() if line.startswith("data:")),
|
|
None,
|
|
)
|
|
if data_line is None:
|
|
return None
|
|
try:
|
|
decoded = json.loads(data_line[len("data:") :].strip())
|
|
except json.JSONDecodeError:
|
|
return None
|
|
return decoded if isinstance(decoded, dict) else None
|