From eee41a7bf344c0bd059a60a14e434e7d4d7f84ff Mon Sep 17 00:00:00 2001 From: adithya-s-k Date: Sun, 2 Aug 2026 19:22:51 +0000 Subject: [PATCH 01/74] capture: dialect-agnostic token capture for agent rollouts An OpenAI-spec proxy that sits between a coding agent and an inference endpoint and records the exact token ids and per-token logprobs of every model call, so a rollout is trainable. Nothing is tokenised locally: the engine returns prompt_token_ids, so turn k+1's prompt is the canonical tokenisation of everything before it and turns link by exact token prefix. Re-rendering a prompt offline drifts from what the model saw, and a drifted prompt silently fragments one conversation into several. Four wire dialects (chat-completions, OpenAI Responses, Anthropic Messages, Google generateContent), adapted from the Polar gateway (Apache-2.0); provenance in dialects/README.md. Vendored because the package named polar on PyPI is unrelated. Two ideas are borrowed from verifiers: aux routes, so a count_tokens call is answered without becoming a model turn, and per-dialect streaming detection, since Google signals streaming in the URL. Includes engine certification, which refuses an endpoint that cannot return token ids, and port forwarding for sandboxes that cannot reach localhost. --- .gitignore | 3 + src/openenv/core/harness/capture/__init__.py | 35 + src/openenv/core/harness/capture/contract.py | 190 +++ src/openenv/core/harness/capture/detection.py | 61 + .../core/harness/capture/dialects/README.md | 31 + .../core/harness/capture/dialects/__init__.py | 23 + .../harness/capture/dialects/anthropic.py | 693 +++++++++++ .../core/harness/capture/dialects/base.py | 125 ++ .../core/harness/capture/dialects/google.py | 724 +++++++++++ .../core/harness/capture/dialects/images.py | 335 +++++ .../harness/capture/dialects/openai_chat.py | 41 + .../capture/dialects/openai_responses.py | 1090 +++++++++++++++++ .../harness/capture/dialects/reasoning.py | 114 ++ src/openenv/core/harness/capture/export.py | 189 +++ .../core/harness/capture/forwarding.py | 276 +++++ src/openenv/core/harness/capture/graph.py | 284 +++++ src/openenv/core/harness/capture/server.py | 532 ++++++++ src/openenv/core/harness/capture/sessions.py | 147 +++ src/openenv/core/harness/capture/sse.py | 175 +++ src/openenv/core/harness/capture/upstream.py | 182 +++ src/openenv/core/harness/capture/validate.py | 305 +++++ .../core/harness/capture/validate_llm.py | 162 +++ 22 files changed, 5717 insertions(+) create mode 100644 src/openenv/core/harness/capture/__init__.py create mode 100644 src/openenv/core/harness/capture/contract.py create mode 100644 src/openenv/core/harness/capture/detection.py create mode 100644 src/openenv/core/harness/capture/dialects/README.md create mode 100644 src/openenv/core/harness/capture/dialects/__init__.py create mode 100644 src/openenv/core/harness/capture/dialects/anthropic.py create mode 100644 src/openenv/core/harness/capture/dialects/base.py create mode 100644 src/openenv/core/harness/capture/dialects/google.py create mode 100644 src/openenv/core/harness/capture/dialects/images.py create mode 100644 src/openenv/core/harness/capture/dialects/openai_chat.py create mode 100644 src/openenv/core/harness/capture/dialects/openai_responses.py create mode 100644 src/openenv/core/harness/capture/dialects/reasoning.py create mode 100644 src/openenv/core/harness/capture/export.py create mode 100644 src/openenv/core/harness/capture/forwarding.py create mode 100644 src/openenv/core/harness/capture/graph.py create mode 100644 src/openenv/core/harness/capture/server.py create mode 100644 src/openenv/core/harness/capture/sessions.py create mode 100644 src/openenv/core/harness/capture/sse.py create mode 100644 src/openenv/core/harness/capture/upstream.py create mode 100644 src/openenv/core/harness/capture/validate.py create mode 100644 src/openenv/core/harness/capture/validate_llm.py diff --git a/.gitignore b/.gitignore index 51fe1e842..726e19e49 100644 --- a/.gitignore +++ b/.gitignore @@ -141,3 +141,6 @@ docs/source/_env_assets/ # Sphinx-gallery generated output docs/source/auto_getting_started/ docs/source/sg_execution_times.rst + +# Gradio UI build artifacts +.gradio/ diff --git a/src/openenv/core/harness/capture/__init__.py b/src/openenv/core/harness/capture/__init__.py new file mode 100644 index 000000000..f419fa959 --- /dev/null +++ b/src/openenv/core/harness/capture/__init__.py @@ -0,0 +1,35 @@ +"""Token-level capture for agentic rollouts. + +An OpenAI-spec proxy that sits between a coding agent and an inference engine, recording the exact +token ids and logprobs of every model call so a rollout can be trained on. + +The core idea is that nothing is ever tokenised locally. The engine returns `prompt_token_ids`, so +turn k+1's prompt IS the canonical tokenisation of everything before it, and turns are linked by +exact token prefix. Re-rendering a prompt offline drifts from what the model actually saw, and a +drifted prompt silently fragments one long conversation into several short ones. + +Four wire dialects are supported (chat-completions, Responses, Anthropic Messages, Google +generateContent) because coding agents did not agree on one. Validated across 16 harnesses, each +cross-checked against the harness's own trace. +""" + +from .contract import measure_retokenization_skew, to_trace_entries, to_turn_records +from .detection import APIType, detect +from .graph import RolloutGraph, TurnNode +from .upstream import InferenceClient, UpstreamError +from .validate_llm import LLMReport, require_llm, validate_llm + +__all__ = [ + "to_turn_records", + "to_trace_entries", + "measure_retokenization_skew", + "APIType", + "detect", + "RolloutGraph", + "TurnNode", + "InferenceClient", + "UpstreamError", + "LLMReport", + "validate_llm", + "require_llm", +] diff --git a/src/openenv/core/harness/capture/contract.py b/src/openenv/core/harness/capture/contract.py new file mode 100644 index 000000000..f0d0029e4 --- /dev/null +++ b/src/openenv/core/harness/capture/contract.py @@ -0,0 +1,190 @@ +"""What a rollout hands a trainer. + +The contract is one thing, and it is small: + + per turn -> (prompt_token_ids, completion_token_ids, per_token_logps) + +That is everything an on-policy method needs. `prompt_token_ids` is the engine's own tokenisation of +the conversation up to that turn, `completion_token_ids` is what it sampled, and `per_token_logps` +are the behaviour-policy logprobs for exactly those sampled tokens. `to_turn_records` emits it, and +it is lossless because nothing is ever re-derived. + +Everything else in this module is an ADAPTER to some consumer's existing shape, and adapters lose +information. `to_trace_entries` produces TRL's `TraceEntry` +(`{request, response, completion_token_ids, per_token_logps}`), which has no field for the prompt's +token ids — so a consumer of it must re-render the prompt with `apply_chat_template` to recover +them. That re-render is not free: measured against Qwen3.5 it was 0/6 turns exact, off by two tokens +every turn, until thinking was disabled. A prompt that differs by one token from what the model saw +looks like divergence, and a long conversation silently fragments into several short ones. + +So: new consumers should take `to_turn_records`. `to_trace_entries` exists to work with TRL as it is +today, and `measure_retokenization_skew` exists to tell you what that costs for a given +model + harness pair instead of guessing. +""" + +from __future__ import annotations + +from typing import Any + +from .graph import RolloutGraph, TurnNode + + +def _agent_nodes(graph: RolloutGraph, document: dict[str, Any]) -> list[TurnNode]: + """Nodes on the agent's conversation, in arrival order, excluding discarded retries.""" + agent_rows = [r for r in document["sequences"] if r["role"] == "agent"] + if not agent_rows: + return [] + root = agent_rows[0]["root_id"] + keep = set(agent_rows[0]["node_ids"]) + return [ + n + for n in graph.nodes() + if graph.root_of(n.node_id) == root and n.node_id in keep + ] + + +def to_trace_entries( + graph: RolloutGraph, document: dict[str, Any] +) -> list[dict[str, Any]]: + """Rollout graph -> TRL `list[TraceEntry]`. Works with TRL today, at the cost of re-tokenization. + + Auxiliary roots and discarded retries are already excluded here, so the caller does not need an + `agent_turn_fn`. That hook exists because a flat trace cannot tell an aux call from an agent + turn; a graph can, structurally. + """ + entries = [] + for node in _agent_nodes(graph, document): + entries.append( + { + "request": { + "messages": node.request_messages, + "tools": node.request_tools, + }, + "response": { + "choices": [ + { + "message": node.response_message, + "finish_reason": node.finish_reason, + } + ] + }, + "completion_token_ids": node.sampled_ids, + "per_token_logps": node.sampled_logprobs or [], + } + ) + return entries + + +def to_turn_records( + graph: RolloutGraph, document: dict[str, Any] +) -> list[tuple[list[int], list[int], list[float]]]: + """Rollout graph -> `(prompt_ids, output_ids, output_log_probs)` per turn, losslessly. + + Maps 1:1 onto TRL's `TurnRecord`, using the engine's own prompt tokenization. Returned as plain + tuples so this module stays importable without TRL on the path. + """ + return [ + (node.prompt_ids, node.sampled_ids, node.sampled_logprobs or []) + for node in _agent_nodes(graph, document) + ] + + +def measure_retokenization_skew( + graph: RolloutGraph, + document: dict[str, Any], + tokenizer, + *, + chat_template: str | None = None, + chat_template_kwargs: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Compare TRL's re-tokenized prompt against the engine's, per turn. The number nobody has had. + + Reproduces `_turns_from_trace` exactly, including the `_decode_tool_call_arguments` step (the + trace stores tool-call `arguments` as a JSON string; XML-style templates such as Qwen3.5's + iterate it and raise on a string). + + Returns per-turn exact-match, common-prefix length and length delta. `exact_match_frac == 1.0` + means re-tokenization is safe for this model+harness pair and the hook buys nothing. Anything + less means TRL is training on a prompt the model never saw, and `_chain_to_sequences` will fork + the conversation at the first divergence. + """ + import json as _json + + def decode_arguments(messages): + out = [] + for message in messages: + calls = message.get("tool_calls") + if not calls: + out.append(message) + continue + new = [] + for call in calls: + function = call.get("function") + arguments = (function or call).get("arguments") + if not isinstance(arguments, str): + new.append(call) + continue + try: + arguments = _json.loads(arguments) + except _json.JSONDecodeError: + arguments = {} + new.append( + {**call, "function": {**function, "arguments": arguments}} + if function + else {**call, "arguments": arguments} + ) + out.append({**message, "tool_calls": new}) + return out + + def prefix_len(a, b): + n = min(len(a), len(b)) + i = 0 + while i < n and a[i] == b[i]: + i += 1 + return i + + turns = [] + for i, node in enumerate(_agent_nodes(graph, document)): + try: + rebuilt = tokenizer.apply_chat_template( + decode_arguments(node.request_messages), + tools=node.request_tools, + add_generation_prompt=True, + tokenize=True, + return_dict=False, + chat_template=chat_template, + **(chat_template_kwargs or {}), + ) + except Exception as exc: # noqa: BLE001 - a template that raises IS the finding + turns.append( + { + "turn": i, + "error": f"{type(exc).__name__}: {str(exc)[:160]}", + "engine_len": len(node.prompt_ids), + } + ) + continue + turns.append( + { + "turn": i, + "engine_len": len(node.prompt_ids), + "rebuilt_len": len(rebuilt), + "delta": len(rebuilt) - len(node.prompt_ids), + "prefix_match": prefix_len(node.prompt_ids, rebuilt), + "exact": list(rebuilt) == list(node.prompt_ids), + } + ) + + scored = [t for t in turns if "exact" in t] + return { + "n_turns": len(turns), + "n_errors": len(turns) - len(scored), + "exact_match_frac": (sum(t["exact"] for t in scored) / len(scored)) + if scored + else 0.0, + "max_abs_delta": max((abs(t["delta"]) for t in scored), default=0), + "min_prefix_match_frac": min( + (t["prefix_match"] / max(t["engine_len"], 1) for t in scored), default=0.0 + ), + "turns": turns, + } diff --git a/src/openenv/core/harness/capture/detection.py b/src/openenv/core/harness/capture/detection.py new file mode 100644 index 000000000..f20742ca8 --- /dev/null +++ b/src/openenv/core/harness/capture/detection.py @@ -0,0 +1,61 @@ +"""Which wire dialect is this request written in? + +Four dialects reach the intercept, because coding agents did not agree on one: + + openai_chat opencode, qwen-coder, goose, swe-agent, mini-swe-agent, terminus-2, ... + openai_responses codex, trae-agent + anthropic claude-code + google gemini-cli, antigravity-sdk + +Detection is by path first, then headers, then body shape — strongest signal to weakest. Path is +unambiguous when present; a header is a deliberate client declaration; body shape is a last resort +and can be coincidental, so it is only consulted when nothing better exists. + +Getting this wrong is not subtle: a Google request parsed as chat-completions produces a 400 and the +agent silently does nothing, which reads as "captured nothing" rather than as a routing bug. That is +why trae-agent looked like a chat harness for a full night — its config says `provider: openai`, but +the access log showed exactly one `POST /v1/responses` against 465 chat calls. +""" + +from __future__ import annotations + +from enum import Enum +from typing import Any + + +class APIType(str, Enum): + ANTHROPIC = "anthropic" + OPENAI_CHAT = "openai_chat" + OPENAI_RESPONSES = "openai_responses" + GOOGLE = "google" + + +def detect(path: str, headers: dict[str, str], body: dict[str, Any]) -> APIType: + """Classify one request. Defaults to chat-completions, the most common dialect.""" + if "/v1/messages" in path: + return APIType.ANTHROPIC + if "/v1/chat/completions" in path: + return APIType.OPENAI_CHAT + if "/v1/responses" in path: + return APIType.OPENAI_RESPONSES + # Google puts the method in the path: `/v1beta/models/{model}:generateContent`, and its + # streaming variant `:streamGenerateContent`. Both contain this substring. + if "generateContent" in path: + return APIType.GOOGLE + + if "anthropic-version" in {k.lower() for k in headers}: + return APIType.ANTHROPIC + + if "contents" in body: + return APIType.GOOGLE + if "input" in body and "instructions" in body: + return APIType.OPENAI_RESPONSES + + return APIType.OPENAI_CHAT + + +def extract_model(api_type: APIType, body: dict[str, Any]) -> str: + """The model name the client asked for, whatever the dialect calls it.""" + if api_type is APIType.GOOGLE: + return body.get("model", "gemini-pro") + return body.get("model", "unknown") diff --git a/src/openenv/core/harness/capture/dialects/README.md b/src/openenv/core/harness/capture/dialects/README.md new file mode 100644 index 000000000..993090f69 --- /dev/null +++ b/src/openenv/core/harness/capture/dialects/README.md @@ -0,0 +1,31 @@ +# Dialect transformers + +Translation between the four wire dialects coding agents speak and the OpenAI chat-completions +shape a vLLM server understands. + +| dialect | spoken by | +|---|---| +| `openai_chat` | opencode, qwen-coder, goose, swe-agent, mini-swe-agent, terminus-2, vibe, mimo, kimi-cli, hermes, openclaw, openhands-sdk, pi | +| `openai_responses` | codex, trae-agent | +| `anthropic` | claude-code | +| `google` | gemini-cli, antigravity-sdk | + +Without this layer the intercept only works for chat-completions agents, which is 13 of the 16 +validated harnesses — but the three it loses (claude-code, codex, gemini-cli) are the ones whose +capture is hardest to get right, so they are also the ones most worth having covered. + +## Provenance + +Adapted from the Polar gateway (`polar/gateway/transform/`, Apache-2.0). Changes made when +vendoring: + +- import paths rewritten to be relative and self-contained; no dependency on Polar remains +- the internal request marker `_polar_model_served` renamed to `_served_model` +- reasoning-signature wire prefixes `polar:` / `sg_polar_` renamed to `oe:` / `sg_oe_` + (an opaque, symmetric encode/decode pair — the value is arbitrary as long as both sides agree) +- `engine.py` and `proxy.py` were **not** vendored. They carried an SGLang backend that cannot + support token capture at all, so they were replaced by `capture/upstream.py`, a vLLM-only client + in ~160 lines. + +`images.py` and `reasoning.py` are required: the anthropic, google and responses transformers all +import them for multimodal content blocks and thinking-block round-tripping respectively. diff --git a/src/openenv/core/harness/capture/dialects/__init__.py b/src/openenv/core/harness/capture/dialects/__init__.py new file mode 100644 index 000000000..41b9183af --- /dev/null +++ b/src/openenv/core/harness/capture/dialects/__init__.py @@ -0,0 +1,23 @@ +"""Transform manager — dispatches to the right transformer by API type.""" + +from ..detection import APIType +from .anthropic import AnthropicTransformer +from .base import BaseTransformer +from .google import GoogleTransformer +from .openai_chat import OpenAIChatTransformer +from .openai_responses import OpenAIResponsesTransformer + + +class TransformManager: + """Route to the correct transformer based on detected API type.""" + + def __init__(self): + self._transformers: dict[APIType, BaseTransformer] = { + APIType.ANTHROPIC: AnthropicTransformer(), + APIType.OPENAI_CHAT: OpenAIChatTransformer(), + APIType.OPENAI_RESPONSES: OpenAIResponsesTransformer(), + APIType.GOOGLE: GoogleTransformer(), + } + + def get(self, api_type: APIType) -> BaseTransformer: + return self._transformers[api_type] diff --git a/src/openenv/core/harness/capture/dialects/anthropic.py b/src/openenv/core/harness/capture/dialects/anthropic.py new file mode 100644 index 000000000..b5573e532 --- /dev/null +++ b/src/openenv/core/harness/capture/dialects/anthropic.py @@ -0,0 +1,693 @@ +"""Anthropic Messages API transformer. + +Transforms between Anthropic Messages API and OpenAI Chat Completions API. +Aligned with agent-harness-proxy/src/harness_proxy/transform/anthropic.py. +""" + +from __future__ import annotations + +import json +import re +import uuid +from dataclasses import dataclass +from typing import Any, Optional + +from .base import BaseTransformer +from .images import ( + anthropic_content_to_openai_chat, + openai_chat_content_to_anthropic_blocks, +) +from .reasoning import extract_reasoning_from_anthropic_content, make_signature + +# Claude Code SDK leaks `x-anthropic-billing-header: ...cch=;` as the +# first line of the system prompt. The cch= hash changes per request, so +# rendered prompt tokens drift every turn and prefix_merging can't chain +# multi-turn traces. Strip the line before forwarding to SGLang. +_CLAUDE_CODE_BILLING_HEADER_RE = re.compile( + r"^\s*x-anthropic-billing-header:[^\n]*\n?", re.IGNORECASE +) + + +@dataclass +class _AnthropicToolCallState: + id: str + name: str = "" + anthropic_index: int | None = None + buffered_arguments: str = "" + started: bool = False + + +class AnthropicStreamState: + """Per-request Anthropic streaming state. + + Anthropic SSE blocks are stateful across chunks: content blocks must be + explicitly started, optionally receive multiple deltas, and then be closed + before the final message delta. This helper tracks those open blocks for a + single upstream OpenAI/SGLang stream. + """ + + def __init__(self, model: str, finish_to_stop_reason: dict[str, str]): + self.model = model + self.finish_to_stop_reason = finish_to_stop_reason + self.message_id = f"msg_{uuid.uuid4().hex}" + self.next_block_index = 0 + self.text_block_index: int | None = None + self.text_block_started = False + self.thinking_block_index: int | None = None + self.thinking_block_started = False + self.thinking_buffer = "" + self.tool_calls: dict[int, _AnthropicToolCallState] = {} + self.stop_reason = "end_turn" + self.output_tokens = 0 + self.any_block_started = False + self.completed = False + + def process_chunk( + self, chunk: dict[str, Any], is_first: bool = False + ) -> list[dict[str, Any]]: + events: list[dict[str, Any]] = [] + + if is_first: + events.append( + { + "type": "message_start", + "message": { + "id": self.message_id, + "type": "message", + "role": "assistant", + "content": [], + "model": self.model, + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 0, "output_tokens": 0}, + }, + } + ) + + usage = chunk.get("usage", {}) + if usage: + self.output_tokens = usage.get("completion_tokens", self.output_tokens) + + choices = chunk.get("choices", []) + if not choices: + return events + + choice = choices[0] + delta = choice.get("delta", {}) or {} + finish_reason = choice.get("finish_reason") + if finish_reason: + self.stop_reason = self.finish_to_stop_reason.get(finish_reason, "end_turn") + + # Thinking blocks must precede text and tool_use per Anthropic spec. + reasoning = delta.get("reasoning_content") + if reasoning: + if not self.thinking_block_started: + events.append(self._open_thinking_block()) + events.append( + { + "type": "content_block_delta", + "index": self.thinking_block_index, + "delta": {"type": "thinking_delta", "thinking": reasoning}, + } + ) + self.thinking_buffer += reasoning + + content = delta.get("content") + if content: + thinking_stop = self._close_thinking_block() + if thinking_stop: + events.extend(thinking_stop) + if not self.text_block_started: + events.append(self._open_text_block()) + events.append( + { + "type": "content_block_delta", + "index": self.text_block_index, + "delta": {"type": "text_delta", "text": content}, + } + ) + + tool_call_deltas = delta.get("tool_calls") or [] + if not isinstance(tool_call_deltas, list): + tool_call_deltas = [tool_call_deltas] + for tool_call_delta in tool_call_deltas: + if isinstance(tool_call_delta, dict): + events.extend(self._process_tool_call(tool_call_delta)) + + return events + + def finalize(self) -> list[dict[str, Any]]: + if self.completed: + return [] + + events: list[dict[str, Any]] = [] + + thinking_stop = self._close_thinking_block() + if thinking_stop: + events.extend(thinking_stop) + + text_stop = self._close_text_block() + if text_stop: + events.append(text_stop) + + for tool_index in sorted(self.tool_calls): + tool_state = self.tool_calls[tool_index] + if tool_state.started and tool_state.anthropic_index is not None: + events.append( + { + "type": "content_block_stop", + "index": tool_state.anthropic_index, + } + ) + + if not self.any_block_started: + empty_index = self.next_block_index + events.append( + { + "type": "content_block_start", + "index": empty_index, + "content_block": {"type": "text", "text": ""}, + } + ) + events.append({"type": "content_block_stop", "index": empty_index}) + + events.append( + { + "type": "message_delta", + "delta": {"stop_reason": self.stop_reason, "stop_sequence": None}, + "usage": {"output_tokens": self.output_tokens}, + } + ) + events.append({"type": "message_stop"}) + + self.completed = True + return events + + def _open_text_block(self) -> dict[str, Any]: + self.text_block_started = True + self.text_block_index = self.next_block_index + self.next_block_index += 1 + self.any_block_started = True + return { + "type": "content_block_start", + "index": self.text_block_index, + "content_block": {"type": "text", "text": ""}, + } + + def _close_text_block(self) -> dict[str, Any] | None: + if not self.text_block_started or self.text_block_index is None: + return None + + event = {"type": "content_block_stop", "index": self.text_block_index} + self.text_block_started = False + self.text_block_index = None + return event + + def _open_thinking_block(self) -> dict[str, Any]: + self.thinking_block_started = True + self.thinking_block_index = self.next_block_index + self.next_block_index += 1 + self.any_block_started = True + return { + "type": "content_block_start", + "index": self.thinking_block_index, + "content_block": {"type": "thinking", "thinking": "", "signature": ""}, + } + + def _close_thinking_block(self) -> list[dict[str, Any]] | None: + if not self.thinking_block_started or self.thinking_block_index is None: + return None + idx = self.thinking_block_index + events = [ + { + "type": "content_block_delta", + "index": idx, + "delta": { + "type": "signature_delta", + "signature": make_signature(self.thinking_buffer), + }, + }, + {"type": "content_block_stop", "index": idx}, + ] + self.thinking_block_started = False + self.thinking_block_index = None + return events + + def _process_tool_call( + self, tool_call_delta: dict[str, Any] + ) -> list[dict[str, Any]]: + events: list[dict[str, Any]] = [] + + tool_index = tool_call_delta.get("index", 0) + if not isinstance(tool_index, int): + tool_index = 0 + + tool_state = self.tool_calls.get(tool_index) + if tool_state is None: + tool_state = _AnthropicToolCallState( + id=tool_call_delta.get("id", f"toolu_{uuid.uuid4().hex[:24]}"), + ) + self.tool_calls[tool_index] = tool_state + elif tool_call_delta.get("id"): + tool_state.id = tool_call_delta["id"] + + function = tool_call_delta.get("function", {}) + name = function.get("name") + if isinstance(name, str) and name: + tool_state.name += name + + args = function.get("arguments") + args_str = "" + if isinstance(args, str) and args: + args_str = args + elif args not in (None, ""): + args_str = json.dumps(args) + + if args_str: + tool_state.buffered_arguments += args_str + + if tool_state.name and not tool_state.started: + thinking_stop = self._close_thinking_block() + if thinking_stop: + events.extend(thinking_stop) + + text_stop = self._close_text_block() + if text_stop: + events.append(text_stop) + + tool_state.started = True + tool_state.anthropic_index = self.next_block_index + self.next_block_index += 1 + self.any_block_started = True + + events.append( + { + "type": "content_block_start", + "index": tool_state.anthropic_index, + "content_block": { + "type": "tool_use", + "id": tool_state.id, + "name": tool_state.name, + "input": {}, + }, + } + ) + + if tool_state.buffered_arguments: + events.append( + { + "type": "content_block_delta", + "index": tool_state.anthropic_index, + "delta": { + "type": "input_json_delta", + "partial_json": tool_state.buffered_arguments, + }, + } + ) + tool_state.buffered_arguments = "" + elif tool_state.started and args_str and tool_state.anthropic_index is not None: + events.append( + { + "type": "content_block_delta", + "index": tool_state.anthropic_index, + "delta": { + "type": "input_json_delta", + "partial_json": args_str, + }, + } + ) + + return events + + +class AnthropicTransformer(BaseTransformer): + """Transform between Anthropic and OpenAI API formats.""" + + FINISH_TO_STOP_REASON: dict[str, str] = { + "stop": "end_turn", + "length": "max_tokens", + "tool_calls": "tool_use", + "content_filter": "refusal", + "stop_sequence": "stop_sequence", + } + + def transform_request(self, body: dict[str, Any]) -> dict[str, Any]: + messages = [] + + # Handle system message + system = body.get("system") + if system: + system_content = self._flatten_content(system) + # Drop Claude Code's per-request billing header line (breaks + # prefix_merging because cch= changes every turn). + system_content = _CLAUDE_CODE_BILLING_HEADER_RE.sub("", system_content) + if system_content: + messages.append({"role": "system", "content": system_content}) + + # Transform messages + for msg in body.get("messages", []): + transformed = self._transform_message(msg) + if transformed: + if isinstance(transformed, list): + messages.extend(transformed) + else: + messages.append(transformed) + + result: dict[str, Any] = { + "messages": messages, + "max_tokens": body.get("max_tokens", 4096), + } + if "model" in body: + result["model"] = body["model"] + + if "temperature" in body: + result["temperature"] = body["temperature"] + if "top_p" in body: + result["top_p"] = body["top_p"] + if "top_k" in body: + result["top_k"] = body["top_k"] + if "stop_sequences" in body: + result["stop"] = body["stop_sequences"] + if body.get("stream", False): + result["stream"] = True + + # Anthropic `thinking` request param → enable_thinking on chat template. + thinking_cfg = body.get("thinking") + if isinstance(thinking_cfg, dict) and thinking_cfg.get("type") in { + "enabled", + "adaptive", + }: + chat_template_kwargs = dict(result.get("chat_template_kwargs") or {}) + chat_template_kwargs["enable_thinking"] = True + result["chat_template_kwargs"] = chat_template_kwargs + + # Tools. Claude Code sometimes sends tools=[] on compaction/summary + # turns; forwarding tool_choice without a non-empty tools list makes + # SGLang reject with "tool_choice only allowed when tools specified". + if "tools" in body: + tools = self._transform_tools_to_openai(body["tools"]) + if tools: + result["tools"] = tools + result["tool_choice"] = self._transform_tool_choice_to_openai( + body.get("tool_choice", {"type": "auto"}) + ) + + return self._normalize_request( + result, + body.get("_served_model"), + ) + + def transform_response( + self, + response: dict[str, Any], + original_request: dict[str, Any], + ) -> dict[str, Any]: + choices = response.get("choices", []) + if not choices: + return self._error_response("No choices in response") + + choice = choices[0] + message = choice.get("message", {}) + + content = [] + reasoning = message.get("reasoning_content") + if isinstance(reasoning, str) and reasoning: + content.append( + { + "type": "thinking", + "thinking": reasoning, + "signature": make_signature(reasoning), + } + ) + + text = message.get("content") + if text or (isinstance(text, list) and text): + content.extend(openai_chat_content_to_anthropic_blocks(text)) + + for tool_call in message.get("tool_calls") or []: + content.append( + { + "type": "tool_use", + "id": tool_call.get("id", f"toolu_{uuid.uuid4().hex[:24]}"), + "name": tool_call.get("function", {}).get("name", ""), + "input": self._parse_json_safe( + tool_call.get("function", {}).get("arguments", "{}") + ), + } + ) + + finish_reason = choice.get("finish_reason", "stop") + stop_reason = self.FINISH_TO_STOP_REASON.get(finish_reason, "end_turn") + usage = response.get("usage", {}) + anthropic_usage = self._usage_to_anthropic(usage) + + if not content: + content.append({"type": "text", "text": ""}) + + return { + "id": f"msg_{response.get('id', uuid.uuid4().hex)}", + "type": "message", + "role": "assistant", + "content": content, + "model": original_request.get("model", "claude-3"), + "stop_reason": stop_reason, + "stop_sequence": None, + "usage": anthropic_usage, + } + + def create_stream_state( + self, original_request: dict[str, Any] + ) -> AnthropicStreamState: + return AnthropicStreamState( + model=original_request.get("model", "claude-3"), + finish_to_stop_reason=self.FINISH_TO_STOP_REASON, + ) + + def transform_stream_chunk( + self, + chunk: dict[str, Any], + original_request: dict[str, Any], + is_first: bool = False, + ) -> list[dict[str, Any]]: + """Best-effort single-chunk Anthropic transform. + + The server uses `create_stream_state()` for request-scoped streaming. + This fallback keeps direct callers working for simple single-chunk cases. + """ + state = self.create_stream_state(original_request) + events = state.process_chunk(chunk, is_first=is_first) + choices = chunk.get("choices", []) + if choices and choices[0].get("finish_reason"): + events.extend(state.finalize()) + return events + + def _transform_message(self, msg: dict[str, Any]) -> Optional[dict | list]: + """Transform a single Anthropic message to OpenAI format.""" + role = msg.get("role", "user") + content = msg.get("content", "") + + if isinstance(content, str): + return {"role": role, "content": content} + + if not isinstance(content, list): + return {"role": role, "content": str(content)} + + # Check for mixed content: tool_result blocks + other content + tool_results = [ + c for c in content if isinstance(c, dict) and c.get("type") == "tool_result" + ] + tool_uses = [ + c for c in content if isinstance(c, dict) and c.get("type") == "tool_use" + ] + text_blocks = [ + c for c in content if isinstance(c, dict) and c.get("type") == "text" + ] + + messages = [] + + # Assistant `thinking` blocks → reasoning_content (kept for replay). + reasoning_text = "" + if role == "assistant": + reasoning_text = extract_reasoning_from_anthropic_content(content) + + # Handle assistant messages with tool_use blocks + if role == "assistant" and tool_uses: + tool_calls = [] + text_parts = [] + for block in content: + if isinstance(block, dict): + if block.get("type") == "text": + text_parts.append(block.get("text", "")) + elif block.get("type") == "tool_use": + tool_calls.append( + { + "id": block.get("id", f"call_{uuid.uuid4().hex[:24]}"), + "type": "function", + "function": { + "name": block.get("name", ""), + "arguments": json.dumps(block.get("input", {})), + }, + } + ) + msg_dict: dict[str, Any] = { + "role": "assistant", + "content": "\n".join(text_parts) if text_parts else None, + } + if reasoning_text: + msg_dict["reasoning_content"] = reasoning_text + if tool_calls: + msg_dict["tool_calls"] = tool_calls + return msg_dict + + # Handle user messages with tool_result blocks + if role == "user" and tool_results: + # Each tool_result becomes a tool message + for tr in tool_results: + tool_content = tr.get("content", "") + converted_content = anthropic_content_to_openai_chat(tool_content) + text_content = self._flatten_content(converted_content) + # Anthropic marks failed tool results with is_error=true. + # Surface this to the model so it can see the call failed + # rather than treating the payload as normal output. + if tr.get("is_error"): + text_content = ( + f"[Tool Error] {text_content}" + if text_content + else "[Tool Error]" + ) + messages.append( + { + "role": "tool", + "tool_call_id": tr.get("tool_use_id", ""), + "content": text_content, + } + ) + # OpenAI tool messages stay text-only; images are sent as a + # follow-up user message, so mixed text/image order is not preserved. + image_parts = self._image_parts(converted_content) + if image_parts: + messages.append({"role": "user", "content": image_parts}) + + # Any extra user text should come after the tool results. + text_parts = [b.get("text", "") for b in text_blocks if b.get("text")] + if text_parts: + messages.append({"role": "user", "content": "\n".join(text_parts)}) + return messages if messages else None + + # Regular content blocks — keep images when present. + result: dict[str, Any] = { + "role": role, + "content": anthropic_content_to_openai_chat(content), + } + if role == "assistant" and reasoning_text: + result["reasoning_content"] = reasoning_text + return result + + def _flatten_content(self, content: Any) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + parts = [] + for block in content: + if isinstance(block, str): + parts.append(block) + elif isinstance(block, dict): + if block.get("type") == "text": + parts.append(block.get("text", "")) + elif block.get("type") == "tool_result": + parts.append(self._flatten_content(block.get("content", ""))) + return "\n".join(parts) + return str(content) + + def _image_parts(self, content: Any) -> list[dict[str, Any]]: + if not isinstance(content, list): + return [] + return [ + part + for part in content + if isinstance(part, dict) and part.get("type") == "image_url" + ] + + def _transform_tools_to_openai(self, tools: list[dict]) -> list[dict]: + result = [] + for tool in tools: + # Anthropic server tools (web_search_*, code_execution_*) carry an + # explicit `type` and have no `input_schema`. SGLang can't dispatch + # them, so drop rather than forwarding a stub function tool. + tool_type = tool.get("type") + if ( + tool_type + and tool_type not in ("custom", "function") + and "input_schema" not in tool + ): + continue + name = tool.get("name") + if not isinstance(name, str) or not name: + continue + result.append( + { + "type": "function", + "function": { + "name": name, + "description": tool.get("description", ""), + "parameters": tool.get("input_schema", {}), + }, + } + ) + return result + + def _transform_tool_choice_to_openai(self, tool_choice: Any) -> Any: + if isinstance(tool_choice, dict): + tc_type = tool_choice.get("type") + if tc_type == "auto": + return "auto" + elif tc_type == "any": + return "required" + elif tc_type == "none": + return "none" + elif tc_type == "tool": + return { + "type": "function", + "function": {"name": tool_choice.get("name", "")}, + } + return "auto" + + def _parse_json_safe(self, s: str) -> dict: + try: + return json.loads(s) + except (json.JSONDecodeError, TypeError): + return {} + + def _usage_to_anthropic(self, usage: dict[str, Any]) -> dict[str, Any]: + prompt_tokens = usage.get("prompt_tokens", 0) + completion_tokens = usage.get("completion_tokens", 0) + cache_read = self._cached_prompt_tokens(usage) + input_tokens = ( + max(prompt_tokens - cache_read, 0) if cache_read else prompt_tokens + ) + + result: dict[str, Any] = { + "input_tokens": input_tokens, + "output_tokens": completion_tokens, + } + if cache_read: + result["cache_read_input_tokens"] = cache_read + cache_creation = usage.get("cache_creation_input_tokens") + if isinstance(cache_creation, int) and cache_creation: + result["cache_creation_input_tokens"] = cache_creation + return result + + def _cached_prompt_tokens(self, usage: dict[str, Any]) -> int: + details = usage.get("prompt_tokens_details") + if isinstance(details, dict): + cached = details.get("cached_tokens") + if isinstance(cached, int): + return cached + cached = usage.get("cached_tokens") + return cached if isinstance(cached, int) else 0 + + def _error_response(self, message: str) -> dict[str, Any]: + return { + "type": "error", + "error": {"type": "api_error", "message": message}, + } diff --git a/src/openenv/core/harness/capture/dialects/base.py b/src/openenv/core/harness/capture/dialects/base.py new file mode 100644 index 000000000..9da1b9ff7 --- /dev/null +++ b/src/openenv/core/harness/capture/dialects/base.py @@ -0,0 +1,125 @@ +"""Base transformer interface with inference-backend request enhancement.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + + +class BaseTransformer(ABC): + """Abstract base class for API transformers. + + Transforms requests from source API format to OpenAI format (for the + inference backend), and transforms responses back to source API format. + """ + + @abstractmethod + def transform_request(self, body: dict[str, Any]) -> dict[str, Any]: + """Transform request body to OpenAI format for the inference backend.""" + pass + + @abstractmethod + def transform_response( + self, + response: dict[str, Any], + original_request: dict[str, Any], + ) -> dict[str, Any]: + """Transform response back to source API format.""" + pass + + @abstractmethod + def transform_stream_chunk( + self, + chunk: dict[str, Any], + original_request: dict[str, Any], + is_first: bool = False, + ) -> dict[str, Any] | list[dict[str, Any]]: + """Transform a streaming chunk to source API format.""" + pass + + def is_streaming_request(self, body: dict[str, Any]) -> bool: + """Check if request is for streaming response.""" + return body.get("stream", False) + + def create_stream_state(self, original_request: dict[str, Any]) -> Any | None: + """Create per-request stream state when chunk transforms need memory.""" + return None + + @staticmethod + def _is_qwen35_model(model_name: str | None) -> bool: + if not model_name: + return False + return "qwen3.5" in model_name.lower() + + @staticmethod + def _content_to_text(content: Any) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + parts: list[str] = [] + for block in content: + if isinstance(block, str): + parts.append(block) + elif isinstance(block, dict): + text = block.get("text") + if isinstance(text, str): + parts.append(text) + return "\n".join(parts) + return str(content) if content else "" + + @classmethod + def _merge_developer_role(cls, request: dict[str, Any]) -> dict[str, Any]: + """Rename 'developer' role to 'system' and merge all system messages into one.""" + messages = request.get("messages") + if not isinstance(messages, list): + return request + + # Rename developer -> system + normalized = [ + {**msg, "role": "system"} + if isinstance(msg, dict) and msg.get("role") == "developer" + else msg + for msg in messages + ] + + # Merge multiple system messages into one at the top + system_parts: list[str] = [] + non_system: list[Any] = [] + for msg in normalized: + if isinstance(msg, dict) and msg.get("role") == "system": + text = cls._content_to_text(msg.get("content", "")) + if text: + system_parts.append(text) + else: + non_system.append(msg) + + if system_parts: + request["messages"] = [ + {"role": "system", "content": "\n\n".join(system_parts)}, + *non_system, + ] + else: + request["messages"] = non_system + return request + + def _normalize_request( + self, + request: dict[str, Any], + model_name: str | None = None, + ) -> dict[str, Any]: + """Normalize the OpenAI request: drop internal keys, merge system roles, + and apply per-model template fixes. Training-signal params (logprobs, + token ids) are added later by the inference engine. + """ + request.pop("_served_model", None) + + request = self._merge_developer_role(request) + + if self._is_qwen35_model(model_name): + # Qwen3.5 outputs tool calls inside thinking; disable thinking. + # https://www.reddit.com/r/LocalLLaMA/comments/1sccqt2/i_think_i_got_solutions_for_qwen_35_tool_call_in/ + chat_template_kwargs = dict(request.get("chat_template_kwargs") or {}) + chat_template_kwargs.setdefault("enable_thinking", False) + request["chat_template_kwargs"] = chat_template_kwargs + + return request diff --git a/src/openenv/core/harness/capture/dialects/google.py b/src/openenv/core/harness/capture/dialects/google.py new file mode 100644 index 000000000..65ef9fbbb --- /dev/null +++ b/src/openenv/core/harness/capture/dialects/google.py @@ -0,0 +1,724 @@ +"""Google Generative AI API transformer with tool-call support.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import Any + +from .base import BaseTransformer +from .images import ( + google_content_parts_to_openai_chat, + google_part_to_openai_chat, + openai_chat_content_to_google_parts, +) +from .reasoning import extract_reasoning_from_gemini_parts, make_signature + + +@dataclass +class _GoogleToolCallState: + call_id: str = "" + name: str = "" + arguments: str = "" + + +class _GoogleStreamState: + """Accumulate streamed OpenAI tool-call deltas into Google response parts.""" + + def __init__(self, transformer: "GoogleTransformer") -> None: + self._transformer = transformer + self._tool_calls: dict[int, _GoogleToolCallState] = {} + self._finish_reason: str | None = None + self._usage: dict[str, Any] | None = None + self._emitted_tool_calls = False + self._emitted_finish_reason = False + + def process_chunk( + self, + chunk: dict[str, Any], + *, + is_first: bool = False, + ) -> list[dict[str, Any]]: + del is_first + + choices = chunk.get("choices", []) + if not choices: + usage = chunk.get("usage") + if isinstance(usage, dict): + self._usage = usage + return [] + + choice = choices[0] + delta = choice.get("delta", {}) or {} + usage = chunk.get("usage") + if isinstance(usage, dict): + self._usage = usage + + parts: list[dict[str, Any]] = [] + reasoning = delta.get("reasoning_content") + if isinstance(reasoning, str) and reasoning: + parts.append( + { + "thought": True, + "text": reasoning, + "thoughtSignature": make_signature(reasoning), + } + ) + content = delta.get("content") + if isinstance(content, str) and content: + parts.append({"text": content}) + + for tool_call in self._transformer._normalize_tool_call_deltas( + delta.get("tool_calls") + ): + tool_index = tool_call.get("index", 0) + if not isinstance(tool_index, int): + tool_index = 0 + state = self._tool_calls.setdefault(tool_index, _GoogleToolCallState()) + + if isinstance(tool_call.get("id"), str) and tool_call["id"]: + state.call_id = tool_call["id"] + + function = tool_call.get("function", {}) + if isinstance(function, dict): + name = function.get("name") + if isinstance(name, str) and name: + state.name += name + + arguments = function.get("arguments") + if isinstance(arguments, str) and arguments: + state.arguments += arguments + elif arguments not in (None, ""): + state.arguments += json.dumps(arguments) + + finish_reason = choice.get("finish_reason") + current_finish_reason = ( + finish_reason if isinstance(finish_reason, str) and finish_reason else None + ) + if current_finish_reason: + self._finish_reason = current_finish_reason + + if parts: + if current_finish_reason: + self._emitted_finish_reason = True + return [ + self._transformer._build_stream_response( + parts, + finish_reason=current_finish_reason, + usage=self._usage, + ) + ] + + if self._finish_reason and self._tool_calls and not self._emitted_tool_calls: + self._emitted_tool_calls = True + tool_parts = [ + self._transformer._tool_call_part( + name=state.name, + arguments=state.arguments, + call_id=state.call_id, + ) + for _, state in sorted(self._tool_calls.items()) + if state.name + ] + if tool_parts: + if self._finish_reason: + self._emitted_finish_reason = True + return [ + self._transformer._build_stream_response( + tool_parts, + finish_reason=self._finish_reason, + usage=self._usage, + ) + ] + + if current_finish_reason and not self._emitted_finish_reason: + self._emitted_finish_reason = True + return [ + self._transformer._build_stream_response( + [], + finish_reason=current_finish_reason, + usage=self._usage, + ) + ] + + return [] + + def finalize(self) -> list[dict[str, Any]]: + if self._tool_calls and not self._emitted_tool_calls: + self._emitted_tool_calls = True + tool_parts = [ + self._transformer._tool_call_part( + name=state.name, + arguments=state.arguments, + call_id=state.call_id, + ) + for _, state in sorted(self._tool_calls.items()) + if state.name + ] + if tool_parts: + if self._finish_reason: + self._emitted_finish_reason = True + return [ + self._transformer._build_stream_response( + tool_parts, + finish_reason=self._finish_reason, + usage=self._usage, + ) + ] + return [] + + +class GoogleTransformer(BaseTransformer): + """Transform between Google Generative AI and OpenAI API formats.""" + + ROLE_MAP = { + "user": "user", + "model": "assistant", + "system": "system", + "developer": "system", + } + FINISH_REASON_MAP_REVERSE = { + "stop": "STOP", + "length": "MAX_TOKENS", + "content_filter": "SAFETY", + "tool_calls": "STOP", + "stop_sequence": "STOP", + } + + def transform_request(self, body: dict[str, Any]) -> dict[str, Any]: + messages: list[dict[str, Any]] = [] + config = body.get("config") + config_section = config if isinstance(config, dict) else {} + + system_instruction = ( + body.get("systemInstruction") + or body.get("system_instruction") + or config_section.get("systemInstruction") + or config_section.get("system_instruction") + ) + system_text = self._extract_system_instruction_text(system_instruction) + if system_text: + messages.append({"role": "system", "content": system_text}) + + for content in body.get("contents", []): + messages.extend(self._convert_content_to_messages(content)) + + result: dict[str, Any] = {"messages": messages} + if "model" in body: + result["model"] = body["model"] + + gen_config: dict[str, Any] = {} + for source in ( + config_section.get("generationConfig"), + body.get("generationConfig"), + ): + if isinstance(source, dict): + gen_config.update(source) + if not gen_config and config_section: + gen_config = config_section + + if "maxOutputTokens" in gen_config: + result["max_tokens"] = gen_config["maxOutputTokens"] + if "temperature" in gen_config: + result["temperature"] = gen_config["temperature"] + if "topP" in gen_config: + result["top_p"] = gen_config["topP"] + if "topK" in gen_config: + result["top_k"] = gen_config["topK"] + if "stopSequences" in gen_config: + result["stop"] = gen_config["stopSequences"] + if "candidateCount" in gen_config: + result["n"] = gen_config["candidateCount"] + if "presencePenalty" in gen_config: + result["presence_penalty"] = gen_config["presencePenalty"] + if "frequencyPenalty" in gen_config: + result["frequency_penalty"] = gen_config["frequencyPenalty"] + if "seed" in gen_config: + result["seed"] = gen_config["seed"] + if "logprobs" in gen_config: + result["top_logprobs"] = gen_config["logprobs"] + + response_format = self._convert_response_format(gen_config) + if response_format is not None: + result["response_format"] = response_format + + if body.get("_streaming", False): + result["stream"] = True + + # Gemini `thinkingConfig.includeThoughts: true` → enable_thinking. + thinking_cfg = gen_config.get("thinkingConfig") or config_section.get( + "thinkingConfig" + ) + if isinstance(thinking_cfg, dict) and thinking_cfg.get("includeThoughts"): + chat_template_kwargs = dict(result.get("chat_template_kwargs") or {}) + chat_template_kwargs["enable_thinking"] = True + result["chat_template_kwargs"] = chat_template_kwargs + + # SGLang rejects tool_choice without a non-empty tools list; bind + # the pair so one can't be forwarded without the other. + tools = self._convert_tools( + body.get("tools") or config_section.get("tools") or [] + ) + if tools: + result["tools"] = tools + tool_choice = self._convert_tool_choice( + body.get("toolConfig") or config_section.get("toolConfig") or {} + ) + if tool_choice is not None: + result["tool_choice"] = tool_choice + + return self._normalize_request( + result, + body.get("_served_model"), + ) + + def _convert_response_format( + self, gen_config: dict[str, Any] + ) -> dict[str, Any] | None: + response_format_cfg = gen_config.get("responseFormat") + if isinstance(response_format_cfg, dict): + text_format = response_format_cfg.get("text") + if isinstance(text_format, dict): + mime_type = text_format.get("mimeType") + if self._is_json_mime_type(mime_type): + schema = text_format.get("schema") + return self._chat_response_format_from_schema(schema) + + mime_type = gen_config.get("responseMimeType") + if not self._is_json_mime_type(mime_type): + return None + + schema = ( + gen_config.get("responseJsonSchema") + or gen_config.get("_responseJsonSchema") + or gen_config.get("responseSchema") + ) + return self._chat_response_format_from_schema(schema) + + @staticmethod + def _is_json_mime_type(value: Any) -> bool: + if not isinstance(value, str): + return False + return ( + value.lower() == "application/json" or value.upper() == "APPLICATION_JSON" + ) + + def _chat_response_format_from_schema(self, schema: Any) -> dict[str, Any]: + if isinstance(schema, dict) and schema: + return { + "type": "json_schema", + "json_schema": { + "name": "google_response", + "schema": self._normalize_google_schema(schema), + }, + } + return {"type": "json_object"} + + def _normalize_google_schema(self, schema: dict[str, Any]) -> dict[str, Any]: + normalized: dict[str, Any] = {} + for key, value in schema.items(): + if key == "type" and isinstance(value, str): + normalized[key] = value.lower() + elif key == "properties" and isinstance(value, dict): + normalized[key] = { + prop_name: self._normalize_google_schema(prop_schema) + if isinstance(prop_schema, dict) + else prop_schema + for prop_name, prop_schema in value.items() + } + elif key == "items" and isinstance(value, dict): + normalized[key] = self._normalize_google_schema(value) + elif key in {"anyOf", "oneOf", "allOf"} and isinstance(value, list): + normalized[key] = [ + self._normalize_google_schema(item) + if isinstance(item, dict) + else item + for item in value + ] + else: + normalized[key] = value + return normalized + + def transform_response( + self, + response: dict[str, Any], + original_request: dict[str, Any], + ) -> dict[str, Any]: + del original_request + + candidates = [] + for i, choice in enumerate(response.get("choices", [])): + message = choice.get("message", {}) + parts = [] + reasoning = message.get("reasoning_content") + if isinstance(reasoning, str) and reasoning: + parts.append( + { + "thought": True, + "text": reasoning, + "thoughtSignature": make_signature(reasoning), + } + ) + content = message.get("content") + if content or isinstance(content, list): + parts.extend(openai_chat_content_to_google_parts(content)) + parts.extend(self._tool_call_parts_from_message(message)) + + finish_reason = choice.get("finish_reason", "stop") + google_finish = self.FINISH_REASON_MAP_REVERSE.get(finish_reason, "STOP") + + candidates.append( + { + "content": {"parts": parts, "role": "model"}, + "finishReason": google_finish, + "index": i, + "safetyRatings": [], + } + ) + + usage = response.get("usage", {}) + usage_metadata = { + "promptTokenCount": usage.get("prompt_tokens", 0), + "candidatesTokenCount": usage.get("completion_tokens", 0), + "totalTokenCount": usage.get("total_tokens", 0), + } + cached_tokens = self._cached_prompt_tokens(usage) + if cached_tokens: + usage_metadata["cachedContentTokenCount"] = cached_tokens + result = { + "candidates": candidates, + "usageMetadata": usage_metadata, + } + function_calls = self._response_function_calls(candidates) + if function_calls: + result["functionCalls"] = function_calls + return result + + def transform_stream_chunk( + self, + chunk: dict[str, Any], + original_request: dict[str, Any], + is_first: bool = False, + ) -> dict[str, Any]: + del original_request, is_first + + candidates = [] + for choice in chunk.get("choices", []): + delta = choice.get("delta", {}) or {} + parts = [] + reasoning_chunk = delta.get("reasoning_content") + if isinstance(reasoning_chunk, str) and reasoning_chunk: + parts.append( + { + "thought": True, + "text": reasoning_chunk, + "thoughtSignature": make_signature(reasoning_chunk), + } + ) + content = delta.get("content") + if content: + parts.append({"text": content}) + for tool_call in self._normalize_tool_call_deltas(delta.get("tool_calls")): + function = tool_call.get("function", {}) + parts.append( + self._tool_call_part( + name=str(function.get("name") or ""), + arguments=function.get("arguments", ""), + call_id=str(tool_call.get("id") or ""), + ) + ) + + candidate: dict[str, Any] = { + "content": {"parts": parts, "role": "model"}, + "index": choice.get("index", 0), + } + finish_reason = choice.get("finish_reason") + if finish_reason: + candidate["finishReason"] = self.FINISH_REASON_MAP_REVERSE.get( + finish_reason, "STOP" + ) + candidates.append(candidate) + + result: dict[str, Any] = {"candidates": candidates} + usage = chunk.get("usage") + if usage: + result["usageMetadata"] = { + "promptTokenCount": usage.get("prompt_tokens", 0), + "candidatesTokenCount": usage.get("completion_tokens", 0), + "totalTokenCount": usage.get("total_tokens", 0), + } + function_calls = self._response_function_calls(candidates) + if function_calls: + result["functionCalls"] = function_calls + return result + + def create_stream_state( + self, original_request: dict[str, Any] + ) -> _GoogleStreamState: + del original_request + return _GoogleStreamState(self) + + def is_streaming_request(self, body: dict[str, Any]) -> bool: + return body.get("_streaming", False) + + def _convert_content_to_messages(self, content: Any) -> list[dict[str, Any]]: + if not isinstance(content, dict): + return [] + + parts = content.get("parts", []) + role = content.get("role", "user") + openai_role = self.ROLE_MAP.get(role, "user") + messages: list[dict[str, Any]] = [] + user_parts: list[dict[str, Any]] = [] + tool_calls: list[dict[str, Any]] = [] + tool_messages: list[dict[str, Any]] = [] + + for part in parts: + if isinstance(part, str): + user_parts.append({"type": "text", "text": part}) + continue + if not isinstance(part, dict): + continue + if "text" in part and isinstance(part["text"], str): + user_parts.append({"type": "text", "text": part["text"]}) + continue + image_part = google_part_to_openai_chat(part) + if image_part: + user_parts.append(image_part) + continue + if "functionCall" in part and isinstance(part["functionCall"], dict): + function_call = part["functionCall"] + tool_calls.append( + { + "id": function_call.get("id") + or function_call.get("call_id") + or "", + "type": "function", + "function": { + "name": function_call.get("name", ""), + "arguments": json.dumps(function_call.get("args", {})), + }, + } + ) + continue + if "functionResponse" in part and isinstance( + part["functionResponse"], dict + ): + function_response = part["functionResponse"] + tool_messages.append( + { + "role": "tool", + "tool_call_id": function_response.get("id") + or function_response.get("call_id") + or function_response.get("name", ""), + "content": json.dumps(function_response.get("response", {})), + } + ) + + if openai_role == "assistant": + message_content = google_content_parts_to_openai_chat(parts) + reasoning_text = extract_reasoning_from_gemini_parts(parts) + if message_content or tool_calls or reasoning_text: + assistant_message: dict[str, Any] = { + "role": "assistant", + "content": message_content, + } + if reasoning_text: + assistant_message["reasoning_content"] = reasoning_text + if tool_calls: + assistant_message["tool_calls"] = tool_calls + messages.append(assistant_message) + elif openai_role == "system": + system_text = self._extract_text_from_parts(parts) + if system_text: + messages.append({"role": "system", "content": system_text}) + else: + if user_parts: + messages.append( + { + "role": "user", + "content": google_content_parts_to_openai_chat(parts), + } + ) + messages.extend(tool_messages) + + return messages + + def _convert_tools(self, google_tools: list[Any]) -> list[dict[str, Any]]: + openai_tools: list[dict[str, Any]] = [] + for tool in google_tools: + if not isinstance(tool, dict): + continue + declarations = tool.get("functionDeclarations") or tool.get( + "function_declarations" + ) + if not isinstance(declarations, list): + continue + for declaration in declarations: + if not isinstance(declaration, dict): + continue + name = declaration.get("name") + if not isinstance(name, str) or not name: + continue + function: dict[str, Any] = {"name": name} + description = declaration.get("description") + if isinstance(description, str) and description: + function["description"] = description + parameters = ( + declaration.get("parameters") + or declaration.get("parametersJsonSchema") + or declaration.get("parameters_json_schema") + ) + if not isinstance(parameters, dict): + parameters = {"type": "object", "properties": {}} + function["parameters"] = parameters + openai_tools.append({"type": "function", "function": function}) + return openai_tools + + def _convert_tool_choice(self, tool_config: Any) -> Any | None: + if not isinstance(tool_config, dict): + return None + function_calling_config = tool_config.get( + "functionCallingConfig" + ) or tool_config.get("function_calling_config") + if not isinstance(function_calling_config, dict): + return None + + mode = str(function_calling_config.get("mode", "")).upper() + allowed_names = function_calling_config.get( + "allowedFunctionNames" + ) or function_calling_config.get("allowed_function_names") + if mode == "NONE": + return "none" + if mode in {"ANY", "VALIDATED"}: + if isinstance(allowed_names, list) and len(allowed_names) == 1: + allowed_name = allowed_names[0] + if isinstance(allowed_name, str) and allowed_name: + return {"type": "function", "function": {"name": allowed_name}} + return "required" + return None + + def _tool_call_parts_from_message( + self, message: dict[str, Any] + ) -> list[dict[str, Any]]: + parts: list[dict[str, Any]] = [] + for tool_call in message.get("tool_calls", []) or []: + if not isinstance(tool_call, dict): + continue + function = tool_call.get("function", {}) + if not isinstance(function, dict): + continue + name = function.get("name") + if not isinstance(name, str) or not name: + continue + parts.append( + self._tool_call_part( + name=name, + arguments=function.get("arguments", ""), + call_id=str(tool_call.get("id") or ""), + ) + ) + return parts + + def _tool_call_part( + self, *, name: str, arguments: Any, call_id: str + ) -> dict[str, Any]: + function_call: dict[str, Any] = { + "name": name, + "args": self._parse_arguments(arguments), + } + if call_id: + function_call["id"] = call_id + return {"functionCall": function_call} + + def _parse_arguments(self, arguments: Any) -> Any: + if isinstance(arguments, (dict, list, int, float, bool)) or arguments is None: + return arguments if arguments is not None else {} + if not isinstance(arguments, str): + return {"value": arguments} + stripped = arguments.strip() + if not stripped: + return {} + try: + return json.loads(stripped) + except json.JSONDecodeError: + return {"raw": stripped} + + def _normalize_tool_call_deltas(self, tool_calls: Any) -> list[dict[str, Any]]: + if isinstance(tool_calls, list): + return [ + tool_call for tool_call in tool_calls if isinstance(tool_call, dict) + ] + if isinstance(tool_calls, dict): + return [tool_calls] + return [] + + def _build_stream_response( + self, + parts: list[dict[str, Any]], + *, + finish_reason: str | None = None, + usage: dict[str, Any] | None = None, + ) -> dict[str, Any]: + candidate: dict[str, Any] = { + "content": {"parts": parts, "role": "model"}, + "index": 0, + } + if finish_reason: + candidate["finishReason"] = self.FINISH_REASON_MAP_REVERSE.get( + finish_reason, "STOP" + ) + result: dict[str, Any] = {"candidates": [candidate]} + if usage: + result["usageMetadata"] = { + "promptTokenCount": usage.get("prompt_tokens", 0), + "candidatesTokenCount": usage.get("completion_tokens", 0), + "totalTokenCount": usage.get("total_tokens", 0), + } + function_calls = self._response_function_calls(result["candidates"]) + if function_calls: + result["functionCalls"] = function_calls + return result + + def _response_function_calls( + self, candidates: list[dict[str, Any]] + ) -> list[dict[str, Any]]: + function_calls: list[dict[str, Any]] = [] + for candidate in candidates: + content = candidate.get("content", {}) + if not isinstance(content, dict): + continue + for part in content.get("parts", []) or []: + if not isinstance(part, dict): + continue + function_call = part.get("functionCall") + if isinstance(function_call, dict): + function_calls.append(function_call) + return function_calls + + def _extract_text_from_parts(self, parts: list) -> str: + texts = [] + for part in parts: + if isinstance(part, dict) and "text" in part: + texts.append(part["text"]) + elif isinstance(part, str): + texts.append(part) + return "\n".join(texts) + + def _extract_system_instruction_text(self, system_instruction: Any) -> str: + if isinstance(system_instruction, str): + return system_instruction + if isinstance(system_instruction, dict): + return self._extract_text_from_parts(system_instruction.get("parts", [])) + if isinstance(system_instruction, list): + return self._extract_text_from_parts(system_instruction) + return "" + + def _cached_prompt_tokens(self, usage: dict[str, Any]) -> int: + details = usage.get("prompt_tokens_details") + if isinstance(details, dict): + cached = details.get("cached_tokens") + if isinstance(cached, int): + return cached + cached = usage.get("cached_tokens") + return cached if isinstance(cached, int) else 0 diff --git a/src/openenv/core/harness/capture/dialects/images.py b/src/openenv/core/harness/capture/dialects/images.py new file mode 100644 index 000000000..16fdeb1e5 --- /dev/null +++ b/src/openenv/core/harness/capture/dialects/images.py @@ -0,0 +1,335 @@ +"""Image content conversion helpers for gateway API transformers.""" + +from __future__ import annotations + +import re +from typing import Any + +_DATA_URL_RE = re.compile( + r"^data:(?P[^;,]+);base64,(?P.*)$", re.DOTALL +) + + +def is_image_mime_type(mime_type: Any) -> bool: + return isinstance(mime_type, str) and mime_type.lower().startswith("image/") + + +def make_data_url(mime_type: str, data: str) -> str: + if data.startswith("data:"): + return data + return f"data:{mime_type};base64,{data}" + + +def parse_data_url(url: str) -> tuple[str, str] | None: + match = _DATA_URL_RE.match(url) + if not match: + return None + mime_type = match.group("mime_type") + if not is_image_mime_type(mime_type): + return None + return mime_type, match.group("data") + + +# OpenAI's image_url.detail only accepts these values; vLLM rejects anything +# else. Harnesses send their own (e.g. codex's "original"), so drop unknowns +# rather than forward a value that 400s the whole image request. +_VALID_IMAGE_DETAILS = frozenset({"auto", "low", "high"}) + + +def openai_image_url_block(url: str, *, detail: Any = None) -> dict[str, Any]: + image_url: dict[str, Any] = {"url": url} + if isinstance(detail, str) and detail in _VALID_IMAGE_DETAILS: + image_url["detail"] = detail + return {"type": "image_url", "image_url": image_url} + + +def openai_text_block(text: str) -> dict[str, Any]: + return {"type": "text", "text": text} + + +def openai_image_url(block: dict[str, Any]) -> str | None: + image_url = block.get("image_url") + if isinstance(image_url, str) and image_url: + return image_url + if isinstance(image_url, dict): + url = image_url.get("url") + if isinstance(url, str) and url: + return url + return None + + +def openai_image_detail(block: dict[str, Any]) -> str | None: + image_url = block.get("image_url") + if isinstance(image_url, dict): + detail = image_url.get("detail") + if isinstance(detail, str) and detail: + return detail + detail = block.get("detail") + if isinstance(detail, str) and detail: + return detail + return None + + +def openai_content_from_text_and_images( + parts: list[dict[str, Any]], + *, + text_separator: str = "\n", +) -> str | list[dict[str, Any]]: + has_image = any(part.get("type") == "image_url" for part in parts) + if has_image: + return parts + return text_separator.join( + part.get("text", "") + for part in parts + if part.get("type") == "text" and isinstance(part.get("text"), str) + ) + + +def openai_responses_input_content_to_chat(content: Any) -> str | list[dict[str, Any]]: + if isinstance(content, str): + return content + if not isinstance(content, list): + return str(content) if content else "" + + parts: list[dict[str, Any]] = [] + for block in content: + if isinstance(block, str): + parts.append(openai_text_block(block)) + continue + if not isinstance(block, dict): + continue + + block_type = block.get("type") + if block_type in ("input_text", "output_text", "text"): + text = block.get("text") + if isinstance(text, str): + parts.append(openai_text_block(text)) + continue + if block_type in ("input_image", "image_url"): + url = _responses_image_url(block) + if url: + parts.append(openai_image_url_block(url, detail=block.get("detail"))) + + return openai_content_from_text_and_images(parts) + + +def _responses_image_url(block: dict[str, Any]) -> str | None: + image_url = block.get("image_url") + if isinstance(image_url, str) and image_url: + return image_url + if isinstance(image_url, dict): + url = image_url.get("url") + if isinstance(url, str) and url: + return url + return None + + +def anthropic_content_to_openai_chat(content: Any) -> str | list[dict[str, Any]]: + if isinstance(content, str): + return content + if not isinstance(content, list): + return str(content) if content else "" + + parts: list[dict[str, Any]] = [] + for block in content: + if isinstance(block, str): + parts.append(openai_text_block(block)) + continue + if not isinstance(block, dict): + continue + + block_type = block.get("type") + if block_type == "text": + text = block.get("text") + if isinstance(text, str): + parts.append(openai_text_block(text)) + continue + if block_type == "image": + image = anthropic_image_to_openai_chat(block) + if image: + parts.append(image) + continue + if block_type == "document": + text = anthropic_document_to_text(block) + if text: + parts.append(openai_text_block(text)) + + return openai_content_from_text_and_images(parts) + + +def anthropic_document_to_text(block: dict[str, Any]) -> str: + """Extract text from an Anthropic `document` block. + + Handles `source.type == "text"` and `source.type == "content"`. Base64 + PDFs are dropped — SGLang can't render binary docs through the chat + template. + """ + source = block.get("source") + if not isinstance(source, dict): + return "" + source_type = source.get("type") + if source_type == "text": + data = source.get("data") + return data if isinstance(data, str) else "" + if source_type == "content": + inner = source.get("content") + if isinstance(inner, list): + pieces: list[str] = [] + for inner_block in inner: + if isinstance(inner_block, dict) and inner_block.get("type") == "text": + text = inner_block.get("text") + if isinstance(text, str): + pieces.append(text) + return "\n".join(pieces) + return "" + + +def anthropic_image_to_openai_chat(block: dict[str, Any]) -> dict[str, Any] | None: + source = block.get("source") + if not isinstance(source, dict): + return None + + source_type = source.get("type") + if source_type == "base64": + mime_type = source.get("media_type") or source.get("mediaType") + data = source.get("data") + if is_image_mime_type(mime_type) and isinstance(data, str) and data: + return openai_image_url_block(make_data_url(mime_type, data)) + if source_type == "url": + url = source.get("url") + if isinstance(url, str) and url: + return openai_image_url_block(url) + return None + + +def google_content_parts_to_openai_chat(parts: Any) -> str | list[dict[str, Any]]: + if not isinstance(parts, list): + return "" + + openai_parts: list[dict[str, Any]] = [] + for part in parts: + if isinstance(part, str): + openai_parts.append(openai_text_block(part)) + continue + if not isinstance(part, dict): + continue + + # Thought parts are reasoning_content, not user-visible content. + if part.get("thought") is True: + continue + + text = part.get("text") + if isinstance(text, str): + openai_parts.append(openai_text_block(text)) + continue + + image = google_part_to_openai_chat(part) + if image: + openai_parts.append(image) + + return openai_content_from_text_and_images(openai_parts) + + +def google_part_to_openai_chat(part: dict[str, Any]) -> dict[str, Any] | None: + inline_data = part.get("inline_data") or part.get("inlineData") + if isinstance(inline_data, dict): + mime_type = inline_data.get("mime_type") or inline_data.get("mimeType") + data = inline_data.get("data") + if is_image_mime_type(mime_type) and isinstance(data, str) and data: + return openai_image_url_block(make_data_url(mime_type, data)) + + file_data = part.get("file_data") or part.get("fileData") + if isinstance(file_data, dict): + mime_type = file_data.get("mime_type") or file_data.get("mimeType") + uri = file_data.get("file_uri") or file_data.get("fileUri") + if is_image_mime_type(mime_type) and isinstance(uri, str) and uri: + return openai_image_url_block(uri) + + return None + + +def openai_chat_content_to_anthropic_blocks(content: Any) -> list[dict[str, Any]]: + if isinstance(content, str): + return [{"type": "text", "text": content}] + if not isinstance(content, list): + return [{"type": "text", "text": str(content) if content else ""}] + + blocks: list[dict[str, Any]] = [] + for part in content: + if isinstance(part, str): + blocks.append({"type": "text", "text": part}) + continue + if not isinstance(part, dict): + continue + part_type = part.get("type") + if part_type == "text": + text = part.get("text") + if isinstance(text, str): + blocks.append({"type": "text", "text": text}) + continue + if part_type == "image_url": + image = openai_chat_image_to_anthropic(part) + if image: + blocks.append(image) + + return blocks or [{"type": "text", "text": ""}] + + +def openai_chat_image_to_anthropic(part: dict[str, Any]) -> dict[str, Any] | None: + url = openai_image_url(part) + if not url: + return None + + parsed = parse_data_url(url) + if parsed: + mime_type, data = parsed + return { + "type": "image", + "source": { + "type": "base64", + "media_type": mime_type, + "data": data, + }, + } + + return {"type": "image", "source": {"type": "url", "url": url}} + + +def openai_chat_content_to_google_parts(content: Any) -> list[dict[str, Any]]: + if isinstance(content, str): + return [{"text": content}] + if not isinstance(content, list): + return [{"text": str(content) if content else ""}] + + parts: list[dict[str, Any]] = [] + for part in content: + if isinstance(part, str): + parts.append({"text": part}) + continue + if not isinstance(part, dict): + continue + part_type = part.get("type") + if part_type == "text": + text = part.get("text") + if isinstance(text, str): + parts.append({"text": text}) + continue + if part_type == "image_url": + image = openai_chat_image_to_google(part) + if image: + parts.append(image) + + return parts or [{"text": ""}] + + +def openai_chat_image_to_google(part: dict[str, Any]) -> dict[str, Any] | None: + url = openai_image_url(part) + if not url: + return None + + parsed = parse_data_url(url) + if parsed: + mime_type, data = parsed + return {"inline_data": {"mime_type": mime_type, "data": data}} + + return {"file_data": {"mime_type": "image/jpeg", "file_uri": url}} diff --git a/src/openenv/core/harness/capture/dialects/openai_chat.py b/src/openenv/core/harness/capture/dialects/openai_chat.py new file mode 100644 index 000000000..af9297b24 --- /dev/null +++ b/src/openenv/core/harness/capture/dialects/openai_chat.py @@ -0,0 +1,41 @@ +"""OpenAI Chat Completions transformer with SGLang training enhancements.""" + +from __future__ import annotations + +from typing import Any + +from .base import BaseTransformer + + +class OpenAIChatTransformer(BaseTransformer): + """Transform OpenAI Chat requests (passthrough + training params).""" + + def transform_request(self, body: dict[str, Any]) -> dict[str, Any]: + result = body.copy() + if "max_tokens" not in result and "max_completion_tokens" in result: + result["max_tokens"] = result["max_completion_tokens"] + return self._normalize_request( + result, + body.get("_served_model"), + ) + + def transform_response( + self, + response: dict[str, Any], + original_request: dict[str, Any], + ) -> dict[str, Any]: + result = response.copy() + if "model" in original_request: + result["model"] = original_request["model"] + return result + + def transform_stream_chunk( + self, + chunk: dict[str, Any], + original_request: dict[str, Any], + is_first: bool = False, + ) -> dict[str, Any]: + result = chunk.copy() + if "model" in original_request: + result["model"] = original_request["model"] + return result diff --git a/src/openenv/core/harness/capture/dialects/openai_responses.py b/src/openenv/core/harness/capture/dialects/openai_responses.py new file mode 100644 index 000000000..0c2b0fa94 --- /dev/null +++ b/src/openenv/core/harness/capture/dialects/openai_responses.py @@ -0,0 +1,1090 @@ +"""OpenAI Responses API transformer. + +Transforms between OpenAI Responses API (Codex CLI) and OpenAI Chat Completions. +Aligned with agent-harness-proxy/src/harness_proxy/transform/openai_responses.py. +""" + +from __future__ import annotations + +import json +import time +import uuid +from dataclasses import dataclass +from typing import Any, Optional + +from .base import BaseTransformer +from .images import openai_responses_input_content_to_chat +from .reasoning import encrypt_reasoning, extract_reasoning_from_responses_item + + +@dataclass +class _ResponsesToolCallState: + name: str = "" + call_id: str = "" + arguments: str = "" + started: bool = False + fc_id: str = "" + + +class ResponsesStreamState: + """Per-request Responses API streaming state.""" + + def __init__(self, model: str): + self.response_id = f"resp_{uuid.uuid4().hex[:24]}" + self.model = model + self.text_started = False + self.text_content = "" + self.message_output_index = 0 + self.output_index_offset = 0 + self.tool_calls: dict[int, _ResponsesToolCallState] = {} + self.usage = {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0} + self.reasoning_started = False + self.reasoning_closed = False + self.reasoning_content = "" + self.reasoning_id = "" + self.completed = False + + def process_chunk( + self, chunk: dict[str, Any], is_first: bool = False + ) -> list[dict[str, Any]]: + events: list[dict[str, Any]] = [] + + if is_first: + events.append( + { + "type": "response.created", + "response": { + "id": self.response_id, + "object": "response", + "status": "in_progress", + "model": self.model, + "output": [], + "usage": self.usage.copy(), + }, + } + ) + + usage = chunk.get("usage") + if isinstance(usage, dict): + self.usage["input_tokens"] = usage.get( + "prompt_tokens", self.usage["input_tokens"] + ) + self.usage["output_tokens"] = usage.get( + "completion_tokens", self.usage["output_tokens"] + ) + self.usage["total_tokens"] = usage.get( + "total_tokens", self.usage["total_tokens"] + ) + + choices = chunk.get("choices", []) + if not choices: + return events + + choice = choices[0] + delta = choice.get("delta", {}) or {} + + # Reasoning item must come first so the harness sees the chain-of-thought + # before any output_text or function_call items. + reasoning_delta = delta.get("reasoning_content") + if isinstance(reasoning_delta, str) and reasoning_delta: + if not self.reasoning_started: + self.reasoning_started = True + self.reasoning_id = f"rs_{uuid.uuid4().hex[:24]}" + events.append( + { + "type": "response.output_item.added", + "output_index": 0, + "item": { + "type": "reasoning", + "id": self.reasoning_id, + "summary": [], + "content": [], + "status": "in_progress", + }, + } + ) + events.append( + { + "type": "response.reasoning_summary_part.added", + "item_id": self.reasoning_id, + "output_index": 0, + "summary_index": 0, + "part": {"type": "summary_text", "text": ""}, + } + ) + self.reasoning_content += reasoning_delta + events.append( + { + "type": "response.reasoning_summary_text.delta", + "item_id": self.reasoning_id, + "output_index": 0, + "summary_index": 0, + "delta": reasoning_delta, + } + ) + + content = delta.get("content") + if content: + # Close reasoning before opening message. + events.extend(self._close_reasoning()) + if not self.text_started: + self.text_started = True + self.message_output_index = 1 if self.reasoning_started else 0 + self.output_index_offset = self.message_output_index + 1 + message_id = f"msg_{uuid.uuid4().hex[:24]}" + events.append( + { + "type": "response.output_item.added", + "output_index": self.message_output_index, + "item": { + "type": "message", + "id": message_id, + "role": "assistant", + "status": "in_progress", + "content": [], + }, + } + ) + events.append( + { + "type": "response.content_part.added", + "output_index": self.message_output_index, + "content_index": 0, + "part": {"type": "output_text", "text": ""}, + } + ) + + self.text_content += content + events.append( + { + "type": "response.output_text.delta", + "output_index": self.message_output_index, + "content_index": 0, + "delta": content, + } + ) + + tool_calls_delta = delta.get("tool_calls") or [] + if not isinstance(tool_calls_delta, list): + tool_calls_delta = [tool_calls_delta] + + if tool_calls_delta and self.reasoning_started and not self.reasoning_closed: + events.extend(self._close_reasoning()) + if not self.text_started: + # No text — tools come immediately after reasoning. + self.output_index_offset = 1 + + for tool_call in tool_calls_delta: + if not isinstance(tool_call, dict): + continue + + tool_index = tool_call.get("index", 0) + if not isinstance(tool_index, int): + tool_index = 0 + + tool_state = self.tool_calls.get(tool_index) + if tool_state is None: + tool_state = _ResponsesToolCallState( + call_id=tool_call.get("id", f"call_{uuid.uuid4().hex[:24]}"), + ) + self.tool_calls[tool_index] = tool_state + elif tool_call.get("id"): + tool_state.call_id = tool_call["id"] + + function = tool_call.get("function", {}) + name = function.get("name") + if isinstance(name, str) and name: + tool_state.name += name + + arguments = function.get("arguments") + arguments_str = "" + if isinstance(arguments, str) and arguments: + arguments_str = arguments + elif arguments not in (None, ""): + arguments_str = json.dumps(arguments) + if arguments_str: + tool_state.arguments += arguments_str + + output_index = self.output_index_offset + tool_index + if tool_state.name and not tool_state.started: + tool_state.started = True + tool_state.fc_id = f"fc_{uuid.uuid4().hex[:24]}" + events.append( + { + "type": "response.output_item.added", + "output_index": output_index, + "item": { + "type": "function_call", + "id": tool_state.fc_id, + "call_id": tool_state.call_id, + "name": tool_state.name, + "arguments": "", + "status": "in_progress", + }, + } + ) + + if tool_state.arguments: + events.append( + { + "type": "response.function_call_arguments.delta", + "output_index": output_index, + "delta": tool_state.arguments, + } + ) + elif tool_state.started and arguments_str: + events.append( + { + "type": "response.function_call_arguments.delta", + "output_index": output_index, + "delta": arguments_str, + } + ) + + return events + + def finalize(self) -> list[dict[str, Any]]: + if self.completed: + return [] + + events: list[dict[str, Any]] = [] + + # Close reasoning if it never got closed by content/tools. + events.extend(self._close_reasoning()) + + if self.text_started: + events.append( + { + "type": "response.content_part.done", + "output_index": self.message_output_index, + "content_index": 0, + "part": {"type": "output_text", "text": self.text_content}, + } + ) + events.append( + { + "type": "response.output_item.done", + "output_index": self.message_output_index, + "item": { + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": self.text_content}], + }, + } + ) + + for tool_index in sorted(self.tool_calls): + tool_state = self.tool_calls[tool_index] + if not tool_state.started: + continue + + output_index = self.output_index_offset + tool_index + events.append( + { + "type": "response.function_call_arguments.done", + "output_index": output_index, + "arguments": tool_state.arguments, + } + ) + events.append( + { + "type": "response.output_item.done", + "output_index": output_index, + "item": { + "type": "function_call", + "id": tool_state.fc_id or f"fc_{uuid.uuid4().hex[:24]}", + "call_id": tool_state.call_id, + "name": tool_state.name, + "arguments": tool_state.arguments, + "status": "completed", + }, + } + ) + + output: list[dict[str, Any]] = [] + if self.reasoning_started: + output.append( + { + "type": "reasoning", + "id": self.reasoning_id, + "summary": [ + {"type": "summary_text", "text": self.reasoning_content} + ], + "content": [ + {"type": "reasoning_text", "text": self.reasoning_content} + ], + "encrypted_content": encrypt_reasoning(self.reasoning_content), + "status": "completed", + } + ) + if self.text_started: + output.append( + { + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": self.text_content}], + } + ) + for tool_index in sorted(self.tool_calls): + tool_state = self.tool_calls[tool_index] + if tool_state.started: + output.append( + { + "type": "function_call", + "id": tool_state.fc_id or f"fc_{uuid.uuid4().hex[:24]}", + "call_id": tool_state.call_id, + "name": tool_state.name, + "arguments": tool_state.arguments, + "status": "completed", + } + ) + + events.append( + { + "type": "response.completed", + "response": { + "id": self.response_id, + "object": "response", + "created_at": int(time.time()), + "status": "completed", + "model": self.model, + "output": output, + "usage": self.usage.copy(), + }, + } + ) + self.completed = True + return events + + def _close_reasoning(self) -> list[dict[str, Any]]: + if not self.reasoning_started or self.reasoning_closed: + return [] + self.reasoning_closed = True + return [ + { + "type": "response.reasoning_summary_text.done", + "item_id": self.reasoning_id, + "output_index": 0, + "summary_index": 0, + "text": self.reasoning_content, + }, + { + "type": "response.reasoning_summary_part.done", + "item_id": self.reasoning_id, + "output_index": 0, + "summary_index": 0, + "part": {"type": "summary_text", "text": self.reasoning_content}, + }, + { + "type": "response.output_item.done", + "output_index": 0, + "item": { + "type": "reasoning", + "id": self.reasoning_id, + "summary": [ + {"type": "summary_text", "text": self.reasoning_content} + ], + "content": [ + {"type": "reasoning_text", "text": self.reasoning_content} + ], + "encrypted_content": encrypt_reasoning(self.reasoning_content), + "status": "completed", + }, + }, + ] + + +class OpenAIResponsesTransformer(BaseTransformer): + """Transform OpenAI Responses API to/from SGLang chat completions.""" + + def transform_request(self, body: dict[str, Any]) -> dict[str, Any]: + messages: list[dict[str, Any]] = [] + + instructions = body.get("instructions") + if instructions: + messages.append({"role": "system", "content": instructions}) + + input_data = body.get("input", "") + if isinstance(input_data, str): + messages.append({"role": "user", "content": input_data}) + elif isinstance(input_data, list): + messages.extend(self._convert_input_items_to_messages(input_data)) + + result: dict[str, Any] = {"messages": messages} + if "model" in body: + result["model"] = body["model"] + + if "max_tokens" in body: + result["max_tokens"] = body["max_tokens"] + if "max_output_tokens" in body: + result["max_tokens"] = body["max_output_tokens"] + if "temperature" in body: + result["temperature"] = body["temperature"] + if "top_p" in body: + result["top_p"] = body["top_p"] + if "top_logprobs" in body: + result["top_logprobs"] = body["top_logprobs"] + if "parallel_tool_calls" in body: + result["parallel_tool_calls"] = body["parallel_tool_calls"] + if "stream" in body: + result["stream"] = body["stream"] + + text_cfg = body.get("text") + if isinstance(text_cfg, dict): + response_format = self._response_format_from_text_config(text_cfg) + if response_format is not None: + result["response_format"] = response_format + + # Responses `reasoning` request param → enable_thinking. + reasoning_cfg = body.get("reasoning") + if isinstance(reasoning_cfg, dict) and self._reasoning_config_enables_thinking( + reasoning_cfg + ): + chat_template_kwargs = dict(result.get("chat_template_kwargs") or {}) + chat_template_kwargs["enable_thinking"] = True + result["chat_template_kwargs"] = chat_template_kwargs + + # SGLang rejects tool_choice without a non-empty tools list; bind + # the pair so one can't be forwarded without the other. + tools = self._convert_tools(body.get("tools", [])) + if tools: + result["tools"] = tools + if "tool_choice" in body: + result["tool_choice"] = self._tool_choice_to_openai_chat( + body["tool_choice"] + ) + + return self._normalize_request( + result, + body.get("_served_model"), + ) + + def _response_format_from_text_config( + self, + text_cfg: dict[str, Any], + ) -> dict[str, Any] | None: + format_cfg = text_cfg.get("format") + if not isinstance(format_cfg, dict): + return None + + format_type = format_cfg.get("type") + if format_type == "text": + return None + if format_type == "json_object": + return {"type": "json_object"} + if format_type != "json_schema": + return None + + json_schema = format_cfg.get("json_schema") + if isinstance(json_schema, dict): + return {"type": "json_schema", "json_schema": json_schema} + + converted = { + key: format_cfg[key] + for key in ("name", "description", "schema", "strict") + if key in format_cfg + } + if not converted: + return None + return {"type": "json_schema", "json_schema": converted} + + def transform_response( + self, + response: dict[str, Any], + original_request: dict[str, Any], + ) -> dict[str, Any]: + choices = response.get("choices", []) + if not choices: + return self._make_error_response("No choices in response") + + choice = choices[0] + message = choice.get("message", {}) + + output_items: list[dict[str, Any]] = [] + + reasoning = message.get("reasoning_content") + if isinstance(reasoning, str) and reasoning: + output_items.append( + { + "type": "reasoning", + "id": f"rs_{uuid.uuid4().hex[:24]}", + "summary": [{"type": "summary_text", "text": reasoning}], + "content": [{"type": "reasoning_text", "text": reasoning}], + "encrypted_content": encrypt_reasoning(reasoning), + "status": "completed", + } + ) + + content = message.get("content") + if content: + output_items.append( + { + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": content}], + } + ) + + for tc in message.get("tool_calls") or []: + func = tc.get("function", {}) + name = func.get("name", "") + if name in ("shell", "execute", "run_command"): + output_items.append(self._local_shell_call_from_tool_call(tc)) + else: + output_items.append( + { + "type": "function_call", + "id": f"fc_{uuid.uuid4().hex[:24]}", + "call_id": tc.get("id", ""), + "name": name, + "arguments": func.get("arguments", "{}"), + "status": "completed", + } + ) + + usage = response.get("usage", {}) + response_usage = { + "input_tokens": usage.get("prompt_tokens", 0), + "output_tokens": usage.get("completion_tokens", 0), + "total_tokens": usage.get("total_tokens", 0), + } + cached_tokens = self._cached_prompt_tokens(usage) + if cached_tokens: + response_usage["input_tokens_details"] = {"cached_tokens": cached_tokens} + return { + "id": response.get("id", f"resp_{uuid.uuid4().hex}"), + "object": "response", + "created_at": response.get("created", int(time.time())), + "status": "completed", + "model": original_request.get("model", response.get("model", "unknown")), + "output": output_items, + "usage": response_usage, + } + + def create_stream_state( + self, original_request: dict[str, Any] + ) -> ResponsesStreamState: + return ResponsesStreamState( + model=original_request.get("model", "unknown"), + ) + + def transform_stream_chunk( + self, + chunk: dict[str, Any], + original_request: dict[str, Any], + is_first: bool = False, + ) -> list[dict[str, Any]]: + """Best-effort single-chunk Responses transform.""" + state = self.create_stream_state(original_request) + events = state.process_chunk(chunk, is_first=is_first) + choices = chunk.get("choices", []) + if choices and choices[0].get("finish_reason"): + events.extend(state.finalize()) + return events + + def _convert_input_items_to_messages( + self, + items: list[dict[str, Any]], + ) -> list[dict[str, Any]]: + messages: list[dict[str, Any]] = [] + pending_tool_calls: list[dict[str, Any]] = [] + pending_tool_outputs: list[dict[str, Any]] = [] + pending_input_content: list[dict[str, Any]] = [] + pending_reasoning: str = "" + + for item in items: + item_type = item.get("type") + + if item_type == "reasoning": + # A new reasoning item starts a new turn block. If the prior + # block already has its function_call_output, flush it now so + # this reasoning attaches to the NEXT function_call, not the + # previous one. (Otherwise codex's per-fc reasoning gets + # accumulated and dumped onto the wrong assistant message, + # breaking the prefix_merging chain.) + if pending_tool_outputs: + messages.extend( + self._flush_tool_block( + pending_tool_calls, + pending_tool_outputs, + pending_reasoning, + ) + ) + pending_tool_calls = [] + pending_tool_outputs = [] + pending_reasoning = "" + reasoning_text = extract_reasoning_from_responses_item(item) + if reasoning_text: + pending_reasoning = ( + f"{pending_reasoning}\n{reasoning_text}" + if pending_reasoning + else reasoning_text + ) + continue + + if item_type in {"input_text", "input_image"}: + if pending_tool_calls or pending_tool_outputs: + messages.extend( + self._flush_tool_block( + pending_tool_calls, pending_tool_outputs, pending_reasoning + ) + ) + pending_tool_calls = [] + pending_tool_outputs = [] + pending_reasoning = "" + pending_input_content.append(item) + continue + + if item_type == "message": + if pending_input_content: + messages.extend(self._flush_input_content(pending_input_content)) + pending_input_content = [] + if pending_tool_calls or pending_tool_outputs: + messages.extend( + self._flush_tool_block( + pending_tool_calls, pending_tool_outputs, pending_reasoning + ) + ) + pending_tool_calls = [] + pending_tool_outputs = [] + pending_reasoning = "" + + role = item.get("role", "user") + content = openai_responses_input_content_to_chat( + item.get("content", "") + ) + msg: dict[str, Any] = {"role": role, "content": content} + if role == "assistant" and pending_reasoning: + msg["reasoning_content"] = pending_reasoning + pending_reasoning = "" + messages.append(msg) + + elif item_type == "function_call": + if pending_input_content: + messages.extend(self._flush_input_content(pending_input_content)) + pending_input_content = [] + if pending_tool_outputs: + messages.extend( + self._flush_tool_block( + pending_tool_calls, + pending_tool_outputs, + pending_reasoning, + ) + ) + pending_tool_calls = [] + pending_tool_outputs = [] + pending_reasoning = "" + pending_tool_calls.append( + { + "id": item.get("call_id", f"call_{uuid.uuid4().hex[:24]}"), + "type": "function", + "function": { + "name": item.get("name", ""), + "arguments": item.get("arguments", "{}"), + }, + } + ) + + elif item_type in {"local_shell_call", "shell_call"}: + if pending_input_content: + messages.extend(self._flush_input_content(pending_input_content)) + pending_input_content = [] + if pending_tool_outputs: + messages.extend( + self._flush_tool_block( + pending_tool_calls, + pending_tool_outputs, + pending_reasoning, + ) + ) + pending_tool_calls = [] + pending_tool_outputs = [] + pending_reasoning = "" + pending_tool_calls.append(self._local_shell_call_to_tool_call(item)) + + elif item_type == "function_call_output": + if pending_input_content: + messages.extend(self._flush_input_content(pending_input_content)) + pending_input_content = [] + pending_tool_outputs.extend(self._function_call_output_messages(item)) + + elif item_type in {"local_shell_call_output", "shell_call_output"}: + if pending_input_content: + messages.extend(self._flush_input_content(pending_input_content)) + pending_input_content = [] + pending_tool_outputs.extend(self._local_shell_output_messages(item)) + + else: + if pending_input_content: + messages.extend(self._flush_input_content(pending_input_content)) + pending_input_content = [] + if pending_tool_calls or pending_tool_outputs: + messages.extend( + self._flush_tool_block( + pending_tool_calls, + pending_tool_outputs, + pending_reasoning, + ) + ) + pending_tool_calls = [] + pending_tool_outputs = [] + pending_reasoning = "" + converted = self._convert_response_item_to_message(item) + if isinstance(converted, list): + messages.extend(converted) + elif converted: + messages.append(converted) + + if pending_input_content: + messages.extend(self._flush_input_content(pending_input_content)) + + if pending_tool_calls or pending_tool_outputs: + messages.extend( + self._flush_tool_block( + pending_tool_calls, pending_tool_outputs, pending_reasoning + ) + ) + pending_reasoning = "" + + # Trailing reasoning with no following assistant message: synthesize one. + if pending_reasoning: + messages.append( + { + "role": "assistant", + "content": None, + "reasoning_content": pending_reasoning, + } + ) + + return messages + + def _function_call_output_messages( + self, item: dict[str, Any] + ) -> list[dict[str, Any]]: + output = self._function_call_output_content(item.get("output", "")) + converted_content = openai_responses_input_content_to_chat(output) + messages = [ + { + "role": "tool", + "tool_call_id": item.get("call_id", ""), + "content": self._flatten_function_call_output(output), + } + ] + + image_parts = self._image_parts(converted_content) + if image_parts: + messages.append({"role": "user", "content": image_parts}) + return messages + + def _local_shell_call_to_tool_call(self, item: dict[str, Any]) -> dict[str, Any]: + return { + "id": item.get("call_id") + or item.get("id") + or f"call_{uuid.uuid4().hex[:24]}", + "type": "function", + "function": { + "name": "shell", + "arguments": self._local_shell_action_to_arguments(item.get("action")), + }, + } + + def _local_shell_output_messages( + self, item: dict[str, Any] + ) -> list[dict[str, Any]]: + call_id = item.get("call_id") or item.get("id") or "" + return self._function_call_output_messages( + {"call_id": call_id, "output": item.get("output", "")} + ) + + def _local_shell_action_to_arguments(self, action: Any) -> str: + if isinstance(action, str): + return action + if not isinstance(action, dict): + return "{}" + + command = action.get("command") + if isinstance(command, str): + stripped = command.strip() + if stripped.startswith(("{", "[")): + try: + json.loads(stripped) + return stripped + except json.JSONDecodeError: + pass + return json.dumps({"cmd": command}) + + commands = action.get("commands") + if isinstance(commands, list): + command_values = [cmd for cmd in commands if isinstance(cmd, str)] + args: dict[str, Any] + if len(command_values) == 1: + args = {"cmd": command_values[0]} + else: + args = {"commands": command_values} + for key in ("timeout_ms", "max_output_length"): + if key in action: + args[key] = action[key] + return json.dumps(args) + + args = {key: value for key, value in action.items() if key != "type"} + return json.dumps(args) if args else "{}" + + def _local_shell_call_from_tool_call( + self, tool_call: dict[str, Any] + ) -> dict[str, Any]: + function = tool_call.get("function", {}) + arguments = ( + function.get("arguments", "{}") if isinstance(function, dict) else "{}" + ) + call_id = tool_call.get("id", "") + return { + "type": "local_shell_call", + "id": f"lsh_{uuid.uuid4().hex[:24]}", + "call_id": call_id, + "status": "completed", + "action": self._local_shell_action_from_arguments(arguments), + } + + def _local_shell_action_from_arguments(self, arguments: Any) -> dict[str, Any]: + parsed: Any = None + if isinstance(arguments, str): + try: + parsed = json.loads(arguments) + except json.JSONDecodeError: + parsed = None + elif isinstance(arguments, dict): + parsed = arguments + + if isinstance(parsed, dict): + commands = parsed.get("commands") + if isinstance(commands, list): + action = {"commands": [cmd for cmd in commands if isinstance(cmd, str)]} + else: + command = parsed.get("cmd") or parsed.get("command") + action = {"commands": [command]} if isinstance(command, str) else {} + for key in ("timeout_ms", "max_output_length"): + if key in parsed: + action[key] = parsed[key] + if action.get("commands"): + return action + + if isinstance(arguments, str) and arguments: + return {"commands": [arguments]} + return {"commands": []} + + def _function_call_output_content(self, output: Any) -> Any: + if isinstance(output, dict): + if self._is_responses_content_block(output): + return [output] + for key in ("output", "body", "content"): + if key in output: + return self._function_call_output_content(output[key]) + return output + + def _flatten_function_call_output(self, output: Any) -> str: + if isinstance(output, str): + return output + if isinstance(output, list): + parts = [] + for block in output: + if isinstance(block, str): + parts.append(block) + elif isinstance(block, dict): + if block.get("type") in {"input_text", "output_text", "text"}: + text = block.get("text") + if isinstance(text, str): + parts.append(text) + return "\n".join(parts) + if isinstance(output, dict): + return json.dumps(output) + return str(output) if output is not None else "" + + def _image_parts(self, content: Any) -> list[dict[str, Any]]: + if not isinstance(content, list): + return [] + return [ + part + for part in content + if isinstance(part, dict) and part.get("type") == "image_url" + ] + + def _is_responses_content_block(self, block: dict[str, Any]) -> bool: + return block.get("type") in { + "input_text", + "output_text", + "text", + "input_image", + "image_url", + } + + def _flush_input_content( + self, content_parts: list[dict[str, Any]] + ) -> list[dict[str, Any]]: + content = openai_responses_input_content_to_chat(content_parts) + return [{"role": "user", "content": content}] if content else [] + + def _flush_tool_block( + self, + tool_calls: list[dict[str, Any]], + tool_outputs: list[dict[str, Any]], + reasoning: str = "", + ) -> list[dict[str, Any]]: + messages: list[dict[str, Any]] = [] + if tool_calls: + assistant_msg: dict[str, Any] = { + "role": "assistant", + "content": None, + "tool_calls": list(tool_calls), + } + if reasoning: + assistant_msg["reasoning_content"] = reasoning + messages.append(assistant_msg) + messages.extend(tool_outputs) + return messages + + def _convert_response_item_to_message( + self, + item: dict[str, Any], + ) -> Optional[dict[str, Any] | list[dict[str, Any]]]: + item_type = item.get("type", "") + + if item_type == "message": + role = item.get("role", "user") + content = openai_responses_input_content_to_chat(item.get("content", [])) + if content: + return {"role": role, "content": content} + + elif item_type == "function_call_output": + return self._function_call_output_messages(item) + + # Fallback: plain {role, content} dict + if not item_type and "role" in item and "content" in item: + role = item["role"] + content = item["content"] + if isinstance(content, str): + return {"role": role, "content": content} + if isinstance(content, list): + converted = openai_responses_input_content_to_chat(content) + if converted: + return {"role": role, "content": converted} + + return None + + def _convert_tools(self, tools: list[dict[str, Any]]) -> list[dict[str, Any]]: + converted = [] + for tool in tools: + if tool.get("type") == "function" and "function" in tool: + converted.append({"type": "function", "function": tool["function"]}) + continue + + tool_type = tool.get("type") + if tool_type in {"shell", "local_shell"}: + converted.append( + { + "type": "function", + "function": { + "name": "shell", + "description": tool.get( + "description", + "Run shell commands in the local workspace.", + ), + "parameters": { + "type": "object", + "properties": { + "cmd": {"type": "string"}, + "commands": { + "type": "array", + "items": {"type": "string"}, + }, + "timeout_ms": {"type": "number"}, + "max_output_length": {"type": "number"}, + }, + }, + }, + } + ) + continue + + # Drop server-side tool types Polar can't dispatch (web_search, + # file_search, computer_use, mcp, code_interpreter, image_generation, + # custom, etc.). Only client-side functions/shell are convertible. + if tool_type and tool_type != "function": + continue + + name = tool.get("name") or tool.get("id", "") + if not name: + continue + + parameters = tool.get("parameters") + if parameters is None: + input_schema = tool.get("inputSchema") or tool.get("input_schema") + if isinstance(input_schema, dict): + json_schema = input_schema.get("jsonSchema") + parameters = ( + json_schema if isinstance(json_schema, dict) else input_schema + ) + else: + parameters = {} + + func_def: dict[str, Any] = { + "name": name, + "description": tool.get("description", ""), + "parameters": parameters, + } + if "strict" in tool: + func_def["strict"] = tool["strict"] + converted.append({"type": "function", "function": func_def}) + + return converted + + def _tool_choice_to_openai_chat(self, tool_choice: Any) -> Any: + if isinstance(tool_choice, str): + if tool_choice == "shell": + return {"type": "function", "function": {"name": "shell"}} + return tool_choice + + if not isinstance(tool_choice, dict): + return tool_choice + + choice_type = tool_choice.get("type") + if choice_type == "function": + function = tool_choice.get("function") + if isinstance(function, dict): + return tool_choice + name = tool_choice.get("name") + if isinstance(name, str) and name: + return {"type": "function", "function": {"name": name}} + if choice_type in {"shell", "local_shell"}: + return {"type": "function", "function": {"name": "shell"}} + return tool_choice + + def _reasoning_config_enables_thinking(self, reasoning_cfg: dict[str, Any]) -> bool: + if not reasoning_cfg: + return False + effort = reasoning_cfg.get("effort") + if isinstance(effort, str) and effort.lower() == "none": + return False + return True + + def _cached_prompt_tokens(self, usage: dict[str, Any]) -> int: + details = usage.get("prompt_tokens_details") + if isinstance(details, dict): + cached = details.get("cached_tokens") + if isinstance(cached, int): + return cached + cached = usage.get("cached_tokens") + return cached if isinstance(cached, int) else 0 + + def _make_error_response(self, message: str) -> dict[str, Any]: + return { + "type": "response.failed", + "response": { + "id": "resp_error", + "object": "response", + "status": "failed", + "error": {"code": "internal_error", "message": message}, + }, + } diff --git a/src/openenv/core/harness/capture/dialects/reasoning.py b/src/openenv/core/harness/capture/dialects/reasoning.py new file mode 100644 index 000000000..17360b62e --- /dev/null +++ b/src/openenv/core/harness/capture/dialects/reasoning.py @@ -0,0 +1,114 @@ +"""Reasoning round-trip helpers shared across API transformers. + +SGLang's `--reasoning-parser` (split mode: `qwen3`, `minimax`, `deepseek-r1`, +etc.) splits the model's chain-of-thought into the assistant message's +`reasoning_content` field. This module helps each transform convert that +field to / from the API-specific reasoning shape: + +- Anthropic: `thinking` content block with `thinking` + `signature` +- Gemini: part with `thought: true`, `text`, `thoughtSignature` +- Responses: `reasoning` output item with `summary`, `content`, `encrypted_content` +- OAI Chat: `reasoning_content` field on the assistant message (passthrough) + +Signatures and `encrypted_content` only need to round-trip opaquely through +the harness (the gateway is the API server on both ends), so we use +deterministic synthetic tokens — no real cryptography is necessary. +""" + +from __future__ import annotations + +import base64 +import hashlib +from typing import Any + + +def make_signature(reasoning_text: str) -> str: + """Deterministic synthetic signature for an Anthropic/Gemini thought block.""" + if not reasoning_text: + return "" + digest = hashlib.sha256(reasoning_text.encode("utf-8")).digest() + return "sg_oe_" + base64.urlsafe_b64encode(digest).decode("ascii").rstrip("=") + + +def encrypt_reasoning(reasoning_text: str) -> str: + """Pack reasoning into Responses-style `encrypted_content`. + + Base64-encoded so it survives transport. Decoded by `decrypt_reasoning` + when the harness replays it on the next turn. + """ + if not reasoning_text: + return "" + return "oe:" + base64.urlsafe_b64encode(reasoning_text.encode("utf-8")).decode( + "ascii" + ) + + +def decrypt_reasoning(encrypted: str | None) -> str: + """Reverse of `encrypt_reasoning`. Returns empty string on any failure.""" + if not isinstance(encrypted, str) or not encrypted.startswith("oe:"): + return "" + try: + return base64.urlsafe_b64decode(encrypted[len("oe:") :].encode("ascii")).decode( + "utf-8" + ) + except Exception: + return "" + + +def extract_reasoning_from_anthropic_content(content: Any) -> str: + """Extract reasoning_content from Anthropic assistant content blocks.""" + if not isinstance(content, list): + return "" + parts: list[str] = [] + for block in content: + if not isinstance(block, dict): + continue + if block.get("type") == "thinking": + text = block.get("thinking", "") + if isinstance(text, str) and text: + parts.append(text) + return "\n".join(parts) + + +def extract_reasoning_from_gemini_parts(parts: Any) -> str: + """Extract reasoning_content from Gemini content parts (thought:true).""" + if not isinstance(parts, list): + return "" + pieces: list[str] = [] + for part in parts: + if not isinstance(part, dict): + continue + if part.get("thought") is True: + text = part.get("text", "") + if isinstance(text, str) and text: + pieces.append(text) + return "\n".join(pieces) + + +def extract_reasoning_from_responses_item(item: dict[str, Any]) -> str: + """Extract reasoning_content text from a Responses `reasoning` input item. + + Prefers `content[*].text` (full chain), falls back to `summary[*].text`, + finally tries `encrypted_content` (decoded by `decrypt_reasoning`). + """ + content = item.get("content") + if isinstance(content, list): + chunks = [ + b.get("text", "") + for b in content + if isinstance(b, dict) and isinstance(b.get("text"), str) + ] + joined = "\n".join(c for c in chunks if c) + if joined: + return joined + summary = item.get("summary") + if isinstance(summary, list): + chunks = [ + b.get("text", "") + for b in summary + if isinstance(b, dict) and isinstance(b.get("text"), str) + ] + joined = "\n".join(c for c in chunks if c) + if joined: + return joined + return decrypt_reasoning(item.get("encrypted_content")) diff --git a/src/openenv/core/harness/capture/export.py b/src/openenv/core/harness/capture/export.py new file mode 100644 index 000000000..9ba844740 --- /dev/null +++ b/src/openenv/core/harness/capture/export.py @@ -0,0 +1,189 @@ +"""Graph -> the JSON a trainer consumes. + +One document per rollout. Every field a trainer needs is precomputed and validated; nothing +downstream has to re-derive, re-tokenize, or guess. + + { + "session_id": ..., + "stats": {...}, graph shape: turns, roots, forks, discards + "sequences": [ one per root-to-leaf path + {"input_ids", "loss_mask", "logprobs", "prompt_len", "n_turns", + "turn_lengths", sampled tokens per turn: the join key against a harness trace + "role", "agent" | "auxiliary" | "discarded" + "validation": [...]} + ], + "validation": [...], rollout-level findings + "trainable": bool the single gate: did anything survive + } + +Sequences are labelled rather than filtered. A caller that silently drops rows cannot be +distinguished from one that had none to drop, and "the group quietly shrank" is far harder to +diagnose than "three rows were labelled auxiliary". The trainer picks by `role`. + +ROLE ASSIGNMENT is structural first, heuristic only as a tiebreak. A rollout's real work is the +longest path with tool access; a title generator or a summariser is a short toolless root. The old +approach (matching known system-prompt strings) needed a new entry per harness and failed silently +on the harnesses nobody had profiled yet. +""" + +from __future__ import annotations + +from typing import Any + +from .validate import check_rollout, check_sequence + +AGENT, AUXILIARY, DISCARDED = "agent", "auxiliary", "discarded" + + +def _assign_roles(graph, sequences) -> list[str]: + """Label each flattened path. Purely structural. + + - a path ending in a discarded node (a sibling that never continued) is a retry + - a path whose turns carry a TOOL MANIFEST is the agent working + - anything else is auxiliary: title generators, summarisers, classifiers + + **Multiple agent paths are normal and all of them are trainable.** A harness that rewrites its + system prompt mid-run breaks token-prefix continuity and starts a new root, even though the + conversation continued: claude-code does exactly this, swapping a 12118-char system prompt for a + 12541-char one at call 6 while its message list grows 2 -> 26 unbroken. Both roots are the agent + doing real work on the same task and earn the same reward. + + An earlier version kept only the single longest tool-using path. On opencode that was + indistinguishable from correct (its second root really is a title generator), but on claude-code + it silently discarded 6 genuine agent turns. Tool access alone is the honest signal: aux calls + essentially never pass a tool manifest, and coding agents essentially always do. + """ + discarded_ids = {n.node_id for n in graph.discarded_nodes()} + live = [ + (i, s) for i, s in enumerate(sequences) if s.node_ids[-1] not in discarded_ids + ] + + # Tools only DISCRIMINATE when some paths have them and others do not. That is the opencode + # shape: an agent chain with a manifest plus a toolless title generator. + # + # Some harnesses never send a manifest at all. terminus-2 parses tool calls out of raw model + # text, so every one of its paths has n_tools == 0. Applying the tool rule there labels the whole + # rollout auxiliary, and an earlier "keep the longest" fallback then kept exactly ONE of its 13 + # turns -- a harness-trace cross-check caught it as `captured [263] vs trace [167, ..., 136]`. + # + # So: if nothing in the rollout uses tools, tools carry no signal and every live path is agent + # work. If something does, the toolless paths really are auxiliary. + any_tools = any(graph.get(nid).n_tools > 0 for _, s in live for nid in s.node_ids) + + roles = [] + for i, seq in enumerate(sequences): + if seq.node_ids[-1] in discarded_ids: + roles.append(DISCARDED) + continue + if not any_tools: + roles.append(AGENT if seq.n_trainable else AUXILIARY) + continue + has_tools = any(graph.get(nid).n_tools > 0 for nid in seq.node_ids) + roles.append(AGENT if has_tools else AUXILIARY) + return roles + + +def export_session( + session, *, include_discarded: bool = False, include_messages: bool = False +) -> dict[str, Any]: + """Build the training document for one rollout. + + `include_messages` adds each turn's request messages, tools and response message. Off by default + because it multiplies payload size by the full conversation text, on when you need to feed TRL's + `TraceEntry` contract or measure re-tokenization skew. + """ + graph = session.graph + rollout_report = check_rollout(graph) + + sequences = graph.sequences() + roles = _assign_roles(graph, sequences) + + rows: list[dict[str, Any]] = [] + for seq, role in zip(sequences, roles): + if role == DISCARDED and not include_discarded: + continue + report = check_sequence(seq) + rows.append( + { + "role": role, + "root_id": seq.root_id, + "node_ids": seq.node_ids, + "n_turns": seq.n_turns, + "prompt_len": seq.prompt_len, + "n_trainable": seq.n_trainable, + "turn_lengths": seq.turn_lengths(), + "input_ids": seq.input_ids, + "loss_mask": seq.loss_mask, + "logprobs": seq.logprobs, + "trainable": report.ok and role == AGENT, + "validation": [str(f) for f in report.findings], + } + ) + + # Every call in arrival order, including the ones excluded from training. This is what an + # external trace can be reconciled against: the harness logged every LLM call it made, + # so comparing only the surviving path would report a mismatch on any rollout that retried. + discarded_ids = {n.node_id for n in graph.discarded_nodes()} + turns = [ + { + "node_id": node.node_id, + "index": node.index, + "root_id": graph.root_of(node.node_id), + "n_sampled": len(node.sampled_ids), + "n_prompt": len(node.prompt_ids), + "n_tools": node.n_tools, + "finish_reason": node.finish_reason, + "harness_session_id": node.harness_session_id, + "discarded": node.node_id in discarded_ids, + **( + { + "request_messages": node.request_messages, + "request_tools": node.request_tools, + "response_message": node.response_message, + } + if include_messages + else {} + ), + } + for node in graph.nodes() + ] + + trainable_rows = [r for r in rows if r["trainable"]] + return { + "session_id": session.session_id, + "metadata": session.metadata, + "turns": turns, + "stats": { + **graph.stats(), + "n_sequences": len(rows), + "n_trainable_sequences": len(trainable_rows), + "n_trainable_tokens": sum(r["n_trainable"] for r in trainable_rows), + }, + "sequences": rows, + "validation": [str(f) for f in rollout_report.findings] + session.findings, + "trainable": bool(trainable_rows) and rollout_report.ok, + } + + +def summarise(document: dict[str, Any]) -> str: + """One screen of text. What you actually read after a rollout.""" + stats = document["stats"] + lines = [ + f"session {document['session_id']} trainable={document['trainable']}", + f" graph: {stats['n_turns']} turns, {stats['n_roots']} roots, " + f"{stats['n_forks']} forks, {stats['n_discarded']} discarded", + f" training: {stats['n_trainable_sequences']} sequence(s), " + f"{stats['n_trainable_tokens']} trainable tokens", + ] + for row in document["sequences"]: + lines.append( + f" [{row['role']:<10}] turns={row['n_turns']:<3} prompt={row['prompt_len']:<6} " + f"len={len(row['input_ids']):<6} trainable={row['n_trainable']:<5} " + f"turn_lengths={row['turn_lengths']}" + ) + for finding in row["validation"]: + if not finding.startswith("[INFO]"): + lines.append(f" {finding}") + for finding in document["validation"]: + lines.append(f" {finding}") + return "\n".join(lines) diff --git a/src/openenv/core/harness/capture/forwarding.py b/src/openenv/core/harness/capture/forwarding.py new file mode 100644 index 000000000..e6afe4756 --- /dev/null +++ b/src/openenv/core/harness/capture/forwarding.py @@ -0,0 +1,276 @@ +"""Publish the intercept server at a URL the sandbox can reach. + +The agent runs inside a sandbox on the public internet; the intercept runs next to the engine on a +cluster node with no inbound connectivity. Exactly one hop must be forwarded, and it is this one. +The engine itself never needs exposing: it stays on localhost behind the intercept. + +`PortForwarder` is that hop, as a swappable strategy. Three implementations, chosen by what the +sandbox actually is rather than by preference: + + DirectExposure the sandbox can already route to us (local docker, same VPC). No third party, + no expiry, no throughput ceiling. Prefer this whenever it is true. + GradioForwarder frpc via gradio.networking.setup_tunnel. + CloudflareForwarder cloudflared quick forwards, or a named forwards in production. + +MEASURED, not assumed. Over a full day of harness bring-up on one intercept: + + gradio / frpc 521 POSTs, ZERO forwarding errors in the server log, ~370ms health round trip, + still up after 24h. + cloudflared 10765 log lines on the sibling experiment, with repeated + `failed to accept QUIC stream: timeout`, `datagram manager encountered a + failure while serving`, and `lookup region1.v2.argotunnel.com: i/o timeout`. + It always reconnected, so this is churn rather than outage, but it is churn. + +So gradio is the better default at eval scale. Cloudflare earns its place elsewhere: quick forwards +expire in a way named forwards do not, and a named forwards gives a stable hostname, real access +policies, and no shared relay. gradio.live URLs expire at 72h and are a single frpc hop, which is +fine for a sweep and a bottleneck at GRPO group width. + +SHARE TOKENS ARE NOT AUTH. `share_token` identifies the forward to the share server; the resulting +URL is public either way. What protects the GPU behind it is the intercept's own key check, which is +why `SessionRegistry.require_registered` defaults to True. +""" + +from __future__ import annotations + +import re +import secrets +import shutil +import subprocess +import time +from abc import ABC, abstractmethod + + +class ForwardingError(RuntimeError): + """Raised when a forwarder cannot be established. Never returns a half-open forward. + + Failing here is strictly better than returning a URL that does not resolve: a stale URL that + still looks valid produces a rollout that silently captures nothing, which is the exact class of + failure this whole layer exists to make impossible. + """ + + +class PortForwarder(ABC): + """Publish `local_host:local_port` and hand back a URL reachable from the sandbox.""" + + def __init__(self) -> None: + self._url: str | None = None + self._local_port: int | None = None + + @classmethod + def preflight(cls) -> None: + """Raise ForwardingError if this strategy cannot possibly work here. + + Called BEFORE a run starts, so a missing binary or an uninstalled dependency is a startup + error rather than a failure discovered after the first sandbox has been billed. + """ + + @abstractmethod + def start(self, local_port: int, *, local_host: str = "127.0.0.1") -> str: + """Begin forwarding and return the public URL.""" + + @abstractmethod + def stop(self) -> None: + """Tear down. Must be idempotent: teardown runs on both success and failure paths.""" + + @property + def url(self) -> str | None: + return self._url + + @property + def name(self) -> str: + return type(self).__name__ + + def __enter__(self) -> "PortForwarder": + return self + + def __exit__(self, *exc) -> None: + self.stop() + + +class DirectExposure(PortForwarder): + """No forward: hand back the address as-is. + + For local docker sandboxes, or any deployment where the sandbox can already route to the host. + This is the production answer whenever it is available, and it is worth checking before reaching + for a forward: a forward exists because E2B is off-cluster, not because the design needs one. + """ + + def __init__(self, advertise_host: str = "127.0.0.1", scheme: str = "http") -> None: + super().__init__() + self._host = advertise_host + self._scheme = scheme + + def start(self, local_port: int, *, local_host: str = "127.0.0.1") -> str: + self._local_port = local_port + self._url = f"{self._scheme}://{self._host}:{local_port}" + return self._url + + def stop(self) -> None: + self._url = None + + +class GradioForwarder(PortForwarder): + """frpc, via `gradio.networking.setup_tunnel`. + + Preferred over shelling out to a binary for a reason that matters operationally: it RETURNS the + URL, rather than leaving us to grep a subprocess log for it and hope it appeared. It is also + outbound-only, so it needs no inbound firewall rule, and gradio is already an OpenEnv dependency. + Verifiers reached the same conclusion independently and forwards via frpc too. + + Pass `share_server_address` to point at your own frps: stable URLs, no 72h expiry, no third + party, no shared throughput ceiling. + """ + + def __init__( + self, + share_server_address: str | None = None, + share_server_tls_certificate: str | None = None, + ) -> None: + super().__init__() + self._share_server = share_server_address + self._tls_cert = share_server_tls_certificate + + @classmethod + def preflight(cls) -> None: + try: + from gradio.networking import setup_tunnel # noqa: F401 + except Exception as exc: # noqa: BLE001 + raise ForwardingError( + f"gradio is required for GradioForwarder: {exc}" + ) from exc + + def start(self, local_port: int, *, local_host: str = "127.0.0.1") -> str: + from gradio.networking import setup_tunnel + + try: + url = setup_tunnel( + local_host=local_host, + local_port=local_port, + share_token=secrets.token_hex(16), + share_server_address=self._share_server, + share_server_tls_certificate=self._tls_cert, + ) + except Exception as exc: # noqa: BLE001 + raise ForwardingError(f"gradio forward failed to open: {exc}") from exc + self._local_port, self._url = local_port, url + return url + + def stop(self) -> None: + # frpc runs in-process and dies with it. That coupling is deliberate: a forward outliving the + # intercept it points at is a 502 generator. + self._url = None + + +class CloudflareForwarder(PortForwarder): + """`cloudflared`, either a quick forwards or a named one. + + Quick tunnels (no `tunnel_name`) need no account and print a `*.trycloudflare.com` URL on + stderr, which we parse. Named forwards need `cloudflared login` beforehand but give a stable + hostname that survives restarts, which is what you want once this is not a sweep any more. + + The URL arrives asynchronously on stderr, so `start` blocks until it appears or gives up. That + wait is the entire reason this class is more code than GradioForwarder. + """ + + _URL_RE = re.compile(r"https://[-a-z0-9]+\.trycloudflare\.com") + + def __init__( + self, + tunnel_name: str | None = None, + hostname: str | None = None, + binary: str = "cloudflared", + startup_timeout_s: float = 60.0, + ) -> None: + super().__init__() + self._tunnel_name = tunnel_name + self._hostname = hostname + self._binary = binary + self._startup_timeout_s = startup_timeout_s + self._proc: subprocess.Popen | None = None + + @classmethod + def preflight(cls, binary: str = "cloudflared") -> None: + if shutil.which(binary) is None: + raise ForwardingError( + f"`{binary}` not found on PATH. Install it, or use GradioForwarder, which needs no " + "binary because frpc ships with gradio." + ) + + def start(self, local_port: int, *, local_host: str = "127.0.0.1") -> str: + self.preflight(self._binary) + target = f"http://{local_host}:{local_port}" + + if self._tunnel_name: + cmd = [self._binary, "tunnel", "run", "--url", target, self._tunnel_name] + else: + cmd = [self._binary, "forward", "--url", target, "--no-autoupdate"] + + self._proc = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + universal_newlines=True, + ) + self._local_port = local_port + + # A named forwards serves a hostname we already know, so there is nothing to parse. + if self._tunnel_name and self._hostname: + self._url = f"https://{self._hostname}" + return self._url + + url = self._await_url() + if url is None: + self.stop() + raise ForwardingError( + f"cloudflared printed no forward URL within {self._startup_timeout_s:.0f}s. " + "Check that the binary can reach Cloudflare, or use GradioForwarder." + ) + self._url = url + return url + + def _await_url(self) -> str | None: + """Read stderr until the URL appears, the process dies, or we run out of patience.""" + assert self._proc is not None and self._proc.stdout is not None + deadline = time.monotonic() + self._startup_timeout_s + while time.monotonic() < deadline: + if self._proc.poll() is not None: + return None # died during startup + line = self._proc.stdout.readline() + if not line: + time.sleep(0.05) + continue + match = self._URL_RE.search(line) + if match: + return match.group(0) + return None + + def stop(self) -> None: + proc, self._proc, self._url = self._proc, None, None + if proc is None or proc.poll() is not None: + return + proc.terminate() + try: + proc.wait(timeout=10) + except subprocess.TimeoutExpired: + proc.kill() + + +_FORWARDERS: dict[str, type[PortForwarder]] = { + "direct": DirectExposure, + "gradio": GradioForwarder, + "cloudflare": CloudflareForwarder, +} + + +def make_forwarder(kind: str = "gradio", **kwargs) -> PortForwarder: + """Build a forwarder by name, for CLI wiring (`--expose gradio|cloudflare|direct`).""" + try: + cls = _FORWARDERS[kind] + except KeyError: + raise ForwardingError( + f"unknown port forwarder {kind!r}; choose one of {sorted(_FORWARDERS)}" + ) from None + return cls(**kwargs) diff --git a/src/openenv/core/harness/capture/graph.py b/src/openenv/core/harness/capture/graph.py new file mode 100644 index 000000000..22bf8c253 --- /dev/null +++ b/src/openenv/core/harness/capture/graph.py @@ -0,0 +1,284 @@ +"""The rollout graph: every model call a harness made, linked by token prefix. + +A rollout is not a list of turns. Harnesses retry, spawn subagents, generate titles, and compact +context, and all of it arrives on one wire looking identical. A flat list forces you to guess which +turns belong together; a graph records it. + + node one model call: the prompt the server tokenized, the tokens it sampled back + parent the call whose prompt+completion is a token-prefix of this call's prompt + root a call that extends nothing (a new conversation: the agent's, a subagent's, + a title generator's, or the continuation after a context compaction) + path root -> leaf, which is exactly one training sequence + +Everything downstream is a walk. `sequences()` concatenates a path into +`input_ids / loss_mask / logprobs`; branch structure tells you which paths are the agent's work and +which are discards. + +WHY A GRAPH AND NOT PREFIX-MERGED CHAINS. Three things fall out of it that a chain cannot express: + + * **Retries become visible.** A retried turn is a sibling that never continued: same parent, no + children. A proxy cannot ask the harness what it discarded, but the shape of the graph shows it. + * **Subagents separate themselves.** A subagent has its own system prompt, so its first call + extends nothing and starts its own root. No system-prompt keyword matching required. + * **Compaction is representable.** A rewritten history is not a prefix extension, so it opens a + new root instead of corrupting the chain it came from. + +TOKEN FIDELITY. We never tokenize. The inference server tokenizes each prompt as a side effect of +serving it and returns `prompt_token_ids`, so turn k+1's prompt IS the canonical tokenization of +everything up to that point, including the tool results the harness inserted. Assistant bodies come +back as sampled `token_ids` with aligned logprobs. So for any path: + + concat(node.context_ids + node.sampled_ids for node in path) + +reproduces, exactly, the token sequences the model actually saw and produced. Nothing is re-rendered +through a local chat template, which is the single largest source of silent train/inference skew. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Iterator + + +def common_prefix_len(a: list[int], b: list[int]) -> int: + """How many leading tokens `a` and `b` share.""" + limit = min(len(a), len(b)) + i = 0 + while i < limit and a[i] == b[i]: + i += 1 + return i + + +@dataclass +class TurnNode: + """One model call, and where it sits relative to the call before it.""" + + node_id: str + prompt_ids: list[int] + sampled_ids: list[int] + sampled_logprobs: list[float] | None = None + parent_id: str | None = None + + # Provenance, for attribution and for the per-harness notes. Never used in token math. + index: int = 0 # arrival order within the session + model: str | None = None + finish_reason: str | None = None + harness_session_id: str | None = ( + None # the harness's OWN session id, when it sends one + ) + system_digest: str | None = ( + None # cheap identity for the conversation this belongs to + ) + n_tools: int = 0 + request_messages: list[dict[str, Any]] = field(default_factory=list) + # Retained because TRL's `_turns_from_trace` passes `tools` to `apply_chat_template`: the tool + # manifest is part of the rendered prompt, so a re-tokenization without it does not match what + # the engine actually saw. + request_tools: list[dict[str, Any]] | None = None + response_message: dict[str, Any] = field(default_factory=dict) + + @property + def end_ids(self) -> list[int]: + """Cumulative token sequence after this turn: its prompt plus what it sampled.""" + return self.prompt_ids + self.sampled_ids + + def context_ids(self, parent: "TurnNode | None") -> list[int]: + """Tokens this node adds to the sequence BEFORE the model starts generating. + + For a root that is the whole prompt. For a child it is the interstitial span: the tool + results, user turns and template scaffolding the harness inserted since the parent stopped + generating. These are real tokens the model conditioned on, but it did not produce them, so + they are context (mask 0) rather than targets. + """ + if parent is None: + return list(self.prompt_ids) + return self.prompt_ids[len(parent.end_ids) :] + + +@dataclass +class TrainingSequence: + """One path through the graph, flattened. Maps onto TRL's `TrainingSequence` fields.""" + + input_ids: list[int] + loss_mask: list[int] + logprobs: list[float] + node_ids: list[str] + prompt_len: int # tokens before the first sampled token + root_id: str + n_turns: int + + @property + def n_trainable(self) -> int: + return sum(self.loss_mask) + + def turn_lengths(self) -> list[int]: + """Sampled-token count per turn, in order. The join key against a harness trace.""" + lengths, run = [], 0 + for m in self.loss_mask: + if m: + run += 1 + elif run: + lengths.append(run) + run = 0 + if run: + lengths.append(run) + return lengths + + +class RolloutGraph: + """All model calls for one rollout, linked by prefix. + + `add_turn` is the whole ingestion path: it finds the parent by token prefix and appends. Calls + arrive in wire order, but the graph does not depend on that order being meaningful, which matters + because harnesses issue concurrent requests (parallel subagents, background summarisation). + """ + + def __init__(self) -> None: + self._nodes: dict[str, TurnNode] = {} + self._order: list[str] = [] + self._children: dict[str, list[str]] = {} + + # --- construction -------------------------------------------------- + def add_turn(self, node: TurnNode) -> TurnNode: + node.index = len(self._order) + node.parent_id = self._find_parent(node) + self._nodes[node.node_id] = node + self._order.append(node.node_id) + self._children.setdefault(node.node_id, []) + if node.parent_id is not None: + self._children[node.parent_id].append(node.node_id) + return node + + def _find_parent(self, node: TurnNode) -> str | None: + """The existing node whose prompt+completion is the LONGEST exact prefix of this prompt. + + Longest wins so a deep chain attaches to its immediate predecessor rather than to an early + ancestor that also matches. Requiring an exact prefix (not a fuzzy match) is deliberate: a + harness that mutates its history has genuinely produced a different sequence, and quietly + attaching it would fabricate a trajectory the model never saw. Such a call becomes a new root + instead, which `roots()` surfaces rather than hides. + """ + best_id, best_len = None, 0 + for candidate_id in self._order: + candidate = self._nodes[candidate_id] + end = candidate.end_ids + if len(end) > len(node.prompt_ids) or len(end) <= best_len: + continue + if common_prefix_len(end, node.prompt_ids) == len(end): + best_id, best_len = candidate_id, len(end) + return best_id + + # --- structure ----------------------------------------------------- + def nodes(self) -> list[TurnNode]: + return [self._nodes[i] for i in self._order] + + def get(self, node_id: str) -> TurnNode: + return self._nodes[node_id] + + def children(self, node_id: str) -> list[TurnNode]: + return [self._nodes[i] for i in self._children.get(node_id, [])] + + def roots(self) -> list[TurnNode]: + return [self._nodes[i] for i in self._order if self._nodes[i].parent_id is None] + + def leaves(self) -> list[TurnNode]: + return [self._nodes[i] for i in self._order if not self._children.get(i)] + + def root_of(self, node_id: str) -> str: + """Which conversation this node belongs to. Walks parents to the top.""" + node = self._nodes[node_id] + while node.parent_id is not None: + node = self._nodes[node.parent_id] + return node.node_id + + def path_to(self, leaf_id: str) -> list[TurnNode]: + path: list[TurnNode] = [] + node: TurnNode | None = self._nodes[leaf_id] + while node is not None: + path.append(node) + node = self._nodes[node.parent_id] if node.parent_id else None + return list(reversed(path)) + + def paths(self) -> Iterator[list[TurnNode]]: + """Every root-to-leaf path. One per distinct trajectory, including discarded branches.""" + for leaf in self.leaves(): + yield self.path_to(leaf.node_id) + + def forks(self) -> list[tuple[str, list[str]]]: + """Nodes with more than one child: retries, resamples, or parallel branches.""" + return [(pid, kids) for pid, kids in self._children.items() if len(kids) > 1] + + def discarded_nodes(self) -> list[TurnNode]: + """Sampled turns that led nowhere. + + A sibling with no children whose parent has another child that DID continue was generated and + thrown away: a retry after a parse failure, or a resample. Training it with the rollout's + reward credits work that never happened. Detected purely from shape, with no harness + cooperation, which matters because a proxy is otherwise blind to retries. + + The final turn of a real trajectory is also childless, so a sibling is only called discarded + when at least one of its siblings continued. + """ + discarded: list[TurnNode] = [] + for _, kids in self.forks(): + continued = [k for k in kids if self._children.get(k)] + if not continued: + continue # all siblings are terminal: ambiguous, keep them all + discarded.extend(self._nodes[k] for k in kids if not self._children.get(k)) + return discarded + + # --- flattening ---------------------------------------------------- + def sequence_for(self, leaf_id: str) -> TrainingSequence: + """Flatten one root-to-leaf path into token ids, mask and logprobs. + + Invariant enforced here rather than trusted: a turn whose logprobs are missing or misaligned + contributes its tokens as CONTEXT (mask 0), never as targets. A trainable token without a + real behaviour-policy logprob would make GRPO's importance ratio `exp(new - old)` a ratio + against a number we invented. + """ + path = self.path_to(leaf_id) + input_ids: list[int] = [] + loss_mask: list[int] = [] + logprobs: list[float] = [] + prompt_len = 0 + + parent: TurnNode | None = None + for position, node in enumerate(path): + context = node.context_ids(parent) + input_ids.extend(context) + loss_mask.extend([0] * len(context)) + logprobs.extend([0.0] * len(context)) + if position == 0: + prompt_len = len(context) + + usable = node.sampled_logprobs is not None and len( + node.sampled_logprobs + ) == len(node.sampled_ids) + input_ids.extend(node.sampled_ids) + loss_mask.extend([1 if usable else 0] * len(node.sampled_ids)) + logprobs.extend( + node.sampled_logprobs if usable else [0.0] * len(node.sampled_ids) + ) + parent = node + + return TrainingSequence( + input_ids=input_ids, + loss_mask=loss_mask, + logprobs=logprobs, + node_ids=[n.node_id for n in path], + prompt_len=prompt_len, + root_id=path[0].node_id, + n_turns=len(path), + ) + + def sequences(self) -> list[TrainingSequence]: + return [self.sequence_for(leaf.node_id) for leaf in self.leaves()] + + def stats(self) -> dict[str, Any]: + return { + "n_turns": len(self._order), + "n_roots": len(self.roots()), + "n_leaves": len(self.leaves()), + "n_forks": len(self.forks()), + "n_discarded": len(self.discarded_nodes()), + } diff --git a/src/openenv/core/harness/capture/server.py b/src/openenv/core/harness/capture/server.py new file mode 100644 index 000000000..4e59fc47e --- /dev/null +++ b/src/openenv/core/harness/capture/server.py @@ -0,0 +1,532 @@ +"""The intercept server. + + input an OpenAI-spec endpoint you already host (vLLM or SGLang) + the served model name + output per rollout, a JSON document of exact token ids, logprobs and loss masks, ready to train + +In between: a coding agent points at this URL, in whichever wire dialect it speaks, and +nothing about the agent changes except a base URL and an API key. + + agent (in an E2B sandbox, any of ~37) + | OPENAI_BASE_URL / ANTHROPIC_BASE_URL / provider config = this server + | API key = the rollout's session id <- the entire multiplexing scheme + v + THIS --detect dialect--> normalise to chat --inject capture params--> your engine + <--replay in the agent's dialect (SSE if it asked for SSE)---------┘ + | + └─ each call becomes a node in the rollout graph, linked by token prefix + +WHY IT CAPTURES FAITHFULLY: we never tokenize. The engine tokenizes each prompt to serve it and hands +back `prompt_token_ids`, so turn k+1's prompt is the canonical tokenization of everything up to that +point, tool results included. Completions come back as sampled ids with aligned logprobs. Stitching +those along a graph path reproduces exactly what the model saw and produced, with no local chat +template involved. See `graph.py`. + +TWO ASYMMETRIES, both learned from real failures: + + * **Capture is non-streaming, the reply is whatever the client asked for.** One complete response + carries ids and logprobs whole; reassembling them from SSE deltas is error-prone in exactly the + way that silently corrupts training data. But a harness that requested SSE and receives a JSON + body does not error, it yields nothing: opencode reported `step-finish reason:"unknown"`, zero + tokens, no error, having been handed a perfectly valid tool call. See `sse.py`. + * **We validate on ingest, not on export.** A turn whose logprobs are misaligned must be caught + while we still know which turn it was. + +Run: python -m intercept.server --llm-url http://127.0.0.1:8000 --model Qwen3.5-9B +""" + +from __future__ import annotations + +import argparse +import logging +import uuid +from typing import Any + +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse, StreamingResponse + +from . import sse +from .detection import APIType, detect +from .dialects import TransformManager +from .export import export_session +from .graph import TurnNode +from .sessions import extract_harness_session, SessionRegistry +from .upstream import InferenceClient, UpstreamError +from .validate import check_turn + +logger = logging.getLogger("intercept") + + +def _system_digest(messages: list[dict[str, Any]]) -> str | None: + """Cheap identity for 'which conversation is this'. Recorded, never used for routing.""" + import hashlib + + for message in messages or []: + if message.get("role") == "system": + content = message.get("content") + if isinstance(content, list): # anthropic / responses send block lists + content = " ".join( + p.get("text", "") for p in content if isinstance(p, dict) + ) + if isinstance(content, str) and content: + return hashlib.sha256(content.encode()).hexdigest()[:16] + return None + + +# Routes an agent calls that are NOT model turns. They must be answered, but must never become graph +# nodes: recording them adds a bogus root and corrupts the trajectory structure. +# +# Borrowed from verifiers, whose Dialect ABC carries `aux_routes` for exactly this +# (v1/dialects/anthropic.py:273, "relayed as native JSON, never recorded on the trace"). +# claude-code calls count_tokens before sending a turn; without this the catch-all would hand it to +# `transform_request`, forward nonsense upstream, and file the result as a model call. +AUX_ROUTES: tuple[str, ...] = ("/v1/messages/count_tokens",) + + +def is_aux_route(path: str) -> bool: + normalised = "/" + path.lstrip("/") + return any(normalised.endswith(route) for route in AUX_ROUTES) + + +def approximate_token_count(body: dict[str, Any]) -> int: + """Answer a count_tokens request without a tokenizer. + + We deliberately do not load one: the whole design keeps tokenization on the engine, and pulling a + tokenizer in here just to serve a side request would reintroduce the "two sources of truth" + problem this architecture exists to avoid. Agents use this figure for context-budget decisions, + not for anything that reaches training, so a ~4-chars-per-token estimate is sufficient. If a + harness turns out to depend on exactness, forward it to the engine's /tokenize endpoint instead. + """ + text_len = 0 + for message in body.get("messages") or []: + content = message.get("content") + if isinstance(content, str): + text_len += len(content) + elif isinstance(content, list): + for part in content: + if isinstance(part, dict): + text_len += len(str(part.get("text") or part.get("content") or "")) + system = body.get("system") + if isinstance(system, str): + text_len += len(system) + elif isinstance(system, list): + text_len += sum( + len(str(p.get("text", ""))) for p in system if isinstance(p, dict) + ) + text_len += len(str(body.get("tools") or "")) + return max(1, text_len // 4) + + +def wants_stream(path: str, body: dict[str, Any]) -> bool: + """Did the client ask for SSE? Each dialect says so differently. + + OpenAI chat, Responses and Anthropic all set `stream: true` in the body. **Google does not.** It + signals streaming in the URL: `:streamGenerateContent`, usually with `?alt=sse`. gemini-cli calls + + POST /v1beta/models/:streamGenerateContent?alt=sse + + with no `stream` key anywhere in the body, so a body-only check returns False and we answer a + streaming request with a plain JSON document. It arrives as HTTP 200 and the client dies parsing + it: + + Error: Incomplete JSON segment at the end + at ApiClient.processStreamResponse_1 (@google/gemini-cli/...) + + Same failure family as opencode's silent `reason:"unknown"`: a valid-looking response in the + wrong envelope. verifiers models this as a per-dialect `Dialect.streaming(body)`; this is the + same idea kept to one function. + """ + if body.get("stream") is True: + return True + lowered = path.lower() + return "streamgeneratecontent" in lowered or "alt=sse" in lowered + + +_MAX_TOKENS_KEYS = ("max_tokens", "max_completion_tokens", "max_output_tokens") + + +def clamp_output_tokens(chat_request: dict[str, Any], cap: int | None) -> int | None: + """Cap the requested output length so prompt + completion fits the served context window. + + Harnesses ask for absurd output budgets. qwen-coder requests **64000** output tokens, which on a + 65536-token model leaves room for a 1536-token prompt and then fails on the next character: + + maximum context length is 65536 tokens. However, you requested 64000 output tokens and + your prompt contains at least 1537 input tokens, for a total of at least 65537 + + Every call 502s, the agent does nothing, and it presents as "reached the intercept, captured + nothing". Polar caps this too (`proxy_max_tokens_cap = 16384`, noting opencode's ~32000 default + "exceeds some provider limits"), so it is a known hazard rather than one harness misbehaving. + + A fixed cap rather than `context - len(prompt)`: computing the latter needs a tokenizer here, and + keeping tokenization on the engine is the whole design. Agent turns are short (the longest seen + across every validated harness is 874 tokens), so a few thousand is generous. + + Returns the value it replaced, for logging, or None if nothing changed. + """ + if not cap: + return None + for key in _MAX_TOKENS_KEYS: + value = chat_request.get(key) + if isinstance(value, int) and value > cap: + chat_request[key] = cap + return value + return None + + +def normalise_for_capture(chat_request: dict[str, Any]) -> None: + """Force the upstream call into the one shape that yields complete, capturable responses. + + `stream_options` is not cosmetic: vLLM validates it against `stream` and rejects the pair with + "Stream options can only be defined when `stream=True`", which 400s the ENTIRE request. opencode + sends it on every call, so leaving it in is a total outage rather than a degradation. + + Both keys are set here rather than left to the inference client, because the client rewrites + `stream` only after this point: a check against the incoming value sees `true` and leaves + `stream_options` behind, which is precisely the bug this exists to prevent. + + An empty `tools` array is dropped for the same reason. vLLM rejects it outright: + + `tools` must not be an empty array. Either provide at least one tool or omit the field + entirely. + + kimi-cli sends `tools: []` once its agent loop has no tools left to offer, which 400s the call. + The two forms mean the same thing to the model, so dropping the key is lossless and keeps the + rollout alive rather than truncating it mid-trajectory. + """ + chat_request["stream"] = False + chat_request.pop("stream_options", None) + for key in ("tools", "functions"): + if key in chat_request and not chat_request[key]: + chat_request.pop(key) + # `tool_choice` without `tools` is equally invalid, and is meaningless once the list is gone. + if "tools" not in chat_request: + chat_request.pop("tool_choice", None) + + +def normalise_response(response: dict[str, Any]) -> None: + """Fill in usage sub-objects that vLLM leaves null but the OpenAI schema always returns. + + vLLM returns `"prompt_tokens_details": null` when prefix caching is off. OpenAI always returns + the object, so a harness that reads `usage.prompt_tokens_details.cached_tokens` without guarding + gets an AttributeError. trae-agent does exactly that and dies after its FIRST call: + + 'NoneType' object has no attribute 'cached_tokens' + + which produced a clean single-turn capture and a task the agent never attempted. + + This touches ONLY accounting fields. No token id, logprob or message content is altered, so it + cannot affect what gets captured or trained. It is a compatibility shim that makes us MORE + OpenAI-conformant than the engine behind us, which is the safe direction: a client that already + guarded for null sees a zeroed object instead, which reads the same. + """ + usage = response.get("usage") + if not isinstance(usage, dict): + return + if usage.get("prompt_tokens_details") is None: + usage["prompt_tokens_details"] = {"cached_tokens": 0, "audio_tokens": 0} + if usage.get("completion_tokens_details") is None: + usage["completion_tokens_details"] = { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + } + + +def normalise_client_payload(payload: dict[str, Any], api_type: APIType) -> None: + """Fill in usage sub-objects the OUTBOUND dialect promises but the transformer omits. + + Sibling of `normalise_response`, one layer further out. That one repairs the chat-completions + usage we get FROM vLLM; this repairs the usage we hand TO the client after translation. + + Polar's Responses transformer builds usage as exactly + `{"input_tokens", "output_tokens", "total_tokens"}` (transform/openai_responses.py:43), with no + detail sub-objects. The real Responses API always returns them, and trae-agent reads them without + a guard (trae_agent/utils/llm_clients/openai_client.py): + + cache_read_input_tokens=response.usage.input_tokens_details.cached_tokens or 0, + reasoning_tokens=response.usage.output_tokens_details.reasoning_tokens or 0, + + so `input_tokens_details` is None and it dies with + `'NoneType' object has no attribute 'cached_tokens'` after its FIRST call. + + Note this is why trae-agent looked like a chat-completions harness for a whole night: its seam + says `openai_chat`, but the access log shows exactly one `POST /v1/responses` against 465 + chat-completions calls. It speaks Responses. + + Accounting fields only. No token id, logprob or content is touched, so capture is unaffected. + """ + usage = payload.get("usage") + if not isinstance(usage, dict): + return + if api_type is APIType.OPENAI_RESPONSES: + if usage.get("input_tokens_details") is None: + usage["input_tokens_details"] = {"cached_tokens": 0} + if usage.get("output_tokens_details") is None: + usage["output_tokens_details"] = {"reasoning_tokens": 0} + + +def create_app( + *, + llm_url: str, + model: str | None = None, + engine: str = "vllm", + require_registered: bool = True, + max_output_tokens: int | None = 8192, +) -> FastAPI: + app = FastAPI(title="openenv-capture") + + # Identifies this app instance on /health. A caller that binds a port cannot tell "my server is + # up" from "someone else's server already held this port" by connecting alone, and answering the + # wrong process is silent: sessions are minted here and rejected there, so the agent gets 401 and + # the rollout reports no model calls. + app.state.instance_id = uuid.uuid4().hex + app.state.inference = InferenceClient( + base_url=llm_url.rstrip("/"), served_model=model + ) + app.state.transforms = TransformManager() + app.state.registry = SessionRegistry(require_registered=require_registered) + app.state.model = model + app.state.llm_url = llm_url + app.state.max_output_tokens = max_output_tokens + + @app.get("/health") + async def health() -> dict[str, Any]: + return { + "status": "ok", + "instance": app.state.instance_id, + "upstream": llm_url, + "engine": engine, + "model": app.state.model, + "sessions": len(app.state.registry.list_ids()), + "require_registered": app.state.registry.require_registered, + } + + @app.post("/sessions") + async def create_session(payload: dict[str, Any] | None = None) -> dict[str, Any]: + """Mint a rollout id. Hand it to the agent as its API key; that is the whole integration.""" + payload = payload or {} + session = app.state.registry.create( + payload.get("session_id"), **(payload.get("metadata") or {}) + ) + return {"session_id": session.session_id} + + @app.get("/sessions") + async def list_sessions() -> dict[str, Any]: + return {"sessions": app.state.registry.summary()} + + @app.get("/sessions/{session_id}") + async def session_status(session_id: str) -> Any: + """Live progress. `idle_s` is the cheapest wedge detector: turns arriving means progress.""" + session = app.state.registry.get(session_id) + if session is None: + return JSONResponse({"error": "unknown session"}, status_code=404) + return { + "session_id": session_id, + "idle_s": round(session.idle_seconds, 1), + "upstream_errors": session.upstream_errors, + **session.graph.stats(), + } + + @app.get("/sessions/{session_id}/rollout") + async def rollout( + session_id: str, include_discarded: bool = False, include_messages: bool = False + ) -> Any: + """THE training endpoint: stitched, masked, logprob-aligned, validated.""" + session = app.state.registry.get(session_id) + if session is None: + return JSONResponse({"error": "unknown session"}, status_code=404) + return export_session( + session, + include_discarded=include_discarded, + include_messages=include_messages, + ) + + @app.delete("/sessions/{session_id}") + async def delete_session(session_id: str) -> dict[str, Any]: + return {"deleted": app.state.registry.delete(session_id)} + + @app.get("/v1/models") + async def models() -> Any: + return await app.state.inference.list_models() + + @app.post("/{path:path}") + async def proxy(path: str, request: Request) -> Any: + """Catch-all: /v1/chat/completions, /v1/messages, /v1/responses, :generateContent.""" + headers = dict(request.headers) + try: + body = await request.json() + except Exception: # noqa: BLE001 + return JSONResponse( + {"error": {"message": "body must be JSON"}}, status_code=400 + ) + + # Answered, never recorded. Must come before session routing and dialect handling: an aux + # route is not a model turn, so it has no business creating a node or a session. + if is_aux_route(path): + logger.info("aux route %s (answered, not recorded)", path) + return JSONResponse({"input_tokens": approximate_token_count(body)}) + + session = app.state.registry.resolve(headers, body) + if session is None: + # Deliberately 401 rather than serving an unknown caller: this port is public. + return JSONResponse( + { + "error": { + "message": "unknown API key; register a session via POST /sessions", + "type": "invalid_request_error", + } + }, + status_code=401, + ) + + api_type: APIType = detect(f"/{path}", headers, body) + transformer = app.state.transforms.get(api_type) + + original_request = dict(body) + # Include the query string: Google puts `alt=sse` there, not in the body. + full_target = ( + f"/{path}?{request.url.query}" if request.url.query else f"/{path}" + ) + client_wants_stream = wants_stream(full_target, body) + chat_request = transformer.transform_request(dict(body)) + if app.state.model: + chat_request["model"] = app.state.model + normalise_for_capture(chat_request) + clamped = clamp_output_tokens(chat_request, app.state.max_output_tokens) + if clamped: + logger.info( + "clamped requested output tokens %d -> %d", + clamped, + app.state.max_output_tokens, + ) + + try: + response = await app.state.inference.completion(chat_request) + except UpstreamError as exc: + session.upstream_errors += 1 + logger.warning("upstream error [%s]: %s", session.session_id, exc) + return JSONResponse({"error": {"message": str(exc)}}, status_code=502) + + normalise_response(response) + _ingest(session, chat_request, response, api_type) + + if client_wants_stream: + return StreamingResponse( + sse.replay(api_type, transformer, response, original_request), + media_type="text/event-stream", + headers=sse.SSE_HEADERS, + ) + payload = transformer.transform_response(response, original_request) + normalise_client_payload(payload, api_type) + return JSONResponse(payload) + + def _ingest( + session, + chat_request: dict[str, Any], + response: dict[str, Any], + api_type: APIType, + ) -> None: + """Turn one upstream response into a graph node, validating before it lands. + + Never raises. A capture problem must degrade one turn, not kill a rollout that is otherwise + producing usable data, and certainly not take down the server serving every other rollout. + """ + try: + choice = (response.get("choices") or [{}])[0] + logprob_entries = (choice.get("logprobs") or {}).get("content") or [] + logprobs = [e.get("logprob") for e in logprob_entries] or None + sampled_ids = choice.get("token_ids") or [] + prompt_ids = response.get("prompt_token_ids") or [] + index = session.graph.stats()["n_turns"] + + report = check_turn( + prompt_ids, + sampled_ids, + logprobs, + finish_reason=choice.get("finish_reason"), + index=index, + ) + session.findings.extend(str(f) for f in report.findings) + if not report.ok: + logger.warning( + "[%s] turn %d rejected: %s", + session.session_id, + index, + "; ".join(str(f) for f in report.fatal), + ) + # Still recorded, with logprobs dropped: the tokens are real context for later turns, + # and `sequence_for` masks a turn whose logprobs it cannot trust. + logprobs = None + + session.graph.add_turn( + TurnNode( + node_id=uuid.uuid4().hex[:12], + prompt_ids=list(prompt_ids), + sampled_ids=list(sampled_ids), + sampled_logprobs=list(logprobs) if logprobs else None, + model=chat_request.get("model"), + finish_reason=choice.get("finish_reason"), + harness_session_id=extract_harness_session({}, chat_request), + system_digest=_system_digest(chat_request.get("messages") or []), + n_tools=len(chat_request.get("tools") or []), + request_messages=chat_request.get("messages") or [], + request_tools=chat_request.get("tools"), + response_message=choice.get("message") or {}, + ) + ) + session.last_turn_at = __import__("time").time() + session.metadata.setdefault("api_type", api_type.value) + except Exception: # noqa: BLE001 + logger.exception("[%s] ingest failed; turn dropped", session.session_id) + + return app + + +def main() -> None: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument( + "--llm-url", required=True, help="OpenAI-spec endpoint you host" + ) + parser.add_argument( + "--model", default=None, help="served model name to send upstream" + ) + parser.add_argument("--engine", default="vllm", choices=["vllm", "sglang"]) + parser.add_argument("--host", default="0.0.0.0") + parser.add_argument("--port", type=int, default=8100) + parser.add_argument( + "--max-output-tokens", + type=int, + default=8192, + help="cap on requested completion length; 0 disables", + ) + parser.add_argument( + "--allow-unregistered", + action="store_true", + help="serve unknown API keys (local debugging only; this port may be public)", + ) + args = parser.parse_args() + + import uvicorn + + logging.basicConfig( + level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s" + ) + uvicorn.run( + create_app( + llm_url=args.llm_url, + model=args.model, + engine=args.engine, + require_registered=not args.allow_unregistered, + max_output_tokens=args.max_output_tokens or None, + ), + host=args.host, + port=args.port, + log_level="info", + ) + + +if __name__ == "__main__": + main() diff --git a/src/openenv/core/harness/capture/sessions.py b/src/openenv/core/harness/capture/sessions.py new file mode 100644 index 000000000..613bbeeec --- /dev/null +++ b/src/openenv/core/harness/capture/sessions.py @@ -0,0 +1,147 @@ +"""Session routing: one server, one port, N concurrent rollouts. + +The whole multiplexing scheme is one decision: **the API key IS the session id**. We mint a key per +rollout, hand it to the agent as its `OPENAI_API_KEY` (or `ANTHROPIC_API_KEY`, or a provider config +field), and every SDK forwards it unchanged on every request. So the bearer token that arrives is +already the rollout identifier, and no agent needs to know it is being recorded. + +That is what makes one intercept server serve a whole GRPO group. The alternative, a proxy per +sandbox, means N processes, N ports, N forwards, and capture living inside the thing most likely to +die. + +Two rules learned the hard way: + + * **A registered key beats every other hint.** Harnesses inject their own session headers, and + opencode sends `x-session-id: ses_...` from the AI SDK. Letting that win files the trajectory + under an id the caller has never seen, so lookups return nothing and it reads as "the agent made + no model calls" while every turn was in fact captured. That cost real debugging time. + * **The harness's own session id is kept, not discarded.** It is recorded on the node as + `harness_session_id`. Our key scopes the ROLLOUT; theirs identifies the sub-conversation, which + is exactly the ground truth needed to separate a subagent from the main agent. + +Unknown keys are rejected when `require_registered` is on. It defaults to on because this server sits +behind a public forward in front of a GPU, and an open inference endpoint is a real cost, not a +theoretical one. +""" + +from __future__ import annotations + +import re +import secrets +import threading +import time +from dataclasses import dataclass, field +from typing import Any + +from .graph import RolloutGraph + +# Session ids become dict keys, filenames, and URL path segments. +_SESSION_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,127}$") + + +def clean_session_id(value: Any) -> str | None: + if isinstance(value, str) and _SESSION_ID_RE.fullmatch(value.strip()): + return value.strip() + return None + + +def extract_api_key(headers: dict[str, str]) -> str | None: + """Every major SDK puts its key in one of these three places.""" + lower = {k.lower(): v for k, v in headers.items()} + auth = lower.get("authorization", "") + if auth.lower().startswith("bearer "): + return auth[7:].strip() + return lower.get("x-api-key") or lower.get("x-goog-api-key") + + +def extract_harness_session( + headers: dict[str, str], body: dict[str, Any] +) -> str | None: + """The harness's own conversation id, when it volunteers one. Never used for routing.""" + lower = {k.lower(): v for k, v in headers.items()} + return ( + clean_session_id(lower.get("x-session-id")) + or clean_session_id(lower.get("proxy-x-session-id")) + or clean_session_id(body.get("_session_id")) + or clean_session_id(body.get("user")) + ) + + +@dataclass +class Session: + """One rollout's capture buffer.""" + + session_id: str + created_at: float = field(default_factory=time.time) + graph: RolloutGraph = field(default_factory=RolloutGraph) + metadata: dict[str, Any] = field(default_factory=dict) + findings: list[str] = field(default_factory=list) + last_turn_at: float | None = None + upstream_errors: int = 0 + + @property + def idle_seconds(self) -> float: + """Since the last captured turn. The cheapest signal that separates progress from a wedge.""" + return time.time() - (self.last_turn_at or self.created_at) + + +class SessionRegistry: + """Thread-safe. Uvicorn serves concurrently and rollouts are independent.""" + + def __init__(self, *, require_registered: bool = True) -> None: + self._sessions: dict[str, Session] = {} + self._lock = threading.Lock() + self.require_registered = require_registered + + def create(self, session_id: str | None = None, **metadata: Any) -> Session: + sid = clean_session_id(session_id) or f"s{secrets.token_hex(12)}" + with self._lock: + session = self._sessions.get(sid) or Session(session_id=sid) + session.metadata.update(metadata) + self._sessions[sid] = session + return session + + def get(self, session_id: str | None) -> Session | None: + if not session_id: + return None + with self._lock: + return self._sessions.get(session_id) + + def resolve(self, headers: dict[str, str], body: dict[str, Any]) -> Session | None: + """Route a request to its rollout. `None` means reject. + + Order matters and is the opposite of what looks natural: the registered API key wins over any + session header the harness supplies. See the module docstring. + """ + api_key = extract_api_key(headers) + session = self.get(api_key) + if session is not None: + return session + if self.require_registered: + return None + # Open mode (local debugging only): an unknown caller still gets a session, so a + # misconfigured agent shows up as an orphan trajectory instead of vanishing. + return self.create(api_key) + + def list_ids(self) -> list[str]: + with self._lock: + return sorted(self._sessions) + + def delete(self, session_id: str) -> bool: + with self._lock: + return self._sessions.pop(session_id, None) is not None + + def summary(self) -> list[dict[str, Any]]: + with self._lock: + sessions = list(self._sessions.values()) + return [ + { + "session_id": s.session_id, + "turns": s.graph.stats()["n_turns"], + "roots": s.graph.stats()["n_roots"], + "idle_s": round(s.idle_seconds, 1), + "upstream_errors": s.upstream_errors, + **s.metadata, + } + for s in sessions + ] diff --git a/src/openenv/core/harness/capture/sse.py b/src/openenv/core/harness/capture/sse.py new file mode 100644 index 000000000..644a3536d --- /dev/null +++ b/src/openenv/core/harness/capture/sse.py @@ -0,0 +1,175 @@ +"""Synthetic SSE: capture non-streaming, reply streaming. + +Every coding harness streams. That is not a preference we can talk them out of -- opencode, codex and +claude-code all drive their UI off token deltas -- and a harness that asks for SSE and receives a +plain JSON body does not error. Its stream parser simply yields nothing, so opencode reports +`step-finish reason:"unknown"` with zero tokens and no message, having received a *perfectly valid* +tool call. Capture looks flawless from our side and the agent does nothing. That failure cost real +debugging time, hence this module. + +Meanwhile capture wants the opposite: one complete response, because token ids and logprobs arrive +whole and reassembling them from deltas is error-prone in exactly the way that silently corrupts +training data. + +So we do both. Fetch non-streaming upstream, store that for capture, then replay the complete +response to the client as a synthetic SSE stream. The client cannot tell the difference; we never +parse deltas. + +The per-dialect machinery is reused from Polar's transformers (`create_stream_state` / +`transform_stream_chunk`), which are dependency-clean. Only the small formatting helpers are ported +here, because in Polar they live in `server.py` next to its node/dispatcher layer. +""" + +from __future__ import annotations + +import json +from typing import Any + +from .detection import APIType +from .dialects.base import BaseTransformer + +SSE_HEADERS = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + # Without this an intermediate proxy (nginx, cloudflared) may buffer the whole stream and hand + # it over at once, which defeats the point for the client even though the bytes are correct. + "X-Accel-Buffering": "no", +} + + +def _format_typed_events(events: list[dict[str, Any]]) -> str: + """Anthropic and OpenAI-Responses both use named events (`event: `).""" + return "".join( + f"event: {event.get('type', 'unknown')}\ndata: {json.dumps(event, default=str)}\n\n" + for event in events + ) + + +def _format_data_only(chunk: dict[str, Any]) -> str: + """OpenAI chat-completions and Google use bare `data:` lines.""" + return f"data: {json.dumps(chunk, default=str)}\n\n" + + +def format_events(api_type: APIType, events: list[dict[str, Any]]) -> str: + """Format EVERY event. Dropping any of them truncates the stream. + + This used to emit only `events[0]` for the data-only dialects, which is invisible for + chat-completions (we synthesise exactly one chunk) but silently truncates Google. gemini-cli calls + `:streamGenerateContent?alt=sse`, whose stream state emits several events, and receiving only the + first produced: + + Error: Incomplete JSON segment at the end + at ApiClient.processStreamResponse_1 (@google/gemini-cli/...) + + A 200 with a truncated body, which is the failure shape this whole layer keeps running into. + """ + if api_type in (APIType.ANTHROPIC, APIType.OPENAI_RESPONSES): + return _format_typed_events(events) + return "".join(_format_data_only(event) for event in events) + + +def format_chunk( + api_type: APIType, + transformer: BaseTransformer, + chunk: dict[str, Any], + original_request: dict[str, Any], + *, + is_first: bool, +) -> str: + """Fallback for transformers with no stream-state machine: one-shot chunk transform.""" + transformed = transformer.transform_stream_chunk( + chunk, original_request, is_first=is_first + ) + if api_type == APIType.ANTHROPIC: + return _format_typed_events(transformed) + if api_type == APIType.OPENAI_RESPONSES: + events = ( + transformed + if isinstance(transformed, list) + else ([transformed] if transformed else []) + ) + return _format_typed_events(events) + return _format_data_only(transformed) + + +def response_to_chunk(response: dict[str, Any]) -> dict[str, Any]: + """Repackage a complete chat completion as a single `chat.completion.chunk` delta. + + One chunk carrying everything, rather than a plausible-looking token-by-token replay. The client + only needs a well-formed stream, and faking granularity would invent timing information we do not + have. Tool calls have to be re-indexed into delta form: streaming clients accumulate + `tool_calls[i].function.arguments` across chunks, so the `index` field is required even when + there is exactly one chunk to accumulate. + """ + choice = (response.get("choices") or [{}])[0] + message = choice.get("message") or {} + + tool_calls_delta = [ + { + "index": i, + "id": tc.get("id"), + "type": tc.get("type", "function"), + "function": { + "name": (tc.get("function") or {}).get("name", ""), + "arguments": (tc.get("function") or {}).get("arguments", ""), + }, + } + for i, tc in enumerate(message.get("tool_calls") or []) + ] + + delta: dict[str, Any] = {"role": "assistant"} + if message.get("content") is not None: + delta["content"] = message["content"] + # Reasoning models put thinking here; harnesses that render it expect it in the delta. + for key in ("reasoning_content", "reasoning"): + if message.get(key) is not None: + delta["reasoning_content"] = message[key] + break + if tool_calls_delta: + delta["tool_calls"] = tool_calls_delta + + return { + "id": response.get("id"), + "object": "chat.completion.chunk", + "created": response.get("created"), + "model": response.get("model"), + "choices": [ + {"index": 0, "delta": delta, "finish_reason": choice.get("finish_reason")} + ], + # Clients that sent stream_options.include_usage expect this; we dropped the option upstream + # (vLLM rejects it with stream=False) but the response carries usage anyway, so honour it. + "usage": response.get("usage"), + } + + +async def replay( + api_type: APIType, + transformer: BaseTransformer, + response: dict[str, Any], + original_request: dict[str, Any], +): + """Async generator yielding the SSE body for one complete upstream response.""" + chunk = response_to_chunk(response) + stream_state = transformer.create_stream_state(original_request) + + if stream_state is not None: + # Dialects with real state machines (Anthropic, Responses, Google) emit a sequence of + # lifecycle events -- message_start, content_block_delta, message_stop and friends -- and the + # client will reject a stream that skips them, so finalize() is not optional. + events = stream_state.process_chunk(chunk, is_first=True) + if events: + yield format_events(api_type, events) + final_events = stream_state.finalize() + if final_events: + yield format_events(api_type, final_events) + else: + output = format_chunk( + api_type, transformer, chunk, original_request, is_first=True + ) + if output: + yield output + + if api_type == APIType.OPENAI_CHAT: + # Only chat-completions uses this sentinel. The typed-event dialects signal completion with + # their own terminal event, and an extra [DONE] there is a parse error. + yield "data: [DONE]\n\n" diff --git a/src/openenv/core/harness/capture/upstream.py b/src/openenv/core/harness/capture/upstream.py new file mode 100644 index 000000000..ce02d247b --- /dev/null +++ b/src/openenv/core/harness/capture/upstream.py @@ -0,0 +1,182 @@ +"""The upstream leg: one HTTP client to a vLLM OpenAI-compatible server. + +vLLM only, deliberately. SGLang cannot support this layer at all — none of +`return_tokens_as_token_ids`, `logprobs_mode`, `processed_logprobs` or `return_token_ids` exist in +its tree, and its chat route returns token *text* with no ids (sgl-project/sglang#18378 requests +exactly this, for the same train/inference consistency reason). Carrying a two-engine abstraction for +a backend that structurally cannot work would be pretending we have a choice. + +Two request params do all the work, and both are easy to get subtly wrong: + + return_token_ids=True makes vLLM emit `response.prompt_token_ids` and `choice.token_ids`. + `prompt_token_ids` is the load-bearing one: it is the engine's own + tokenisation of the whole conversation so far, which is what lets turn + k+1 be matched against turn k by exact token prefix without us ever + tokenising locally. + top_logprobs=0 must be SET, not omitted. vLLM only populates `logprobs.content[]` when + `top_logprobs` is not None, even with `logprobs=True`. Zero returns just + the sampled token's logprob, which is all training needs. +""" + +from __future__ import annotations + +import json +from typing import Any + +import httpx + + +class UpstreamError(RuntimeError): + """Any failure talking to the engine. Never leaks httpx types to callers.""" + + +class UpstreamHTTPError(UpstreamError): + """Engine answered with a non-2xx status.""" + + def __init__( + self, status_code: int, body: dict[str, Any] | str | None = None + ) -> None: + self.status_code = status_code + self.body = body + detail = body + if isinstance(body, dict): + error = body.get("error") + detail = error.get("message") if isinstance(error, dict) else error or body + super().__init__(f"upstream returned {status_code}: {str(detail)[:400]}") + + +class UpstreamTimeoutError(UpstreamError): + """Engine did not answer within the liveness ceiling.""" + + +class UpstreamTransportError(UpstreamError): + """Connection-level failure: refused, reset, DNS.""" + + +def prepare_request( + request: dict[str, Any], *, served_model: str | None = None +) -> dict[str, Any]: + """Add the params that make a response capturable. Mutates and returns `request`.""" + request["logprobs"] = True + request["return_token_ids"] = True + request.setdefault("top_logprobs", 0) + + # vLLM reads a prior turn's thinking from `reasoning`, while the dialect transformers emit the + # canonical `reasoning_content`. Without this rename an earlier turn's interleaved thinking + # renders as an empty `` and the prompt silently differs from what the model + # actually produced — which breaks prefix matching for the turn after it. + for message in request.get("messages") or []: + if isinstance(message, dict) and message.get("reasoning_content") is not None: + message["reasoning"] = message.pop("reasoning_content") + + if served_model: + # Read by the transformers for per-model request fixes (e.g. Qwen3.5 emits tool calls inside + # thinking, so thinking has to be disabled for tool use). Stripped again before the request + # leaves, in `BaseTransformer._normalize_request`. + request["_served_model"] = served_model + return request + + +def normalize_response(response: dict[str, Any]) -> dict[str, Any]: + """Canonicalise vLLM's response shape in place.""" + choices = response.get("choices") + if not isinstance(choices, list): + return response + + for choice in choices: + if not isinstance(choice, dict): + continue + + message = choice.get("message") + if isinstance(message, dict): + if ( + message.get("reasoning_content") is None + and message.get("reasoning") is not None + ): + message["reasoning_content"] = message.pop("reasoning") + + # Copy each token id onto its logprob entry. Not load-bearing — capture reads + # `choice.token_ids` and the per-entry `logprob` — but it keeps a stored trace one shape, + # so a consumer never has to know which engine produced it. Guarded on equal length because + # a mismatch means the two lists are not describing the same tokens, and pairing them anyway + # would silently attach the wrong id to every logprob. + token_ids = choice.get("token_ids") + entries = ((choice.get("logprobs") or {}).get("content")) or [] + if isinstance(token_ids, list) and len(token_ids) == len(entries): + for token_id, entry in zip(token_ids, entries): + if isinstance(entry, dict): + entry.setdefault("token_id", token_id) + return response + + +class InferenceClient: + """Async client to one engine. One instance per server, shared across sessions.""" + + # A high ceiling, not a per-request budget. Callers impose their own deadline; this exists only + # so a wedged engine cannot pin a connection forever. + _LIVENESS_TIMEOUT_S = 900.0 + _CONNECT_TIMEOUT_S = 30.0 + + def __init__(self, base_url: str, *, served_model: str | None = None) -> None: + self.base_url = base_url.rstrip("/") + self.served_model = served_model + self._client: httpx.AsyncClient | None = None + + async def _get_client(self) -> httpx.AsyncClient: + if self._client is None or self._client.is_closed: + self._client = httpx.AsyncClient( + base_url=self.base_url, + timeout=httpx.Timeout( + self._LIVENESS_TIMEOUT_S, connect=self._CONNECT_TIMEOUT_S + ), + ) + return self._client + + async def aclose(self) -> None: + if self._client is not None and not self._client.is_closed: + await self._client.aclose() + self._client = None + + async def completion(self, request: dict[str, Any]) -> dict[str, Any]: + """One non-streaming chat completion, prepared for capture and normalised on the way back.""" + body = prepare_request(dict(request), served_model=self.served_model) + payload = await self._post("/v1/chat/completions", body) + return normalize_response(payload) + + async def list_models(self) -> dict[str, Any]: + client = await self._get_client() + try: + response = await client.get("/v1/models") + except httpx.RequestError as exc: + raise self._transport_error(exc) from exc + await self._raise_for_status(response) + return response.json() + + async def _post(self, path: str, body: dict[str, Any]) -> dict[str, Any]: + client = await self._get_client() + try: + response = await client.post(path, json=body) + except httpx.RequestError as exc: + raise self._transport_error(exc) from exc + await self._raise_for_status(response) + return response.json() + + async def _raise_for_status(self, response: httpx.Response) -> None: + if response.is_success: + return + content = await response.aread() + await response.aclose() + body: dict[str, Any] | str | None = None + text = content.decode("utf-8", errors="replace").strip() + if text: + try: + body = json.loads(text) + except json.JSONDecodeError: + body = text + raise UpstreamHTTPError(response.status_code, body) + + @staticmethod + def _transport_error(exc: httpx.RequestError) -> UpstreamError: + if isinstance(exc, httpx.TimeoutException): + return UpstreamTimeoutError(f"engine timed out: {exc}") + return UpstreamTransportError(f"could not reach engine: {exc}") diff --git a/src/openenv/core/harness/capture/validate.py b/src/openenv/core/harness/capture/validate.py new file mode 100644 index 000000000..1c81fe015 --- /dev/null +++ b/src/openenv/core/harness/capture/validate.py @@ -0,0 +1,305 @@ +"""Token-in / token-out validation. Nothing leaves this system unchecked. + +Capture failures in this stack are silent by construction. Every bug found so far returned a +perfectly well-formed payload and reported success: a missing `--return-tokens-as-token-ids` yields +text with no ids and trains on nothing; a harness that rewrites its history yields chains that stitch +into a trajectory the model never saw; an SSE client handed a JSON body yields zero tokens and no +error anywhere. None of these raise. All of them produce plausible JSON. + +So validation is not a debug aid here, it is the only thing standing between a clean-looking run and +weeks of training on corrupted sequences. Checks are graded: + + FATAL the row is not trainable. Drop it. Training on it is worse than dropping it. + WARN the row is trainable but something is off and should be understood. + INFO recorded for the per-harness notes. + +`check_upstream` runs before a rollout is spent (endpoint capability), `check_turn` runs per model +call (token/logprob alignment), `check_sequence` runs on the flattened output (mask/logprob +invariants), and `check_rollout` runs on the whole graph (structure). +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +FATAL, WARN, INFO = "FATAL", "WARN", "INFO" + + +@dataclass +class Finding: + level: str + code: str + detail: str + + def __str__(self) -> str: + return f"[{self.level}] {self.code}: {self.detail}" + + +@dataclass +class Report: + findings: list[Finding] = field(default_factory=list) + # Nodes the harness's own trace does not count as agent steps (auxiliary calls). Populated only + # by a harness-trace reconciler. Captured correctly, but excluded from training so they cannot be credited + # with the rollout's reward. Empty for every harness whose counts agree exactly. + aux_node_ids: list[str] = field(default_factory=list) + + def add(self, level: str, code: str, detail: str) -> None: + self.findings.append(Finding(level, code, detail)) + + @property + def fatal(self) -> list[Finding]: + return [f for f in self.findings if f.level == FATAL] + + @property + def ok(self) -> bool: + return not self.fatal + + def merge(self, other: "Report") -> "Report": + self.findings.extend(other.findings) + return self + + def text(self) -> str: + if not self.findings: + return "all checks passed" + return "\n".join(str(f) for f in self.findings) + + +# --- per model call ---------------------------------------------------- +def check_turn( + prompt_ids, sampled_ids, logprobs, *, finish_reason=None, index=0 +) -> Report: + """Validate one captured call before it becomes a graph node.""" + report = Report() + tag = f"turn {index}" + + if not prompt_ids: + report.add( + FATAL, + "no_prompt_ids", + f"{tag}: server returned no prompt_token_ids. The endpoint " + "is missing --return-tokens-as-token-ids, or the engine does not support it.", + ) + if not sampled_ids: + # Legitimate when the model is cut off at zero tokens, but nothing is trainable either way. + report.add( + WARN, + "no_sampled_ids", + f"{tag}: no completion token ids (finish={finish_reason})", + ) + + if logprobs is None: + if sampled_ids: + report.add( + FATAL, + "no_logprobs", + f"{tag}: {len(sampled_ids)} sampled tokens with no " + "logprobs. GRPO's importance ratio needs the behaviour-policy logprob for " + "every trainable token; without them the turn can only be context.", + ) + elif len(logprobs) != len(sampled_ids): + report.add( + FATAL, + "logprob_misalign", + f"{tag}: {len(logprobs)} logprobs vs {len(sampled_ids)} sampled ids. An off-by-one " + "here shifts credit onto the wrong tokens and still trains.", + ) + return report + + +# --- per flattened sequence ------------------------------------------- +def check_sequence(seq, *, min_trainable: int = 1) -> Report: + """Validate a flattened path. These are the invariants the trainer assumes and never re-checks.""" + report = Report() + n = len(seq.input_ids) + + if len(seq.loss_mask) != n or len(seq.logprobs) != n: + report.add( + FATAL, + "length_mismatch", + f"input_ids={n} loss_mask={len(seq.loss_mask)} logprobs={len(seq.logprobs)}", + ) + return report # every later check would be meaningless + + trainable = sum(seq.loss_mask) + if trainable < min_trainable: + report.add( + FATAL, + "nothing_trainable", + f"{trainable} trainable tokens in a {n}-token sequence: this row contributes no " + "gradient and only shrinks the effective group", + ) + + # A masked position carrying a logprob means context was scored; a trainable position without one + # means a target was invented. Both are silent corruption, in opposite directions. + masked_with_lp = sum( + 1 for m, lp in zip(seq.loss_mask, seq.logprobs) if m == 0 and lp != 0.0 + ) + if masked_with_lp: + report.add( + FATAL, + "masked_has_logprob", + f"{masked_with_lp} context positions carry a non-zero logprob", + ) + + if seq.prompt_len >= n: + report.add( + FATAL, + "empty_response", + f"prompt_len={seq.prompt_len} covers the whole sequence", + ) + + # Not an error: a token with logprob exactly 0.0 has probability 1.0, which is common for + # structural tokens in a constrained tool-call grammar (``, closing brackets). + # Recorded so a *sudden* change in the rate is visible. + zero_lp_trainable = sum( + 1 for m, lp in zip(seq.loss_mask, seq.logprobs) if m == 1 and lp == 0.0 + ) + if zero_lp_trainable: + report.add( + INFO, + "certain_tokens", + f"{zero_lp_trainable}/{trainable} trainable tokens have logprob 0.0 (p=1.0)", + ) + + positive = [lp for lp in seq.logprobs if lp > 0.0] + if positive: + report.add( + FATAL, + "positive_logprob", + f"{len(positive)} logprobs > 0 (max {max(positive):.4f}); log-probabilities cannot " + "be positive, so these are not logprobs", + ) + return report + + +# --- per rollout ------------------------------------------------------- +def check_rollout(graph, *, expect_single_root: bool = False) -> Report: + """Validate graph structure. This is where harness-specific weirdness shows up first.""" + report = Report() + stats = graph.stats() + + if stats["n_turns"] == 0: + report.add( + FATAL, + "no_turns", + "the intercept saw no model calls: the agent never reached it " + "(wrong base URL, unresolved model, or auth rejected)", + ) + return report + + if stats["n_roots"] == stats["n_turns"] and stats["n_turns"] > 1: + # One root PER TURN: the harness re-renders its prompt each turn instead of appending, so no + # token prefix is shared. terminus-2 does this (its message list grows 1,3,5..15 while the + # rendered tokens never line up). + # + # WARN, not FATAL. Each turn is still an exact prompt with exact sampled tokens and real + # logprobs, which is perfectly good SINGLE-turn training data. What is lost is cross-turn + # structure: the interstitial tool results are never masked-in as context, so credit cannot + # flow across turns. Calling that unusable would throw away correct data; calling it clean + # would hide a real degradation. So: trainable, and labelled. + report.add( + WARN, + "per_turn_capture_only", + f"every turn is its own root ({stats['n_turns']}). This harness re-renders its " + "prompt rather than appending, so rows are single-turn. Tokens and logprobs are " + "exact; multi-turn credit assignment is not available.", + ) + elif stats["n_roots"] > 1: + # Normal: aux calls (title generation), subagents, or a harness that rewrote its system + # prompt partway (claude-code) and so continued under a second prefix family. + report.add( + WARN, + "multiple_roots", + f"{stats['n_roots']} roots across {stats['n_turns']} turns. Each root is a separate " + "conversation (subagent, aux call, or a rewritten prompt that broke the chain).", + ) + elif expect_single_root and stats["n_roots"] != 1: + report.add(FATAL, "root_count", f"expected 1 root, got {stats['n_roots']}") + + if stats["n_turns"] == 1: + # ONE call for an entire agentic task. Capture is trivially self-consistent here (a single + # turn has nothing to stitch to and no prefix to disagree with), so every other check in this + # file passes and the rollout reads as clean. It is not: an agent that made one model call + # and stopped did not attempt the task. + # + # Found the hard way. swe-agent, trae-agent, nemo-agent and antigravity-sdk each passed 5/5 + # while producing exactly one turn per task and solving 0/5, for four unrelated harness-side + # reasons (litellm cost registry, a null `prompt_tokens_details`, a tool-less prompt format, + # and an SDK loop that exits after the first tool call). The capture layer was right every + # time and the rollouts were still worthless. + # + # FATAL because the whole point of this layer is refusing to hand over data we cannot stand + # behind, and a one-turn agentic rollout is a harness failure wearing a clean capture. + report.add( + FATAL, + "degenerate_rollout", + "exactly 1 model call for the whole task: the agent stopped after its first " + "response. Capture is self-consistent because there is nothing to stitch, so the " + "other checks cannot see this. Read the trial's agent stdout for the real cause.", + ) + + if stats["n_discarded"]: + report.add( + WARN, + "discarded_turns", + f"{stats['n_discarded']} sampled turn(s) led nowhere (retries or resamples). They " + "are excluded from training paths; the tokens were still generated and billed.", + ) + if stats["n_forks"]: + report.add(INFO, "forks", f"{stats['n_forks']} fork point(s) in the graph") + return report + + +# --- endpoint capability, before spending a sandbox -------------------- +def check_upstream_response(payload: dict[str, Any]) -> Report: + """Assert a raw chat-completions reply actually carries what capture needs. + + Run this against the endpoint before booting anything. The failure it catches (an endpoint served + without the capture flags) otherwise surfaces only as empty training rows, hours later. + """ + report = Report() + choices = payload.get("choices") or [] + if not choices: + report.add(FATAL, "no_choices", "response has no choices") + return report + choice = choices[0] + + if not payload.get("prompt_token_ids"): + report.add( + FATAL, + "no_prompt_token_ids", + "no top-level prompt_token_ids. Serve with --return-tokens-as-token-ids; without " + "it multi-turn stitching is impossible because turn k+1's prompt is unknown.", + ) + if not choice.get("token_ids"): + report.add( + FATAL, + "no_completion_token_ids", + "choices[0].token_ids missing. Send return_token_ids=True and serve with " + "--return-tokens-as-token-ids.", + ) + + content = (choice.get("logprobs") or {}).get("content") + if not content: + report.add( + FATAL, + "no_logprobs", + "choices[0].logprobs.content missing. Send logprobs=True.", + ) + else: + if choice.get("token_ids") and len(content) != len(choice["token_ids"]): + report.add( + FATAL, + "logprob_misalign", + f"{len(content)} logprobs vs {len(choice['token_ids'])} token ids", + ) + token = (content[0] or {}).get("token", "") + if not str(token).startswith("token_id:"): + report.add( + WARN, + "token_strings", + f"logprob tokens are strings ({token!r}) not 'token_id:N'. Ids are being " + "recovered from a side channel rather than the tokenizer's own numbering.", + ) + return report diff --git a/src/openenv/core/harness/capture/validate_llm.py b/src/openenv/core/harness/capture/validate_llm.py new file mode 100644 index 000000000..ccfe22e1e --- /dev/null +++ b/src/openenv/core/harness/capture/validate_llm.py @@ -0,0 +1,162 @@ +"""Certify that an inference engine can actually support token-level capture. + +This runs BEFORE a server binds a port, and refusing here is the entire point. + +An engine missing `--return-tokens-as-token-ids --logprobs-mode processed_logprobs` still answers +every request perfectly well: it returns text, a `200`, and plausible-looking usage. What it does not +return is token ids. Every row rebuilt downstream is then empty, training silently does nothing, and +the first symptom is a loss curve that never moves days later. That failure has no loud edge, so the +check has to be up front and fatal. + +Only vLLM implements the contract today. SGLang has none of the four capture knobs +(`return_tokens_as_token_ids`, `logprobs_mode`, `processed_logprobs`, `return_token_ids` all match +zero files in its tree) and its chat route returns token *text* only — see sgl-project/sglang#18378, +which requests exactly this and is motivated by the same train/inference consistency problem. A +hosted alternative exists but is narrow: fireworks-ai via the HF router honours vLLM's +`return_token_ids`, though every one of its live models is a reasoning model whose reasoning tokens +are dropped from history, so multi-turn stitching degrades to per-turn. +""" + +from __future__ import annotations + +import json +import urllib.error +import urllib.request +from dataclasses import dataclass, field + +from .validate import check_upstream_response + + +@dataclass +class LLMReport: + """Outcome of certification. `ok` gates whether a server may start.""" + + ok: bool + llm_url: str + model: str + findings: list[str] = field(default_factory=list) + n_prompt_ids: int = 0 + n_completion_ids: int = 0 + served_models: list[str] = field(default_factory=list) + + def summary(self) -> str: + if self.ok: + return ( + f"engine OK: {self.n_completion_ids} completion ids, " + f"{self.n_prompt_ids} prompt ids" + ) + return "engine NOT usable for capture:\n " + "\n ".join(self.findings) + + +def _post(url: str, body: dict, timeout: float) -> dict: + request = urllib.request.Request( + url, + data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(request, timeout=timeout) as response: + return json.loads(response.read()) + + +def list_models(llm_url: str, timeout: float = 30.0) -> list[str]: + """Served model ids, or [] if the endpoint is unreachable.""" + try: + with urllib.request.urlopen( + f"{llm_url.rstrip('/')}/v1/models", timeout=timeout + ) as r: + return [m.get("id", "") for m in json.loads(r.read()).get("data", [])] + except Exception: # noqa: BLE001 - unreachable is reported by the caller, not raised here + return [] + + +def validate_llm(llm_url: str, model: str, *, timeout: float = 120.0) -> LLMReport: + """Send one real completion and assert the response carries what capture needs. + + Deliberately a live probe rather than a flag inspection: launch flags are not readable over the + API, and an engine can be started with the right arguments and still not behave (wrong version, + a proxy in between that strips fields). The only trustworthy check is asking for a completion and + looking at what comes back. + """ + base = llm_url.rstrip("/") + served = list_models(base, timeout=min(timeout, 30.0)) + + if not served: + return LLMReport( + ok=False, + llm_url=base, + model=model, + findings=[ + f"GET {base}/v1/models returned nothing; the engine is unreachable" + ], + ) + + if model not in served: + # Worth failing on rather than warning: a mismatched name is silently accepted by some + # servers and then every request 404s at rollout time instead of at startup. + return LLMReport( + ok=False, + llm_url=base, + model=model, + served_models=served, + findings=[f"model {model!r} is not served here; available: {served}"], + ) + + body = { + "model": model, + "messages": [{"role": "user", "content": "Reply with the single word: ok"}], + "max_tokens": 8, + "temperature": 0.0, + "logprobs": True, + "top_logprobs": 0, + # vLLM >= 0.10.2 exposes this on the OpenAI route. Harmless where unsupported; its ABSENCE + # from the response is exactly the signal we are testing for. + "return_token_ids": True, + } + try: + payload = _post(f"{base}/v1/chat/completions", body, timeout) + except urllib.error.HTTPError as exc: + detail = exc.read()[:300].decode(errors="replace") + return LLMReport( + ok=False, + llm_url=base, + model=model, + served_models=served, + findings=[f"probe failed: HTTP {exc.code}: {detail}"], + ) + except Exception as exc: # noqa: BLE001 + return LLMReport( + ok=False, + llm_url=base, + model=model, + served_models=served, + findings=[f"probe failed: {type(exc).__name__}: {str(exc)[:300]}"], + ) + + report = check_upstream_response(payload) + choice = (payload.get("choices") or [{}])[0] + return LLMReport( + ok=report.ok, + llm_url=base, + model=model, + served_models=served, + findings=[str(f) for f in report.findings], + n_prompt_ids=len(payload.get("prompt_token_ids") or []), + n_completion_ids=len(choice.get("token_ids") or []), + ) + + +def require_llm(llm_url: str, model: str, *, timeout: float = 120.0) -> LLMReport: + """`validate_llm`, but raises instead of returning a failed report. + + For the startup path, where continuing past a bad engine is never the right behaviour. + """ + report = validate_llm(llm_url, model, timeout=timeout) + if not report.ok: + raise RuntimeError( + report.summary() + + "\n\nA vLLM server must be started with:" + + "\n --return-tokens-as-token-ids --logprobs-mode processed_logprobs" + + "\nWithout them the engine returns text with no token ids, every rebuilt training row" + + " is empty, and nothing downstream reports an error." + ) + return report From 5976ac92e45f245021f6916174c06c3478fabf25 Mon Sep 17 00:00:00 2001 From: adithya-s-k Date: Sun, 2 Aug 2026 19:23:03 +0000 Subject: [PATCH 02/74] harbor: run Harbor tasks as an OpenEnv environment Serves Harbor's task datasets over the Task API and runs a rollout through one long-running MCP tool, with the agent and the sandbox chosen per call rather than baked into the deployment. A failed rollout returns a result, never an exception. That is the reason this layer exists: in the in-process predecessor a rollout exception reached the trainer and hung every rank at the NCCL barrier, which is why trl.experimental.harbor wraps nearly every environment call individually. Behind an HTTP boundary that failure class cannot occur. Rewards are forwarded, never recomputed. Harbor's dict travels verbatim and the scalar is chosen by an explicit rule, refusing rather than guessing when several keys exist. reward=None is not zero: it means the verifier never ran, and conflating them makes a dead sandbox look like a wrong answer. Sandbox availability is asked for rather than assumed. A backend counts as usable only if its class imports, its SDK is present, and Harbor's own preflight passes; checking credentials alone reports a backend available and then fails at rollout time. Hosted deployments mount the capture proxy on the env server's own app, since a Space has one port and one public URL and nothing needs forwarding there. --- src/openenv/harbor/__init__.py | 18 + src/openenv/harbor/atif.py | 342 +++++++++ src/openenv/harbor/capabilities.py | 254 +++++++ src/openenv/harbor/client.py | 173 +++++ src/openenv/harbor/environment.py | 232 ++++++ src/openenv/harbor/install_fixes.py | 479 ++++++++++++ src/openenv/harbor/models.py | 253 +++++++ src/openenv/harbor/rollout.py | 352 +++++++++ src/openenv/harbor/runner.py | 280 +++++++ src/openenv/harbor/seams.py | 766 +++++++++++++++++++ src/openenv/harbor/serving.py | 243 ++++++ src/openenv/harbor/startup.py | 159 ++++ src/openenv/harbor/tasks.py | 301 ++++++++ src/openenv/harbor/ui.py | 1090 +++++++++++++++++++++++++++ 14 files changed, 4942 insertions(+) create mode 100644 src/openenv/harbor/__init__.py create mode 100644 src/openenv/harbor/atif.py create mode 100644 src/openenv/harbor/capabilities.py create mode 100644 src/openenv/harbor/client.py create mode 100644 src/openenv/harbor/environment.py create mode 100644 src/openenv/harbor/install_fixes.py create mode 100644 src/openenv/harbor/models.py create mode 100644 src/openenv/harbor/rollout.py create mode 100644 src/openenv/harbor/runner.py create mode 100644 src/openenv/harbor/seams.py create mode 100644 src/openenv/harbor/serving.py create mode 100644 src/openenv/harbor/startup.py create mode 100644 src/openenv/harbor/tasks.py create mode 100644 src/openenv/harbor/ui.py diff --git a/src/openenv/harbor/__init__.py b/src/openenv/harbor/__init__.py new file mode 100644 index 000000000..edf7c0ca2 --- /dev/null +++ b/src/openenv/harbor/__init__.py @@ -0,0 +1,18 @@ +"""Harbor integration: run Harbor tasks as OpenEnv environments, for eval and training. + +Harbor owns what it is good at — task datasets, sandbox backends, coding agents, verifiers, trial +concurrency, pass@k. This package adds the OpenEnv side and nothing more: + + tasks.py dataset discovery over the Task API (HF repo | local dir | Harbor registry) + seams.py how each agent is pointed at the capture proxy — the only per-agent knowledge + install_fixes.py subclasses for agents whose Harbor wrapper cannot be configured as shipped + atif.py cross-check captured tokens against Harbor's own ATIF trajectory + models.py wire types + +Generic capture lives in `openenv.core.harness.capture` and knows nothing about Harbor. The +dependency runs one way only: `openenv.harbor` imports capture, never the reverse. ATIF is here +rather than there because it is Harbor's trace format, not a general one. + +Harbor itself is an optional dependency (`pip install openenv[harbor]`), imported lazily so that +importing `openenv` never requires it. +""" diff --git a/src/openenv/harbor/atif.py b/src/openenv/harbor/atif.py new file mode 100644 index 000000000..2772fa4ce --- /dev/null +++ b/src/openenv/harbor/atif.py @@ -0,0 +1,342 @@ +"""Reconcile our token capture against Harbor's ATIF trajectory. The independent cross-check. + +Harbor agents write `agent/trajectory.json` in ATIF (Agent Trajectory Interchange Format, currently +v1.7), a published spec for logging agent interaction histories across debugging, SFT and RL. ~27 of +Harbor's agents build ATIF trajectories, which is far more than the six the docs list. + +Why this matters more than it sounds. ATIF and the intercept measure the same rollout through +completely independent paths: the harness counts its own tokens and reports them to Harbor, while we +derive ours from engine-returned token ids stitched along a graph. If they agree turn by turn, the +masking, the turn segmentation and the prefix stitching are all correct simultaneously. Nothing else +we can run gives that assurance, because every internal check shares our own assumptions. + +Validated on opencode + Qwen3.5-4B, 8 agent steps: + + ATIF completion_tokens : [37, 36, 104, 264, 255, 119, 32, 27] total 874 + intercept turn_lengths : [37, 36, 104, 264, 255, 119, 32, 27] total 874 + ATIF step-1 prompt_tokens 7990 == intercept prompt_len 7990 + +ATIF also carries three things a proxy structurally cannot see, which is the other half of the value: + + llm_call_count >1 means the harness burned several model calls on one logical step, + i.e. it retried. A proxy sees the calls but not that they were retries. + subagent_trajectories nested trajectories (v1.7). Ground truth for which turns are a subagent, + instead of inferring it from graph roots. + tool_call_id <-> observation.source_call_id + which tool result answered which call. + +`Metrics` in ATIF has optional `logprobs` and `completion_token_ids` fields, which harnesses leave +empty. So the end state is not two formats to reconcile: it is ATIF with our token fields filled in, +one artifact that is trace, SFT dataset and RL data at once. `merge_into_atif` does that. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +from openenv.core.harness.capture.validate import FATAL, INFO, Report, WARN + + +def load_atif(trial_dir: str | Path) -> dict[str, Any] | None: + """Read `agent/trajectory.json` from a Harbor trial dir. None if the agent emitted none.""" + path = Path(trial_dir) / "agent" / "trajectory.json" + if not path.is_file(): + return None + try: + return json.loads(path.read_text()) + except Exception: # noqa: BLE001 - a malformed trace must not break a good rollout + return None + + +def agent_steps(atif: dict[str, Any]) -> list[dict[str, Any]]: + """Steps the agent produced. `source` is one of user | agent | system.""" + return [s for s in (atif.get("steps") or []) if s.get("source") == "agent"] + + +def atif_turn_lengths(atif: dict[str, Any]) -> list[int]: + return [ + int((s.get("metrics") or {}).get("completion_tokens") or 0) + for s in agent_steps(atif) + ] + + +def _subsequence_gap(needle: list[int], haystack: list[int]) -> list[int] | None: + """Indices of `haystack` skipped when `needle` is matched as a subsequence, else None. + + Greedy two-pointer, which is exact for subsequence membership. Returns None the moment a needle + element cannot be found, so a genuine disagreement (a value we never captured, or counts that + differ) falls through to the FATAL path rather than being explained away as auxiliary calls. + + Requires the needle to be strictly shorter; equal lists are handled by the exact-match path and + a LONGER needle means ATIF logged calls we never saw, which is a capture failure, not aux calls. + """ + if len(needle) >= len(haystack): + return None + skipped: list[int] = [] + j = 0 + for i, value in enumerate(haystack): + if j < len(needle) and value == needle[j]: + j += 1 + else: + skipped.append(i) + return skipped if j == len(needle) else None + + +def reconcile(document: dict[str, Any], atif: dict[str, Any] | None) -> Report: + """Compare the exported rollout against ATIF. Disagreement is the signal. + + Deliberately FATAL on a per-turn mismatch. Two independent measurements of the same rollout + disagreeing means one is wrong, and we cannot tell which. Training on data we cannot corroborate + is exactly the failure this whole layer exists to prevent. + """ + report = Report() + if atif is None: + report.add( + INFO, "no_atif", "agent emitted no ATIF trajectory; cross-check unavailable" + ) + return report + + report.add( + INFO, + "atif_version", + f"schema {atif.get('schema_version')} " + f"agent {(atif.get('agent') or {}).get('name')}", + ) + + agent_rows = [r for r in document["sequences"] if r["role"] == "agent"] + if not agent_rows: + # When the intercept saw NO calls at all, `check_rollout` already says so plainly. Adding a + # second FATAL here just buries the real cause: I misread this as a distinct failure mode + # twice tonight before noticing it was always downstream of `no_turns`. + if not document.get("turns"): + report.add( + INFO, + "no_turns_upstream", + "no model calls were captured, so there is nothing to reconcile; see the " + "rollout's own no_turns finding for the cause", + ) + else: + report.add( + FATAL, + "no_agent_sequence", + f"{len(document['turns'])} calls captured but none labelled 'agent'", + ) + return report + + # Compare against EVERY call the agent made, in arrival order, across ALL agent roots. + # + # Two reasons this is the right comparison rather than the surviving training path: + # * ATIF logs every LLM call including retries, so a rollout that retried would report a false + # mismatch. Comparing the full list also validates the discard decisions: get one wrong and + # the equality breaks. + # * A harness that rewrites its system prompt mid-run (claude-code) splits one conversation + # across several roots. ATIF still sees one flat list, so the union of agent roots is what + # lines up with it. + agent_roots = {r["root_id"] for r in agent_rows} + turns = [t for t in document.get("turns", []) if t["root_id"] in agent_roots] + ours_all = [t["n_sampled"] for t in turns] + theirs = atif_turn_lengths(atif) + + # A converter that never fills in token counts gives us nothing to compare against. vibe reports + # completion_tokens=0 on every step while capturing perfectly (reward 1.0, 6 turns, 1197 + # trainable tokens), so calling that a MISMATCH would fail a harness for its trace converter's + # laziness rather than for anything wrong with the rollout. Downgrade to "no cross-check + # available", which is what `atif=none` already means for harnesses that emit nothing at all. + if not any(theirs): + # Covers BOTH shapes of an unhelpful converter: all-zero counts (vibe) and no agent steps at + # all (antigravity-sdk reported `[]` while we captured a 145-token call). Either way there is + # nothing to compare against, and failing the rollout would punish it for the trace + # converter's gaps rather than for anything wrong with the capture. + detail = ( + f"{len(theirs)} steps but all completion_tokens are 0" + if theirs + else "no agent steps at all" + ) + report.add( + WARN, + "atif_no_token_counts", + f"ATIF has {detail}, so no independent cross-check is possible. Intercept " + f"captured {sum(ours_all)} tokens across {len(ours_all)} calls; those numbers " + "stand unverified.", + ) + return report + + # ATIF being a strict SUBSEQUENCE of our calls is a real and benign shape, distinct from a + # disagreement. It means we saw calls the harness did not log as agent steps, which is what an + # auxiliary call is: gemini-cli fires a "next speaker" check between steps, mini-swe-agent does + # something similar. Both measurements are individually right. + # + # intercept: [58, 132, 266, 370, 105, 32, 54, 164, 33] + # ATIF: [58, 132, 266, 370, 105, 32, 33] + # + # The old code called any inequality FATAL, which failed a rollout for correctly capturing MORE + # than the harness chose to log. Note the asymmetry that makes this safe: extra calls on OUR side + # are explainable, whereas MISSING calls would mean we lost something, and that still fails. + # + # The aux calls are identified by position and excluded from training, because they are not the + # agent working on the task and must not carry the rollout's reward. + aux_idx = _subsequence_gap(theirs, ours_all) if ours_all != theirs else None + + # A subsequence match only carries evidence when ATIF accounts for MOST of what we captured. + # The shorter the ATIF list relative to ours, the more likely a match is coincidence: 5 values + # will embed in 49 almost by construction, so "the other 44 are auxiliary" is an inference the + # data does not support. + # + # Real example that forced this. mimo on one task: 49 captured calls, ATIF logged 5, and the + # matcher happily demoted 44 to auxiliary under a WARN, discarding 90% of a rollout without + # failing anything. Contrast the cases where the inference IS sound, where ATIF covers the large + # majority: gemini-cli 8/12, 14/19, 6/10, mini-swe-agent 5/6. + # + # Below the floor we refuse rather than guess. Silently training on a tenth of a rollout is a + # worse outcome than an explicit failure, which is the whole premise of this layer. + if aux_idx is not None and len(theirs) < 0.5 * len(ours_all): + report.add( + FATAL, + "atif_coverage_too_low", + f"ATIF accounts for only {len(theirs)} of {len(ours_all)} captured calls " + f"({len(theirs) / len(ours_all):.0%}). They embed as a subsequence, but at this " + "ratio that is as likely coincidence as signal, so the extra calls cannot be " + "called auxiliary with any confidence. Refusing rather than discarding " + f"{len(aux_idx)} calls on a guess.\n" + f" intercept: {ours_all}\n ATIF : {theirs}", + ) + return report + + if aux_idx is not None: + aux_nodes = [turns[i]["node_id"] for i in aux_idx] + report.add( + WARN, + "atif_aux_calls", + f"{len(aux_idx)} of {len(ours_all)} captured calls are absent from ATIF, so the " + f"harness did not consider them agent steps (auxiliary calls such as " + f"gemini-cli's next-speaker check). The remaining {len(theirs)} agree " + f"token-for-token. Sizes: {[ours_all[i] for i in aux_idx]}. These are excluded " + f"from training rather than credited with the rollout's reward.", + ) + report.aux_node_ids = aux_nodes + return report + + if ours_all == theirs: + n_discarded = sum(1 for t in turns if t["discarded"]) + n_trained = sum(len(r["turn_lengths"]) for r in agent_rows) + detail = ( + f"all {len(ours_all)} calls agree token-for-token (total {sum(ours_all)}); " + f"{n_discarded} discarded, {n_trained} trained" + ) + if len(agent_rows) > 1: + detail += ( + f"; across {len(agent_rows)} agent sequences (the harness rewrote its prompt " + "mid-run, so the conversation spans several token-prefix families)" + ) + report.add(INFO, "turns_match", detail) + else: + report.add( + FATAL, + "turn_mismatch", + f"per-call completion tokens disagree.\n" + f" intercept (all calls): {ours_all}\n" + f" ATIF : {theirs}\n" + f" intercept (trained) : {[r['turn_lengths'] for r in agent_rows]}", + ) + + steps = agent_steps(atif) + if steps: + atif_prompt = int((steps[0].get("metrics") or {}).get("prompt_tokens") or 0) + # The first agent sequence in arrival order holds the rollout's opening prompt. + first_prompt_len = agent_rows[0]["prompt_len"] + if atif_prompt and atif_prompt != first_prompt_len: + report.add( + WARN, + "prompt_len_mismatch", + f"first-turn prompt: intercept {first_prompt_len} vs ATIF {atif_prompt}", + ) + + retried = [s for s in steps if int(s.get("llm_call_count") or 1) > 1] + if retried: + report.add( + WARN, + "atif_retries", + f"{len(retried)} ATIF step(s) report llm_call_count>1: the harness retried. " + f"Graph found {document['stats']['n_discarded']} discarded turn(s).", + ) + + subagents = atif.get("subagent_trajectories") or [] + if subagents: + report.add( + WARN, + "atif_subagents", + f"ATIF reports {len(subagents)} subagent trajectory(ies); graph has " + f"{document['stats']['n_roots']} roots. Subagent turns must not be trained with " + "the parent rollout's reward.", + ) + return report + + +def merge_into_atif(atif: dict[str, Any], document: dict[str, Any]) -> dict[str, Any]: + """Fill ATIF's empty `completion_token_ids` / `logprobs` with our captured values. + + Produces one artifact that is both the human-readable trace and the training data, rather than + two formats a consumer has to join. Only attempted when the per-turn counts already agree; a + mismatch means we cannot map our tokens onto their steps, and guessing would be worse than + leaving the fields empty. + """ + agent_rows = [r for r in document["sequences"] if r["role"] == "agent"] + if not agent_rows: + return atif + + # ATIF logs every call the harness made; our training rows contain only the ones that survived, + # possibly split across several roots. Align on the FULL call list in arrival order and skip the + # discarded positions, rather than zipping lists of different length and silently shifting every + # token onto the wrong step. + agent_roots = {r["root_id"] for r in agent_rows} + turns_meta = [t for t in document.get("turns", []) if t["root_id"] in agent_roots] + steps_all = [s for s in atif.get("steps") or [] if s.get("source") == "agent"] + if len(turns_meta) != len(steps_all): + return atif + if [t["n_sampled"] for t in turns_meta] != atif_turn_lengths(atif): + return atif + + merged = json.loads(json.dumps(atif)) # never mutate Harbor's artifact in place + steps = [s for s in merged["steps"] if s.get("source") == "agent"] + + # Sampled tokens from every agent sequence, in the order their nodes arrived. + by_node: dict[str, list[tuple[int, float]]] = {} + for row in agent_rows: + sampled = [ + (tid, lp) + for tid, m, lp in zip(row["input_ids"], row["loss_mask"], row["logprobs"]) + if m + ] + # One mask-run per node that contributed trainable tokens. If a node was masked out (its + # logprobs could not be trusted) the counts diverge and the node->span mapping is no longer + # reliable, so we decline to merge rather than attach tokens to the wrong step. + if len(row["node_ids"]) != len(row["turn_lengths"]): + return atif + cursor = 0 + for node_id, length in zip(row["node_ids"], row["turn_lengths"]): + by_node[node_id] = sampled[cursor : cursor + length] + cursor += length + + for meta, step in zip(turns_meta, steps): + metrics = step.setdefault("metrics", {}) + if meta["discarded"]: + # Generated, then abandoned by the harness. Recorded so the trace stays complete, but + # flagged so nobody trains it with the rollout's reward. + metrics["discarded"] = True + continue + span = by_node.get(meta["node_id"]) + if span is None: + continue + metrics["completion_token_ids"] = [t for t, _ in span] + metrics["logprobs"] = [lp for _, lp in span] + + first = agent_rows[0] + merged.setdefault("extra", {})["intercept"] = { + "prompt_ids": first["input_ids"][: first["prompt_len"]], + "n_trainable": sum(r["n_trainable"] for r in agent_rows), + "n_agent_sequences": len(agent_rows), + "session_id": document["session_id"], + } + return merged diff --git a/src/openenv/harbor/capabilities.py b/src/openenv/harbor/capabilities.py new file mode 100644 index 000000000..bf108365c --- /dev/null +++ b/src/openenv/harbor/capabilities.py @@ -0,0 +1,254 @@ +"""What can this server actually run, right now? + +A client should not have to guess. `capabilities()` answers three questions in one call — which +harnesses exist and how well each is trusted, which sandboxes have working credentials, and which +datasets are served — so the failure "you asked for a sandbox whose API key is missing" happens at +discovery time instead of 90 seconds into a rollout. + +Credential checks reuse Harbor's own `preflight()` per backend rather than reimplementing key +lookups. Two things about that call have to be handled and are easy to miss: + + * it raises **`SystemExit`**, a `BaseException`, so a bare `except Exception` will not catch it and + the server process dies instead of reporting an unavailable sandbox + * `EnvironmentFactory.run_preflight` is never called by `Job.create()` — only Harbor's CLI calls + it — so a library caller gets no credential validation at all unless it asks +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from . import seams + +# Harbor agents that run HOST-SIDE, in this server's process, rather than being installed into the +# sandbox. The distinction is not cosmetic: their LLM traffic never leaves the host, so they reach +# the capture proxy on localhost and need no public URL at all. +HOST_SIDE_AGENTS = frozenset({"terminus-2", "computer-1", "oracle", "nop", "dspy-rlm"}) + +# Backends worth advertising. Harbor registers 23; these are the ones with a credential story we +# check and have exercised. Others still work via `--sandbox `, just unadvertised. +KNOWN_SANDBOXES = ("docker", "e2b", "modal", "daytona") + + +@dataclass +class SandboxStatus: + name: str + available: bool + detail: str = "" + + +@dataclass +class HarnessStatus: + name: str + dialect: str + kind: str # "installed" (runs in the sandbox) | "base" (runs host-side) + status: str # "validated" | "untested" + needs_subclass: bool # True when Harbor's wrapper cannot be configured as shipped + notes: str = "" + + +@dataclass +class Capabilities: + harnesses: list[HarnessStatus] = field(default_factory=list) + sandboxes: list[SandboxStatus] = field(default_factory=list) + datasets: list[dict[str, Any]] = field(default_factory=list) + llm: dict[str, Any] = field(default_factory=dict) + + def to_dict(self) -> dict[str, Any]: + return { + "harnesses": [vars(h) for h in self.harnesses], + "sandboxes": [vars(s) for s in self.sandboxes], + "datasets": self.datasets, + "llm": self.llm, + } + + @property + def available_sandboxes(self) -> list[str]: + return [s.name for s in self.sandboxes if s.available] + + @property + def validated_harnesses(self) -> list[str]: + return [h.name for h in self.harnesses if h.status == "validated"] + + def render(self, *, verbose: bool = False) -> str: + """Human-readable summary, printed at server start. + + Deliberately shows what is NOT available and why, not just what is. A sandbox missing its + API key is the single most common reason a rollout dies 90 seconds in, and Harbor's own + preflight message names the exact variable — so it is worth surfacing at startup rather than + making someone read a traceback later. + + Args: + verbose (`bool`, *optional*, defaults to `False`): + List every harness instead of only the validated ones. + + Returns: + `str`: A multi-line report. + """ + out: list[str] = [] + + if self.llm: + model = self.llm.get("model", "?") + ok = self.llm.get("ok") + mark = "ok" if ok else ("FAILED" if ok is False else "unchecked") + out.append(f"llm {model} [{mark}]") + if self.llm.get("url"): + out.append(f" {self.llm['url']}") + + usable = [s for s in self.sandboxes if s.available] + out.append(f"\nsandboxes {len(usable)} of {len(self.sandboxes)} usable") + for s in self.sandboxes: + if s.available: + out.append(f" [ok] {s.name}") + else: + out.append(f" [--] {s.name:<10} {s.detail[:88]}") + + if self.datasets: + total = sum(d.get("num_tasks", 0) for d in self.datasets) + out.append(f"\ndatasets {len(self.datasets)} split(s), {total} tasks") + for d in self.datasets: + if d.get("error"): + out.append(f" [--] {d['name']:<48} {d['error'][:60]}") + else: + out.append( + f" [ok] {d['name']:<48} {d.get('num_tasks', 0):>6} tasks" + ) + + validated = [h for h in self.harnesses if h.status == "validated"] + out.append( + f"\nharnesses {len(validated)} validated of {len(self.harnesses)} known" + ) + shown = self.harnesses if verbose else validated + by_dialect: dict[str, list[str]] = {} + for h in shown: + label = h.name + ("*" if h.kind == "base" else "") + by_dialect.setdefault(h.dialect, []).append(label) + for dialect, names in sorted(by_dialect.items()): + out.append(f" {dialect:<18} {', '.join(sorted(names))}") + out.append(" (* runs host-side in this process, so it needs no public URL)") + + if not usable: + out.append( + "\nWARNING: no sandbox has working credentials; every rollout will fail." + ) + return "\n".join(out) + + +def _missing_sdk(environment_class: Any) -> str: + """Report the backend's own verdict on whether its SDK is importable. + + Loading the class is not enough. Every Harbor backend guards its provider SDK with a + module-level `try: import ... except ImportError: _HAS_X = False`, and raises `MissingExtraError` + from `__init__` rather than at import. So the module imports cleanly, the class loads cleanly, + the backend reports available, and the failure arrives only once a rollout tries to build a + sandbox, by which time it reads as a broken rollout rather than a missing dependency. + + Reading the flag asks the backend the same question its constructor will ask, before offering it. + + Returns: + `str`: A message naming the missing extra, or `""` when the SDK is present. + """ + import sys + + module = sys.modules.get(getattr(environment_class, "__module__", ""), None) + if module is None: + return "" + absent = sorted( + flag + for flag, value in vars(module).items() + if flag.startswith("_HAS_") and value is False + ) + if not absent: + return "" + extras = ", ".join(flag.removeprefix("_HAS_").lower() for flag in absent) + return ( + f"SDK not installed ({extras}). Install it with: pip install 'harbor[cloud]', " + "or `openenv[harbor]`, which pulls every cloud backend." + ) + + +def check_sandbox(name: str) -> SandboxStatus: + """Ask Harbor whether this backend's credentials are present. + + `SystemExit` is caught explicitly: Harbor's preflight raises it as its failure signal, and it + does not inherit from `Exception`. + """ + try: + from harbor.environments.factory import ( + _load_environment_class, + EnvironmentFactory, + ) + from harbor.models.environment_type import EnvironmentType + except ImportError as exc: + return SandboxStatus(name, False, f"harbor not installed: {exc}") + + try: + env_type = EnvironmentType(name) + except ValueError: + return SandboxStatus(name, False, f"unknown environment type {name!r}") + + # Credentials are only half of it. Harbor's preflight checks env vars but never imports the + # backend, so a provider with perfect credentials and no SDK installed reports "available" and + # then fails at rollout time with MissingExtraError. Load the class first. + try: + environment_class = _load_environment_class(env_type) + except Exception as exc: # noqa: BLE001 - Harbor raises its own MissingExtraError here + hint = str(exc).replace("\n", " ")[:160] + return SandboxStatus( + name, False, hint or f"backend {name!r} could not be loaded" + ) + + missing = _missing_sdk(environment_class) + if missing: + return SandboxStatus(name, False, missing) + + try: + EnvironmentFactory.run_preflight(env_type) + except SystemExit as exc: + # Harbor's own message names the missing variable; pass it through rather than paraphrase. + return SandboxStatus(name, False, str(exc) or "credentials missing") + except ImportError as exc: + return SandboxStatus( + name, False, f"extra not installed: pip install 'harbor[{name}]' ({exc})" + ) + except Exception as exc: # noqa: BLE001 - an unexpected failure is still "not available" + return SandboxStatus(name, False, f"{type(exc).__name__}: {str(exc)[:160]}") + return SandboxStatus(name, True) + + +def list_harnesses() -> list[HarnessStatus]: + """Every harness with a seam, and how much to trust it.""" + out: list[HarnessStatus] = [] + for name in sorted(set(seams.SEAMS)): + seam = seams.get(name) + out.append( + HarnessStatus( + name=name, + dialect=seam.dialect, + kind="base" if name in HOST_SIDE_AGENTS else "installed", + status=seam.status, + needs_subclass=seam.import_path is not None, + notes=seam.notes, + ) + ) + return out + + +def capabilities( + *, + datasets: list[str] | None = None, + sandboxes: tuple[str, ...] = KNOWN_SANDBOXES, + llm: dict[str, Any] | None = None, +) -> Capabilities: + """Full picture. Safe to call on a fresh instance; does no I/O beyond credential checks.""" + result = Capabilities( + harnesses=list_harnesses(), + sandboxes=[check_sandbox(s) for s in sandboxes], + llm=dict(llm or {}), + ) + if datasets: + from .tasks import HarborTaskProvider + + result.datasets = HarborTaskProvider(datasets).list_splits() + return result diff --git a/src/openenv/harbor/client.py b/src/openenv/harbor/client.py new file mode 100644 index 000000000..f2f87f3d2 --- /dev/null +++ b/src/openenv/harbor/client.py @@ -0,0 +1,173 @@ +"""Typed client for a deployed harbor_env server. + +Two surfaces, matching the server: the Task API for discovery (plain HTTP, cheap, side-effect free) +and one MCP tool for execution (a single long call). + +```python +from openenv.harbor.client import HarborEnv + +with HarborEnv(base_url="http://localhost:8000") as env: + split = env.splits()[0]["name"] + print(env.num_tasks(split)) + result = env.run_rollout(split=split, task_index=0, harness="opencode", sandbox="e2b") + print(result.reward, result.turns[0].completion_token_ids[:8]) +``` +""" + +from __future__ import annotations + +import inspect +import json +from typing import Any + +import httpx +from openenv.core.env_server.mcp_types import CallToolAction, CallToolObservation +from openenv.core.mcp_client import MCPToolClient +from openenv.core.utils import run_async_safely + +from .models import HarborRolloutResult, HarborTaskRef + +# A rollout is minutes. The base client defaults to 60s, which fires mid-run. +_DEFAULT_MESSAGE_TIMEOUT_S = 1800.0 + + +class HarborEnv(MCPToolClient): + """Client for `harbor_env`. + + Args: + base_url (`str`): + Server root, e.g. `http://localhost:8000`. + message_timeout_s (`float`, *optional*, defaults to `1800.0`): + Websocket message timeout. Raised well above the base client's 60s because a single + rollout runs for minutes and the default would time out mid-call. + """ + + def __init__( + self, + base_url: str, + *, + message_timeout_s: float = _DEFAULT_MESSAGE_TIMEOUT_S, + **kwargs: Any, + ) -> None: + super().__init__( + base_url=base_url, message_timeout_s=message_timeout_s, **kwargs + ) + self._http = httpx.Client(base_url=base_url.rstrip("/"), timeout=120.0) + + # --- discovery (Task API) -------------------------------------------- + def splits(self) -> list[dict[str, Any]]: + """Datasets this server offers, with task counts.""" + return self._http.get("/harbor_env/splits").raise_for_status().json() + + def num_tasks(self, split: str = "") -> int: + payload = self._post("/harbor_env/num_tasks", {"split": split}) + return int(payload.get("num_tasks", 0)) + + def get_task(self, split: str, index: int) -> HarborTaskRef: + payload = self._post("/harbor_env/task", {"split": split, "index": index}) + return HarborTaskRef.model_validate(payload.get("task", payload)) + + def get_task_range( + self, split: str, start: int = 0, stop: int = 20 + ) -> list[HarborTaskRef]: + payload = self._post( + "/harbor_env/task_range", {"split": split, "start": start, "stop": stop} + ) + return [HarborTaskRef.model_validate(t) for t in payload.get("tasks", [])] + + # --- execution (MCP) --------------------------------------------------- + def run_rollout( + self, + *, + split: str = "", + task_index: int = 0, + harness: str = "opencode", + sandbox: str = "e2b", + reward_key: str = "", + keep_sandbox: bool = False, + force_build: bool = False, + ) -> HarborRolloutResult: + """Run one rollout and return its result. + + `harness` and `sandbox` are per-call, so consecutive rollouts can use different agents and + different backends against the same server. + + Args: + split (`str`, *optional*): + Dataset spec. Defaults to the server's first. + task_index (`int`, *optional*, defaults to `0`): + Index into the split. + harness (`str`, *optional*, defaults to `"opencode"`): + A validated seam name, or a `module:Class` import path for your own agent. + sandbox (`str`, *optional*, defaults to `"e2b"`): + Harbor environment type, e.g. `e2b` or `modal`. + reward_key (`str`, *optional*): + Which reward key is the training signal, for multi-reward tasks. + + Returns: + [`HarborRolloutResult`]: Reward, per-turn token ids and logprobs, and findings. + """ + raw = self._call( + "run_rollout", + split=split, + task_index=task_index, + harness=harness, + sandbox=sandbox, + reward_key=reward_key, + keep_sandbox=keep_sandbox, + force_build=force_build, + ) + return HarborRolloutResult.model_validate_json(_as_text(raw)) + + def capabilities(self) -> dict[str, Any]: + """Harnesses, sandboxes, datasets and LLM status for this server.""" + return json.loads(_as_text(self._call("capabilities"))) + + # --- internals --------------------------------------------------------- + def _call(self, name: str, **kwargs: Any) -> Any: + """Call an MCP tool from synchronous code. + + `MCPToolClient.call_tool` cannot be used here. It is a coroutine that internally does + `await self.step(action)`, but `EnvClient.step` dispatches on execution mode and returns a + concrete `StepResult` in sync mode, so awaiting it raises `TypeError: object StepResult + can't be used in 'await' expression`. Driving `step` directly works in both modes. + """ + result = self.step(CallToolAction(tool_name=name, arguments=kwargs)) + if inspect.isawaitable(result): # async mode returns an awaitable instead + result = run_async_safely(result) + + observation = result.observation + if isinstance(observation, CallToolObservation): + if observation.error is not None: + raise RuntimeError( + f"tool {name!r} failed: {observation.error.message} " + f"({observation.error.error_type.value})" + ) + return observation.result + return observation + + def _post(self, path: str, body: dict[str, Any]) -> dict[str, Any]: + return self._http.post(path, json=body).raise_for_status().json() + + def close(self) -> None: + try: + self._http.close() + finally: + super().close() + + +def _as_text(raw: Any) -> str: + """MCP tool results arrive as text content; unwrap whatever shape the transport used.""" + if isinstance(raw, str): + return raw + if isinstance(raw, dict): + content = raw.get("content") + if isinstance(content, list) and content: + first = content[0] + if isinstance(first, dict) and "text" in first: + return str(first["text"]) + return json.dumps(raw) + if isinstance(raw, list) and raw: + first = raw[0] + return str(getattr(first, "text", first)) + return str(raw) diff --git a/src/openenv/harbor/environment.py b/src/openenv/harbor/environment.py new file mode 100644 index 000000000..dc03c70b5 --- /dev/null +++ b/src/openenv/harbor/environment.py @@ -0,0 +1,232 @@ +"""The OpenEnv environment: Task API for discovery, one MCP tool for execution. + +`run_rollout` is a single long tool call rather than `reset`/`step`, because OpenEnv's HTTP handlers +construct and close an environment per request while a Harbor rollout is one stateful 60-600s run +that the harness drives. There is no meaningful `step()` to expose while opencode drives itself. +""" + +from __future__ import annotations + +import json +from typing import Any +from uuid import uuid4 + +from openenv.core.env_server.mcp_environment import MCPEnvironment +from openenv.core.env_server.types import Observation + +from .models import HarborState + +# A rollout is minutes, not seconds. OpenEnv's MCP tools default to 30s and there is no config knob, +# so the only way to raise it is to shadow `step`/`step_async` and inject a default. +_ROLLOUT_TIMEOUT_S = 1800.0 + + +class HarborEnvironment(MCPEnvironment): + """Per-session environment exposing `run_rollout` and `capabilities` over MCP.""" + + SUPPORTS_CONCURRENT_SESSIONS = True + + # Server-wide config, set once by `serving.build_app`. Class-level because OpenEnv builds a + # throwaway instance per `/metadata` and `/schema` request, and `__init__` must stay cheap and + # credential-free or the docs page cannot load. + _datasets: list[str] = [] + _llm_url: str = "" + _model: str = "" + # The validated LLM report from startup. Rebuilding it per request would mean re-probing the + # endpoint on every `capabilities()` call, so the startup verdict (including `ok`) is carried. + _llm: dict[str, Any] = {} + + @classmethod + def configure( + cls, + *, + datasets: list[str], + llm_url: str = "", + model: str = "", + llm: dict[str, Any] | None = None, + ) -> None: + cls._datasets = list(datasets) + cls._llm_url = llm_url + cls._model = model + cls._llm = dict(llm or {}) + + def __init__(self) -> None: + from fastmcp import FastMCP + + from .capabilities import capabilities as _capabilities + from .tasks import HarborTaskProvider + + self._provider = HarborTaskProvider(self._datasets) + self._state = HarborState(episode_id=str(uuid4()), llm_url=self._llm_url) + + mcp = FastMCP("harbor_env") + + @mcp.tool + def run_rollout( + split: str = "", + task_index: int = 0, + harness: str = "opencode", + sandbox: str = "e2b", + reward_key: str = "", + keep_sandbox: bool = False, + force_build: bool = False, + ) -> str: + """Run one Harbor rollout and return a JSON `HarborRolloutResult`. + + `harness` and `sandbox` are per-call, so consecutive rollouts can use different agents + and different backends against the same server. + """ + return self._run_rollout( + split, + task_index, + harness, + sandbox, + reward_key, + keep_sandbox, + force_build, + ) + + @mcp.tool + def capabilities() -> str: + """Harnesses, sandboxes, datasets and LLM status for this server.""" + caps = _capabilities( + datasets=self._datasets, + llm=self._llm or {"url": self._llm_url, "model": self._model}, + ) + return json.dumps(caps.to_dict()) + + @mcp.tool + def list_tasks(split: str = "", start: int = 0, stop: int = 20) -> str: + """A window of tasks in a split, for browsing without pulling all of them.""" + return json.dumps(self._provider.get_task_range(split, start, stop)) + + super().__init__(mcp) + + # --- Task API (OpenEnv discovers these by duck typing) ---------------- + def list_splits(self) -> list[dict[str, Any]]: + return self._provider.list_splits() + + def num_tasks(self, split: str) -> int: + return self._provider.num_tasks(split) + + def list_tasks(self, split: str) -> list[dict[str, Any]]: + return self._provider.list_tasks(split) + + def get_task(self, split: str, index: int) -> dict[str, Any]: + return self._provider.get_task(split, index) + + def get_task_range( + self, split: str, start: int | None = None, stop: int | None = None + ) -> list[dict[str, Any]]: + return self._provider.get_task_range(split, start, stop) + + # --- Environment ------------------------------------------------------ + def reset( + self, seed: int | None = None, episode_id: str | None = None, **_: Any + ) -> Observation: + """New episode. Boots nothing — a sandbox is created per `run_rollout`, not per reset.""" + service = self._service() + self._state = HarborState( + episode_id=episode_id or str(uuid4()), + llm_url=self._llm_url, + intercept_url=getattr(service, "public_url", "") if service else "", + ) + return Observation( + done=False, + reward=None, + metadata={ + "status": "ready", + "message": "Call run_rollout(split=..., task_index=..., harness=..., sandbox=...)", + "datasets": self._datasets, + }, + ) + + def _step_impl( + self, action: Any, timeout_s: float | None = None, **_: Any + ) -> Observation: + return Observation( + done=False, + reward=None, + metadata={ + "error": f"Unknown action {type(action).__name__}; " + "use CallToolAction(name='run_rollout', ...)" + }, + ) + + def step( + self, action: Any, timeout_s: float | None = None, **kwargs: Any + ) -> Observation: + return super().step(action, timeout_s=timeout_s or _ROLLOUT_TIMEOUT_S, **kwargs) + + async def step_async( + self, action: Any, timeout_s: float | None = None, **kwargs: Any + ) -> Observation: + return await super().step_async( + action, timeout_s=timeout_s or _ROLLOUT_TIMEOUT_S, **kwargs + ) + + @property + def state(self) -> HarborState: + return self._state + + # --- internals -------------------------------------------------------- + @staticmethod + def _service() -> Any: + from .serving import HarborService + + return HarborService.current() + + def _run_rollout( + self, + split: str, + task_index: int, + harness: str, + sandbox: str, + reward_key: str, + keep_sandbox: bool, + force_build: bool, + ) -> str: + import asyncio + from pathlib import Path + + from .models import HarborRolloutResult + from .rollout import run_rollout as _run + + service = self._service() + if service is None: + return HarborRolloutResult( + ok=False, + error="server not initialised: no capture proxy is running", + harness=harness, + sandbox=sandbox, + ).model_dump_json() + + split = split or (self._datasets[0] if self._datasets else "") + try: + task_dir = self._provider.task_dir(split, int(task_index)) + except Exception as exc: # noqa: BLE001 + return HarborRolloutResult( + ok=False, error=str(exc)[:400], harness=harness, sandbox=sandbox + ).model_dump_json() + + result = asyncio.run( + _run( + task_dir=task_dir, + harness=harness, + sandbox=sandbox, + registry=service.capture.registry, + intercept_url=service.public_url, + model=service.model, + trials_dir=Path("/tmp/openenv-harbor-trials"), + dataset=split, + reward_key=reward_key, + keep_sandbox=keep_sandbox, + force_build=force_build, + ) + ) + + self._state.rollouts_completed += 1 + self._state.last_reward = result.reward + self._state.last_task_id = result.task_id + self._state.last_trial_name = result.trial_name + return result.model_dump_json() diff --git a/src/openenv/harbor/install_fixes.py b/src/openenv/harbor/install_fixes.py new file mode 100644 index 000000000..98e3e6764 --- /dev/null +++ b/src/openenv/harbor/install_fixes.py @@ -0,0 +1,479 @@ +"""Local subclasses that fix agent INSTALL failures, registered via `AgentConfig.import_path`. + +Same escape hatch as `pi_agent.py`, different reason. These agents' seams are fine; they simply never +get far enough to use them, because installing the CLI into the DataAgent sandbox fails. Both are +upstream bugs in Harbor's wrappers rather than anything about the intercept, and both are fixed here +without touching Harbor. + +Each override is deliberately minimal: call Harbor's own `install()` and change only the one thing +that is wrong, so we inherit every future upstream fix instead of forking the install logic. +""" + +from __future__ import annotations + +import json +import shlex +from pathlib import Path +from typing import Any + +from harbor.agents.installed.cline.cline import ClineCli, ExecInput +from harbor.agents.installed.gemini_cli import GeminiCli +from harbor.agents.installed.hermes import Hermes +from harbor.agents.installed.kimi_cli import KimiCli +from harbor.agents.installed.openclaw import OpenClaw +from harbor.agents.installed.openhands import OpenHands +from harbor.agents.installed.pi import Pi +from harbor.agents.installed.swe_agent import SweAgent +from harbor.environments.base import BaseEnvironment + + +class InterceptGeminiCli(GeminiCli): + """gemini-cli, with `bash` guaranteed before nvm runs. + + Harbor's `GeminiCli.install` declares only `("curl",)` as a system dependency + (gemini_cli.py:111), but `nvm_node_install_snippet()` pipes the installer into **bash**: + + curl -o- .../install.sh | env -u NODE_VERSION bash + + In an image without bash that pipe fails, nvm never lands, and the snippet's own guard reports: + + Error: NVM failed to load + + which reads like an nvm problem rather than a missing shell. opencode does the same nvm install + successfully in the same image because it declares `("curl", "bash")` (opencode.py:91). So this + is simply a missing dependency in gemini-cli's declaration. + """ + + async def install(self, environment: BaseEnvironment) -> None: + await self.ensure_system_dependencies(environment, ("curl", "bash")) + await super().install(environment) + + +class InterceptOpenHands(OpenHands): + """OpenHands pinned to the last V0 release, which still has `openhands.core.main`. + + Harbor's wrapper installs `openhands-ai` unpinned and then verifies with: + + /opt/openhands-venv/bin/python -m openhands.core.main --version + + OpenHands V1 restructured the package: the agent core moved out into `openhands-sdk` / + `openhands-agent-server`, so `openhands.core` no longer exists and install fails with + + ModuleNotFoundError: No module named 'openhands.core' + + V0 is scheduled for removal on 2026-04-01, so this pin is a stopgap: once Harbor's wrapper is + updated for the V1 entry point, drop the pin and this subclass. 0.49.0 is the newest 0.x on PyPI. + """ + + DEFAULT_V0_VERSION = "0.49.0" + # Harbor defaults to `uv python install 3.13`, but 0.49.0 shipped in July 2025 and does not + # resolve there, so the version pin alone is not enough. + DEFAULT_PYTHON = "3.12" + + # Transitive dependencies openhands-ai 0.49.0 imports but does not declare. Installing 0.49.0 + # succeeds, and then the import check dies: + # File ".../openhands/events/event_store_abc.py", line 5, in + # from deprecated import deprecated + # ModuleNotFoundError: No module named 'deprecated' + # Pinning an old release means living with whatever its metadata got wrong at the time. + MISSING_DEPS = ("Deprecated",) + VENV = "/opt/openhands-venv" + + def __init__(self, *args: Any, **kwargs: Any): + kwargs.setdefault("version", self.DEFAULT_V0_VERSION) + kwargs.setdefault("python_version", self.DEFAULT_PYTHON) + super().__init__(*args, **kwargs) + + async def install(self, environment: BaseEnvironment) -> None: + """Install via Harbor, and repair the missing deps if its verify step trips over them. + + Harbor's install ends with `python -m openhands.core.main --version`, so an undeclared + dependency surfaces as a failed install rather than a failed run. We cannot pre-empt it + without forking Harbor's install command, so instead: let it run, and if it fails, add the + known-missing packages to the venv it already built and re-run the same verification. If that + passes, the install is genuinely fine. + """ + try: + await super().install(environment) + return + except Exception as exc: # noqa: BLE001 - remediate, then re-verify honestly + self.logger.warning( + "openhands install failed (%s); attempting dependency repair", + str(exc)[:160], + ) + + packages = " ".join(self.MISSING_DEPS) + await self.exec_as_agent( + environment, + command=( + f"set -euo pipefail; {self.VENV}/bin/python -m ensurepip --upgrade || true; " + f"{self.VENV}/bin/python -m pip install {packages}" + ), + ) + await self._install_poetry_shim(environment) + # Same check Harbor uses. If this still fails it raises, and the failure is real. + await self.exec_as_agent( + environment, + command=f"{self.VENV}/bin/python -m openhands.core.main --version", + ) + + async def _install_poetry_shim(self, environment: BaseEnvironment) -> None: + """Make `poetry run python ...` work in a venv that poetry never created. + + OpenHands' LocalRuntime starts its action-execution server with: + + ['poetry', 'run', 'python', '-u', '-m', 'openhands.runtime.action_execution_server', ...] + + which assumes a poetry-managed source checkout. Installed from PyPI into a uv venv there is no + pyproject.toml anywhere above site-packages, so poetry refuses: + + server: Poetry could not find a pyproject.toml file in + /opt/openhands-venv/lib/python3.12/site-packages or its parents + server process exited + + The agent then waits for a server that will never come up, and tenacity converts that into + `RetryError[]` with the actual cause nowhere in the traceback. + + Rather than fabricate a pyproject.toml (which makes poetry resolve and possibly reinstall a + dependency tree), shim the one invocation OpenHands makes: drop the `run` verb and exec the + venv's own interpreter. Everything the server needs is already installed there. + """ + shim = ( + "#!/bin/sh\n" + '[ "$1" = "run" ] && shift\n' + f'[ "$1" = "python" ] && {{ shift; exec {self.VENV}/bin/python "$@"; }}\n' + f'exec {self.VENV}/bin/"$@"\n' + ) + # Installed as ROOT into /usr/local/bin, and over any existing poetry. + # + # A first attempt wrote it to ~/.local/bin and changed nothing: the error was + # "Poetry could not find a pyproject.toml", not "poetry: command not found", so a real poetry + # is already on PATH ahead of ~/.local/bin. Shadowing it is the only way the shim is reached. + # /usr/local/bin precedes ~/.local/bin on every image we run. + for directory in ("/usr/local/bin", "/usr/bin"): + await self.exec_as_root( + environment, + command=( + f"mkdir -p {directory} && cat > {directory}/poetry <<'SHIM'\n{shim}SHIM\n" + f"chmod 0755 {directory}/poetry" + ), + ) + + +class InterceptSweAgent(SweAgent): + """swe-agent, given a git repo to work in so Harbor's working code path is taken. + + Harbor builds the repo argument as: + + "$(if [ -d /testbed ]; then echo '--env.repo.type=preexisting --env.repo.repo_name=/testbed'; " + "else echo '--env.repo.path=$(pwd)'; fi)" + + The else-branch is broken: `$(pwd)` sits inside SINGLE quotes, so it is never expanded and the + literal string is passed through. swe-agent then resolves it relative to its cwd and dies: + + git.exc.NoSuchPathError: /workdir/$(pwd) + + Underneath that is a second problem: swe-agent is a SWE-bench agent and requires a git repository, + while DataAgent tasks are a CSV and a question. + + Both are solved by satisfying the `[ -d /testbed ]` test that Harbor already checks. We `git init` + the task's own /workdir and expose it as /testbed, so Harbor takes its preexisting-repo branch + (which has no quoting bug) and the agent still works where the task data and /workdir/answer.txt + live. Nothing in Harbor changes. + """ + + async def setup(self, environment: BaseEnvironment) -> None: + await super().setup(environment) + await self.exec_as_root( + environment, + command=( + "set -eu; mkdir -p /workdir; cd /workdir; " + # A repo with no commit still fails some checks, so make one. + "git rev-parse --git-dir >/dev/null 2>&1 || { " + " git init -q .; " + " git config user.email harbor@example.com; git config user.name harbor; " + " touch .harbor-keep; git add -A; git commit -qm 'harbor: initial' || true; }; " + "[ -e /testbed ] || ln -s /workdir /testbed" + ), + ) + + +class InterceptOpenClaw(OpenClaw): + """openclaw, with its config actually present inside the sandbox. + + Harbor writes the merged config to the HOST trial dir and then copies it from a CONTAINER path: + + upload_path = self.logs_dir / "openclaw.upload.json" # host + "mkdir -p ~/.openclaw && cp /logs/agent/openclaw.upload.json ~/.openclaw/openclaw.json" + + with the comment "trial mounts logs here as /logs/agent". That holds for a bind-mounted docker + runtime. **E2B has no bind mounts**, so the file exists on the host and nowhere in the sandbox: + + cp: cannot stat '/logs/agent/openclaw.upload.json' + + Supplying `openclaw_config` does not help, because the problem is not that the config is empty. + + Fix: build the same config Harbor would and upload it to the container path during setup, so the + copy in `run()` finds it. Harbor's own `_build_full_openclaw_config` is reused, so the content + stays whatever Harbor intended, merges included. + """ + + async def setup(self, environment: BaseEnvironment) -> None: + await super().setup(environment) + + payload = json.dumps(self._build_full_openclaw_config(), indent=2) + "\n" + local = Path(self.logs_dir) / self._UPLOAD_CONFIG_FILENAME + local.parent.mkdir(parents=True, exist_ok=True) + local.write_text(payload, encoding="utf-8") + + target = f"{self._CONTAINER_LOGS_AGENT}/{self._UPLOAD_CONFIG_FILENAME}" + await self.exec_as_root( + environment, + command=f"mkdir -p {self._CONTAINER_LOGS_AGENT} && " + f"chmod 777 {self._CONTAINER_LOGS_AGENT}", + ) + await environment.upload_file(local, target) + + +class InterceptCline(ClineCli): + """cline-cli, given a base URL through the only channel it has: its settings store. + + Harbor forwards exactly `{PROVIDER, API_KEY, MODELID}` and runs + + cline -P -k $API_KEY -m $MODELID --json --yolo + + There is no base-URL flag and no base-URL env var, so cline resolves the provider's REAL endpoint + and dies with the session id as a bearer token: + + Incorrect API key provided: s55f2f5a… You can find your API key at + https://platform.openai.com/account/api-keys + + Cline's OpenAI-Compatible provider takes Base URL + key + model id from its settings store + (`~/.cline/data/globalState.json`), not from the CLI. Harbor writes that file itself at the start + of `create_run_agent_commands`, so anything written earlier is overwritten. Instead we let + Harbor's command run and INSERT a merge step between it and the agent invocation, which keeps + Harbor's own keys (`welcomeViewCompleted`, `isNewUser`) intact. + """ + + def __init__( + self, *args: Any, intercept_config: dict[str, str] | None = None, **kwargs: Any + ): + self._intercept_config = intercept_config or {} + super().__init__(*args, **kwargs) + + def create_run_agent_commands(self, instruction: str): + commands = list(super().create_run_agent_commands(instruction)) + base_url = self._intercept_config.get("base_url") + api_key = self._intercept_config.get("api_key") + model = self._intercept_config.get("model") + if not (base_url and api_key and model) or not commands: + return commands + + # Merge rather than replace: Harbor's globalState keys must survive. + settings = { + "openAiBaseUrl": f"{base_url}/v1", + "openAiApiKey": api_key, + "openAiModelId": model, + "apiProvider": "openai", + } + merge = ( + "python3 - <<'__HARBOR_CLINE_SETTINGS__'\n" + "import json, pathlib\n" + "p = pathlib.Path.home() / '.cline' / 'data' / 'globalState.json'\n" + "p.parent.mkdir(parents=True, exist_ok=True)\n" + "try:\n" + " cfg = json.loads(p.read_text())\n" + "except Exception:\n" + " cfg = {}\n" + f"cfg.update({json.dumps(settings)})\n" + "p.write_text(json.dumps(cfg))\n" + "__HARBOR_CLINE_SETTINGS__" + ) + commands.insert(1, ExecInput(command=merge)) + return commands + + +class InterceptHermes(Hermes): + """hermes, pointed at our endpoint via its `custom_providers` config, compression off. + + Harbor generates hermes' config.yaml with `provider: "auto"`, so hermes routes itself and never + honours OPENAI_BASE_URL. It authenticates somewhere else entirely and the intercept records zero + calls: + + HTTP 401: Missing Authentication header + + Pinning `provider: "openai"` does NOT work -- hermes rejects it outright + (`Unknown provider 'openai'`), because its table maps the openai prefix to a `None` CLI flag, + i.e. "use the full provider/model name and let auto-routing resolve it". + + The supported route is a **custom provider** entry + (https://hermes-agent.nousresearch.com/docs/integrations/providers): + + custom_providers: + - name: intercept + base_url: http://host:port/v1 + api_key: ... + + Compression is also disabled. Harbor enables it (`threshold: 0.85`), and compaction rewrites the + conversation, which breaks token-prefix stitching: the rollout would fragment into per-turn rows + exactly as terminus-2 does. + + `_build_config_yaml` is a staticmethod upstream but is invoked as `self._build_config_yaml(...)`, + so overriding it as an instance method is safe and gives access to the injected endpoint. + """ + + PROVIDER = "intercept" + + def __init__( + self, *args: Any, intercept_config: dict[str, str] | None = None, **kwargs: Any + ): + self._intercept_config = intercept_config or {} + super().__init__(*args, **kwargs) + + def _build_config_yaml(self, model: str) -> str: # type: ignore[override] + import yaml + + cfg = yaml.safe_load(Hermes._build_config_yaml(model)) or {} + + base_url = self._intercept_config.get("base_url") + api_key = self._intercept_config.get("api_key") + if base_url: + entry = { + "name": self.PROVIDER, + "base_url": f"{base_url}/v1", + "model": model, + } + if api_key: + entry["api_key"] = api_key + cfg["custom_providers"] = [entry] + cfg["provider"] = self.PROVIDER + + # Compaction is incompatible with prefix stitching; see docs/HARNESS_NOTES.md. + cfg.setdefault("compression", {})["enabled"] = False + return yaml.safe_dump(cfg, sort_keys=False) + + +class InterceptKimi(KimiCli): + """kimi-cli, surviving the stream reset its own teardown causes. + + Harbor runs kimi as (kimi_cli.py:379-394): + + (echo $PROMPT; sleep 86400) | kimi --wire --yolo --afk ... | ( + while IFS= read -r line; do ... case "$line" in *'"id":"1"'*) break ;; esac; done + ...; kill 0) + + `sleep 86400` holds stdin open for a day, and `kill 0` tears down the whole process group once + the terminating wire event arrives. Harbor already expects part of the fallout and swallows + `NonZeroAgentExitCodeError` for "exit 143" (SIGTERM). What it does not expect is that killing the + group also kills the E2B exec stream mid-flight, so the HTTP/2 connection on the HOST side dies: + + httpcore.RemoteProtocolError: + (raised in .venv312/site-packages/httpcore/_async/http2.py, i.e. OUR process, not the sandbox) + + That propagates out of `run`, so Harbor abandons the trial and never runs the verifier. Every one + of 11 kimi trials died this way, each AFTER completing real work (one had 37 captured turns), and + no other harness has ever produced this error on the same E2B backend, which is what identifies + it as kimi's teardown rather than transport flakiness. + + By the time it fires, kimi has already written its wire output to /logs/agent/, so the trajectory + and answer are on disk and the trial can be graded normally. Swallowing it here is the same + judgement Harbor already made for exit 143, applied to the other half of the same teardown. + + Deliberately narrow: ONLY httpcore/httpx RemoteProtocolError. Any other failure still raises, + because a rollout that broke for an unknown reason must not be quietly graded. + """ + + async def run(self, instruction, environment, context) -> None: # type: ignore[override] + try: + await super().run(instruction, environment, context) + except Exception as exc: # noqa: BLE001 - re-raised below unless it is the known teardown + if type(exc).__name__ != "RemoteProtocolError" or "StreamReset" not in str( + exc + ): + raise + # Expected: `kill 0` took the exec stream down with the process group. + + +# ------------------------------------------------------------------------------------------------ +# pi +# ------------------------------------------------------------------------------------------------ +# pi, taught to talk to our intercept. Registered via `AgentConfig.import_path`, no Harbor patch. +# +# THE PROBLEM. Harbor's `pi` wrapper forwards a fixed list of API-key variables and nothing else +# (installed/pi.py:100-137). It has no base-URL handling at all, and unlike opencode there is no +# `_build_register_config_command` hook to write a provider config. Pointed at our intercept, pi ignores +# `OPENAI_BASE_URL`, calls api.openai.com with our session id as the key, and dies: +# +# OpenAI API error (401): Incorrect API key provided: sc6f2e5a… +# You can find your API key at https://platform.openai.com/account/api-keys +# +# THE FIX. pi reads custom providers from `~/.pi/agent/models.json` +# (https://pi.dev/docs/latest/custom-provider). Harbor gives us no hook to write it, but it does let an +# agent be supplied by `import_path`, so we subclass `Pi`, write the file in `setup()`, and register +# the subclass. Harbor is untouched. +# +# AgentConfig(import_path="harnesses.pi_agent:InterceptPi", ...) +# +# This is the general escape hatch for any harness whose config Harbor does not know how to write: +# subclass locally, override `setup()`, register by import path. +# +# TWO DETAILS THAT MATTER. +# +# `api` must be `openai-completions`. Left to its own devices pi picks `openai-responses` for a provider +# named `openai` (we watched it do exactly that: `"api":"openai-responses"` in its session log). The +# intercept handles both, but chat-completions is the dialect with the least translation and by far the +# most mileage on it. +# +# The provider is NOT named `openai`. A distinct name keeps pi off its built-in OpenAI defaults, the +# same reason opencode's provider is called `intercepted`. +PROVIDER = "intercept" +MODELS_JSON = "~/.pi/agent/models.json" + + +def build_models_json(base_url: str, api_key: str, model: str) -> str: + return json.dumps( + { + "providers": { + PROVIDER: { + "baseUrl": f"{base_url}/v1", + "api": "openai-completions", + "apiKey": api_key, + "models": [{"id": model}], + } + } + }, + indent=2, + ) + + +class InterceptPi(Pi): + """`Pi` that writes a custom-provider config into the sandbox before running. + + Config arrives through `AgentConfig.kwargs` as `intercept_config`, mirroring how opencode + receives `opencode_config`, so the seam table stays uniform across harnesses. + """ + + def __init__( + self, *args: Any, intercept_config: dict[str, str] | None = None, **kwargs: Any + ): + self._intercept_config = intercept_config or {} + super().__init__(*args, **kwargs) + + async def setup(self, environment: BaseEnvironment) -> None: + await super().setup(environment) + + base_url = self._intercept_config.get("base_url") + api_key = self._intercept_config.get("api_key") + model = self._intercept_config.get("model") + if not (base_url and api_key and model): + # Refuse quietly rather than run against api.openai.com with a session id as the key, + # which is what happens by default and costs a sandbox to discover. + raise ValueError( + "InterceptPi requires intercept_config with base_url, api_key, model" + ) + + payload = shlex.quote(build_models_json(base_url, api_key, model)) + await self.exec_as_agent( + environment, + command=f"mkdir -p ~/.pi/agent && printf '%s' {payload} > {MODELS_JSON}", + ) diff --git a/src/openenv/harbor/models.py b/src/openenv/harbor/models.py new file mode 100644 index 000000000..b66645589 --- /dev/null +++ b/src/openenv/harbor/models.py @@ -0,0 +1,253 @@ +"""Wire types for harbor_env. + +Two shapes matter. `HarborTaskRef` is what the Task API hands out during discovery, one per dataset +item. `HarborRolloutResult` is what one `run_rollout` returns: the reward, and enough token detail to +train on. + +Everything here is JSON-serialisable by construction — `run_rollout` returns +`result.model_dump_json()` and the client re-validates, matching how `opencode_env` and `pi_env` do +it. There is no shared memory between server and client. +""" + +from __future__ import annotations + +from typing import Any + +from openenv.core.env_server.types import State +from pydantic import BaseModel, Field + + +class HarborTaskRef(BaseModel): + """One task, as returned by the Task API. What a trainer's dataset holds per row.""" + + index: int + task_id: str + task_name: str + dataset: str = "" + instruction: str = "" + + +class HarborTurn(BaseModel): + """One model call, captured exactly. The unit a trainer consumes. + + `prompt_token_ids` is the engine's own tokenisation of everything before this turn, not a local + re-render. That distinction is the whole point of the capture layer: re-tokenising a prompt + offline drifts from what the model actually saw (measured at 0/6 exact on Qwen3.5 until thinking + was disabled), and a drifted prompt silently forks a long conversation into short fragments. + """ + + turn: int + role: str = "agent" + finish_reason: str | None = None + prompt_token_ids: list[int] = Field(default_factory=list) + completion_token_ids: list[int] = Field(default_factory=list) + per_token_logps: list[float] = Field(default_factory=list) + n_tools: int = 0 + discarded: bool = False + + # What the model actually produced, in readable form. Only the assistant's own output is kept, + # never the prompt side: the prompt is already present as token ids and repeating it as text + # would roughly double the payload for no new information. This is what makes a result + # inspectable without a tokenizer, and it is what a reward function keys on when it needs to + # know which tool was called rather than how many tokens were spent. + text: str = "" + tool_calls: list[dict[str, Any]] = Field(default_factory=list) + + +class HarborConversation(BaseModel): + """One complete conversation from a rollout, exactly as the harness assembled it. + + A rollout can contain several: a root is a conversation that started from a fresh prompt, so + subagents and auxiliary calls each get their own. `messages` is the full list including the + system prompt and every tool result, which is what makes a finished rollout readable rather than + a column of token counts. + """ + + root_id: str = "" + role: str = "agent" # agent | auxiliary | discarded + n_turns: int = 0 + messages: list[dict[str, Any]] = Field(default_factory=list) + + +class HarborStepResult(BaseModel): + """One step of a multi-step task. Harbor gates progression on `min_reward`.""" + + name: str = "" + rewards: dict[str, float] = Field(default_factory=dict) + passed: bool = True + + +class HarborRolloutResult(BaseModel): + """Everything one rollout produced. + + A failed rollout is still a valid result: `ok=False`, `error` set, `reward=None`. Nothing raises + across the server boundary, because a rollout exception reaching a trainer is what hangs every + rank at the NCCL barrier forever. + """ + + # identity + task_id: str = "" + task_name: str = "" + dataset: str = "" + harness: str = "" + sandbox: str = "" + trial_name: str = "" + session_id: str = "" + + # outcome — forwarded from Harbor's verifier, never recomputed here + reward: float | None = None + rewards: dict[str, float] = Field(default_factory=dict) + reward_key: str = "" + step_results: list[HarborStepResult] = Field(default_factory=list) + + # capture + turns: list[HarborTurn] = Field(default_factory=list) + conversations: list[HarborConversation] = Field(default_factory=list) + n_turns: int = 0 + n_roots: int = 0 + n_trainable_tokens: int = 0 + multi_turn: bool = False + atif: str = "none" + findings: list[str] = Field(default_factory=list) + + # timings and diagnostics. There is no metrics endpoint and no structured logging in the env + # server, so observability has to ride back inside the payload or it does not exist. + wall_s: float = 0.0 + phase_timings: dict[str, float] = Field(default_factory=dict) + agent_log_tail: str = "" + + # failure + ok: bool = True + error: str | None = None + exception_type: str | None = None + + @property + def solved(self) -> bool: + """Graded AND positive. `reward is None` means the verifier never ran, which is not a zero.""" + return self.reward is not None and self.reward > 0 + + +class HarborState(State): + """Per-session counters. Mutated inside the tool, since `step` only dispatches.""" + + rollouts_completed: int = 0 + last_reward: float | None = None + last_task_id: str | None = None + last_trial_name: str | None = None + llm_url: str = "" + intercept_url: str = "" + + +def _assistant_text(response: dict[str, Any]) -> str: + """The assistant's own words, flattened across the shapes the four dialects produce.""" + content = response.get("content") + if isinstance(content, list): # anthropic / responses send block lists + return " ".join( + part.get("text", "") + for part in content + if isinstance(part, dict) and part.get("text") + ) + return content if isinstance(content, str) else "" + + +def _tool_calls(response: dict[str, Any]) -> list[dict[str, Any]]: + """Tool calls as `{name, arguments}`, normalised across dialects. + + Kept as data rather than a rendered string: a reward function that wants to check which tool ran + should not have to parse a display format. + """ + out: list[dict[str, Any]] = [] + for call in response.get("tool_calls") or []: + function = call.get("function") or {} + name = function.get("name") or call.get("name") + if not name: + continue + out.append( + { + "name": str(name), + "arguments": function.get("arguments", call.get("arguments", "")), + } + ) + return out + + +def conversations_from_document(document: dict[str, Any]) -> list[HarborConversation]: + """Rebuild the full conversations, system prompt and tool results included. + + The deepest node of a chain already carries the whole conversation in `request_messages`, since + each call replays everything before it. So the last node per root plus its own response is the + complete transcript, with no stitching and no risk of drifting from what was actually sent. + """ + by_node = {t["node_id"]: t for t in document.get("turns", [])} + # One per root, not one per sequence. A fork produces several paths through the same root, and + # each replays the same conversation up to the branch point, so emitting one per sequence shows + # the reader near-identical transcripts and calls both of them the main conversation. The + # longest path is the complete one. + best: dict[str, HarborConversation] = {} + + for sequence in document.get("sequences", []): + node_ids = sequence.get("node_ids") or [] + if not node_ids: + continue + last = by_node.get(node_ids[-1], {}) + messages = list(last.get("request_messages") or []) + response = last.get("response_message") or {} + if response: + messages.append({**response, "role": response.get("role", "assistant")}) + if not messages: + continue + root_id = str(sequence.get("root_id", "")) or node_ids[0] + candidate = HarborConversation( + root_id=root_id, + role=str(sequence.get("role", "agent")), + n_turns=int(sequence.get("n_turns", len(node_ids))), + messages=messages, + ) + current = best.get(root_id) + if current is None or len(candidate.messages) > len(current.messages): + best[root_id] = candidate + return list(best.values()) + + +def turns_from_document(document: dict[str, Any]) -> list[HarborTurn]: + """Flatten a capture document into per-turn training rows. + + Only `agent` sequences become turns. Auxiliary calls are dropped here rather than marked, + because they are not the agent working on the task and must never carry its reward — a + next-speaker classification credited with solving a task is a reward-hacking gift. + """ + by_node = {t["node_id"]: t for t in document.get("turns", [])} + rows: list[HarborTurn] = [] + index = 0 + + for sequence in document.get("sequences", []): + if sequence.get("role") != "agent": + continue + cursor = 0 + input_ids = sequence["input_ids"] + mask = sequence["loss_mask"] + logprobs = sequence["logprobs"] + for node_id, length in zip(sequence["node_ids"], sequence["turn_lengths"]): + node = by_node.get(node_id, {}) + response = node.get("response_message") or {} + sampled = [(i, m) for i, m in zip(input_ids, mask) if m] + span = sampled[cursor : cursor + length] + lp = [p for p, m in zip(logprobs, mask) if m][cursor : cursor + length] + rows.append( + HarborTurn( + turn=index, + finish_reason=node.get("finish_reason"), + prompt_token_ids=input_ids[: sequence["prompt_len"]] + if index == 0 + else [], + completion_token_ids=[i for i, _ in span], + per_token_logps=lp, + n_tools=node.get("n_tools", 0), + discarded=bool(node.get("discarded")), + text=_assistant_text(response), + tool_calls=_tool_calls(response), + ) + ) + cursor += length + index += 1 + return rows diff --git a/src/openenv/harbor/rollout.py b/src/openenv/harbor/rollout.py new file mode 100644 index 000000000..1b0dafa90 --- /dev/null +++ b/src/openenv/harbor/rollout.py @@ -0,0 +1,352 @@ +"""Run one Harbor trial through the capture proxy and return trainable output. + +This is where every piece meets: a task from the dataset, a harness from the seam table, a sandbox +from Harbor, and a capture session whose id doubles as the agent's API key. + + task dir ──┐ + harness ───┼──> TrialConfig ──> Trial.run() ──> TrialResult (reward) + sandbox ───┘ │ + └──> agent talks to the intercept + (api key == session id) + │ + RolloutGraph ──> HarborRolloutResult + +**Nothing raises out of `run_rollout`.** A failed rollout returns `ok=False` with `reward=None`. That +is not defensive habit, it is the reason this layer exists: in the white-box predecessor a rollout +exception reached the trainer and hung every rank at the NCCL barrier forever, so every method there +had to be individually wrapped. Behind a result object that failure mode cannot occur. +""" + +from __future__ import annotations + +import asyncio +import os +import time +import uuid +from pathlib import Path +from typing import Any + +from openenv.core.harness.capture.export import export_session + +from . import seams +from .atif import load_atif, reconcile +from .models import ( + conversations_from_document, + HarborRolloutResult, + HarborStepResult, + HarborTurn, + turns_from_document, +) + +# Harbor's own retry is deliberately off. A Harbor-level retry re-runs the agent against the SAME +# capture session, so two attempts merge into one graph and the trace cross-check compares a +# two-attempt capture against a one-attempt trajectory. One attempt = one session = one rollout; +# retries belong to the caller, with a fresh session each time. +_NO_HARBOR_RETRY = 0 + + +# Some agents read os.environ at CONSTRUCTION, before any sandbox exists — Harbor's claude-code +# wrapper decides there which credentials to forward. For those, `agent_env` alone is too late, so +# the seam also carries a process-env channel. +# +# That channel is global, so two concurrent rollouts of the same harness would overwrite each +# other's session key. The lock is held only across agent construction (`Trial.create`), which is +# brief; the rollout itself then runs unserialised. +_PROC_ENV_LOCK = asyncio.Lock() + + +def apply_process_env(seam_name: str, env: dict[str, str]) -> None: + """Set the seam's process-level env vars, warning when one would break the grader. + + OPENAI_API_KEY is shared with the task grader: Harbor forwards it into the sandbox and the + DataAgent grader's LLM-judge tier fires on `if os.environ.get("OPENAI_API_KEY")`. Overwriting it + with a session id makes the judge 401, so every semantically-correct-but-not-exact answer scores + 0 — which reads as a weak model rather than a broken harness, and poisons the RL baseline. + """ + grader_key = os.environ.get("OPENAI_API_KEY") + for key, value in env.items(): + if key == "OPENAI_API_KEY" and grader_key and value != grader_key: + print( + f"[{seam_name}] WARNING: this seam overwrites OPENAI_API_KEY, which the grader " + "uses for its LLM-judge tier. Judging is disabled for this run (exact-match and " + "numeric tolerance still apply)." + ) + os.environ[key] = value + + +def _pick_reward( + rewards: dict[str, float], reward_key: str = "" +) -> tuple[float | None, str]: + """Collapse Harbor's reward dict to the one scalar an RL trainer consumes. + + Refuses rather than guesses. Harbor lets a task emit any keys it likes, and inventing a rule to + combine them is inventing reward semantics — which is exactly how a previous run got reward + hacked: a `+0.2 for submitting anything` term made the policy learn to quick-submit, training + reward looked healthy, and eval collapsed from 0.740 to 0.178. + + Args: + rewards (`dict[str, float]`): + The verifier's reward dict, as Harbor produced it. + reward_key (`str`, *optional*): + Force a specific key. Required when the dict has several and none is named `reward`. + + Returns: + `tuple[float | None, str]`: The chosen value and the key it came from. + + Raises: + ValueError: If the key is ambiguous and none was given. + """ + if not rewards: + return None, "" + if reward_key: + if reward_key not in rewards: + raise ValueError(f"reward_key {reward_key!r} not in {sorted(rewards)}") + return float(rewards[reward_key]), reward_key + if len(rewards) == 1: + key = next(iter(rewards)) + return float(rewards[key]), key + if "reward" in rewards: + return float(rewards["reward"]), "reward" + raise ValueError( + f"task produced several rewards {sorted(rewards)} and none is named 'reward'. " + "Pass reward_key= to say which one is the training signal." + ) + + +def build_trial_config( + *, + task_dir: Path | str, + harness: str, + sandbox: str, + intercept_url: str, + session_id: str, + model: str, + trial_name: str, + trials_dir: Path | str, + keep_sandbox: bool = False, + agent_timeout_sec: float | None = None, + force_build: bool = False, +) -> Any: + """Assemble Harbor's `TrialConfig` for one rollout. + + The seam decides how this particular agent learns about the proxy — an env var, a constructor + kwarg, or a config file written by a subclass. That is the only per-agent knowledge involved; + everything downstream is identical. + """ + from harbor.models.trial.config import ( + AgentConfig, + EnvironmentConfig, + TaskConfig, + TrialConfig, + VerifierConfig, + ) + + seam = seams.get(harness) + model_name, kwargs, agent_env, proc_env = seam.resolve( + base_url=intercept_url, session=session_id, model=model + ) + apply_process_env(seam.name, proc_env) + + agent = AgentConfig( + name=seam.import_path or harness, + import_path=seam.import_path, + model_name=model_name, + kwargs=kwargs, + env=agent_env, + override_timeout_sec=agent_timeout_sec, + ) + return TrialConfig( + task=TaskConfig(path=Path(str(task_dir))), + agent=agent, + environment=EnvironmentConfig( + type=sandbox, delete=not keep_sandbox, force_build=force_build + ), + verifier=VerifierConfig(), + trial_name=trial_name, + trials_dir=Path(str(trials_dir)), + ) + + +async def run_rollout( + *, + task_dir: Path | str, + harness: str, + sandbox: str, + registry: Any, + intercept_url: str, + model: str, + trials_dir: Path | str, + dataset: str = "", + reward_key: str = "", + keep_sandbox: bool = False, + agent_timeout_sec: float | None = None, + force_build: bool = False, + session_prefix: str = "oe", +) -> HarborRolloutResult: + """Run one rollout end to end. Never raises. + + Args: + task_dir (`Path` or `str`): + The Harbor task directory to run. + harness (`str`): + A seam name (`opencode`, `claude-code`, ...) or a `module:Class` import path. + sandbox (`str`): + A Harbor `EnvironmentType`, e.g. `e2b` or `modal`. + registry (`SessionRegistry`): + The live capture registry; a session is minted here and its id becomes the agent's key. + intercept_url (`str`): + Public URL of the capture proxy, as the sandbox must reach it. + model (`str`): + Served model id. Normalised per harness by the seam. + trials_dir (`Path` or `str`): + Where Harbor writes trial artifacts. + reward_key (`str`, *optional*): + Which reward key is the training signal, for multi-reward tasks. + keep_sandbox (`bool`, *optional*, defaults to `False`): + Leave the sandbox alive after the run, for debugging. + + Returns: + [`HarborRolloutResult`]: Reward, per-turn token ids and logprobs, and validation findings. + """ + task_dir = Path(str(task_dir)) + started = time.monotonic() + + # The session id lands in Harbor's E2B sandbox metadata (it is derived from trial_name), which + # is what later makes it possible to reap only the sandboxes this server created. + trial_name = f"{session_prefix}-{task_dir.name[:28]}-{uuid.uuid4().hex[:8]}" + session = registry.create( + session_id=None, harness=harness, sandbox=sandbox, task=task_dir.name + ) + + result = HarborRolloutResult( + task_id=str(task_dir), + task_name=task_dir.name, + dataset=dataset, + harness=harness, + sandbox=sandbox, + trial_name=trial_name, + session_id=session.session_id, + ) + + trial_result = None + try: + from harbor.trial.trial import Trial + + config_kwargs = dict( + task_dir=task_dir, + harness=harness, + sandbox=sandbox, + intercept_url=intercept_url, + session_id=session.session_id, + model=model, + trial_name=trial_name, + trials_dir=trials_dir, + keep_sandbox=keep_sandbox, + agent_timeout_sec=agent_timeout_sec, + force_build=force_build, + ) + + async with _PROC_ENV_LOCK: + config = build_trial_config(**config_kwargs) + trial = await Trial.create(config) + trial_result = await trial.run() + except Exception as exc: # noqa: BLE001 - a rollout failure is a RESULT, never an exception + result.ok = False + result.error = str(exc)[:600] + result.exception_type = type(exc).__name__ + + result.wall_s = round(time.monotonic() - started, 2) + + # --- reward, forwarded verbatim ------------------------------------- + if trial_result is not None: + verifier = getattr(trial_result, "verifier_result", None) + rewards = dict(getattr(verifier, "rewards", None) or {}) if verifier else {} + result.rewards = {k: float(v) for k, v in rewards.items()} + try: + result.reward, result.reward_key = _pick_reward(result.rewards, reward_key) + except ValueError as exc: + result.ok = False + result.error = str(exc) + for step in getattr(trial_result, "step_results", None) or []: + step_rewards = dict( + getattr(getattr(step, "verifier_result", None), "rewards", None) or {} + ) + result.step_results.append( + HarborStepResult( + name=getattr(step, "name", "") or "", + rewards={k: float(v) for k, v in step_rewards.items()}, + ) + ) + info = getattr(trial_result, "exception_info", None) + if info is not None and result.error is None: + result.ok = False + result.exception_type = getattr(info, "exception_type", None) + result.error = str(getattr(info, "exception_message", ""))[:600] + + # --- capture -------------------------------------------------------- + try: + # `include_messages` is what puts the assistant's own output in the result. + # Only the response side is kept downstream (see `turns_from_document`), so + # the payload grows by the completion text, not by the whole conversation. + document = export_session(session, include_messages=True) + stats = document.get("stats", {}) + result.n_turns = stats.get("n_turns", 0) + result.n_roots = stats.get("n_roots", 0) + result.n_trainable_tokens = stats.get("n_trainable_tokens", 0) + result.multi_turn = result.n_turns > result.n_roots + result.findings = [ + f for f in document.get("validation", []) if not f.startswith("[INFO]") + ] + + trial_dir = _trial_dir(trial_result, trials_dir, trial_name) + atif = load_atif(trial_dir) if trial_dir else None + report = reconcile(document, atif) + result.atif = "none" if atif is None else ("match" if report.ok else "MISMATCH") + result.findings += [ + str(f) for f in report.findings if not str(f).startswith("[INFO]") + ] + + # Calls the harness's own trace does not count as agent steps are auxiliary; drop them so + # they cannot be credited with the reward earned by solving the task. + if report.aux_node_ids: + aux = set(report.aux_node_ids) + for sequence in document["sequences"]: + if sequence["role"] == "agent" and set(sequence["node_ids"]) <= aux: + sequence["role"] = "auxiliary" + + result.turns = turns_from_document(document) + # Every conversation, not only the trainable ones: an auxiliary call that + # went wrong is exactly what someone reading a bad rollout needs to see. + result.conversations = conversations_from_document(document) + if not report.ok: + result.ok = False + result.error = result.error or "capture failed validation" + except Exception as exc: # noqa: BLE001 + result.ok = False + result.error = ( + result.error or f"capture export failed: {type(exc).__name__}: {exc}" + ) + finally: + registry.delete(session.session_id) + + # A rollout that produced no reward is not a zero: the verifier never ran. Keeping the two + # distinct is what stops a dead sandbox being scored as a wrong answer. + if result.ok and result.reward is None and trial_result is not None: + result.findings.append( + "[WARN] ungraded: the verifier produced no reward for this trial" + ) + return result + + +def _trial_dir( + trial_result: Any, trials_dir: Path | str, trial_name: str +) -> Path | None: + """Where Harbor wrote this trial's artifacts, including its trajectory.""" + uri = getattr(trial_result, "trial_uri", None) if trial_result is not None else None + if uri: + return Path(str(uri).replace("file://", "")) + candidate = Path(str(trials_dir)) / trial_name + return candidate if candidate.is_dir() else None + + +__all__ = ["run_rollout", "build_trial_config", "HarborRolloutResult", "HarborTurn"] diff --git a/src/openenv/harbor/runner.py b/src/openenv/harbor/runner.py new file mode 100644 index 000000000..d3f3ececd --- /dev/null +++ b/src/openenv/harbor/runner.py @@ -0,0 +1,280 @@ +"""Drive rollouts without a server: boot capture, run tasks, tear down. + +`openenv harbor rollout` uses this. It exists so the whole path (LLM, capture proxy, forwarding, +seam, Harbor trial, sandbox, verifier, reconciliation) can be exercised with no env server in +the way. When something breaks, that halves the search space immediately: if this works and `serve` does +not, the problem is the serving layer and nothing below it. +""" + +from __future__ import annotations + +import contextlib +import socket +import threading +import time +from pathlib import Path +from typing import Any + +from openenv.core.harness.capture.server import create_app + +from .models import HarborRolloutResult +from .rollout import run_rollout +from .tasks import resolve_task_dirs + + +def _require_free_port(port: int) -> None: + """Raise if anything is already listening on `port`, naming the holder when we can find it.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: + probe.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + try: + probe.bind(("0.0.0.0", port)) + except OSError as exc: + raise RuntimeError( + f"capture port :{port} is already in use ({exc.strerror}). {_port_holder(port)}" + " Stop it or pass a different port: a second server on this port cannot bind, and " + "the agent would silently talk to the older one." + ) from exc + + +def _port_holder(port: int) -> str: + """Best-effort description of the process holding `port`, for the error message only.""" + import shutil + import subprocess + + if not shutil.which("ss"): + return "" + with contextlib.suppress(Exception): + out = subprocess.run( + ["ss", "-ltnp"], capture_output=True, text=True, timeout=5 + ).stdout + for line in out.splitlines(): + if f":{port} " in line and "users:" in line: + return f"Held by {line.split('users:', 1)[1].strip()}." + return "" + + +def _health_instance(port: int) -> str | None: + """Instance id reported by whatever is serving `port`, or `None` if nothing answers yet.""" + import httpx + + with contextlib.suppress(Exception): + resp = httpx.get(f"http://127.0.0.1:{port}/health", timeout=2.0) + if resp.status_code == 200: + return str(resp.json().get("instance") or "unknown") + return None + + +class CaptureServer: + """The capture proxy, running in a background thread for the life of a batch. + + A thread rather than a subprocess because the rollout path needs the live `SessionRegistry` — it + mints a session, then reads the graph back out of it directly. Going through HTTP for that would + add a serialisation round trip and a failure mode for no benefit. + """ + + def __init__( + self, + *, + llm_url: str, + model: str, + port: int = 8100, + max_output_tokens: int = 8192, + ) -> None: + self.app = create_app( + llm_url=llm_url, model=model, max_output_tokens=max_output_tokens + ) + self.port = port + self._thread: threading.Thread | None = None + self._server: Any = None + + @property + def registry(self) -> Any: + return self.app.state.registry + + def start(self, timeout_s: float = 30.0) -> None: + """Bind the port and confirm that the process answering on it is *this* one. + + Raises: + RuntimeError: + If the port is already held, or if the server that comes up on it is not ours. + """ + import uvicorn + + # Fail before uvicorn does. Its bind error surfaces on a background thread, where nothing + # observes it, and the port stays served by whoever holds it. + _require_free_port(self.port) + + config = uvicorn.Config( + self.app, host="0.0.0.0", port=self.port, log_level="warning" + ) + self._server = uvicorn.Server(config) + self._thread = threading.Thread(target=self._server.run, daemon=True) + self._thread.start() + + # Reachability is not identity. A stale process on this port answers every probe, so the + # check is that /health reports our own instance id. + want = self.app.state.instance_id + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + if not self._thread.is_alive(): + raise RuntimeError( + f"capture server thread exited while starting on :{self.port} " + "(most likely the port was taken between the check and the bind)" + ) + got = _health_instance(self.port) + if got == want: + return + if got is not None: + raise RuntimeError( + f"port :{self.port} is served by a different capture server (instance {got}, " + f"expected {want}). Stop the process holding it, or pass a different port; " + "sessions minted here would be rejected there and every rollout would see " + "no model calls." + ) + time.sleep(0.1) + raise RuntimeError( + f"capture server did not come up on :{self.port} within {timeout_s:.0f}s" + ) + + def stop(self) -> None: + if self._server is not None: + self._server.should_exit = True + if self._thread is not None: + self._thread.join(timeout=10) + + +async def run_batch( + *, + llm_url: str, + dataset: str, + task_indices: list[int], + harness: str = "opencode", + sandbox: str = "e2b", + model: str | None = None, + port: int = 8100, + expose: str = "gradio", + trials_dir: Path | None = None, + reward_key: str = "", + keep_sandbox: bool = False, + force_build: bool = False, + env_file: str | None = None, +) -> list[HarborRolloutResult]: + """Run `task_indices` from `dataset` and print a per-rollout report. + + Args: + llm_url (`str`): + OpenAI-spec inference endpoint. + dataset (`str`): + Dataset spec (HF repo id, local dir, or Harbor `name@version`). + task_indices (`list[int]`): + Which tasks to run, by index into the resolved dataset. + harness (`str`, *optional*, defaults to `"opencode"`): + Seam name or `module:Class`. + sandbox (`str`, *optional*, defaults to `"e2b"`): + Harbor environment type. + expose (`str`, *optional*, defaults to `"gradio"`): + How the sandbox reaches the capture proxy: `gradio`, `cloudflare` or `direct`. + + Returns: + `list[HarborRolloutResult]`: One per index, in order. + """ + # Imported here, not at module scope: a hosted deployment mounts the capture proxy on its + # own app and never forwards, so `forwarding` is not shipped there. `serving` imports + # `CaptureServer` from this module, and a module-level import would break that. + from openenv.core.harness.capture.forwarding import make_forwarder + + from .startup import prepare + + caps = prepare( + llm_url=llm_url, + model=model, + datasets=[dataset], + env_file=env_file, + require_llm=True, + quiet=False, + ) + model = caps.llm.get("model") or model or "" + + if sandbox not in caps.available_sandboxes: + detail = next( + (s.detail for s in caps.sandboxes if s.name == sandbox), "not checked" + ) + raise RuntimeError(f"sandbox {sandbox!r} is not usable here: {detail}") + + task_dirs = resolve_task_dirs(dataset) + trials_dir = trials_dir or Path("/tmp/openenv-harbor-trials") + trials_dir.mkdir(parents=True, exist_ok=True) + + capture = CaptureServer(llm_url=llm_url, model=model, port=port) + capture.start() + forwarder = make_forwarder(expose) + public_url = forwarder.start(port) + print(f"\ncapture :{port} -> {public_url} ({forwarder.name})") + print(f"trials {trials_dir}\n") + + results: list[HarborRolloutResult] = [] + try: + for i in task_indices: + if not 0 <= i < len(task_dirs): + print(f" skip index {i}: out of range (dataset has {len(task_dirs)})") + continue + task_dir = task_dirs[i] + print(f"[{harness} / {sandbox}] task {i}: {task_dir.name} ...", flush=True) + result = await run_rollout( + task_dir=task_dir, + harness=harness, + sandbox=sandbox, + registry=capture.registry, + intercept_url=public_url, + model=model, + trials_dir=trials_dir, + dataset=dataset, + reward_key=reward_key, + keep_sandbox=keep_sandbox, + force_build=force_build, + ) + results.append(result) + print(" " + _summarise(result)) + for finding in result.findings[:3]: + print(f" {finding[:150]}") + finally: + forwarder.stop() + capture.stop() + + print("\n" + _report(results)) + return results + + +def _summarise(r: HarborRolloutResult) -> str: + reward = "None" if r.reward is None else f"{r.reward:.2f}" + mode = "multi-turn" if r.multi_turn else "per-turn" + status = "ok" if r.ok else f"FAILED ({r.exception_type or 'error'})" + return ( + f"{status:<26} reward={reward:<6} turns={r.n_turns:<3} roots={r.n_roots:<3} " + f"{mode:<11} tokens={r.n_trainable_tokens:<6} atif={r.atif:<9} {r.wall_s:.0f}s" + + (f"\n {r.error[:180]}" if r.error else "") + ) + + +def _report(results: list[HarborRolloutResult]) -> str: + if not results: + return "no rollouts ran" + ok = sum(1 for r in results if r.ok) + graded = [r for r in results if r.reward is not None] + solved = sum(1 for r in graded if r.reward and r.reward > 0) + lines = [ + "=" * 78, + f"capture {ok}/{len(results)} usable", + f"solved {solved}/{len(graded)} graded" + + ( + f" ({len(results) - len(graded)} ungraded — the verifier never ran)" + if len(graded) != len(results) + else "" + ), + f"tokens {sum(r.n_trainable_tokens for r in results)} trainable across " + f"{sum(r.n_turns for r in results)} turns", + ] + # Capture quality and task success are independent, and conflating them has burned us before: + # a perfectly captured rollout can score 0 because the model was wrong. + lines.append("NOTE: capture and reward are independent measurements.") + return "\n".join(lines) diff --git a/src/openenv/harbor/seams.py b/src/openenv/harbor/seams.py new file mode 100644 index 000000000..2ae80cb02 --- /dev/null +++ b/src/openenv/harbor/seams.py @@ -0,0 +1,766 @@ +"""How each Harbor harness is pointed at the intercept server. + +This is the ONLY per-agent knowledge in the stack, and it is deliberately data. For every Harbor +agent the only thing that differs is which env var or config key carries the base URL and the API +key. Sandbox, capture, stitching, masking and validation are identical downstream. + +A seam has: + env env vars set in OUR process. Harbor's installed agents read os.environ at + construction and forward the provider-relevant subset into the sandbox. + model_fmt what to pass as Harbor's `model_name`. Several agents derive their provider from + the prefix, so this is load-bearing rather than cosmetic. + dialect the wire format we expect. Informational: the server detects per request. Recorded + so a surprise in capture is checkable against what we predicted. + kwargs optional (base_url, session, model) -> dict merged into AgentConfig.kwargs, for + agents needing more than env vars. + status "validated" only after an end-to-end run passed the capture contract AND the ATIF + cross-check. Everything else is "untested", however plausible it looks. + +The API key is always the intercept session id. That is the multiplexing scheme: one server, one +port, N concurrent rollouts, each identified by the key its agent was handed. + +Per-harness findings live in README.md (per-harness findings). Add to it as each agent is brought up. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Callable + +# opencode dispatches on the provider *id*, not the npm package: a provider named `openai` is routed +# to `provider.responses(model)` (the OpenAI Responses API), which @ai-sdk/openai-compatible does not +# implement, and the run dies with `Z.responses is not a function`. Any other name stays on +# chat-completions. Must match the prefix in the opencode seam's model_fmt. +OPENCODE_PROVIDER = "intercepted" + +# Must match `install_fixes.PROVIDER`. Duplicated as a literal rather than imported: `pi_agent` imports +# harbor, and pulling that into this module would make the seam table unusable without Harbor +# installed. Kept honest by `tests/test_seams.py`. +PI_PROVIDER = "intercept" + + +def agent_facing_model(served_model: str) -> str: + """The model name to hand a harness, derived from what the engine actually serves. + + THE PROBLEM. A vLLM started without `--served-model-name` serves under its full repo id, e.g. + `Qwen/Qwen3.5-9B`. Every seam then formats that into its own provider prefix and produces a + two-slash name: + + model_fmt="openai/{model}" -> "openai/Qwen/Qwen3.5-9B" + + Harnesses disagree about what that means. gemini-cli requires exactly `provider/model_name` and + rejects anything else ("Model name must be in the format provider/model_name"); cline-cli wants + `provider:model` with a colon; several split on the FIRST slash and several on the LAST, so the + same string resolves to different models depending on the agent. All of them are "working as + documented" — the string is just ambiguous. + + WHY STRIPPING IS SAFE. The harness-facing name and the upstream name are already decoupled: the + intercept overwrites `chat_request["model"]` with the configured served id on every request + (`capture/server.py`), so whatever a harness puts on the wire is replaced before it reaches the + engine. The harness-facing name only has to be something the harness can parse and route to us; + it never has to match the engine. + + So: take the leaf. `Qwen/Qwen3.5-9B` -> `Qwen3.5-9B`, and a name with no slash is unchanged. + This is preferred over relaunching the engine with `--served-model-name` because it works against + an engine you do not control, including a shared or hosted one. + """ + if not served_model or not served_model.strip(): + raise ValueError( + "served model name is empty; the engine reported no model to route to" + ) + leaf = served_model.strip().rstrip("/").rsplit("/", 1)[-1] + if not leaf: + raise ValueError( + f"cannot derive a harness-facing model name from {served_model!r}" + ) + return leaf + + +def _pi(base_url: str, session: str, model: str) -> dict[str, Any]: + """pi's provider config, carried to our `InterceptPi` subclass via AgentConfig.kwargs.""" + return { + "intercept_config": {"base_url": base_url, "api_key": session, "model": model} + } + + +@dataclass(frozen=True) +class Seam: + """One harness's wiring. Three channels, in order of preference. + + `agent_env` is the RIGHT one and should be the default choice. It becomes `AgentConfig.env`, + which Harbor injects into the sandbox via `agent_environment.scoped_exec_env(agent.extra_env)` + (trial.py:469 for run, :1212 for setup). Two properties make it strictly better than `env`: + + * it works for ANY installed agent, even one whose own code never forwards that variable. + Harbor's `pi`, for instance, forwards only API keys and has no base-URL handling at all; + `agent_env` reaches it anyway because Harbor sets it on every exec in the agent phase. + * it is scoped to the AGENT PHASES ONLY. The verifier runs outside that `with` block, so + setting OPENAI_API_KEY here cannot reach the grader. + + `env` sets variables in OUR process instead. Needed only where the agent reads them at + construction time (before any sandbox exists). It is dangerous: OPENAI_API_KEY set this way + leaks into the grader, whose LLM-judge tier then 401s and silently scores 0 on every answer that + is correct but not an exact string match. + + `kwargs` becomes `AgentConfig.kwargs`, for agents needing structured config (opencode's provider + block). + """ + + name: str + dialect: str + model_fmt: str = "{model}" + # When set, Harbor builds the agent from this class rather than its registered name. The escape + # hatch for a harness whose config Harbor has no seam to write (see pi). + import_path: str | None = None + agent_env: dict[str, str] = field(default_factory=dict) + env: dict[str, str] = field(default_factory=dict) + kwargs: Callable[[str, str, str], dict[str, Any]] | None = None + status: str = "untested" + notes: str = "" + + def resolve( + self, *, base_url: str, session: str, model: str + ) -> tuple[str, dict[str, Any], dict[str, str], dict[str, str]]: + """-> (model_name, AgentConfig.kwargs, AgentConfig.env, os.environ vars). + + `model` is normalised through `agent_facing_model` first, so a served id like + `Qwen/Qwen3.5-9B` cannot leak an extra slash into a harness's model string. + """ + model = agent_facing_model(model) + fmt = {"base_url": base_url, "session": session, "model": model} + agent_env = {k: v.format(**fmt) for k, v in self.agent_env.items()} + proc_env = {k: v.format(**fmt) for k, v in self.env.items()} + extra = self.kwargs(base_url, session, model) if self.kwargs else {} + return self.model_fmt.format(model=model), extra, agent_env, proc_env + + +def _opencode(base_url: str, session: str, model: str) -> dict[str, Any]: + """opencode needs a full provider block; Harbor exposes exactly the right seam for it. + + Harbor writes `provider..options.baseURL` and nothing else (opencode.py:440-447), leaving + opencode to resolve the model through its built-in provider against the models.dev registry. A + locally-served model is not in that registry, so opencode emits step-start/step-finish with zero + tokens and never issues a request: a silent no-op, the worst kind to debug. + + So we supply the provider block outright. `opencode_config` deep-merges LAST (opencode.py:453), + overriding Harbor's generated block without patching Harbor. + """ + return { + "opencode_config": { + "provider": { + OPENCODE_PROVIDER: { + "npm": "@ai-sdk/openai-compatible", + "name": "Harbor Intercept", + "options": { + "baseURL": f"{base_url}/v1", + "apiKey": session, + "timeout": 600_000, + }, + "models": {model: {"name": model}}, + } + } + } + } + + +def _terminus(base_url: str, session: str, model: str) -> dict[str, Any]: + """Terminus runs host-side, in OUR process, so there is no sandbox-side CLI to configure. + + It drives a TmuxSession over environment.exec and calls litellm directly: `api_base` is a + constructor argument. The API key is not; LiteLLM collects surplus kwargs into `_llm_kwargs` and + splats them into the call (lite_llm.py:77,649), so `llm_kwargs` is the channel. + + `model_info` is mandatory, not optional: litellm refuses to route a model name it has no + cost/context metadata for, and a locally-served name is never in its registry. + + Summarisation off: `Terminus2.run` appends subagent rollout details to the main ones + (terminus_2.py:1615), and a summariser is a different task. Training its turns with this + rollout's reward is exactly the contamination to avoid. Compaction also breaks prefix stitching. + """ + return { + "api_base": f"{base_url}/v1", + "llm_kwargs": {"api_key": session}, + "model_info": { + "max_input_tokens": 65536, + "max_output_tokens": 8192, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + "collect_rollout_details": True, # native capture, as a second cross-check + "enable_summarize": False, + "proactive_summarization_threshold": 0, + } + + +SEAMS: dict[str, Seam] = { + # --- priority order for bring-up --------------------------------------- + "opencode": Seam( + name="opencode", + dialect="openai_chat", + model_fmt=OPENCODE_PROVIDER + "/{model}", + # NO env vars. The provider block carries baseURL and apiKey, and setting OPENAI_API_KEY + # here actively breaks grading: Harbor forwards it into the sandbox, where the DataAgent + # grader's tier-3 LLM judge runs `if os.environ.get("OPENAI_API_KEY")` and gets our session + # id instead of a real key. It 401s, the judge is skipped, and every answer that is right + # but not exact-match silently scores 0. Reward corruption, not a crash. + # Naming the provider something other than `openai` also stops Harbor forwarding the key at + # all (agents/installed/opencode.py:516), which is the general trick: alias the provider and + # the grader's key survives untouched. + env={}, + kwargs=_opencode, + status="validated", + notes="Needed 3 global server fixes: SSE replay, stream_options strip, session-id priority.", + ), + "pi": Seam( + name="pi", + status="validated", + dialect="openai_chat", + # Harbor's pi wrapper cannot express a custom endpoint at all, so this seam supplies the + # agent CLASS instead: a local subclass that writes ~/.pi/agent/models.json in setup(). + # See harnesses/pi_agent.py. Harbor stays unmodified. + import_path="openenv.harbor.install_fixes:InterceptPi", + # pi requires `provider/model`; the provider must match the one in models.json. + model_fmt=PI_PROVIDER + "/{model}", + kwargs=_pi, + notes="No base-URL seam in Harbor's wrapper; needs a models.json written into the sandbox. " + "Defaults to the Responses API unless api=openai-completions is pinned.", + ), + "claude-code": Seam( + name="claude-code", + status="validated", + dialect="anthropic", + # No /v1 suffix: the Anthropic SDK appends its own path. + # ANTHROPIC_* does not collide with the grader (which reads OPENAI_API_KEY), so the process + # channel is safe here; agent_env is set too because it is what actually reaches the sandbox. + agent_env={ + "ANTHROPIC_BASE_URL": "{base_url}", + "ANTHROPIC_API_KEY": "{session}", + }, + env={"ANTHROPIC_BASE_URL": "{base_url}", "ANTHROPIC_API_KEY": "{session}"}, + notes="Calls /v1/messages/count_tokens (handled as an aux route). Injects an env/time block " + "in its system prompt: watch n_roots for nonce breakage.", + ), + "codex": Seam( + name="codex", + status="validated", + dialect="openai_responses", + # Key goes ONLY through agent_env: in the process env it would overwrite the grader's + # OPENAI_API_KEY and silently disable the LLM-judge tier. + agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"}, + notes="Responses dialect: exercises a different transform than chat-completions.", + ), + "gemini-cli": Seam( + name="gemini-cli", + status="validated", + dialect="google", + # Requires `provider/model` (gemini_cli.py:781) or it raises before install even matters: + # ValueError: Model name must be in the format provider/model_name + # The first failure here LOOKED like an nvm install problem because the job log echoes the + # install command; the run never got that far. Read result.json's exception_info, not the log. + model_fmt="google/{model}", + # Harbor's wrapper declares only `curl` as a system dep, but nvm's installer pipes into + # bash, so in an image without bash it fails with a misleading "NVM failed to load". + # The subclass adds bash and otherwise defers to Harbor. See harnesses/install_fixes.py. + import_path="openenv.harbor.install_fixes:InterceptGeminiCli", + agent_env={ + "GOOGLE_GEMINI_BASE_URL": "{base_url}", + "GEMINI_API_KEY": "{session}", + "GOOGLE_API_KEY": "{session}", + }, + env={ + "GOOGLE_GEMINI_BASE_URL": "{base_url}", + "GEMINI_API_KEY": "{session}", + "GOOGLE_API_KEY": "{session}", + }, + notes="generateContent dialect; key arrives as x-goog-api-key. Model is carried in the URL " + "path rather than the body: expect that to be the first thing to break.", + ), + # --- second wave -------------------------------------------------------- + "terminus-2": Seam( + name="terminus-2", + status="validated", + dialect="openai_chat", + model_fmt="hosted_vllm/{model}", + kwargs=_terminus, + notes="Host-side agent: no sandbox involved in LLM traffic. Also emits RolloutDetail.", + ), + "openhands": Seam( + name="openhands", + dialect="openai_chat", + model_fmt="openai/{model}", + # Harbor installs `openhands-ai` unpinned and verifies with `python -m openhands.core.main`, + # but V1 moved the core out to openhands-sdk, so latest fails with + # `ModuleNotFoundError: No module named 'openhands.core'`. The subclass pins the last V0 + # release (0.49.0). See harnesses/install_fixes.py. + import_path="openenv.harbor.install_fixes:InterceptOpenHands", + # OpenHands does NOT use OPENAI_*. It reads LLM_MODEL / LLM_BASE_URL / LLM_API_KEY, all via + # `_get_env` (openhands.py:873,920,931), so agent_env reaches them. LLM_MODEL is taken from + # model_name directly, and litellm needs the provider prefix to route. + agent_env={"LLM_BASE_URL": "{base_url}/v1", "LLM_API_KEY": "{session}"}, + notes="LLM_* env vars, not OPENAI_*. Harbor has a 'dummy-key-for-local-vllm' fallback.", + ), + "mini-swe-agent": Seam( + name="mini-swe-agent", + status="validated", + dialect="openai_chat", + model_fmt="openai/{model}", + # Needs `provider/model`. It reads MSWEA_API_KEY, and otherwise derives the key variable from + # the model name via litellm (`openai/` -> OPENAI_API_KEY). Base URL goes to litellm, which + # accepts OPENAI_API_BASE or OPENAI_BASE_URL; set both, the unused one is harmless. + # All through agent_env so the grader's OPENAI_API_KEY is never touched. + agent_env={ + "MSWEA_API_KEY": "{session}", + "OPENAI_API_KEY": "{session}", + "OPENAI_API_BASE": "{base_url}/v1", + "OPENAI_BASE_URL": "{base_url}/v1", + }, + notes="MSWEA_API_KEY plus a model-derived key var; litellm under the hood.", + ), + "qwen-coder": Seam( + name="qwen-coder", + status="validated", + dialect="openai_chat", + # Cleanest seam of the lot: declarative EnvVar descriptors for api_key -> OPENAI_API_KEY and + # base_url -> OPENAI_BASE_URL with env_fallback (qwen_code.py:39-45), then passed explicitly + # as --openai-api-key / --openai-base-url. agent_env feeds `_get_env` directly. + agent_env={"OPENAI_API_KEY": "{session}", "OPENAI_BASE_URL": "{base_url}/v1"}, + ), + "swe-agent": Seam( + name="swe-agent", + status="validated", + dialect="openai_chat", + model_fmt="openai/{model}", + agent_env={"OPENAI_API_KEY": "{session}", "OPENAI_BASE_URL": "{base_url}/v1"}, + # Harbor's repo argument has a quoting bug in its else-branch (`'--env.repo.path=$(pwd)'` + # inside single quotes), giving `git.exc.NoSuchPathError: /workdir/$(pwd)`. The subclass makes + # /workdir a git repo and exposes it as /testbed so Harbor takes its WORKING branch. + import_path="openenv.harbor.install_fixes:InterceptSweAgent", + # Second bug, and the one that made swe-agent produce exactly ONE turn per task while every + # capture check passed. swe-agent asks litellm to cost each call; litellm has no pricing row + # for a locally-served name and raises, which swe-agent treats as fatal: + # + # sweagent.exceptions.ModelConfigurationError: Error calculating cost: + # This model isn't mapped yet. model=openai/Qwen3.5-9B ... + # please make sure you set `per_instance_cost_limit` and `total_cost_limit` to 0 + # + # Harbor already sets exactly these to 0, but only under `if is_hosted_vllm:` (swe_agent.py + # :437-443), and our model_fmt is `openai/` so that branch never runs. They are declared + # CLI_FLAGS, so passing them as kwargs reaches the same flags. `build_cli_flags` skips only + # None (base.py:651), so "0" is emitted rather than dropped as falsy. + kwargs=lambda base_url, session, model: { + "per_instance_cost_limit": "0", + "total_cost_limit": "0", + "max_input_tokens": "0", + }, + notes="SWE-bench agent: requires a git repo, which DataAgent tasks are not.", + ), + "goose": Seam( + name="goose", + status="validated", + dialect="openai_chat", + model_fmt="openai/{model}", + # goose is the one that reads `os.environ.get("OPENAI_API_KEY")` DIRECTLY (goose.py:678) and + # raises if unset, so agent_env cannot reach it and the process env is forced. That disables + # the grader's LLM-judge tier for goose runs; exact-match and numeric tolerance still apply. + # `runner.apply_seam_env` warns when this happens so it is never a silent reward change. + env={"OPENAI_API_KEY": "{session}", "OPENAI_BASE_URL": "{base_url}/v1"}, + agent_env={"OPENAI_BASE_URL": "{base_url}/v1"}, + notes="Reads os.environ directly, so the process env is unavoidable -> LLM judge disabled.", + ), + # --- tier 2: seams derived from each wrapper, none validated yet ---------- + # These follow the same two shapes seen everywhere: a key + base URL pair, delivered through + # `_get_env` (so agent_env works) or `os.environ` (so the process env is forced). Where the + # wrapper reads os.environ directly, both channels are set and the grader warning fires. + "hermes": Seam( + name="hermes", + status="validated", + dialect="openai_chat", + model_fmt="openai/{model}", + # Harbor's generated config sets `provider: "auto"`, so hermes self-routes and ignores + # OPENAI_BASE_URL (401 from an endpoint that is not ours). The subclass pins the provider and + # disables compression, which would otherwise break prefix stitching. + import_path="openenv.harbor.install_fixes:InterceptHermes", + kwargs=lambda base_url, session, model: { + "intercept_config": {"base_url": base_url, "api_key": session} + }, + # hermes.py:365 reads os.environ.get("OPENAI_BASE_URL") directly. + env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"}, + agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"}, + ), + "vibe": Seam( + name="vibe", + status="validated", + dialect="openai_chat", + model_fmt="openai/{model}", + # vibe defaults to Mistral's API unless VIBE_API_BASE or OPENAI_BASE_URL is set (vibe.py:76-96), + # and resolves its key from the var named by VIBE_API_KEY_ENV. + # Its own error names the valid values: + # ValueError: Unknown Vibe backend 'openai'; valid backends are 'mistral' and 'generic' + # (use 'generic' for any OpenAI-compatible endpoint) + agent_env={ + "VIBE_API_BASE": "{base_url}/v1", + "OPENAI_BASE_URL": "{base_url}/v1", + "VIBE_API_KEY_ENV": "OPENAI_API_KEY", + "OPENAI_API_KEY": "{session}", + "VIBE_BACKEND": "generic", + }, + ), + "openclaw": Seam( + name="openclaw", + status="validated", + dialect="openai_chat", + model_fmt="openai/{model}", + agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"}, + # Harbor writes a merged config to openclaw.upload.json and copies it into the sandbox, but + # only when there IS a config. With none, setup dies on + # cp: cannot stat '/logs/agent/openclaw.upload.json' + # A provider block gives it something to write AND points openclaw at us. + # Schema is `models.providers.`, NOT a top-level `providers`. A top-level one is + # accepted by Harbor's merge and then rejected by openclaw itself: + # OpenClaw config is invalid + # Problem: - : Invalid input + # Harbor's own `_align_provider_models` reads cfg["models"]["providers"][provider] and fills + # in a `models` array beside `baseUrl`, which is the shape mirrored here. The provider name + # must match the model prefix so `_model_provider()` resolves it. + kwargs=lambda base_url, session, model: { + "openclaw_config": { + "models": { + "providers": { + "openai": { + "baseUrl": f"{base_url}/v1", + "apiKey": session, + "api": "openai-completions", + "models": [{"id": f"openai/{model}", "name": model}], + } + } + } + }, + # Harbor defaults `--thinking high`, which openclaw rejects for a custom provider: + # Error: Thinking level "high" is not supported for openai/Qwen3.5-9B. Use one of: off. + # It is a CliFlag, so kwargs can override it. Also correct for us: we serve with thinking + # disabled, so anything else would be asking for a mode the model is not running in. + "thinking": "off", + }, + # Harbor writes the config to the HOST logs dir then copies it from the CONTAINER path, + # which assumes a bind mount. E2B has none, so the subclass uploads it into the sandbox. + import_path="openenv.harbor.install_fixes:InterceptOpenClaw", + notes="Harbor assumes /logs/agent is bind-mounted; E2B is not, so the config is uploaded.", + ), + "kimi-cli": Seam( + name="kimi-cli", + status="validated", + dialect="openai_chat", + model_fmt="openai/{model}", + # Harbor DELIBERATELY unsets OPENAI_BASE_URL / OPENAI_API_KEY before spawning kimi + # (`_KIMI_ENV_OVERRIDES_TO_NEUTRALIZE`), because kimi-cli's own + # `augment_provider_with_env_vars` silently overrides its config file from those vars — a + # globally-injected OpenAI key would hijack an OpenRouter run (MoonshotAI/kimi-cli#1165). + # So an env-var seam is not merely ignored here, it is actively erased. That is why every + # attempt showed 0 turns. + # + # The config file is what wins, and Harbor exposes both values as plain constructor kwargs: + # base_url = self._base_url or pcfg["base_url"] (kimi_cli.py:208) + # api_key = self._api_key or (kimi_cli.py:169) + kwargs=lambda base_url, session, model: { + "base_url": f"{base_url}/v1", + "api_key": session, + }, + # Second, separate bug. Harbor's run command ends with `kill 0`, which takes the E2B exec + # stream down along with the process group and raises RemoteProtocolError(StreamReset) in + # OUR process. That killed all 11 kimi trials AFTER the agent had finished its work, so the + # verifier never ran and every reward came back None. The subclass swallows only that one + # error, mirroring Harbor's own handling of the exit-143 half of the same teardown. + import_path="openenv.harbor.install_fixes:InterceptKimi", + ), + "mimo": Seam( + # mimo is an opencode fork and inherits its provider dispatch, so it inherits the same trap: + # a provider literally named `openai` routes to `provider.responses(model)`, which + # @ai-sdk/openai-compatible does not implement, and the run dies with + # TypeError: Z.responses is not a function + # having never reached the intercept (0 turns). + # + # Harbor tries to avoid this: it writes `npm: @ai-sdk/openai-compatible` into the provider + # block, but only when it finds a base URL, and it looks in OUR PROCESS env + # (`os.environ.get(f"{provider.upper()}_BASE_URL")`, mimo.py:389) rather than in the sandbox + # env we set. `agent_env` never reaches os.environ, so that branch does not run and mimo + # falls back to its built-in openai provider. Same shape as the goose problem. + # + # Fixed the way opencode is fixed, and for the same reason: a provider id that is NOT + # "openai", with the block supplied outright via the `mimo_config` kwarg. That deep-merges + # LAST (mimo.py:402), so it wins without patching Harbor and without putting OPENAI_API_KEY + # in our process env where the verifier's LLM judge would inherit it. + name="mimo", + dialect="openai_chat", + model_fmt=OPENCODE_PROVIDER + "/{model}", + kwargs=lambda base_url, session, model: { + "mimo_config": { + "provider": { + OPENCODE_PROVIDER: { + "npm": "@ai-sdk/openai-compatible", + "name": "Harbor Intercept", + "options": { + "baseURL": f"{base_url}/v1", + "apiKey": session, + "timeout": 600_000, + }, + "models": {model: {"name": model}}, + } + } + } + }, + agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"}, + ), + "trae-agent": Seam( + # RESPONSES, not chat. Mislabelled `openai_chat` for a night because the config says + # `provider: openai` and everything else with that label speaks chat. The access log settled + # it: exactly one `POST /v1/responses` against 465 chat-completions calls, and that one was + # trae. Its client reads `usage.input_tokens_details.cached_tokens`, a Responses-only field. + name="trae-agent", + status="validated", + dialect="openai_responses", + model_fmt="openai/{model}", + agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"}, + ), + # --- remaining ATIF agents ------------------------------------------------ + # Goal is every ATIF-capable agent Harbor registers (25). These nine had no seam; the ones with a + # recognisable LLM config get one, and the vendor-service ones are attempted anyway so their block + # reason is RECORDED rather than assumed. + "computer-1": Seam( + name="computer-1", + dialect="openai_chat", + model_fmt="hosted_vllm/{model}", + # Host-side like terminus-2: drives litellm from our process with an `api_base` kwarg, so + # there is no sandbox-side CLI to configure. + kwargs=lambda base_url, session, model: { + "api_base": f"{base_url}/v1", + "llm_kwargs": {"api_key": session}, + "model_info": { + "max_input_tokens": 262144, + "max_output_tokens": 8192, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + }, + notes="Host-side agent, litellm. The other RolloutDetail emitter besides terminus-2.", + ), + "eve": Seam( + name="eve", + dialect="openai_chat", + model_fmt="openai/{model}", + # NOT a general coding agent. `_validate_path` requires a local Eve PROJECT directory with + # package.json plus an agent/ dir (or flat agent.ts / instructions.md), and raises before any + # sandbox work -- the 7s failure with empty Harbor logs. Nothing to do with the intercept. + status="blocked:needs-eve-project", + # Reads OPENAI_BASE_URL / OPENAI_ENDPOINT / OPENAI_API_KEY. + agent_env={ + "OPENAI_BASE_URL": "{base_url}/v1", + "OPENAI_ENDPOINT": "{base_url}/v1", + "OPENAI_API_KEY": "{session}", + }, + env={"OPENAI_BASE_URL": "{base_url}/v1"}, + ), + "cursor-cli": Seam( + name="cursor-cli", + dialect="openai_chat", + model_fmt="openai/{model}", + agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"}, + status="blocked:credentials", + # CONFIRMED by running it, not assumed: + # ValueError: CURSOR_API_KEY environment variable is required. + # It authenticates against Cursor's own service before any model call, so pointing it at a + # local endpoint cannot help. Nothing about the intercept is implicated. + notes="BLOCKED: requires CURSOR_API_KEY (Cursor account). Verified, not assumed.", + ), + "acp": Seam( + name="acp", + dialect="openai_chat", + model_fmt="openai/{model}", + # A LAUNCHER for ACP-speaking agents, not an agent. Needs `registry_entry` / + # `registry_entry_path` describing a distribution ("ACP registry entry must define at least + # one distribution"), and raises immediately without one. + status="blocked:needs-registry-entry", + agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"}, + notes="Agent Client Protocol runner; needs an ACP-speaking agent configured underneath.", + ), + "devin": Seam( + name="devin", + dialect="openai_chat", + model_fmt="openai/{model}", + agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"}, + notes="Cognition hosted service; expected to need vendor credentials.", + ), + "copilot-cli": Seam( + name="copilot-cli", + dialect="openai_chat", + model_fmt="openai/{model}", + agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"}, + notes="Needs GITHUB_TOKEN / COPILOT_GITHUB_TOKEN and a Copilot subscription.", + ), + "antigravity-cli": Seam( + name="antigravity-cli", + dialect="openai_chat", + model_fmt="openai/{model}", + agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"}, + notes="Google Antigravity; auth via AGY_AUTH_JSON_PATH.", + ), + "grok-build": Seam( + name="grok-build", + dialect="openai_chat", + model_fmt="openai/{model}", + agent_env={ + "OPENAI_BASE_URL": "{base_url}/v1", + "OPENAI_API_KEY": "{session}", + "XAI_API_KEY": "{session}", + }, + notes="xAI; expected to need an xAI key.", + ), + "rovodev-cli": Seam( + name="rovodev-cli", + dialect="openai_chat", + model_fmt="openai/{model}", + agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"}, + notes="Atlassian; needs ROVODEV_USER_API_TOKEN + ROVODEV_USER_EMAIL.", + ), + # Found only via Harbor's SUPPORTS_ATIF flag; an import grep missed all four. + "cline-cli": Seam( + name="cline-cli", + dialect="openai_chat", + # COLON, not slash. Its own error is explicit: + # ValueError: model_name must be in format 'provider:model-id', got: 'openai/Qwen3.5-9B' + # The only harness so far that does not use `provider/model`. + model_fmt="openai:{model}", + agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"}, + # Harbor gives cline NO base-URL channel (only PROVIDER/API_KEY/MODELID), so it calls the + # real OpenAI. The subclass merges Cline's own settings store between Harbor's write and the + # run. See harnesses/install_fixes.py. + import_path="openenv.harbor.install_fixes:InterceptCline", + kwargs=lambda base_url, session, model: { + "intercept_config": { + "base_url": base_url, + "api_key": session, + "model": model, + } + }, + ), + "nemo-agent": Seam( + name="nemo-agent", + dialect="openai_chat", + model_fmt="openai/{model}", + # Defaults to `llm_type: "nim"` (NVIDIA NIM), so it never reads OPENAI_* and the OpenAI seam + # was silently ignored. Its own module docstring gives the recipe: + # harbor run --agent nemo-agent --model openai/gpt-4o --ak llm_type=openai + # `--ak` is AgentConfig.kwargs, so this selects the OpenAI-compatible provider. + kwargs=lambda base_url, session, model: {"llm_type": "openai"}, + agent_env={ + "OPENAI_BASE_URL": "{base_url}/v1", + "OPENAI_API_KEY": "{session}", + "OPENAI_API_BASE": "{base_url}/v1", + }, + ), + "openhands-sdk": Seam( + name="openhands-sdk", + status="validated", + dialect="openai_chat", + model_fmt="openai/{model}", + # The V1 SDK packaging, so it should NOT need the V0 pin the classic wrapper does. + agent_env={ + "LLM_BASE_URL": "{base_url}/v1", + "LLM_API_KEY": "{session}", + "OPENAI_BASE_URL": "{base_url}/v1", + "OPENAI_API_KEY": "{session}", + }, + ), + "antigravity-sdk": Seam( + name="antigravity-sdk", + dialect="google", + model_fmt="google/{model}", + # It said so itself: `GEMINI_API_KEY environment variable must be set`. My first seam gave it + # only OPENAI_*, which is why it looked like a vendor-credential block when it is really a + # Google-family agent. Same seam as gemini-cli, which is validated. + agent_env={ + "GOOGLE_GEMINI_BASE_URL": "{base_url}", + "GEMINI_API_KEY": "{session}", + "GOOGLE_API_KEY": "{session}", + }, + env={ + "GOOGLE_GEMINI_BASE_URL": "{base_url}", + "GEMINI_API_KEY": "{session}", + "GOOGLE_API_KEY": "{session}", + }, + notes="Google Antigravity SDK; takes the gemini-cli seam, not an OpenAI one.", + ), +} + +# Every ATIF-capable agent Harbor registers. This IS the goal. +# +# Source of truth is Harbor's own `SUPPORTS_ATIF` class flag, not a grep for trajectory imports. +# The grep undercounted (25 vs 29): it missed antigravity-sdk, cline-cli, nemo-agent and +# openhands-sdk, which build ATIF without importing the models in a way grep could see. Re-derive with: +# +# getattr(import_class(AgentFactory._AGENT_MAP[a], label="agent"), "SUPPORTS_ATIF", False) +# +# computer-1 is deliberately excluded: not needed for this goal. +# Every agent whose Harbor class sets SUPPORTS_ATIF, minus two deliberate exclusions. +# +# `computer-1` is out of scope by request. +# +# `openhands` (V0) is out because Harbor registers it and `openhands-sdk` as two SEPARATE agents, +# not as old and new names for one. V0 bundles the full `openhands-ai` with its own Docker runtime; +# the SDK runs directly in the container. Running a Docker runtime inside a sandbox that is already +# a container is the wrong shape for this work, and the SDK validated 5/5 with zero fixes while V0 +# needed four packaging patches and still did not come through. Its Seam and InterceptOpenHands +# subclass are kept so `--agents openhands` still works; it is simply not a target. +# `nemo-agent` is out as well. Harbor generates its NAT config as +# `workflow: {_type: chat_completion}`: a single-shot LLM call with no tools and no loop, so ONE +# turn is correct behaviour, not a failure. Making it agentic needs a react_agent workflow plus +# NAT's `code_execution` tool, which requires a separate sandbox service (docker or Piston) that +# would not even have the task's CSV. Bring-your-own-workflow, like eve and acp. +# `antigravity-sdk` is out too. Its Go `localharness` binary receives a correctly-translated +# functionCall (verified by direct probe) and then ends the conversation without executing it, the +# SDK's "Received tool call %s but no tool runner is configured. Yielding to user." path. That is +# inside google-antigravity, and Harbor's runner sets the root logger to ERROR so the warning that +# would confirm it never reaches a log. Seam kept; not a target. +ATIF_AGENTS = [ + "acp", + "antigravity-cli", + "claude-code", + "cline-cli", + "codex", + "copilot-cli", + "cursor-cli", + "devin", + "eve", + "gemini-cli", + "goose", + "grok-build", + "hermes", + "kimi-cli", + "mimo", + "mini-swe-agent", + "openclaw", + "opencode", + "openhands-sdk", + "qwen-coder", + "rovodev-cli", + "swe-agent", + "terminus-2", + "trae-agent", + "vibe", +] + +PRIORITY = [ + "opencode", + "pi", + "claude-code", + "codex", + "gemini-cli", + "terminus-2", + "openhands-sdk", + "mini-swe-agent", +] + + +def get(name: str) -> Seam: + if name not in SEAMS: + raise KeyError(f"no seam for {name!r}. Known: {sorted(SEAMS)}") + return SEAMS[name] diff --git a/src/openenv/harbor/serving.py b/src/openenv/harbor/serving.py new file mode 100644 index 000000000..8fbd5b508 --- /dev/null +++ b/src/openenv/harbor/serving.py @@ -0,0 +1,243 @@ +"""Serve Harbor tasks: Task API, MCP rollouts, and a human UI. + +Locally this is two ports: + + :port env server Task API + MCP + UI (faces the trainer / a browser) + :capture_port capture proxy (faces the sandbox, published) + +Two ports on purpose. The sandbox is off-cluster and must reach the capture proxy over a public URL; +the env server has no business being publicly reachable, and sharing one port would expose it as +soon as the capture proxy became reachable. + +On a hosted platform that inverts. A Space gets exactly one public URL and exposes one port, so +there is no second port to publish and nothing to forward. The capture app is mounted onto the env +server's own app at `CAPTURE_MOUNT` instead, and the sandbox reaches it at `/capture`. +The proxy still refuses unregistered callers, which is what keeps a public mount from becoming an +open relay. +""" + +from __future__ import annotations + +import os +import threading +from typing import Any + +from .runner import CaptureServer + +# Where the capture app is mounted when the env server hosts it directly. +CAPTURE_MOUNT = "/capture" + + +def space_public_url() -> str: + """The public URL of the Space this process is running in, or `""` when it is not on one. + + Returns: + `str`: e.g. `https://owner-name.hf.space`, with no trailing slash. + """ + host = os.environ.get("SPACE_HOST", "").strip() + if host: + return "https://" + host.rstrip("/").removeprefix("https://").removeprefix( + "http://" + ) + # SPACE_HOST is the direct answer, but SPACE_ID is the variable that is always set, so derive + # the hostname the same way `auto.auto_env` does. + space_id = os.environ.get("SPACE_ID", "").strip() + if space_id and "/" in space_id: + slug = space_id.replace("/", "-").replace("_", "-").replace(".", "-").lower() + return f"https://{slug}.hf.space" + return "" + + +class HarborService: + """Long-lived state for a serving process: capture proxy, forwarding, datasets. + + Held at module scope by `serve_harbor` so that the environment instances OpenEnv builds + per-request can reach it. They must not own it: `/metadata` and `/schema` construct a throwaway + environment on every call, so anything expensive on `__init__` would be paid per docs hit. + """ + + _instance: "HarborService | None" = None + + def __init__( + self, + *, + llm_url: str, + model: str, + datasets: list[str], + capture_port: int = 8100, + expose: str = "gradio", + ) -> None: + self.llm_url = llm_url + self.model = model + self.datasets = datasets + self.capture = CaptureServer(llm_url=llm_url, model=model, port=capture_port) + self._expose_kind = expose + self.public_url = "" + self.mounted = False + self._forwarder: Any = None + self._lock = threading.Lock() + + def start(self) -> str: + """Make the capture proxy reachable from the sandbox and return its public URL. + + Two different situations, and conflating them is what makes the hosted case awkward: + + - **Hosted** (a Space). The platform already gives this process one public URL and exposes + exactly one port. So the capture app is mounted onto the env server's own app under + `CAPTURE_MOUNT` and reached at `/capture`. No second port, no forwarding, nothing + for the platform to object to. + - **Local.** The sandbox runs off-cluster and cannot reach `127.0.0.1`, so the capture port + is published by whichever forwarder was selected. + """ + public = space_public_url() + if public: + # The env server's app serves it; `build_app` performs the mount. + self.mounted = True + self.public_url = f"{public}{CAPTURE_MOUNT}" + return self.public_url + + from openenv.core.harness.capture.forwarding import make_forwarder + + self.capture.start() + self._forwarder = make_forwarder(self._expose_kind) + self.public_url = self._forwarder.start(self.capture.port) + return self.public_url + + def stop(self) -> None: + if self._forwarder is not None: + self._forwarder.stop() + self.capture.stop() + + @classmethod + def current(cls) -> "HarborService | None": + return cls._instance + + @classmethod + def set_current(cls, service: "HarborService") -> None: + cls._instance = service + + +def serve_harbor( + *, + llm_url: str, + datasets: list[str], + model: str | None = None, + host: str = "0.0.0.0", + port: int = 8000, + capture_port: int = 8100, + expose: str = "gradio", + env_file: str | None = None, +) -> None: + """Boot the capture proxy, then serve the env server with the UI mounted. + + Args: + llm_url (`str`): + OpenAI-spec inference endpoint. + datasets (`list[str]`): + Dataset specs to serve as splits. + port (`int`, *optional*, defaults to `8000`): + Env server port. + capture_port (`int`, *optional*, defaults to `8100`): + Capture proxy port. This is the one published to the sandbox. Ignored on a hosted + platform, where the proxy is mounted on the env server's own app instead. + expose (`str`, *optional*, defaults to `"gradio"`): + How the sandbox reaches the capture proxy locally: `gradio`, `cloudflare` or `direct`. + """ + import uvicorn + + from .startup import prepare + + caps = prepare( + llm_url=llm_url, + model=model, + datasets=datasets, + env_file=env_file, + require_llm=True, + quiet=False, + ) + model = caps.llm.get("model") or model or "" + + service = HarborService( + llm_url=llm_url, + model=model, + datasets=datasets, + capture_port=capture_port, + expose=expose, + ) + public = service.start() + HarborService.set_current(service) + + where = "mounted on this app" if service.mounted else f":{capture_port}" + print(f"\ncapture {where} -> {public}") + print( + f"server http://{host}:{port} (UI at /web, Task API at /{{env}}/splits)" + ) + print("Ctrl-C to stop\n") + + # The UI is the whole point of this entry point, so turn it on rather than making the operator + # discover an env var. + os.environ.setdefault("ENABLE_WEB_INTERFACE", "true") + + app = build_app(datasets=datasets, llm_url=llm_url, model=model, llm=caps.llm) + try: + uvicorn.run(app, host=host, port=port, log_level="info") + finally: + service.stop() + + +def build_app( + *, + datasets: list[str], + llm_url: str = "", + model: str = "", + llm: dict[str, Any] | None = None, +) -> Any: + """The FastAPI app: Task API + MCP + the Gradio UI.""" + from openenv.core.env_server.http_server import create_app + from openenv.core.env_server.mcp_environment import ( + CallToolAction, + CallToolObservation, + ) + + from .environment import HarborEnvironment + from .ui import harbor_gradio_builder + + HarborEnvironment.configure( + datasets=datasets, llm_url=llm_url, model=model, llm=llm + ) + + def gradio_builder( + _web_manager: Any = None, + _action_fields: Any = None, + _metadata: Any = None, + _is_chat: Any = None, + display_title: str = "", + _quick_start: Any = None, + ) -> Any: + """OpenEnv calls this positionally with six web-interface arguments. + + Only the title is useful here: the Harbor UI drives rollouts through its own handlers rather + than the generic action-field form, because a rollout is one long tool call, not a step. + """ + return harbor_gradio_builder(datasets=datasets, title=display_title or "Harbor") + + app = create_app( + HarborEnvironment, + CallToolAction, + CallToolObservation, + env_name="harbor_env", + max_concurrent_envs=int(os.getenv("MAX_CONCURRENT_ENVS", "4")), + gradio_builder=gradio_builder, + custom_tab_name="Harbor", + custom_tab_primary=True, + show_default_tab=False, + ) + + # When the platform gives us a single port, the capture proxy rides on this app instead of + # being published separately. Mounting strips the prefix, so the proxy's own catch-all still + # sees `/v1/chat/completions` and every dialect keeps working unchanged. + service = HarborService.current() + if service is not None and service.mounted: + app.mount(CAPTURE_MOUNT, service.capture.app) + + return app diff --git a/src/openenv/harbor/startup.py b/src/openenv/harbor/startup.py new file mode 100644 index 000000000..48fa8d92f --- /dev/null +++ b/src/openenv/harbor/startup.py @@ -0,0 +1,159 @@ +"""Everything that must be true before the server accepts a request. + +Four checks, in the order that fails cheapest first: + +1. **LLM** — can it return token ids at all? A vLLM without + `--return-tokens-as-token-ids --logprobs-mode processed_logprobs` answers every request perfectly + well and returns no ids, so every rebuilt training row is empty and nothing reports an error. This + is the one failure with no loud edge, so it is checked first and is fatal by default. +2. **sandbox credentials** — via Harbor's own `preflight()`, so the message names the exact missing + variable rather than us guessing at one. +3. **datasets** — resolved and downloaded up front. A 2000-task repo takes real time to fetch, and a + mistyped dataset name should fail here rather than on the first rollout. +4. **report** — print what is usable and, more usefully, what is not and why. + +A caller gets the same `Capabilities` object the server serves over the wire, so what is printed at +startup and what a client can query are the same data. +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any + +from .capabilities import Capabilities, capabilities + +# Load once, at import, so that Harbor's per-backend preflight sees the keys. Harbor reads +# credentials from the process environment and never from a file, so a `.env` that is not exported +# is invisible to it. +_ENV_LOADED = False + + +def load_env_file(path: str | Path | None = None) -> list[str]: + """Export `KEY=value` pairs from a dotenv file into the process environment. + + Existing variables win: an operator who exported something deliberately should not have it + silently replaced by a checked-in file. + + Args: + path (`str` or `Path`, *optional*): + The file to read. Defaults to `$OPENENV_ENV_FILE`, then `./.env`. + + Returns: + `list[str]`: Names of the variables that were set (values are never returned or logged). + """ + global _ENV_LOADED + candidate = Path(path or os.environ.get("OPENENV_ENV_FILE") or ".env").expanduser() + if not candidate.is_file(): + return [] + + applied: list[str] = [] + for raw in candidate.read_text(errors="replace").splitlines(): + line = raw.strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, _, value = line.partition("=") + key, value = key.strip(), value.strip().strip('"').strip("'") + if key and key not in os.environ: + os.environ[key] = value + applied.append(key) + _ENV_LOADED = True + return applied + + +def prepare( + *, + llm_url: str | None = None, + model: str | None = None, + datasets: list[str] | None = None, + sandboxes: tuple[str, ...] | None = None, + env_file: str | Path | None = None, + require_llm: bool = True, + quiet: bool = False, +) -> Capabilities: + """Run every startup check and return what this server can do. + + Args: + llm_url (`str`, *optional*): + OpenAI-spec inference endpoint. No default: callers pass it explicitly. + model (`str`, *optional*): + Served model id. Defaults to `$OPENENV_MODEL`, else the LLM's only served model. + datasets (`list[str]`, *optional*): + Dataset specs to serve. Defaults to `$OPENENV_DATASETS` (comma-separated). + sandboxes (`tuple[str, ...]`, *optional*): + Backends to check. Defaults to `$OPENENV_SANDBOXES`, else the known set. + env_file (`str` or `Path`, *optional*): + Dotenv file to load before checking credentials. + require_llm (`bool`, *optional*, defaults to `True`): + Raise if the LLM cannot support capture. Set `False` to report and continue. + quiet (`bool`, *optional*, defaults to `False`): + Suppress the printed report. + + Returns: + [`Capabilities`]: Harnesses, sandboxes, datasets and LLM status. + + Raises: + RuntimeError: If `require_llm` and no URL was given, or the LLM cannot return token ids. + """ + load_env_file(env_file) + + model = model or os.environ.get("OPENENV_MODEL", "") + if datasets is None: + raw = os.environ.get("OPENENV_DATASETS", "") + datasets = [d.strip() for d in raw.split(",") if d.strip()] + if sandboxes is None: + raw = os.environ.get("OPENENV_SANDBOXES", "") + sandboxes = tuple(s.strip() for s in raw.split(",") if s.strip()) or None + + if require_llm and not llm_url: + raise RuntimeError( + "no LLM URL given. Pass --llm-url (or llm_url=) explicitly; there is no default, " + "because an unset endpoint yields rollouts that look fine and carry no token ids." + ) + + llm: dict[str, Any] = {} + if llm_url: + # Imported from the module rather than the package: the package re-exports a *function* + # named `validate_llm`, which shadows the same-named submodule. + from openenv.core.harness.capture.validate_llm import list_models, validate_llm + + # With no model given, ask the LLM what it serves. Convenient, and it also removes the + # commonest startup mistake: guessing a short alias for a server that publishes its full + # repo id. + if not model: + served = list_models(llm_url) + model = served[0] if len(served) == 1 else "" + + report = validate_llm(llm_url, model) if model else None + if report is None: + llm = { + "url": llm_url, + "model": "", + "ok": False, + "findings": ["no model given and the LLM does not serve exactly one"], + } + else: + llm = { + "url": llm_url, + "model": report.model, + "ok": report.ok, + "findings": report.findings, + "served_models": report.served_models, + } + + kwargs: dict[str, Any] = {"datasets": datasets, "llm": llm} + if sandboxes: + kwargs["sandboxes"] = sandboxes + caps = capabilities(**kwargs) + + if not quiet: + print(caps.render()) + + if require_llm and llm and not llm.get("ok"): + raise RuntimeError( + "LLM cannot support token capture:\n " + + "\n ".join(llm.get("findings") or ["unreachable"]) + + "\n\nStart vLLM with: --return-tokens-as-token-ids --logprobs-mode processed_logprobs" + ) + return caps diff --git a/src/openenv/harbor/tasks.py b/src/openenv/harbor/tasks.py new file mode 100644 index 000000000..e422c26c7 --- /dev/null +++ b/src/openenv/harbor/tasks.py @@ -0,0 +1,301 @@ +"""Dataset discovery: resolve Harbor task sets and serve them over OpenEnv's Task API. + +`HarborTaskProvider` satisfies `openenv.core.env_server.interfaces.TaskProvider`, so the HTTP routes +(`/{env}/splits`, `/{env}/tasks`, `/{env}/task`, ...) come for free once an environment exposes it. + +Two constraints from that Protocol shape the design, and both are easy to violate: + + * **it must be side-effect free.** Discovery must not boot a sandbox or start a job. + * **it must work on a freshly constructed instance.** The route handlers build a throwaway + environment per request purely to answer, so anything expensive has to be cached at module level + rather than on `self`, or every `/task` call re-downloads a dataset. + +Three source kinds, resolved by shape: + + AdithyaSK/data_agent_rl_environment_train HF dataset repo -> snapshot_download + /path/to/tasks local directory + terminal-bench@1.0 Harbor registry name@version + +Harbor has no HuggingFace path of its own — its datasets are git repos or Harbor Hub packages. But an +HF dataset laid out as `tasks//` is already a directory of Harbor task dirs, so downloading it +and pointing Harbor's local mode at the result needs no new concepts. +""" + +from __future__ import annotations + +import os +import threading +from pathlib import Path +from typing import Any + +from .models import HarborTaskRef + +# Resolution is expensive (a download on first use) and the Task API constructs a throwaway +# environment per HTTP request, so the cache has to outlive the instance. +_CACHE: dict[str, list[Path]] = {} +_LOCK = threading.Lock() + + +def _is_hf_repo(spec: str) -> bool: + """`org/name`, not a path and not `name@version`.""" + return ( + "/" in spec + and "@" not in spec + and not spec.startswith((".", "/", "~")) + and len(spec.split("/")) == 2 + ) + + +# Validating a task means reading and parsing several files inside it. That is 0.8 ms per task on +# local SSD, so 2s for a 2238-task suite, and a mounted bucket is an order of magnitude slower per +# read: listing a dataset then costs minutes and the Task API times out before answering. +# +# So discovery lists directories and does not validate. A task that is malformed surfaces as a failed +# rollout, with Harbor's own error, instead of being silently absent from the listing. That is also +# the more honest behaviour: filtering during discovery makes a broken task look like it was never +# in the dataset, and shifts the index of every task after it. +_VALIDATE_TASKS = os.environ.get("OPENENV_VALIDATE_TASKS", "").lower() in ( + "1", + "true", + "yes", +) + + +def _task_dirs_from_directory( + root: Path, *, validate: bool | None = None +) -> list[Path]: + """Task dirs under `root`, preferring a `tasks/` subdir when present. + + Args: + root (`Path`): + Dataset root, either containing `tasks/` or being the task directory itself. + validate (`bool`, *optional*): + Check each directory with Harbor's `Task.is_valid_dir`. Defaults to + `$OPENENV_VALIDATE_TASKS`, off, because it costs a file read per task and discovery is + on the latency path for every `/splits` and `/task` call. + """ + base = root / "tasks" if (root / "tasks").is_dir() else root + candidates = sorted( + p for p in base.iterdir() if p.is_dir() and not p.name.startswith(".") + ) + if not (_VALIDATE_TASKS if validate is None else validate): + return candidates + try: + from harbor.models.task.task import Task + except ImportError: + return candidates + return [p for p in candidates if Task.is_valid_dir(p, disable_verification=True)] + + +def resolve_task_dirs(spec: str, *, refresh: bool = False) -> list[Path]: + """Resolve a dataset spec to an ordered list of Harbor task directories. + + Order is stable (sorted by directory name) because a task's *index* is its identity everywhere + downstream — a trainer's dataset row, a `run_rollout` argument, a result. An unstable order would + silently change which task an index refers to between runs. + """ + with _LOCK: + if not refresh and spec in _CACHE: + return _CACHE[spec] + + path = Path(spec).expanduser() + if path.is_dir(): + dirs = _task_dirs_from_directory(path) + elif _is_hf_repo(spec): + dirs = _task_dirs_from_directory(_materialise_hf_dataset(spec)) + else: + dirs = _registry_task_dirs(spec) + + if not dirs: + raise ValueError( + f"no Harbor tasks found for {spec!r}. Expected an HF dataset repo laid out as " + "`tasks//`, a local directory of task dirs, or a Harbor registry `name@version`." + ) + + with _LOCK: + _CACHE[spec] = dirs + return dirs + + +# Task dirs must contain REAL FILES, not symlinks. +# +# The default HF cache is a symlink farm: every file under `snapshots//` points at +# `../../blobs/`. Harbor uploads a task's `tests/` directory to the sandbox by tarring it, and +# tar faithfully preserves symlinks — so the sandbox receives `test.sh -> ../../../blobs/09b32…`, +# pointing at a path that does not exist there. Bash reports a dangling symlink as +# "No such file or directory", which makes it look like the upload failed when the entry is right +# there in `ls`. +# +# That cost a long debugging session: `upload_dir` appeared to work, `chmod +x` as root succeeded, +# and only `ls -la` revealed the arrow. Backends differ in whether they hit it — E2B's upload path +# does not preserve symlinks, Modal's tar-based one does — so it presents as "Modal is broken". +# +# `local_dir=` makes huggingface_hub write real files instead of populating the symlink cache, which +# fixes it for every backend at once and needs no Harbor change. +_DATASET_ROOT = Path( + os.environ.get("OPENENV_DATASET_CACHE") + or (Path.home() / ".cache" / "openenv" / "harbor-datasets") +) + + +# A Harbor task suite is thousands of tiny files (a `task.toml`, a Dockerfile, a test script per +# task), so wall clock is dominated by per-file round trips rather than bytes. Raising concurrency is +# the lever that matters; `hf_transfer` optimises large-file throughput and does comparatively little +# here, but costs nothing when it is installed. +_DOWNLOAD_WORKERS = int(os.environ.get("OPENENV_DATASET_WORKERS", "32")) + + +def _materialise_hf_dataset(spec: str) -> Path: + """Download an HF dataset as real files and return its local root. + + Mounting beats downloading where it is available: a deployed Space can attach the dataset repo + as a read-only volume and pass the mount path as the dataset spec, which skips this entirely. + `openenv harbor push` does that automatically. This path is for local runs. + """ + from huggingface_hub import snapshot_download + + target = _DATASET_ROOT / spec.replace("/", "__") + target.mkdir(parents=True, exist_ok=True) + snapshot_download( + spec, + repo_type="dataset", + allow_patterns=["tasks/**"], + local_dir=str(target), + max_workers=_DOWNLOAD_WORKERS, + ) + return target + + +def has_symlinks(task_dir: Path) -> list[Path]: + """Any symlinks under `task_dir`. Non-empty means uploads to a tar-based backend will break.""" + return [p for p in task_dir.rglob("*") if p.is_symlink()] + + +def _registry_task_dirs(spec: str) -> list[Path]: + """A Harbor registry dataset, e.g. `terminal-bench@1.0`. Downloads on first use.""" + import asyncio + + from harbor.models.job.config import DatasetConfig + + name, _, version = spec.partition("@") + config = DatasetConfig(name=name, version=version or None) + task_configs = asyncio.run(config.get_task_configs(disable_verification=True)) + return [Path(str(t.get_local_path())) for t in task_configs] + + +def read_instruction(task_dir: Path, *, limit: int = 4000) -> str: + """The task's prompt, for previewing in discovery. Truncated: this is not the authoritative copy. + + The sandbox gets the real instruction from Harbor at run time. Serving a huge prompt over the + Task API for every listed task would make `list_tasks` enormous for no benefit. + """ + path = task_dir / "instruction.md" + if not path.is_file(): + return "" + text = path.read_text(errors="replace").strip() + return text if len(text) <= limit else text[:limit] + "\n…" + + +def prefetch(datasets: list[str]) -> dict[str, Any]: + """Resolve every dataset up front, downloading if needed. + + Called before the server accepts traffic. Two reasons it is worth doing eagerly rather than on + first use: a 2000-task HF repo takes real time to fetch, and a caller who mistypes a dataset + name should learn at startup rather than when the first rollout 404s. Failures are collected + rather than raised, so one bad dataset does not stop the server serving the good ones. + + Args: + datasets (`list[str]`): + Dataset specs — HF repo id, local path, or Harbor `name@version`. + + Returns: + `dict` mapping each spec to `{"num_tasks": int}` or `{"error": str}`. + """ + report: dict[str, Any] = {} + for spec in datasets: + try: + report[spec] = {"num_tasks": len(resolve_task_dirs(spec))} + except Exception as exc: # noqa: BLE001 - one broken dataset must not hide the others + report[spec] = {"error": f"{type(exc).__name__}: {str(exc)[:200]}"} + return report + + +class HarborTaskProvider: + """Serves one or more Harbor datasets as OpenEnv splits. + + A split IS a dataset spec: start the server with two datasets and you get two splits. That keeps + the mapping obvious in both directions — a split name is something you can paste back into + `--dataset` — rather than inventing a train/test split Harbor does not have. + """ + + def __init__(self, datasets: list[str] | None = None) -> None: + self._datasets = list(datasets or []) + + # --- TaskProvider protocol ------------------------------------------ + def list_splits(self) -> list[dict[str, Any]]: + splits = [] + for spec in self._datasets: + try: + n = len(resolve_task_dirs(spec)) + splits.append({"name": spec, "num_tasks": n}) + except Exception as exc: # noqa: BLE001 - a broken dataset must not hide the good ones + splits.append({"name": spec, "num_tasks": 0, "error": str(exc)[:200]}) + return splits + + def num_tasks(self, split: str) -> int: + return len(resolve_task_dirs(self._check(split))) + + def list_tasks(self, split: str) -> list[dict[str, Any]]: + spec = self._check(split) + return [ + self._ref(spec, i, d).model_dump() + for i, d in enumerate(resolve_task_dirs(spec)) + ] + + def get_task(self, split: str, index: int) -> dict[str, Any]: + spec = self._check(split) + dirs = resolve_task_dirs(spec) + if not 0 <= index < len(dirs): + raise IndexError( + f"task index {index} out of range for {spec!r} ({len(dirs)} tasks)" + ) + return self._ref(spec, index, dirs[index]).model_dump() + + def get_task_range( + self, split: str, start: int | None = None, stop: int | None = None + ) -> list[dict[str, Any]]: + spec = self._check(split) + dirs = resolve_task_dirs(spec) + return [ + self._ref(spec, i, d).model_dump() + for i, d in list(enumerate(dirs))[start:stop] + ] + + # --- internals ------------------------------------------------------- + def task_dir(self, split: str, index: int) -> Path: + """The on-disk task dir for an index. Used by the rollout path, not by discovery.""" + dirs = resolve_task_dirs(self._check(split)) + if not 0 <= index < len(dirs): + raise IndexError(f"task index {index} out of range ({len(dirs)} tasks)") + return dirs[index] + + def _check(self, split: str) -> str: + if not self._datasets: + raise ValueError("this server was started with no datasets; pass --dataset") + if not split: + return self._datasets[0] + if split not in self._datasets: + raise ValueError( + f"unknown split {split!r}; served splits are {self._datasets}" + ) + return split + + @staticmethod + def _ref(spec: str, index: int, task_dir: Path) -> HarborTaskRef: + return HarborTaskRef( + index=index, + task_id=str(task_dir), + task_name=task_dir.name, + dataset=spec, + instruction=read_instruction(task_dir), + ) diff --git a/src/openenv/harbor/ui.py b/src/openenv/harbor/ui.py new file mode 100644 index 000000000..0b5a1bf7b --- /dev/null +++ b/src/openenv/harbor/ui.py @@ -0,0 +1,1090 @@ +"""Human-facing UI for a Harbor env server. + +Two columns: the LLM on the left, the task on the right. Validate, pick, run. + +Status text is deliberately terse. The long explanations belong in docs — what a person needs on +screen is whether it will work, what got rewritten, and which sandboxes are usable. + +Validation is a gate, not a hint: an LLM endpoint without token-id capture answers every request +normally and returns nothing trainable, so a rollout looks perfect and is worthless. + +Rich output (the rollout graph, per-turn tokens) is rendered as HTML rather than Gradio widgets, +because a conversation tree with branches and discarded retries is a shape, and a dataframe cannot +show a shape. +""" + +from __future__ import annotations + +import html +import json +import re +from typing import Any + +import gradio as gr + +_UNVALIDATED = "_Enter your LLM URL and press Validate._" + +_CSS = """ +.hb-wrap { max-width: 1400px; margin: 0 auto; } +.hb-card { border: 1px solid var(--border-color-primary); border-radius: 10px; padding: 14px 16px; } +.hb-dim { opacity: .6; } +.hb-kv { display: flex; gap: 22px; flex-wrap: wrap; margin: 4px 0 2px; } +.hb-kv b { font-variant-numeric: tabular-nums; } + +/* The two panels read as one undifferentiated wall of controls without a boundary; the border is + what makes "pick a model" and "pick a task" look like two separate decisions. */ +.hb-cell { border: 1px solid var(--border-color-primary); border-radius: 10px; + padding: 14px 16px; } +.hb-cell + .hb-cell { margin-left: 12px; } + +/* Live conversation. Roles are colour-coded down the left edge so the shape of the loop + (assistant calls a tool, tool answers, assistant calls again) is readable at a glance. */ +/* No max-height here. A fixed-height scroll box nests a second scroller inside the page: + the wheel gets captured while the pointer is over the conversation, and the page stops + growing so there is nothing left to scroll to. Let it run at natural height and let the + page do the scrolling. Length is bounded by the message cap, not by CSS. */ +.hb-tx { margin-top: 10px; } +.hb-msg { border-left: 3px solid var(--border-color-primary); padding: 6px 0 6px 10px; + margin: 8px 0; font-size: 13px; line-height: 1.45; } +.hb-msg pre { white-space: pre-wrap; word-break: break-word; margin: 4px 0 0; + font-size: 12px; opacity: .85; } +.hb-role { display: inline-block; font-size: 11px; text-transform: uppercase; + letter-spacing: .04em; opacity: .65; margin-bottom: 2px; } +.hb-assistant { border-left-color: #22c55e; } +.hb-tool { border-left-color: #38bdf8; } +.hb-user { border-left-color: #a78bfa; } +.hb-system { border-left-color: #94a3b8; opacity: .75; } +.hb-tc { margin-top: 4px; padding: 4px 8px; border-radius: 6px; + background: var(--background-fill-secondary); } +.hb-tc { display: block; } +.hb-tc b { font-family: ui-monospace, SFMono-Regular, Menlo, monospace; font-size: 12px; } +.hb-arrow { opacity: .5; margin-right: 6px; } +.hb-tr { margin-top: 4px; padding: 4px 8px; border-radius: 6px; border-left: 2px solid #38bdf8; + background: var(--background-fill-secondary); } +/* Tool output can be thousands of lines, so this one stays capped, but `contain` stops it + swallowing the page scroll once it reaches its own end. */ +.hb-tr pre{ margin: 0; font-size: 11.5px; opacity: .8; max-height: 220px; + overflow-y: auto; overscroll-behavior: contain; } + +/* A run in flight should look like one. */ +.hb-live { display: flex; align-items: center; gap: 10px; margin-bottom: 6px; } +.hb-pulse { width: 8px; height: 8px; border-radius: 50%; background: #22c55e; + animation: hb-blink 1.2s ease-in-out infinite; } +@keyframes hb-blink { 0%, 100% { opacity: 1; } 50% { opacity: .25; } } +.hb-drop-msg { opacity: .5; border-left-color: #ef4444; } + +/* Verdict. The outcome should be legible from across the room; the numbers behind it should not + compete with it for attention. */ +.hb-verdict { border-left-width: 4px; } +.hb-head { font-size: 17px; font-weight: 650; margin-bottom: 8px; } +.hb-good { border-left-color: #22c55e; } +.hb-warn { border-left-color: #f59e0b; } +.hb-bad { border-left-color: #ef4444; } +.hb-err { white-space: pre-wrap; word-break: break-word; font-size: 12px; margin: 10px 0 0; + padding: 8px 10px; border-radius: 6px; background: var(--background-fill-secondary); } + +/* Findings carry severity: a FATAL means unusable, a WARN means read before training on it. */ +.hb-find { font-size: 12.5px; margin: 5px 0; line-height: 1.45; } +.hb-tag { display: inline-block; min-width: 46px; margin-right: 8px; padding: 1px 6px; + border-radius: 4px; font-size: 10px; font-weight: 700; letter-spacing: .04em; + text-align: center; vertical-align: 1px; } +.hb-fatal .hb-tag { background: #ef4444; color: #fff; } +.hb-warn2 .hb-tag { background: #f59e0b; color: #1f2937; } +.hb-info .hb-tag { background: var(--background-fill-secondary); opacity: .7; } +.hb-info { opacity: .7; } + +/* Turn table: dense, aligned, and the numbers read as numbers. */ +.hb-tbl { width: 100%; border-collapse: collapse; margin-top: 8px; font-size: 13px; } +.hb-tbl th{ text-align: left; font-weight: 600; font-size: 11px; text-transform: uppercase; + letter-spacing: .04em; opacity: .55; padding: 4px 10px 6px 0; + border-bottom: 1px solid var(--border-color-primary); } +.hb-tbl td{ padding: 7px 10px 7px 0; border-bottom: 1px solid var(--border-color-primary); + vertical-align: top; } +.hb-tbl code { font-size: 12px; padding: 1px 6px; border-radius: 4px; + background: var(--background-fill-secondary); } +.hb-num { font-variant-numeric: tabular-nums; text-align: right; white-space: nowrap; + padding-right: 14px !important; } +.hb-prev { margin-top: 3px; font-size: 12px; } +.hb-drop-row { opacity: .45; } +.hb-drop-tag { background: #ef4444; color: #fff; } +.hb-conf { display: inline-block; width: 76px; height: 7px; border-radius: 4px; + background: var(--background-fill-secondary); overflow: hidden; vertical-align: middle; } +.hb-conf span { display: block; height: 100%; } + +/* Each conversation folds away; the main one starts open. */ +.hb-convo { margin-top: 10px; border-top: 1px solid var(--border-color-primary); padding-top: 8px; } +.hb-convo summary { cursor: pointer; padding: 4px 0; } + +/* Setup, before the agent has said anything. */ +.hb-steps { margin: 8px 0 0; } +.hb-step { display: flex; align-items: center; gap: 9px; padding: 3px 0; font-size: 13px; } +.hb-dot { width: 7px; height: 7px; border-radius: 50%; background: var(--border-color-primary); } +.hb-step.done .hb-dot { background: #22c55e; } +.hb-step.now .hb-dot { background: #f59e0b; animation: hb-blink 1.2s ease-in-out infinite; } +.hb-step.todo { opacity: .45; } + +/* The outcome, at a glance. */ +.hb-hero { display: flex; align-items: center; justify-content: space-between; gap: 20px; + padding-bottom: 12px; margin-bottom: 4px; + border-bottom: 1px solid var(--border-color-primary); } +.hb-badge { display: inline-flex; align-items: center; gap: 9px; font-size: 19px; + font-weight: 700; letter-spacing: -.01em; } +.hb-mark { display: inline-flex; align-items: center; justify-content: center; + width: 30px; height: 30px; border-radius: 50%; font-size: 15px; color: #fff; } +.hb-b-good .hb-mark { background: #22c55e; } +.hb-b-warn .hb-mark { background: #f59e0b; } +.hb-b-bad .hb-mark { background: #ef4444; } +.hb-score { text-align: right; line-height: 1.05; } +.hb-score-v { font-size: 42px; font-weight: 700; font-variant-numeric: tabular-nums; + letter-spacing: -.02em; } +.hb-score-c { font-size: 11px; text-transform: uppercase; letter-spacing: .06em; opacity: .55; } +.hb-kv-big span { font-size: 11px; text-transform: uppercase; letter-spacing: .04em; + opacity: .55; } +.hb-kv-big b { display: block; font-size: 19px; margin-top: 3px; text-transform: none; + letter-spacing: normal; opacity: 1; } +.hb-kv-big .hb-key b { color: var(--body-text-color); } +.hb-kv-big .hb-key { opacity: .85; } + +footer { display: none !important; } +""" + + +def _clip(text: Any, limit: int = 400) -> str: + """Escape and shorten a value for display, keeping the head where the meaning usually is.""" + body = text if isinstance(text, str) else json.dumps(text, default=str) + body = body.strip() + return html.escape(body[:limit]) + ("…" if len(body) > limit else "") + + +def _tool_calls(message: dict[str, Any]) -> list[dict[str, Any]]: + """Tool calls on a message, normalised across all four dialects. + + Chat-completions puts them in `tool_calls`; Anthropic puts them in the content block list as + `tool_use`. Reading only the former shows claude-code as a stream of text with no visible + actions, which is exactly the case the live view exists to make visible. + """ + out: list[dict[str, Any]] = [] + for call in message.get("tool_calls") or []: + function = call.get("function") or {} + name = function.get("name") or call.get("name") + if name: + out.append( + { + "name": str(name), + "arguments": function.get("arguments", call.get("arguments", "")), + } + ) + content = message.get("content") + if isinstance(content, list): + for block in content: + if ( + isinstance(block, dict) + and block.get("type") == "tool_use" + and block.get("name") + ): + out.append( + {"name": str(block["name"]), "arguments": block.get("input", "")} + ) + return out + + +def _message_text(message: dict[str, Any]) -> str: + """Readable text of a message, ignoring tool-call and tool-result blocks.""" + content = message.get("content") + if isinstance(content, str): + return content + if isinstance(content, list): + parts = [] + for block in content: + if not isinstance(block, dict): + continue + if block.get("type") in ("tool_use", "tool_result"): + continue + if block.get("text"): + parts.append(str(block["text"])) + return " ".join(parts) + return "" + + +def _tool_results(message: dict[str, Any]) -> list[str]: + """What came back from a tool, in either the chat-completions or the Anthropic shape.""" + if message.get("role") == "tool": + return [ + _message_text(message) or json.dumps(message.get("content"), default=str) + ] + content = message.get("content") + if not isinstance(content, list): + return [] + out = [] + for block in content: + if isinstance(block, dict) and block.get("type") == "tool_result": + body = block.get("content") + if isinstance(body, list): + body = " ".join(b.get("text", "") for b in body if isinstance(b, dict)) + out.append(str(body if body is not None else "")) + return out + + +def _render_calls(calls: list[dict[str, Any]]) -> str: + return "".join( + f'
▸' + f"{html.escape(str(c.get('name', 'tool')))}" + f"
{_clip(c.get('arguments', ''), 600)}
" + for c in calls + ) + + +def _render_message(message: dict[str, Any], *, label: str = "") -> str: + """One row of the conversation: who spoke, what they said, what they invoked or returned.""" + role = str(message.get("role", "?")) + calls = _tool_calls(message) + results = _tool_results(message) + text = _message_text(message) + + # A user message carrying only tool results is the tool speaking, not the user; labelling it + # "user" makes the agent look like it is being prompted between every action. + shown_role = "tool" if results and role != "assistant" else role + # For a `role: tool` message the content IS the result, so rendering both duplicates it. + if shown_role == "tool": + text = "" + body = _clip(text, 700 if shown_role in ("user", "system") else 450) if text else "" + blocks = "".join( + f'
{_clip(r, 500)}
' for r in results + ) + if not body and not blocks and not calls: + return "" + return ( + f'
' + f'{html.escape(label or shown_role)}' + + (f"
{body}
" if body else "") + + _render_calls(calls) + + blocks + + "
" + ) + + +def _transcript_html(session: Any) -> str: + """The conversation as it stands right now: what the agent said, called, and got back. + + Counters answer "is it alive"; this answers "is it doing the right thing", which is the question + worth asking while a rollout is still running. The newest turn's `request_messages` already holds + the whole conversation the harness assembled, tool results included, so rendering that plus the + latest response needs no reconstruction from deltas. + """ + nodes = sorted(session.graph.nodes(), key=lambda n: n.index) + if not nodes: + return "" + latest = nodes[-1] + + rows = [ + row + for row in (_render_message(m) for m in (latest.request_messages or [])) + if row + ] + + response = latest.response_message or {} + tail = _render_message( + {**response, "role": "assistant"}, + label=f"assistant · turn {latest.index} · generating", + ) + if tail: + rows.append(tail) + + # Only the tail is ever new, so cap from the front and say what was dropped. + shown = rows[-18:] + elided = ( + f'
… {len(rows) - len(shown)} earlier message(s)
' + if len(rows) > len(shown) + else "" + ) + # Count across the conversation, not just the response messages: Anthropic carries tool use in + # the assistant content blocks the harness replays back, so a response-only tally reads 0. + calls_so_far = sum( + len(_tool_calls(m)) for m in (latest.request_messages or []) + ) + len(_tool_calls(latest.response_message or {})) + return ( + f'
' + f'Live conversation' + f'turn {latest.index} · {calls_so_far} tool call(s) so far · ' + f"{latest.n_tools} tool(s) offered
{elided}{''.join(shown)}
" + ) + + +# What happens before the agent's first model call, in order. Harbor exposes no progress hook, so +# the stage is inferred from what capture has seen: no session means the trial has not reached the +# agent yet, a session with no turns means the agent is installed and starting up. +_SETUP_STEPS = ( + "creating the sandbox", + "uploading the task", + "installing the agent", + "waiting for the first model call", +) + + +def _steps_html(stage: int) -> str: + """The setup sequence, with the current stage marked.""" + rows = [] + for i, label in enumerate(_SETUP_STEPS): + cls = "done" if i < stage else ("now" if i == stage else "todo") + rows.append( + f'
' + f"{html.escape(label)}
" + ) + return f'
{"".join(rows)}
' + + +def _live_html( + harness: str, + sandbox: str, + phase: str, + elapsed: float, + stats: dict[str, Any] | None, + stage: int = -1, +) -> str: + """The running header: what is running, how far in, and what it has produced so far.""" + bits = [ + f'
' + f'
' + f"Running {html.escape(harness)} on " + f"{html.escape(sandbox)}" + f'{html.escape(phase)} · {elapsed:.0f}s
' + ] + # Before the first call there are no numbers worth showing, so show progress instead. A row of + # zeros for a minute reads as "stuck" when the sandbox is simply still booting. + if stage >= 0: + bits.append(_steps_html(stage)) + if stats: + bits.append( + '
' + + "".join(f"{k}
{v}" for k, v in stats.items()) + + "
" + ) + bits.append("
") + return "".join(bits) + + +# `warn` is already a verdict tone; the finding variant needs its own class name. +_FINDING_CLASS = {"FATAL": "fatal", "WARN": "warn2", "INFO": "info"} + + +def _findings_html(findings: list[str]) -> str: + """Findings, grouped by how much they should worry you. + + They were previously all rendered the same dim grey and truncated to 220 characters, which put + "the intercept saw no model calls" and "3 roots across 7 turns" at equal weight. A FATAL means + the rollout is unusable; a WARN means read it before training on it. + """ + if not findings: + return "" + buckets: dict[str, list[str]] = {"FATAL": [], "WARN": [], "INFO": []} + for raw in findings: + level = ( + "FATAL" + if raw.startswith("[FATAL") + else "WARN" + if raw.startswith("[WARN") + else "INFO" + ) + buckets[level].append( + raw.split("]", 1)[-1].strip() if raw.startswith("[") else raw + ) + + out = [] + for level, items in buckets.items(): + for item in items: + out.append( + f'
' + f'{level}{html.escape(item[:400])}
' + ) + return "".join(out) + + +def _result_html(r: dict[str, Any]) -> str: + """The verdict, the numbers behind it, and anything that qualifies it. + + The outcome is the one thing every reader wants first, so the reward is set at display size and + the supporting counts are deliberately quieter. Getting that hierarchy wrong is how a failed + rollout reads as a successful one at a glance. + """ + reward = r.get("reward") + if not r.get("ok"): + tone, mark, label = "bad", "✕", "Failed" + value, caption = "—", str(r.get("exception_type") or "error") + elif reward is None: + # Not a zero. The verifier never ran, so this says nothing about the model. + tone, mark, label = "warn", "!", "Not graded" + value, caption = "—", "the verifier never ran" + elif reward > 0: + tone, mark, label = "good", "✓", "Solved" + value, caption = f"{reward:.2f}", "reward" + else: + tone, mark, label = "warn", "○", "Not solved" + value, caption = f"{reward:.2f}", "reward" + + turns = r.get("turns") or [] + generated = sum(len(t.get("completion_token_ids") or []) for t in turns) + dropped = sum( + len(t.get("completion_token_ids") or []) for t in turns if t.get("discarded") + ) + tools = sum(len(t.get("tool_calls") or []) for t in turns) + atif = r.get("atif", "none") + + # `key` marks the figures that decide whether this rollout is usable, as opposed to describing it. + # The initial prompt: task instruction plus the harness's system prompt and tool manifest. + # Constant across turns, so it is a property of the rollout rather than a per-row column. + context = len((turns[0].get("prompt_token_ids") or [])) if turns else 0 + + kv = [ + ("trainable tokens", f"{r.get('n_trainable_tokens', 0):,}", True), + ("context", f"{context:,}", False), + ("trace check", atif, atif != "match"), + ("model calls", r.get("n_turns", 0), False), + ("tool calls", tools, False), + ("conversations", r.get("n_roots", 0), False), + ( + "generated", + f"{generated:,}" + (f" · {dropped:,} discarded" if dropped else ""), + False, + ), + ("wall", f"{r.get('wall_s', 0):.0f}s", False), + ] + + out = [ + f'
', + '
', + f'
{mark}' + f"{html.escape(label)}
", + f'
{html.escape(value)}
' + f'
{html.escape(caption)}
', + "
", + '
' + + "".join( + f'{k}
{v}' + for k, v, key in kv + ) + + "
", + ] + + rewards = r.get("rewards") or {} + if len(rewards) > 1: + chosen = r.get("reward_key", "") + parts = [ + f"{html.escape(k)} {v:.3f}" + (" ←" if k == chosen else "") + for k, v in sorted(rewards.items()) + ] + out.append(f'
{"   ".join(parts)}
') + + for step in r.get("step_results") or []: + vals = ", ".join(f"{k}={v:.2f}" for k, v in (step.get("rewards") or {}).items()) + out.append( + f'
step {html.escape(step.get("name", ""))} {vals}
' + ) + + if r.get("error"): + out.append(f'
{html.escape(str(r["error"])[:1200])}
') + if r.get("agent_log_tail"): + out.append( + '
agent log' + f"
{html.escape(str(r['agent_log_tail'])[:4000])}
" + ) + + out.append(_findings_html(r.get("findings") or [])) + out.append( + '
Capture quality and reward are ' + "independent: a perfectly captured rollout can still score 0 because the model was " + "wrong, and reward — means the verifier never ran at all.
" + ) + return "".join(out) + + +def _conversation_html(r: dict[str, Any]) -> str: + """The whole conversation as it was actually sent: system prompt, tools, results, replies. + + Rebuilt from the result rather than the live session, so it survives the run. Several are + possible: each root is a separate conversation, and an auxiliary one (a next-speaker check, a + summariser) is labelled as such so it is not mistaken for the agent working on the task. + """ + conversations = r.get("conversations") or [] + if not conversations: + return "" + + agents = [c for c in conversations if c.get("role", "agent") == "agent"] + blocks = [] + seen_agents = 0 + for i, convo in enumerate(conversations): + role = convo.get("role", "agent") + if role == "agent": + seen_agents += 1 + # Numbered when there is more than one, so two blocks are never both "main". + badge = ( + "main conversation" + if len(agents) == 1 + else f"conversation {seen_agents} of {len(agents)}" + ) + else: + badge = { + "auxiliary": "auxiliary call", + "discarded": "discarded branch", + }.get(role, role) + rows = [ + row + for row in (_render_message(m) for m in convo.get("messages") or []) + if row + ] + if not rows: + continue + blocks.append( + f'
' + f"{html.escape(badge)} " + f'{convo.get("n_turns", 0)} model call(s), ' + f"{len(rows)} message(s){''.join(rows)}
" + ) + if not blocks: + return "" + return ( + f'
Conversation ' + f'everything the model saw and produced' + f"{''.join(blocks)}
" + ) + + +def _confidence(mean_logp: float) -> str: + """A bar for mean logprob. Closer to 0 is more confident; -1.0 is the practical floor here.""" + pct = max(0.0, min(1.0, 1.0 + mean_logp)) # -0 -> 1.0, -1 -> 0.0 + hue = 8 + int(112 * pct) # red through amber to green + return ( + f'' + f'' + ) + + +def _turns_html(r: dict[str, Any]) -> str: + """Turn by turn: what it did, how much it wrote, how sure it was. + + Replaces a table whose most prominent column was "tools", meaning the number of tools *offered* + to the model. That number is a property of the harness, identical on every row, and told nobody + anything. What varies per turn, and is worth reading, is the action taken, the tokens spent on + it, and the model's confidence while producing them. + """ + turns = r.get("turns") or [] + if not turns: + return '
No model calls were captured.
' + + used: dict[str, int] = {} + for t in turns: + for call in t.get("tool_calls") or []: + name = str(call.get("name", "?")) + used[name] = used.get(name, 0) + 1 + + rows = [] + for t in turns: + lp = t.get("per_token_logps") or [] + mean = sum(lp) / len(lp) if lp else 0.0 + gen = len(t.get("completion_token_ids") or []) + calls = t.get("tool_calls") or [] + if calls: + action = " ".join( + f"{html.escape(str(c.get('name', 'tool')))}" for c in calls + ) + elif t.get("finish_reason") == "stop": + action = 'final answer' + else: + action = 'text only' + note = ( + ' discarded' + if t.get("discarded") + else "" + ) + preview = _clip(t.get("text") or "", 160) + rows.append( + f'' + f'{t.get("turn")}' + f"{action}{note}" + + (f'
{preview}
' if preview else "") + + f'{gen:,}' + f"{_confidence(mean) if lp else ''}" + f'{html.escape(str(t.get("finish_reason") or ""))}' + ) + + histogram = "" + if used: + top = sorted(used.items(), key=lambda kv: -kv[1]) + histogram = ( + '
tools used: ' + + "   ".join(f"{html.escape(k)}×{v}" for k, v in top) + + "
" + ) + + return ( + '
Turn by turn' + '' + "" + + "".join(rows) + + "
#actiontokensconfidencestopped because
" + + histogram + + '
Confidence is the mean logprob of the ' + "sampled tokens: full bar means the model was near-certain, short means it was " + "guessing. Discarded turns were generated and billed but lead nowhere, so they are " + "excluded from training paths.
" + ) + + +def _write_contract(r: dict[str, Any]) -> str | None: + """Write `contract.json`: exactly what a trainer consumes, nothing else. + + Per turn, `(prompt_token_ids, completion_token_ids, per_token_logps)` plus the reward. The + logprobs are the load-bearing part and the reason this is a separate file: they are the + behaviour policy's, recorded at sampling time, and cannot be recovered afterwards by re-running + the prompt. Discarded turns are kept but flagged, because they were generated and billed and a + trainer must be able to see them in order to exclude them deliberately. + """ + import tempfile + from pathlib import Path as _Path + + turns = r.get("turns") or [] + if not turns: + return None + contract = { + "task_id": r.get("task_id", ""), + "task_name": r.get("task_name", ""), + "dataset": r.get("dataset", ""), + "harness": r.get("harness", ""), + "sandbox": r.get("sandbox", ""), + "trial_name": r.get("trial_name", ""), + "reward": r.get("reward"), + "rewards": r.get("rewards") or {}, + "reward_key": r.get("reward_key", ""), + "n_trainable_tokens": r.get("n_trainable_tokens", 0), + "turns": [ + { + "turn": t.get("turn"), + "prompt_token_ids": t.get("prompt_token_ids") or [], + "completion_token_ids": t.get("completion_token_ids") or [], + "per_token_logps": t.get("per_token_logps") or [], + "finish_reason": t.get("finish_reason"), + "discarded": bool(t.get("discarded")), + } + for t in turns + ], + } + name = re.sub(r"[^A-Za-z0-9_.-]", "_", str(r.get("task_name") or "rollout")) + target = ( + _Path(tempfile.mkdtemp(prefix="harbor-contract-")) / f"{name}.contract.json" + ) + target.write_text(json.dumps(contract, indent=2)) + return str(target) + + +def _summary_json(r: dict[str, Any]) -> str: + """The result with the token arrays summarised, which is the part anyone actually reads. + + The full document stays available below; printing 8000 integers first buries the fields that + carry meaning. + """ + compact = {k: v for k, v in r.items() if k not in ("turns", "conversations")} + compact["turns"] = [ + { + "turn": t.get("turn"), + "action": [c.get("name") for c in (t.get("tool_calls") or [])] or "text", + "prompt_token_ids": f"<{len(t.get('prompt_token_ids') or [])} ids>", + "completion_token_ids": f"<{len(t.get('completion_token_ids') or [])} ids>", + "per_token_logps": f"<{len(t.get('per_token_logps') or [])} floats>", + "finish_reason": t.get("finish_reason"), + "discarded": t.get("discarded"), + "text": (t.get("text") or "")[:200], + } + for t in (r.get("turns") or [])[:200] + ] + compact["conversations"] = [ + { + "role": c.get("role"), + "n_turns": c.get("n_turns"), + "messages": f"<{len(c.get('messages') or [])} messages>", + } + for c in (r.get("conversations") or []) + ] + return json.dumps(compact, indent=2)[:200_000] + + +def _read(path: Any, limit: int = 20000) -> str: + try: + text = path.read_text(errors="replace") + except Exception: # noqa: BLE001 + return "" + return text if len(text) <= limit else text[:limit] + "\n…truncated…" + + +def harbor_gradio_builder( + *, + datasets: list[str] | None = None, + title: str | None = None, +) -> gr.Blocks: + """Build the Harbor UI. + + Args: + datasets (`list[str]`, *optional*): + Dataset specs served by this server; each becomes a selectable split. + + Returns: + `gr.Blocks`: The interface. + """ + from .tasks import HarborTaskProvider, resolve_task_dirs + + datasets = list(datasets or []) + + def on_validate(url: str, model: str): + from openenv.core.harness.capture.validate_llm import list_models, validate_llm + + from .capabilities import capabilities + from .seams import agent_facing_model + + url = (url or "").strip().rstrip("/") + if not url: + return ( + _UNVALIDATED, + gr.update(), + gr.update(), + {}, + gr.update(interactive=False), + ) + + if not model: + served = list_models(url) + if len(served) != 1: + return ( + f"**Pick a model** — this endpoint serves " + f"`{', '.join(served) or 'nothing reachable'}`.", + gr.update(), + gr.update(), + {}, + gr.update(interactive=False), + ) + model = served[0] + + report = validate_llm(url, model) + if not report.ok: + why = "; ".join(report.findings) or "unreachable" + return ( + f"**Not usable** — {why}\n\n" + "Start vLLM with `--return-tokens-as-token-ids " + "--logprobs-mode processed_logprobs`.", + gr.update(), + gr.update(), + {}, + gr.update(interactive=False), + ) + + caps = capabilities( + datasets=datasets, llm={"url": url, "model": model, "ok": True} + ) + sandboxes = caps.available_sandboxes + by_dialect: dict[str, list[str]] = {} + for h in caps.harnesses: + if h.status == "validated": + by_dialect.setdefault(h.dialect, []).append(h.name) + choices = [ + (f"{n} ({d})", n) + for d, names in sorted(by_dialect.items()) + for n in sorted(names) + ] + values = [v for _, v in choices] + + leaf = agent_facing_model(model) + lines = [f"**Ready** · `{model}` · logprobs + token ids ✓"] + if leaf != model: + lines.append(f"Sent to agents as `{leaf}`, rewritten back on the way out.") + lines.append( + f"Sandboxes: {', '.join(f'`{s}`' for s in sandboxes) or '**none usable**'}" + ) + blocked = [s.name for s in caps.sandboxes if not s.available] + if blocked: + lines.append( + f"unavailable: {', '.join(blocked)}" + ) + + return ( + " \n".join(lines), + gr.update( + choices=choices, + value="opencode" + if "opencode" in values + else (values[0] if values else None), + ), + gr.update(choices=sandboxes, value=sandboxes[0] if sandboxes else None), + {"url": url, "model": model, "ok": True}, + gr.update(interactive=bool(sandboxes)), + ) + + def on_dataset(spec: str): + if not spec: + return gr.update(), "" + try: + n = len(resolve_task_dirs(spec)) + except Exception as exc: # noqa: BLE001 + return gr.update(value=0), f"Cannot load `{spec}` — {exc}" + return gr.update(value=0), f"**{n}** tasks · 0–{n - 1}" + + def on_task(spec: str, index: int): + if not spec: + return "", "", "", "", "" + try: + task_dir = HarborTaskProvider([spec]).task_dir(spec, int(index)) + except Exception as exc: # noqa: BLE001 + return f"_{exc}_", "", "", "", "" + env_dir, tests_dir = task_dir / "environment", task_dir / "tests" + return ( + f"`{task_dir.name}`", + _read(task_dir / "instruction.md"), + _read(env_dir / "Dockerfile"), + _read(task_dir / "task.toml"), + _read(tests_dir / "test.sh"), + ) + + def on_run(engine: dict, spec: str, index: int, harness: str, sandbox: str): + """Stream progress while the rollout runs, then the result and its graph.""" + import asyncio + import queue + import threading + import time + from pathlib import Path + + from .rollout import run_rollout as _run + from .serving import HarborService + + if not engine.get("ok"): + yield _UNVALIDATED, "", "", "{}", None, gr.update(interactive=True) + return + service = HarborService.current() + if service is None: + yield ( + "Server not initialised — no capture proxy running.", + "", + "", + "{}", + None, + gr.update(interactive=True), + ) + return + try: + task_dir = HarborTaskProvider([spec]).task_dir(spec, int(index)) + except Exception as exc: # noqa: BLE001 + yield ( + f"Bad task — {html.escape(str(exc))}", + "", + "", + "{}", + None, + gr.update(interactive=True), + ) + return + + done: queue.Queue = queue.Queue(maxsize=1) + + def worker() -> None: + try: + res = asyncio.run( + _run( + task_dir=task_dir, + harness=harness, + sandbox=sandbox, + registry=service.capture.registry, + intercept_url=service.public_url, + model=service.model, + trials_dir=Path("/tmp/openenv-harbor-trials"), + dataset=spec, + ) + ) + done.put(("ok", res.model_dump())) + except Exception as exc: # noqa: BLE001 - show it, never take the server down + done.put(("err", f"{type(exc).__name__}: {exc}")) + + before = set(service.capture.registry.list_ids()) + thread = threading.Thread(target=worker, daemon=True) + started = time.monotonic() + thread.start() + session_id = None + + while thread.is_alive(): + if session_id is None: + session_id = next( + iter(set(service.capture.registry.list_ids()) - before), None + ) + stats, phase, stage = None, "starting up", 0 + if session_id: + session = service.capture.registry.get(session_id) + if session is not None: + st = session.graph.stats() + # n_trainable_tokens only exists after export; mid-run we can count only what + # has been sampled, before masking and discards. + sampled = sum( + len(n.sampled_ids or []) for n in session.graph.nodes() + ) + turns = st.get("n_turns", 0) + stats = { + "calls": turns, + "roots": st.get("n_roots", 0), + "sampled tokens": sampled, + "discarded": st.get("n_discarded", 0), + } + if turns: + stage = -1 # past setup; the numbers mean something now + phase = "agent working" + # Only meaningful once a call has landed. Before that `idle_seconds` counts + # from session creation, which renders as a stall during a normal boot. + stats["since last call"] = f"{session.idle_seconds:.0f}s" + else: + # The session exists, so the trial reached the agent: the sandbox is up and + # the task is uploaded. What is left is the agent's own startup. + stage = 2 + phase = "sandbox ready, starting the agent" + # The transcript rides in the graph slot: it is empty until the run finishes anyway, + # and the two answer the same question at different times. + transcript = "" + if session_id: + live = service.capture.registry.get(session_id) + if live is not None: + transcript = _transcript_html(live) + yield ( + _live_html( + harness, sandbox, phase, time.monotonic() - started, stats, stage + ), + transcript, + "", + "{}", + None, + gr.update(interactive=False), + ) + time.sleep(2.0) + + kind, payload = done.get() + if kind == "err": + yield ( + f'
Run failed
' + f'
{html.escape(payload)}
', + "", + "", + "{}", + None, + gr.update(interactive=True), + ) + return + yield ( + _result_html(payload), + _conversation_html(payload), + _turns_html(payload), + _summary_json(payload), + _write_contract(payload), + gr.update(interactive=True), + ) + + with gr.Blocks(title=title or "Harbor") as app: + gr.HTML(f"") + state = gr.State({}) + + with gr.Column(elem_classes="hb-wrap"): + gr.Markdown( + "## Harbor\nRun a coding agent on a Harbor task and capture every token " + "and logprob it produced." + ) + + with gr.Row(equal_height=False): + # left — the model + with gr.Column(scale=1, elem_classes="hb-cell"): + gr.Markdown("### LLM") + # Deliberately empty. Prefilling meant the box already held whatever URL the + # server was started with, so Validate confirmed a value nobody chose and a + # stale endpoint could be used without anyone noticing it was stale. + url_in = gr.Textbox( + label="LLM URL", + placeholder="https://… vLLM, OpenAI-compatible", + ) + model_in = gr.Textbox( + label="Model (optional)", placeholder="read from the endpoint" + ) + validate_btn = gr.Button("Validate", variant="secondary") + engine_md = gr.Markdown(_UNVALIDATED) + gr.Markdown("### Agent") + harness_in = gr.Dropdown(label="Agent", choices=[]) + sandbox_in = gr.Dropdown(label="Sandbox", choices=[]) + + # right — the task. Everything but the picker is folded away: the instruction alone + # runs to a screenful, and it pushed the Run button below the fold. + with gr.Column(scale=2, elem_classes="hb-cell"): + gr.Markdown("### Task") + with gr.Row(): + ds_in = gr.Dropdown( + label="Dataset", + choices=datasets, + value=datasets[0] if datasets else None, + scale=3, + ) + idx_in = gr.Number( + label="Index", value=0, precision=0, minimum=0, scale=1 + ) + count_md = gr.Markdown() + task_md = gr.Markdown() + with gr.Accordion("Task details", open=False): + with gr.Accordion("Instruction", open=True): + instruction_box = gr.Code( + label="", language="markdown", lines=16 + ) + with gr.Accordion("Dockerfile", open=False): + dockerfile_box = gr.Code( + label="", language="dockerfile", lines=12 + ) + with gr.Accordion("task.toml", open=False): + toml_box = gr.Code(label="", language="python", lines=12) + with gr.Accordion("Grader", open=False): + tests_box = gr.Code(label="", language="shell", lines=12) + + # Full width, under both columns: the action belongs to the pair, not to either one. + run_btn = gr.Button( + "Run rollout", variant="primary", interactive=False, scale=1 + ) + + gr.Markdown("---") + result_html = gr.HTML() + # Live transcript while running, then the full conversation once finished. + convo_html = gr.HTML() + # Per-turn analysis plus the token-flow graph. + analysis_html = gr.HTML() + # The training contract as a file: token ids and the behaviour-policy logprobs, which + # are the part that cannot be reconstructed after the fact. + contract_file = gr.File( + label="contract.json — token ids, logprobs and reward", + interactive=False, + visible=True, + ) + with gr.Accordion("Result JSON", open=False): + raw_json = gr.Code(language="json", lines=22) + + validate_btn.click( + on_validate, + [url_in, model_in], + [engine_md, harness_in, sandbox_in, state, run_btn], + ) + ds_in.change(on_dataset, [ds_in], [idx_in, count_md]) + ds_in.change( + on_task, + [ds_in, idx_in], + [task_md, instruction_box, dockerfile_box, toml_box, tests_box], + ) + idx_in.change( + on_task, + [ds_in, idx_in], + [task_md, instruction_box, dockerfile_box, toml_box, tests_box], + ) + run_btn.click( + on_run, + [state, ds_in, idx_in, harness_in, sandbox_in], + [result_html, convo_html, analysis_html, raw_json, contract_file, run_btn], + ) + + if datasets: + app.load(on_dataset, [ds_in], [idx_in, count_md]) + app.load( + on_task, + [ds_in, idx_in], + [task_md, instruction_box, dockerfile_box, toml_box, tests_box], + ) + return app From f1ae0e4d6d2f636c64d2fcca566574dcc58b8fde Mon Sep 17 00:00:00 2001 From: adithya-s-k Date: Sun, 2 Aug 2026 19:23:14 +0000 Subject: [PATCH 03/74] cli: openenv harbor info / rollout / serve / push info reports what this machine can actually run. rollout runs one end to end with no server involved, which halves the search space when something breaks: if rollout works and serve does not, the fault is in the serving layer. serve is the env server; push deploys the same thing to a Space. --llm-url is required with no default and no environment fallback, because an unset endpoint produces rollouts that look completely normal and carry no token ids. push attaches the task suites as a bucket volume mounted at /data instead of downloading them: a Harbor suite is thousands of small files and Space disk is ephemeral, so a download is re-paid on every restart. Copies are server side, by xet hash. The mount is verified before the server is pointed at it, and it falls back to downloading rather than reading paths that may not exist. The harbor extra installs every sandbox backend. Not harbor[cloud], which is unsatisfiable: it pulls langsmith[sandbox] and tensorlake, which demand incompatible websockets ranges. --- pyproject.toml | 11 + src/openenv/cli/__main__.py | 6 + src/openenv/cli/commands/harbor.py | 662 +++++++++++++++++++++++++++++ 3 files changed, 679 insertions(+) create mode 100644 src/openenv/cli/commands/harbor.py diff --git a/pyproject.toml b/pyproject.toml index 5d7277ae8..0664bf68c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,17 @@ modal = [ inspect = [ "inspect-ai>=0.3.0", ] +harbor = [ + # Every sandbox backend. Bare `harbor` gives one that lists all 23 backends and can + # instantiate none of them: each raises MissingExtraError from its constructor, which + # surfaces as a failed rollout rather than as a missing dependency. + # + # Not `harbor[cloud]`, which is unsatisfiable: it pulls `langsmith[sandbox]` + # (websockets>=15) and `tensorlake` (websockets>=13,<14) together. + # + # Harbor requires Python >= 3.12, so this extra is excluded from the 3.11 CI leg. + "harbor[e2b,modal,daytona,gke,ec2,runloop,novita,blaxel,beam,islo,opensandbox,cwsandbox,use-computer,cua]>=0.20.0; python_version >= '3.12'", +] [project.scripts] openenv = "openenv.cli.__main__:main" diff --git a/src/openenv/cli/__main__.py b/src/openenv/cli/__main__.py index 42ebbdbe9..e5c86096c 100644 --- a/src/openenv/cli/__main__.py +++ b/src/openenv/cli/__main__.py @@ -14,6 +14,7 @@ build, collect, fork, + harbor, import_env, init, push, @@ -37,6 +38,11 @@ app.command(name="build", help="Build Docker images for OpenEnv environments")( build.build ) +app.add_typer( + harbor.app, + name="harbor", + help="Run Harbor tasks with token-level capture (requires: pip install openenv[harbor])", +) app.command( name="validate", help="Validate environment structure and deployment readiness" )(validate.validate) diff --git a/src/openenv/cli/commands/harbor.py b/src/openenv/cli/commands/harbor.py new file mode 100644 index 000000000..69b6dc2e5 --- /dev/null +++ b/src/openenv/cli/commands/harbor.py @@ -0,0 +1,662 @@ +"""`openenv harbor` — run Harbor tasks with token-level capture. + +Four commands, in the order you would use them: + + openenv harbor info what can this machine run right now? + openenv harbor rollout --task-index 0 one rollout, end to end, no server + openenv harbor serve the env server, for trainers and clients + openenv harbor push the same server, deployed to a Space + +`info` and `rollout` exist so the whole path can be exercised without standing up a server, which +makes the failure surface much smaller when something is wrong: if `rollout` works and `serve` does +not, the problem is the serving layer, not Harbor, the sandbox, the agent or capture. + +Examples: + +```bash +# what is usable, given the credentials on this machine +openenv harbor info --llm-url $LLM --dataset AdithyaSK/data_agent_rl_environment_eval + +# one rollout on E2B with opencode +openenv harbor rollout \\ + --llm-url $LLM \\ + --dataset AdithyaSK/data_agent_rl_environment_eval \\ + --task-index 0 --harness opencode --sandbox e2b + +# the same task on Modal with codex — harness and sandbox are per-rollout +openenv harbor rollout --llm-url $LLM --dataset $DS \\ + --task-index 0 --harness codex --sandbox modal +``` +""" + +from __future__ import annotations + +import asyncio +import contextlib +import json +import os +import shutil +import tempfile +from pathlib import Path +from typing import Annotated, Any, Optional + +import typer + +app = typer.Typer( + name="harbor", + help="Run Harbor tasks with token-level capture", + no_args_is_help=True, +) + +_DATASET_HELP = ( + "Dataset spec: HF repo id, local dir, or Harbor `name@version`. Repeatable." +) +_LLM_HELP = "OpenAI-spec inference endpoint (vLLM). Required: there is no default, because a\nwrong or stale endpoint produces rollouts that look fine and carry no token ids." +_LLM_HELP_OPTIONAL = ( + "OpenAI-spec inference endpoint (vLLM). Optional here: without it, `info` " + "still reports sandboxes, datasets and harnesses." +) + + +def _split(values: Optional[list[str]]) -> list[str]: + """Accept both `--dataset a --dataset b` and `--dataset a,b`.""" + out: list[str] = [] + for value in values or []: + out.extend(v.strip() for v in value.split(",") if v.strip()) + return out + + +@app.command("info") +def info( + llm_url: Annotated[str, typer.Option("--llm-url", help=_LLM_HELP_OPTIONAL)] = "", + model: Annotated[ + str, + typer.Option( + "--model", + help="Served model id. Auto-detected if the engine serves exactly one.", + ), + ] = "", + dataset: Annotated[ + Optional[list[str]], typer.Option("--dataset", help=_DATASET_HELP) + ] = None, + env_file: Annotated[ + str, typer.Option("--env-file", help="dotenv file with provider credentials.") + ] = "", + verbose: Annotated[ + bool, + typer.Option("--verbose", help="List every harness, not only validated ones."), + ] = False, + json_output: Annotated[ + bool, typer.Option("--json", help="Emit machine-readable JSON.") + ] = False, +) -> None: + """Report engine, sandboxes, datasets and harnesses available here.""" + from openenv.harbor.startup import prepare + + caps = prepare( + llm_url=llm_url or None, + model=model or None, + datasets=_split(dataset) or None, + env_file=env_file or None, + require_llm=False, + quiet=True, + ) + print( + json.dumps(caps.to_dict(), indent=2) + if json_output + else caps.render(verbose=verbose) + ) + + +@app.command("rollout") +def rollout( + llm_url: Annotated[str, typer.Option("--llm-url", help=_LLM_HELP)], + model: Annotated[str, typer.Option("--model", help="Served model id.")] = "", + dataset: Annotated[ + Optional[list[str]], typer.Option("--dataset", help=_DATASET_HELP) + ] = None, + task_index: Annotated[ + int, typer.Option("--task-index", help="Index into the split.") + ] = 0, + harness: Annotated[ + str, typer.Option("--harness", help="Seam name, or `module:Class`.") + ] = "opencode", + sandbox: Annotated[ + str, + typer.Option( + "--sandbox", help="Harbor environment type, e.g. e2b | modal | docker." + ), + ] = "e2b", + n: Annotated[ + int, typer.Option("-n", "--n-tasks", help="Run this many consecutive tasks.") + ] = 1, + port: Annotated[ + int, typer.Option("--port", help="Local port for the capture proxy.") + ] = 8100, + expose: Annotated[ + str, + typer.Option( + "--expose", + help="How the sandbox reaches the capture proxy: gradio | cloudflare | direct.", + ), + ] = "gradio", + trials_dir: Annotated[ + str, typer.Option("--trials-dir", help="Where Harbor writes trial artifacts.") + ] = "", + reward_key: Annotated[ + str, + typer.Option("--reward-key", help="Which reward key is the training signal."), + ] = "", + keep_sandbox: Annotated[ + bool, + typer.Option("--keep-sandbox", help="Leave sandboxes alive for debugging."), + ] = False, + force_build: Annotated[ + bool, + typer.Option( + "--force-build", + help="Rebuild the sandbox image, bypassing the content-hash cache. Needed when a task pins deps loosely and its cached image has drifted.", + ), + ] = False, + env_file: Annotated[ + str, typer.Option("--env-file", help="dotenv file with provider credentials.") + ] = "", + out: Annotated[str, typer.Option("--out", help="Write the result JSON here.")] = "", +) -> None: + """Run one or more rollouts without starting a server.""" + from openenv.harbor.runner import run_batch + + datasets = _split(dataset) + if not datasets: + raise typer.BadParameter("--dataset is required") + + results = asyncio.run( + run_batch( + llm_url=llm_url, + model=model or None, + dataset=datasets[0], + task_indices=list(range(task_index, task_index + max(1, n))), + harness=harness, + sandbox=sandbox, + port=port, + expose=expose, + trials_dir=Path(trials_dir) if trials_dir else None, + reward_key=reward_key, + keep_sandbox=keep_sandbox, + force_build=force_build, + env_file=env_file or None, + ) + ) + + if out: + Path(out).write_text(json.dumps([r.model_dump() for r in results], indent=2)) + print(f"\nwrote {out}") + raise typer.Exit(0 if all(r.ok for r in results) else 1) + + +@app.command("serve") +def serve( + llm_url: Annotated[str, typer.Option("--llm-url", help=_LLM_HELP)], + model: Annotated[str, typer.Option("--model", help="Served model id.")] = "", + dataset: Annotated[ + Optional[list[str]], typer.Option("--dataset", help=_DATASET_HELP) + ] = None, + host: Annotated[str, typer.Option("--host")] = "0.0.0.0", + port: Annotated[ + int, typer.Option("--port", help="Env server port (faces the trainer).") + ] = 8000, + capture_port: Annotated[ + int, + typer.Option("--capture-port", help="Capture proxy port (faces the sandbox)."), + ] = 8100, + expose: Annotated[ + str, + typer.Option( + "--expose", + help="How the sandbox reaches the capture proxy: gradio | cloudflare | direct.", + ), + ] = "gradio", + env_file: Annotated[str, typer.Option("--env-file")] = "", +) -> None: + """Serve Harbor tasks over the OpenEnv Task API and MCP. + + Two ports on purpose. The env server faces the trainer on an internal network; the capture proxy + faces the sandbox and is the only thing published. Sharing one port would expose the env + server as soon as the capture proxy became reachable. + """ + from openenv.harbor.serving import serve_harbor + + serve_harbor( + llm_url=llm_url, + model=model or None, + datasets=_split(dataset), + host=host, + port=port, + capture_port=capture_port, + expose=expose, + env_file=env_file or None, + ) + + +@app.command("push") +def push( + llm_url: Annotated[str, typer.Option("--llm-url", help=_LLM_HELP)], + repo_id: Annotated[ + str, typer.Option("--repo-id", help="Target, e.g. your-org/harbor-env.") + ] = "", + model: Annotated[str, typer.Option("--model", help="Served model id.")] = "", + dataset: Annotated[ + Optional[list[str]], typer.Option("--dataset", help=_DATASET_HELP) + ] = None, + private: Annotated[ + bool, + typer.Option( + "--private", + help="Create the Space private. The sandbox then cannot reach the capture proxy, so rollouts are not possible; use it only to park a deployment.", + ), + ] = False, + hardware: Annotated[ + str, typer.Option("--hardware", help="Space hardware, e.g. cpu-basic.") + ] = "", + env_file: Annotated[ + str, + typer.Option( + "--env-file", help="dotenv whose provider keys become Space SECRETS." + ), + ] = "", + bucket: Annotated[ + str, + typer.Option( + "--bucket", + help="Storage bucket holding the task suites. Defaults to a bucket named after the Space. Pass `none` to skip the bucket and let the Space download datasets instead.", + ), + ] = "", + recreate: Annotated[ + bool, + typer.Option( + "--recreate", + help="Delete the Space first, then deploy fresh. A Space keeps variables, secrets, volumes and any file a previous push wrote, so an incremental deploy is not a clean test of what this bundle produces.", + ), + ] = False, + dry_run: Annotated[ + bool, typer.Option("--dry-run", help="Show what would be pushed and stop.") + ] = False, +) -> None: + """Deploy this environment to a Hugging Face Space. + + Takes the same arguments as `serve`, because a deployed Space needs exactly the same + configuration — they are forwarded as Space variables, while provider credentials from + `--env-file` are forwarded as Space *secrets* so they are not readable from the repo. + """ + from pathlib import Path + + from openenv.cli.commands.push import push as _push + from openenv.harbor.startup import load_env_file + + if not repo_id: + raise typer.BadParameter("--repo-id is required, e.g. your-org/harbor-env") + + datasets = _split(dataset) + if not llm_url: + raise typer.BadParameter( + "--llm-url is required: a Space with no engine cannot run anything, and finding " + "that out after deploying is worse than finding it out now." + ) + + if private: + # A hosted deployment serves the capture proxy at /capture. On a private Space + # that URL demands an auth header the agent inside the sandbox does not send, so every model + # call 401s and the rollout records nothing. Worth a warning rather than a hard failure: + # parking a private deployment is legitimate, running rollouts against one is not. + print( + "WARNING: a private Space is not reachable from a sandbox. The capture proxy is " + "served at /capture, and a private Space requires an auth header the " + "agent will not send, so rollouts will capture no model calls. Deploy public for " + "rollouts. The proxy still refuses callers without a registered session id, which " + "is what keeps a public deployment from being an open relay." + ) + + # HF datasets are attached as read-only volumes rather than downloaded. A Harbor suite is + # thousands of small files (13k+ for a 2.2k-task dataset), so downloading on first request takes + # minutes and burns the Space's ephemeral disk; a mount is instant and survives restarts. The + # server needs no special case for it, because a mount path is just a local directory and + # `resolve_task_dirs` already accepts one. + # Every suite lives under one bucket, mounted at /data, named after the Space so the two stay + # obviously paired. `--bucket none` opts out and falls back to downloading. + hf_specs = [spec for spec in datasets if _looks_like_hf_repo(spec)] + bucket = "" if bucket.lower() == "none" else (bucket or repo_id) + mounts = {spec: f"{_MOUNT_ROOT}/{spec.replace('/', '__')}" for spec in hf_specs} + + # Non-secret configuration travels as plain Space variables. + variables = {"OPENENV_LLM_URL": llm_url, "ENABLE_WEB_INTERFACE": "true"} + if datasets: + variables["OPENENV_DATASETS"] = ",".join(datasets) + if model: + variables["OPENENV_MODEL"] = model + + # Provider credentials travel as secrets. Only the keys the sandboxes need — never the whole + # dotenv, which usually holds unrelated tokens. + load_env_file(env_file or None) + # Sandbox credentials, plus the keys a task's own verifier may need. A grader that cannot run + # returns no reward at all, which is reported as `reward=None` rather than 0, so the rollout is + # correctly not scored as a wrong answer, but it is also not usable for training. The DataAgent + # grader reads OPENAI_API_KEY for its LLM-judge tier, and without it every semantically correct + # answer that is not an exact string match goes ungraded. + wanted = ( + "E2B_API_KEY", + "MODAL_TOKEN_ID", + "MODAL_TOKEN_SECRET", + "DAYTONA_API_KEY", + "HF_TOKEN", + "OPENAI_API_KEY", + "ANTHROPIC_API_KEY", + "GEMINI_API_KEY", + ) + secrets = {k: os.environ[k] for k in wanted if os.environ.get(k)} + + # src/openenv/cli/commands/harbor.py -> repo root is parents[4]. + # An installed wheel has no sibling envs/ dir, so fall back to $OPENENV_HARBOR_ENV_DIR. + env_dir = Path( + os.environ.get("OPENENV_HARBOR_ENV_DIR") + or Path(__file__).resolve().parents[4] / "envs" / "harbor_env" + ) + if not (env_dir / "openenv.yaml").is_file(): + raise typer.BadParameter( + f"no harbor_env package at {env_dir}. Set OPENENV_HARBOR_ENV_DIR to its location " + "(an installed openenv wheel does not ship the envs/ directory)." + ) + # `openenv.harbor` does not exist in any released wheel, so a Space that pip-installs `openenv` + # imports the release and dies on `No module named 'openenv.harbor'`. When pushing from a source + # checkout, bundle the working tree instead; the Dockerfile puts /app/env ahead of site-packages. + source_pkg = Path(__file__).resolve().parents[2] # .../src/openenv + bundle = source_pkg if (source_pkg / "harbor").is_dir() else None + + print(f"env {env_dir}") + print(f"repo {repo_id}{' (private)' if private else ''}") + print(f"llm {llm_url}") + print( + f"datasets {', '.join(datasets) or '(none, set OPENENV_DATASETS on the Space)'}" + ) + print(f"variables {sorted(variables)}") + print(f"secrets {sorted(secrets)} (values never printed)") + print(f"openenv {'bundled from ' + str(source_pkg) if bundle else 'from PyPI'}") + if bucket: + print(f"bucket {bucket} -> {_MOUNT_ROOT} ({len(hf_specs)} suite(s))") + for spec, path in mounts.items(): + print( + f"mount {spec} -> {path}" + + (" (via bucket)" if bucket else " (read-only, not downloaded)") + ) + if dry_run: + print("\ndry run: nothing pushed") + return + + if recreate: + _delete_space(repo_id) + + if bucket and hf_specs: + _fill_bucket(bucket, hf_specs) + + with tempfile.TemporaryDirectory(prefix="openenv-harbor-push-") as tmp: + staged = Path(tmp) / "env" + shutil.copytree( + env_dir, + staged, + ignore=shutil.ignore_patterns("__pycache__", "*.pyc", ".venv"), + ) + if bundle is not None: + # A hosted Space mounts the capture proxy on its own app and reaches it at the Space's + # public URL, so the forwarding backends are dead code there. They are excluded rather + # than merely unused: shipping code that shells out to `cloudflared` into a Space is + # both pointless and the kind of thing platform abuse checks reject. `cli` goes for the + # same reason, it is 36 files the server never imports. + shutil.copytree( + bundle, + staged / "openenv", + ignore=shutil.ignore_patterns( + "__pycache__", "*.pyc", "forwarding.py", "cli" + ), + ) + _prune_removed_files(repo_id, staged) + _push( + directory=str(staged), + repo_id=repo_id, + private=private, + hardware=hardware or None, + env_vars=[f"{k}={v}" for k, v in variables.items()], + secrets=[f"{k}={v}" for k, v in secrets.items()], + ) + + # After the push, because volumes attach to a Space that already exists and `--recreate` has + # just deleted it. Setting them triggers one more rebuild, which is why this is last. + attached = ( + _attach_bucket(repo_id, bucket, mounts) + if bucket + else _mount_datasets(repo_id, mounts) + ) + if attached: + # Only now is it safe to point the server at mount paths. Until the mount is confirmed, + # `OPENENV_DATASETS` holds repo ids, so an unattached volume degrades to downloading rather + # than to a server pointed at directories that do not exist. + from huggingface_hub import HfApi + + HfApi().add_space_variable( + repo_id=repo_id, + key="OPENENV_DATASETS", + value=",".join(mounts.get(d, d) for d in datasets), + ) + print("mount OPENENV_DATASETS switched to mount paths") + + +def _prune_removed_files(repo_id: str, staged: Path) -> None: + """Delete files on the Space that this push no longer produces. + + `push` uploads but never deletes, so a file dropped from the bundle keeps running in the + deployment forever. That is not a tidiness point: the first version of this command shipped the + port-forwarding backends, and removing them locally left the deployed Space still carrying code + that shells out to `cloudflared`, which is exactly what a platform abuse check objects to. A + deployment has to reflect the bundle, not the union of every bundle ever pushed. + + Only the bundled `openenv/` subtree is pruned. Everything else in the Space may legitimately have + been added out of band (a README edit through the web UI, a `.gitattributes`), and deleting a + file this command never wrote is not its business. + """ + from huggingface_hub import CommitOperationDelete, HfApi + + api = HfApi() + try: + remote = api.list_repo_files(repo_id, repo_type="space") + except Exception as exc: # noqa: BLE001 - a new Space has nothing to prune + print(f"prune skipped ({type(exc).__name__}); the Space may not exist yet") + return + + local = {str(p.relative_to(staged)) for p in staged.rglob("*") if p.is_file()} + stale = sorted(f for f in remote if f.startswith("openenv/") and f not in local) + if not stale: + return + + print(f"prune {len(stale)} file(s) no longer in the bundle, e.g. {stale[0]}") + api.create_commit( + repo_id=repo_id, + repo_type="space", + operations=[CommitOperationDelete(path_in_repo=f) for f in stale], + commit_message="Remove files no longer part of the harbor_env bundle", + ) + + +# Where dataset volumes are attached inside the Space container. +_MOUNT_ROOT = "/data" + +# Mirrors `openenv.harbor.tasks._DATASET_ROOT`; kept local so the CLI does not import the +# harbor extra just to compute a path. +_DATASET_CACHE = Path( + os.environ.get("OPENENV_DATASET_CACHE") + or (Path.home() / ".cache" / "openenv" / "harbor-datasets") +) + + +def _looks_like_hf_repo(spec: str) -> bool: + """True for `owner/name`, false for a local path or a Harbor `name@version`.""" + return ( + spec.count("/") == 1 + and "@" not in spec + and not spec.startswith((".", "/", "~")) + ) + + +def _mount_datasets(repo_id: str, mounts: dict[str, str]) -> bool: + """Attach each dataset repo to the Space as a read-only volume. + + Downloading a Harbor suite inside a Space is the slow path twice over: thousands of small files + fetched one round trip at a time, onto a disk that is wiped on restart, so the cost is paid again + on every rebuild. A mounted repo is available as ordinary files immediately. + + Volumes are replaced wholesale by the API, so anything already attached is read first and kept. + + Returns: + `bool`: Whether the volumes are confirmed attached. `False` means the caller must keep using + repo ids and let the Space download, which is slower but works. + """ + if not mounts: + return False + try: + from huggingface_hub import HfApi, Volume + except ImportError: + print( + "mount skipped: this huggingface_hub has no Volume support; " + "the Space will download datasets instead" + ) + return False + + api = HfApi() + existing: list[Any] = [] + with contextlib.suppress(Exception): + existing = [ + v + for v in _attached_volumes(api, repo_id) + if getattr(v, "mount_path", None) not in set(mounts.values()) + ] + + volumes = existing + [ + Volume(type="dataset", source=spec, mount_path=path, read_only=True) + for spec, path in sorted(mounts.items()) + ] + try: + api.set_space_volumes(repo_id=repo_id, volumes=volumes) + except Exception as exc: # noqa: BLE001 - a Space that cannot mount still works by downloading + print( + f"mount failed ({type(exc).__name__}: {str(exc)[:160]}). The Space will download " + "datasets instead, which is slow but functional." + ) + return False + + # Accepting the call is not evidence that the volume exists. Read it back, because the failure + # mode of trusting it is a server configured to read directories that were never mounted. + attached = _attached_mount_paths(api, repo_id) + if not set(mounts.values()) <= attached: + print( + "mount not confirmed: the Space reports no attached volumes, so the datasets will " + "be downloaded instead. Attach them from the Space settings if you want the mount." + ) + return False + print(f"mount attached {len(mounts)} dataset volume(s)") + return True + + +def _delete_space(repo_id: str) -> None: + """Delete the Space so the next push is a clean deployment. + + A Space accumulates state a push does not own: variables and secrets set by earlier runs, mounted + volumes, and every file any previous push wrote. That makes an incremental deploy a poor test, + because it can succeed on leftovers the bundle no longer produces. Deleting first means what runs + is exactly what this command uploaded. + + Deliberately destructive, so it only ever happens behind `--recreate`. + """ + from huggingface_hub import HfApi + + api = HfApi() + try: + api.delete_repo(repo_id=repo_id, repo_type="space") + print(f"recreate deleted {repo_id}") + except Exception as exc: # noqa: BLE001 - nothing to delete is the expected first-run case + print(f"recreate nothing to delete ({type(exc).__name__})") + + +def _fill_bucket(bucket: str, specs: list[str]) -> None: + """Create `bucket` if missing and copy each task suite into it, server side. + + `copy_files` copies by xet hash: the Hub moves the references, nothing is downloaded here and + nothing is re-uploaded. That is the difference between seconds and the ~47k-file upload a local + sync performs, and it is why the bucket is filled before the Space exists rather than after. + + Suites already present are skipped, so adding a dataset to a later `push` copies only the new + one and leaves the rest untouched. + """ + from huggingface_hub import HfApi + + api = HfApi() + api.create_bucket(bucket, private=False, exist_ok=True) + + try: + present = { + entry.path.split("/", 1)[0] + for entry in api.list_bucket_tree(bucket) + if getattr(entry, "path", "") + } + except Exception: # noqa: BLE001 - a brand new bucket may not be listable yet + present = set() + + for spec in specs: + prefix = spec.replace("/", "__") + if prefix in present: + print(f"bucket {spec} already present, skipped") + continue + print( + f"copy hf://datasets/{spec} -> hf://buckets/{bucket}/{prefix} (server side)" + ) + api.copy_files(f"hf://datasets/{spec}/", f"hf://buckets/{bucket}/{prefix}/") + + +def _attach_bucket(repo_id: str, bucket: str, mounts: dict[str, str]) -> bool: + """Mount `bucket` on the Space and confirm it attached. + + Returns: + `bool`: Whether the mount is confirmed. `False` leaves the caller on repo ids so the Space + downloads rather than reading a mount that may not be there. + """ + from huggingface_hub import HfApi, Volume + + api = HfApi() + api.set_space_volumes( + repo_id=repo_id, + volumes=[Volume(type="bucket", source=bucket, mount_path=_MOUNT_ROOT)], + ) + if _MOUNT_ROOT not in _attached_mount_paths(api, repo_id): + print( + f"mount not confirmed: no volume at {_MOUNT_ROOT}. Datasets will be downloaded " + "instead. Attach the bucket from the Space settings to use the mount." + ) + return False + print(f"mount {bucket} attached at {_MOUNT_ROOT}") + return bool(mounts) + + +def _attached_volumes(api: Any, repo_id: str) -> list[Any]: + """Volumes currently mounted on `repo_id`. + + Read through `space_info().runtime`, not `get_space_runtime()`. The latter is served by an + endpoint that does not carry a `volumes` key at all, so it always answers `None` and a check + built on it reports every mount as missing. That false negative is worse than no check: it makes + a working mount look broken and sends the caller down the slow path forever. + """ + with contextlib.suppress(Exception): + runtime = api.space_info(repo_id).runtime + if runtime is not None: + return list(runtime.volumes or []) + return [] + + +def _attached_mount_paths(api: Any, repo_id: str) -> set[str]: + """Mount paths currently attached to `repo_id`.""" + return {getattr(v, "mount_path", None) for v in _attached_volumes(api, repo_id)} From 341eebecf4256f2f34fab0adcee09760e38e9d69 Mon Sep 17 00:00:00 2001 From: adithya-s-k Date: Sun, 2 Aug 2026 19:23:29 +0000 Subject: [PATCH 04/74] harbor_env: deployment packaging for a Space Manifest, Dockerfile and ASGI entry point only; the logic lives in openenv.harbor so the capture layer can be shared rather than duplicated per agent environment. The Dockerfile pins UV_PYTHON_INSTALL_DIR and copies it across the stage boundary. Harbor needs Python >= 3.12 while openenv-base ships 3.11, so uv downloads its own interpreter and the venv's bin/python is a symlink into it; copying only .venv leaves a dangling link and the container dies with 'not found'. A build-time assertion now catches that at build rather than at startup. The entry point resolves and validates the served model the way harbor serve does. Without it the proxy has no served model id and forwards whatever name the harness used straight to the engine. --- envs/harbor_env/README.md | 59 ++++++++++++++++++++ envs/harbor_env/__init__.py | 9 ++++ envs/harbor_env/client.py | 5 ++ envs/harbor_env/models.py | 17 ++++++ envs/harbor_env/openenv.yaml | 6 +++ envs/harbor_env/pyproject.toml | 29 ++++++++++ envs/harbor_env/server/Dockerfile | 36 +++++++++++++ envs/harbor_env/server/__init__.py | 0 envs/harbor_env/server/app.py | 87 ++++++++++++++++++++++++++++++ 9 files changed, 248 insertions(+) create mode 100644 envs/harbor_env/README.md create mode 100644 envs/harbor_env/__init__.py create mode 100644 envs/harbor_env/client.py create mode 100644 envs/harbor_env/models.py create mode 100644 envs/harbor_env/openenv.yaml create mode 100644 envs/harbor_env/pyproject.toml create mode 100644 envs/harbor_env/server/Dockerfile create mode 100644 envs/harbor_env/server/__init__.py create mode 100644 envs/harbor_env/server/app.py diff --git a/envs/harbor_env/README.md b/envs/harbor_env/README.md new file mode 100644 index 000000000..275a5e35a --- /dev/null +++ b/envs/harbor_env/README.md @@ -0,0 +1,59 @@ +--- +title: Harbor +emoji: ⚓ +colorFrom: blue +colorTo: indigo +sdk: docker +app_port: 8000 +--- + +# harbor_env + +Run a Harbor task with a coding agent and capture every token id and per-token logprob it produced, +ready to train on. + +Harbor supplies the task datasets, sandbox backends, agents and verifiers. This environment adds the +OpenEnv surface: dataset discovery over the Task API, one `run_rollout` MCP tool, and a capture proxy +between the agent and your model. + +## Configure + +| variable | meaning | +|---|---| +| `OPENENV_LLM_URL` | OpenAI-spec endpoint (vLLM). **Required.** | +| `OPENENV_DATASETS` | comma-separated dataset specs — HF repo id, local dir, or Harbor `name@version` | +| `OPENENV_MODEL` | served model id; read from the engine when it serves exactly one | +| `E2B_API_KEY` | offer the `e2b` sandbox | +| `MODAL_TOKEN_ID`, `MODAL_TOKEN_SECRET` | offer the `modal` sandbox | + +The engine **must** be started with: + +``` +--return-tokens-as-token-ids --logprobs-mode processed_logprobs +``` + +Without them it answers every request normally and returns no token ids, so captured rollouts are +empty and nothing reports an error. The server refuses to start rather than let that happen. + +## Use + +```python +from harbor_env import HarborEnv + +with HarborEnv(base_url="http://localhost:8000") as env: + split = env.splits()[0]["name"] + result = env.run_rollout(split=split, task_index=0, harness="opencode", sandbox="e2b") + print(result.reward, len(result.turns)) +``` + +`harness` and `sandbox` are per-call, so consecutive rollouts can use different agents and different +backends against the same server. + +## Notes + +Two ports: the env server faces trainers and browsers, the capture proxy faces the sandbox and is the +only thing forwarded publicly. + +Capture quality and task success are independent. A rollout can be captured perfectly and score 0 +because the model was wrong; a reward of 1 with unusable capture is worse than useless for training. +Both are reported separately. diff --git a/envs/harbor_env/__init__.py b/envs/harbor_env/__init__.py new file mode 100644 index 000000000..d174ded73 --- /dev/null +++ b/envs/harbor_env/__init__.py @@ -0,0 +1,9 @@ +"""harbor_env — run Harbor tasks with token-level capture. + +The implementation lives in `openenv.harbor` and `openenv.core.harness.capture`; this package is +deployment packaging only (manifest, Dockerfile, ASGI entry point). +""" + +from openenv.harbor.models import HarborRolloutResult, HarborTaskRef, HarborTurn + +__all__ = ["HarborRolloutResult", "HarborTaskRef", "HarborTurn"] diff --git a/envs/harbor_env/client.py b/envs/harbor_env/client.py new file mode 100644 index 000000000..b07299548 --- /dev/null +++ b/envs/harbor_env/client.py @@ -0,0 +1,5 @@ +"""Typed client for a deployed harbor_env.""" + +from openenv.harbor.client import HarborEnv + +__all__ = ["HarborEnv"] diff --git a/envs/harbor_env/models.py b/envs/harbor_env/models.py new file mode 100644 index 000000000..b325e4a3c --- /dev/null +++ b/envs/harbor_env/models.py @@ -0,0 +1,17 @@ +"""Wire types, re-exported so `from harbor_env.models import ...` works like other envs.""" + +from openenv.harbor.models import ( + HarborRolloutResult, + HarborState, + HarborStepResult, + HarborTaskRef, + HarborTurn, +) + +__all__ = [ + "HarborRolloutResult", + "HarborState", + "HarborStepResult", + "HarborTaskRef", + "HarborTurn", +] diff --git a/envs/harbor_env/openenv.yaml b/envs/harbor_env/openenv.yaml new file mode 100644 index 000000000..be39cba77 --- /dev/null +++ b/envs/harbor_env/openenv.yaml @@ -0,0 +1,6 @@ +spec_version: 1 +name: harbor_env +type: space +runtime: fastapi +app: server.app:app +port: 8000 diff --git a/envs/harbor_env/pyproject.toml b/envs/harbor_env/pyproject.toml new file mode 100644 index 000000000..7d23bf7c8 --- /dev/null +++ b/envs/harbor_env/pyproject.toml @@ -0,0 +1,29 @@ +[project] +name = "openenv-harbor-env" +version = "0.1.0" +description = "Run Harbor tasks with a coding agent and capture token-level training data" +requires-python = ">=3.12" +dependencies = [ + "openenv", + # Every sandbox backend, listed individually rather than via `harbor[cloud]`. + # `harbor[cloud]` cannot be installed at all: it pulls both `langsmith[sandbox]`, which + # requires `websockets>=15`, and `tensorlake`, which requires `websockets>=13,<14`. uv + # reports the pair as unsatisfiable and the image build fails. Neither is a sandbox backend + # we offer, so both are dropped and everything else kept. Re-check on a Harbor upgrade. + "harbor[e2b,modal,daytona,gke,ec2,runloop,novita,blaxel,beam,islo,opensandbox,cwsandbox,use-computer,cua]>=0.20.0", + "huggingface_hub>=1.12", + "fastapi>=0.104", + "uvicorn[standard]>=0.24", + "httpx>=0.27", + "gradio>=5", +] + +[project.scripts] +server = "server.app:main" + +[build-system] +requires = ["setuptools>=61"] +build-backend = "setuptools.build_meta" + +[tool.setuptools] +packages = ["server"] diff --git a/envs/harbor_env/server/Dockerfile b/envs/harbor_env/server/Dockerfile new file mode 100644 index 000000000..4c405529f --- /dev/null +++ b/envs/harbor_env/server/Dockerfile @@ -0,0 +1,36 @@ +ARG BASE_IMAGE=ghcr.io/huggingface/openenv-base:latest +FROM ${BASE_IMAGE} AS builder + +# Harbor requires Python >= 3.12 while openenv-base ships 3.11, so `uv sync` downloads its own +# interpreter and the venv's bin/python becomes a symlink into uv's install dir. Pinning that dir +# (and creating it up front, so the COPY below cannot fail when uv reuses a system interpreter) +# is what lets the runtime stage carry the interpreter the venv actually points at. Without it the +# venv arrives with a dangling bin/python and the container dies with "not found". +ENV UV_PYTHON_INSTALL_DIR=/opt/uv-python +RUN mkdir -p /opt/uv-python + +WORKDIR /app/env +COPY . /app/env +RUN --mount=type=cache,target=/root/.cache/uv \ + if [ -f uv.lock ]; then uv sync --frozen --no-editable; else uv sync --no-editable; fi + +FROM ${BASE_IMAGE} +COPY --from=builder /opt/uv-python /opt/uv-python +COPY --from=builder /app/env/.venv /app/.venv +COPY --from=builder /app/env /app/env + +# Fail at build time rather than at startup if the interpreter did not survive the stage boundary. +RUN /app/.venv/bin/python -c "import sys; print('venv python', sys.version)" + +ENV PATH="/app/.venv/bin:$PATH" +# `harbor push` bundles the working tree's openenv/ into /app/env when pushing from a source +# checkout; PYTHONPATH puts it ahead of the released wheel in site-packages, which has no +# `openenv.harbor` until this lands upstream. +ENV PYTHONPATH="/app/env:$PYTHONPATH" +ENV ENABLE_WEB_INTERFACE=true + +HEALTHCHECK --interval=30s --timeout=5s --start-period=20s --retries=3 \ + CMD /app/.venv/bin/python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')" || exit 1 + +EXPOSE 8000 +CMD ["sh", "-c", "cd /app/env && exec /app/.venv/bin/python -m uvicorn server.app:app --host 0.0.0.0 --port 8000"] diff --git a/envs/harbor_env/server/__init__.py b/envs/harbor_env/server/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/envs/harbor_env/server/app.py b/envs/harbor_env/server/app.py new file mode 100644 index 000000000..b575364d9 --- /dev/null +++ b/envs/harbor_env/server/app.py @@ -0,0 +1,87 @@ +"""ASGI entry point for a deployed harbor_env. + +Everything is read from the environment so the same image serves any dataset and any engine without +a rebuild — which is what makes this deployable to a Space: + + OPENENV_DATASETS comma-separated dataset specs (HF repo id, local dir, harbor name@version) + OPENENV_LLM_URL OpenAI-spec inference endpoint + OPENENV_MODEL served model id; read from the engine when it serves exactly one + E2B_API_KEY / MODAL_TOKEN_ID+MODAL_TOKEN_SECRET whichever sandboxes you want offered + +The capture proxy rides on this same app rather than on a second port. A Space publishes exactly one +port and one URL, so the proxy is mounted at `/capture` and the sandbox reaches it at +`https://.hf.space/capture`. Nothing is forwarded and no second listener is opened. +""" + +from __future__ import annotations + +import os + +from openenv.harbor.serving import HarborService, build_app + +_DATASETS = [ + d.strip() for d in os.environ.get("OPENENV_DATASETS", "").split(",") if d.strip() +] +_LLM_URL = os.environ.get("OPENENV_LLM_URL", "") +_MODEL = os.environ.get("OPENENV_MODEL", "") +_LLM: dict = {} + +# Ask the endpoint what it serves when `OPENENV_MODEL` was not set, the same way `harbor serve` does. +# Without this the proxy has no served model id and stops rewriting `model` on the way upstream, so +# whatever name the harness happened to use is forwarded verbatim and the engine rejects it. The +# report is kept so `capabilities()` can state whether capture is actually supported here. +if _LLM_URL: + try: + from openenv.core.harness.capture.validate_llm import list_models, validate_llm + + if not _MODEL: + served = list_models(_LLM_URL) + _MODEL = served[0] if len(served) == 1 else "" + if _MODEL: + _report = validate_llm(_LLM_URL, _MODEL) + _LLM = { + "url": _LLM_URL, + "model": _report.model, + "ok": _report.ok, + "findings": _report.findings, + "served_models": _report.served_models, + } + except Exception as exc: # noqa: BLE001 - a Space must still boot so the UI can show the fault + _LLM = { + "url": _LLM_URL, + "model": _MODEL, + "ok": False, + "findings": [ + f"could not reach the LLM at startup: {type(exc).__name__}: {exc}" + ], + } + +# Resolve capture before the app is built. A Space gives no separate boot hook, the UI needs the +# proxy's public URL to exist by the time anyone presses Run, and `build_app` has to see the service +# in order to mount it. +if _LLM_URL: + _service = HarborService( + llm_url=_LLM_URL, + model=_MODEL, + datasets=_DATASETS, + capture_port=int(os.environ.get("OPENENV_CAPTURE_PORT", "8100")), + expose=os.environ.get("OPENENV_EXPOSE", "gradio"), + ) + # On a Space this only computes the public URL and flags the app for mounting; off one it + # publishes the capture port the usual way. + _service.start() + HarborService.set_current(_service) + +os.environ.setdefault("ENABLE_WEB_INTERFACE", "true") + +app = build_app(datasets=_DATASETS, llm_url=_LLM_URL, model=_MODEL, llm=_LLM) + + +def main() -> None: + import uvicorn + + uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", "8000"))) + + +if __name__ == "__main__": + main() From 1b9b32469f71c3c53c22736d3b67556d8cc69b24 Mon Sep 17 00:00:00 2001 From: adithya-s-k Date: Sun, 2 Aug 2026 19:23:29 +0000 Subject: [PATCH 05/74] tests: port ownership, request normalisation, hosted serving Each of these pins a failure that was silent in production and cheap to reintroduce. No credentials or network needed. Port ownership: a capture server used to report healthy on a port another process owned, because the liveness probe connected to the incumbent while its own bind error died unobserved on a background thread. Sessions were then minted in one registry and rejected by another, producing a 401 and a rollout with zero model calls. Request normalisation: kimi-cli sends tools: [] once its loop has no tools left, and vLLM rejects an empty array outright, truncating the trajectory while leaving a well-formed graph behind. Hosted serving: a Space must mount the capture proxy rather than forward it. One test monkeypatches make_forwarder to raise, so a hosted deployment that ever tries to forward fails the suite. --- tests/envs/test_harbor_capture_normalise.py | 78 ++++++++++++ tests/envs/test_harbor_capture_server.py | 125 ++++++++++++++++++++ tests/envs/test_harbor_hosted_serving.py | 97 +++++++++++++++ 3 files changed, 300 insertions(+) create mode 100644 tests/envs/test_harbor_capture_normalise.py create mode 100644 tests/envs/test_harbor_capture_server.py create mode 100644 tests/envs/test_harbor_hosted_serving.py diff --git a/tests/envs/test_harbor_capture_normalise.py b/tests/envs/test_harbor_capture_normalise.py new file mode 100644 index 000000000..5ba6db61e --- /dev/null +++ b/tests/envs/test_harbor_capture_normalise.py @@ -0,0 +1,78 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Request shapes that vLLM rejects outright, normalised before they reach it. + +Each case here is a real harness sending something an OpenAI-spec engine 400s on. A 400 does not +degrade a rollout, it truncates it: the agent loses the call, and the captured trajectory ends early +while still looking structurally valid. These are cheap to assert and expensive to rediscover. +""" + +from __future__ import annotations + +import pytest + +server = pytest.importorskip("openenv.core.harness.capture.server") + +normalise_for_capture = server.normalise_for_capture + + +def test_stream_is_forced_off(): + """Capture needs one whole response; reassembling ids from SSE deltas corrupts silently.""" + request = {"messages": [], "stream": True} + normalise_for_capture(request) + assert request["stream"] is False + + +def test_stream_options_is_dropped(): + """vLLM 400s on `stream_options` once `stream` is False. opencode sends it on every call.""" + request = { + "messages": [], + "stream": True, + "stream_options": {"include_usage": True}, + } + normalise_for_capture(request) + assert "stream_options" not in request + + +def test_empty_tools_is_dropped(): + """kimi-cli sends `tools: []`, which vLLM rejects: 'must not be an empty array'.""" + request = {"messages": [], "tools": []} + normalise_for_capture(request) + assert "tools" not in request + + +def test_empty_functions_is_dropped(): + """The legacy spelling fails the same way.""" + request = {"messages": [], "functions": []} + normalise_for_capture(request) + assert "functions" not in request + + +def test_tool_choice_is_dropped_with_the_tools_it_referenced(): + """`tool_choice` without `tools` is invalid, and means nothing once the list is gone.""" + request = {"messages": [], "tools": [], "tool_choice": "auto"} + normalise_for_capture(request) + assert "tools" not in request + assert "tool_choice" not in request + + +def test_populated_tools_are_left_alone(): + """The guard must be narrow: dropping real tools would change what the model can do.""" + tools = [{"type": "function", "function": {"name": "bash"}}] + request = {"messages": [], "tools": tools, "tool_choice": "auto"} + normalise_for_capture(request) + assert request["tools"] == tools + assert request["tool_choice"] == "auto" + + +def test_absent_tool_keys_are_not_invented(): + """A request with no tool keys must stay that way rather than gain empty ones.""" + request = {"messages": []} + normalise_for_capture(request) + assert "tools" not in request + assert "functions" not in request + assert "tool_choice" not in request diff --git a/tests/envs/test_harbor_capture_server.py b/tests/envs/test_harbor_capture_server.py new file mode 100644 index 000000000..1697fef19 --- /dev/null +++ b/tests/envs/test_harbor_capture_server.py @@ -0,0 +1,125 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Port-ownership guarantees for the capture proxy. + +A capture server that reports healthy while a *different* process owns its port is the worst +failure this layer has: sessions are minted in one registry and validated against another, so the +agent is rejected with 401, every rollout reports zero model calls, and the UI shows a live view +that can never advance. Nothing in that chain names the port, so these tests pin the invariant that +`start()` refuses rather than proceeds. + +No credentials and no engine are needed: `start()` binds a socket and never contacts `llm_url`. +""" + +from __future__ import annotations + +import socket + +import pytest + +harbor_runner = pytest.importorskip("openenv.harbor.runner") + +CaptureServer = harbor_runner.CaptureServer + +# Never contacted. Discard port, so a stray request would fail loudly rather than reach a real host. +UNUSED_ENGINE = "http://127.0.0.1:9/v1" + + +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]) + + +@pytest.fixture +def capture(): + """Yield a factory that tears every server it built back down.""" + built = [] + + def make(port: int) -> CaptureServer: + server = CaptureServer(llm_url=UNUSED_ENGINE, model="test-model", port=port) + built.append(server) + return server + + yield make + for server in reversed(built): + server.stop() + + +def test_health_reports_this_instance(capture): + """`/health` must identify which app answered, so a probe can check identity not reachability.""" + import httpx + + server = capture(_free_port()) + server.start() + + payload = httpx.get(f"http://127.0.0.1:{server.port}/health", timeout=5.0).json() + assert payload["instance"] == server.app.state.instance_id + assert payload["status"] == "ok" + + +def test_instance_ids_are_distinct(capture): + """Two apps must never share an id, or the identity probe cannot tell them apart.""" + first, second = capture(_free_port()), capture(_free_port()) + assert first.app.state.instance_id != second.app.state.instance_id + + +def test_start_refuses_a_port_another_server_holds(capture): + """The regression: a second server on a held port used to report healthy. + + Its own uvicorn fails to bind on a background thread where nothing observes the error, while the + liveness probe connects successfully to the *incumbent*. Reachability is not ownership. + """ + port = _free_port() + incumbent = capture(port) + incumbent.start() + + intruder = capture(port) + with pytest.raises(RuntimeError, match=f"{port}"): + intruder.start() + + # The incumbent must be untouched: a failed start may not disturb a working server. + import httpx + + payload = httpx.get(f"http://127.0.0.1:{port}/health", timeout=5.0).json() + assert payload["instance"] == incumbent.app.state.instance_id + + +def test_start_refuses_a_port_held_by_a_non_capture_listener(capture): + """Any listener counts, not just another capture server.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as squatter: + squatter.bind(("127.0.0.1", 0)) + squatter.listen(1) + port = int(squatter.getsockname()[1]) + + with pytest.raises(RuntimeError, match="already in use"): + capture(port).start() + + +def test_start_succeeds_on_a_free_port_after_a_refusal(capture): + """A refusal must leave no state behind that breaks the next attempt.""" + port = _free_port() + capture(port).start() + + with pytest.raises(RuntimeError): + capture(port).start() + + recovered = capture(_free_port()) + recovered.start() + assert recovered._thread is not None and recovered._thread.is_alive() + + +def test_stop_releases_the_port(capture): + """Otherwise a restart in the same process hits the new guard and looks like a collision.""" + port = _free_port() + server = capture(port) + server.start() + server.stop() + + successor = capture(port) + successor.start() + assert successor.app.state.instance_id != server.app.state.instance_id diff --git a/tests/envs/test_harbor_hosted_serving.py b/tests/envs/test_harbor_hosted_serving.py new file mode 100644 index 000000000..78113609a --- /dev/null +++ b/tests/envs/test_harbor_hosted_serving.py @@ -0,0 +1,97 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Serving on a hosted platform, where there is one port and one URL. + +Locally the capture proxy runs on its own port and is published to the sandbox. A Space exposes +exactly one port and already has a public URL, so the proxy is mounted onto the env server's app +instead and nothing is forwarded. These tests pin that split, because getting it wrong is not a +crash: it is a deployment that quietly opens a second listener it cannot publish. +""" + +from __future__ import annotations + +import pytest + +serving = pytest.importorskip("openenv.harbor.serving") + +CAPTURE_MOUNT = serving.CAPTURE_MOUNT +HarborService = serving.HarborService +space_public_url = serving.space_public_url + +# Never contacted: nothing here reaches an engine. +UNUSED_LLM = "http://127.0.0.1:9/v1" + + +@pytest.fixture(autouse=True) +def _clear_space_env(monkeypatch): + """Tests must not inherit a Space identity from the developer's shell.""" + monkeypatch.delenv("SPACE_HOST", raising=False) + monkeypatch.delenv("SPACE_ID", raising=False) + + +def test_no_space_means_no_public_url(): + assert space_public_url() == "" + + +def test_space_host_is_used_verbatim(monkeypatch): + monkeypatch.setenv("SPACE_HOST", "owner-env.hf.space") + assert space_public_url() == "https://owner-env.hf.space" + + +def test_space_host_tolerates_a_scheme_already_present(monkeypatch): + monkeypatch.setenv("SPACE_HOST", "https://owner-env.hf.space/") + assert space_public_url() == "https://owner-env.hf.space" + + +def test_space_id_is_slugged_when_host_is_absent(monkeypatch): + """`SPACE_ID` is always set; the hostname lowercases and dash-separates it.""" + monkeypatch.setenv("SPACE_ID", "AdithyaSK/harbor_data.agent-env") + assert space_public_url() == "https://adithyask-harbor-data-agent-env.hf.space" + + +def test_hosted_start_mounts_and_never_forwards(monkeypatch): + """The regression that got a Space flagged: a hosted deployment must not forward.""" + monkeypatch.setenv("SPACE_ID", "owner/env") + + def explode(*_args, **_kwargs): + raise AssertionError("a hosted deployment must not create a forwarder") + + monkeypatch.setattr( + "openenv.core.harness.capture.forwarding.make_forwarder", explode, raising=False + ) + + service = HarborService(llm_url=UNUSED_LLM, model="m", datasets=[]) + url = service.start() + + assert service.mounted is True + assert url == f"https://owner-env.hf.space{CAPTURE_MOUNT}" + # No port was bound, so stop() must be safe even though start() never launched a server. + service.stop() + + +def test_mounted_capture_answers_under_the_prefix(): + """Mounting strips the prefix, so every dialect route keeps working unchanged.""" + fastapi = pytest.importorskip("fastapi") + from fastapi.testclient import TestClient + from openenv.core.harness.capture.server import create_app + + host = fastapi.FastAPI() + host.mount(CAPTURE_MOUNT, create_app(llm_url=UNUSED_LLM, model="m")) + client = TestClient(host) + + assert client.get(f"{CAPTURE_MOUNT}/health").json()["status"] == "ok" + + # The proxy's catch-all must match /v1/chat/completions, not /capture/v1/chat/completions. + response = client.post( + f"{CAPTURE_MOUNT}/v1/chat/completions", + json={"model": "m", "messages": [{"role": "user", "content": "hi"}]}, + headers={"Authorization": "Bearer not-a-session"}, + ) + # 401 rather than 404 proves it routed to the proxy and was rejected on identity, which is also + # what stops a publicly mounted proxy being an open relay. + assert response.status_code == 401 + assert "unknown API key" in response.json()["error"]["message"] From d44f42c462fe030a76c9b749f7da4fd4c648c75b Mon Sep 17 00:00:00 2001 From: adithya-s-k Date: Sun, 2 Aug 2026 19:26:03 +0000 Subject: [PATCH 06/74] docs: add the harbor_env stub and link it check-env-docs generates a docs stub per environment README and fails when one is missing, which it was. The README line about the capture proxy being the only thing forwarded publicly predated the hosted path and was wrong for a Space, where there is one port and one public URL and the proxy is mounted rather than forwarded. Fixed in the README so the generated stub follows. _toctree.yml is maintained by hand, so the generated page needs an entry there or it exists without being reachable from the sidebar. --- docs/source/_toctree.yml | 2 ++ docs/source/environments/harbor.md | 54 ++++++++++++++++++++++++++++++ envs/harbor_env/README.md | 7 ++-- 3 files changed, 61 insertions(+), 2 deletions(-) create mode 100644 docs/source/environments/harbor.md diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index 85b64b7d7..71d34e885 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -71,6 +71,8 @@ title: Terminus - local: environments/coding_tools title: Coding Tools + - local: environments/harbor + title: Harbor - local: environments/chat title: Chat - local: environments/atari diff --git a/docs/source/environments/harbor.md b/docs/source/environments/harbor.md new file mode 100644 index 000000000..de20cb4f5 --- /dev/null +++ b/docs/source/environments/harbor.md @@ -0,0 +1,54 @@ + +# harbor_env + +Run a Harbor task with a coding agent and capture every token id and per-token logprob it produced, +ready to train on. + +Harbor supplies the task datasets, sandbox backends, agents and verifiers. This environment adds the +OpenEnv surface: dataset discovery over the Task API, one `run_rollout` MCP tool, and a capture proxy +between the agent and your model. + +## Configure + +| variable | meaning | +|---|---| +| `OPENENV_LLM_URL` | OpenAI-spec endpoint (vLLM). **Required.** | +| `OPENENV_DATASETS` | comma-separated dataset specs — HF repo id, local dir, or Harbor `name@version` | +| `OPENENV_MODEL` | served model id; read from the engine when it serves exactly one | +| `E2B_API_KEY` | offer the `e2b` sandbox | +| `MODAL_TOKEN_ID`, `MODAL_TOKEN_SECRET` | offer the `modal` sandbox | + +The engine **must** be started with: + +``` +--return-tokens-as-token-ids --logprobs-mode processed_logprobs +``` + +Without them it answers every request normally and returns no token ids, so captured rollouts are +empty and nothing reports an error. The server refuses to start rather than let that happen. + +## Use + +```python +from harbor_env import HarborEnv + +with HarborEnv(base_url="http://localhost:8000") as env: + split = env.splits()[0]["name"] + result = env.run_rollout(split=split, task_index=0, harness="opencode", sandbox="e2b") + print(result.reward, len(result.turns)) +``` + +`harness` and `sandbox` are per-call, so consecutive rollouts can use different agents and different +backends against the same server. + +## Notes + +Two ports when run locally: the env server faces trainers and browsers, the capture proxy faces the +sandbox and is the only thing published. On a hosted platform there is one port and one public URL, +so the proxy is mounted on the env server's own app at `/capture` and nothing is forwarded. It still +refuses callers without a registered session id, which is what keeps a public mount from being an +open relay. + +Capture quality and task success are independent. A rollout can be captured perfectly and score 0 +because the model was wrong; a reward of 1 with unusable capture is worse than useless for training. +Both are reported separately. diff --git a/envs/harbor_env/README.md b/envs/harbor_env/README.md index 275a5e35a..93459b1d0 100644 --- a/envs/harbor_env/README.md +++ b/envs/harbor_env/README.md @@ -51,8 +51,11 @@ backends against the same server. ## Notes -Two ports: the env server faces trainers and browsers, the capture proxy faces the sandbox and is the -only thing forwarded publicly. +Two ports when run locally: the env server faces trainers and browsers, the capture proxy faces the +sandbox and is the only thing published. On a hosted platform there is one port and one public URL, +so the proxy is mounted on the env server's own app at `/capture` and nothing is forwarded. It still +refuses callers without a registered session id, which is what keeps a public mount from being an +open relay. Capture quality and task success are independent. A rollout can be captured perfectly and score 0 because the model was wrong; a reward of 1 with unusable capture is worse than useless for training. From c684348ad80c72eb3c4bd66d2fe59994822c1d55 Mon Sep 17 00:00:00 2001 From: adithya-s-k Date: Sun, 2 Aug 2026 19:35:55 +0000 Subject: [PATCH 07/74] tests: cover the graph, rewards, seams, discovery and rendering Takes harbor coverage from 19 tests to 112, no credentials or network needed. The areas chosen are the ones that fail silently rather than loudly: a bug in any of them produces plausible training data instead of an error. graph prefix linking, roots, forks, discarded branches, loss masking rewards the explicit selection rule, including 0.0 vs None seams model-name normalisation, session threading, dialect coverage discovery ordering stability, the symlink regression, spec classification validation ingest checks and sandbox SDK detection rendering result models, verdict states, contract.json Two real bugs surfaced while writing them. Google streaming requests were misclassified. `detect` tested `"generateContent" in path`, but the streaming variant capitalises the G, so every `:streamGenerateContent` call fell through to chat-completions and would have been parsed by the wrong transformer. `wants_stream` already lowercased the path; `detect` did not. gemini-cli passed the sweep because it used the non-streaming route. Anthropic tool calls were absent from results. `models._tool_calls` read only the chat-completions `tool_calls` key, so claude-code's `tool_use` content blocks never reached `HarborTurn`, leaving `contract.json` and the rendered conversation showing an agent that produced text and took no actions. Two expectations of mine were wrong rather than the code, and are now pinned as behaviour: an empty served model raises instead of returning an empty string, and a turn that sampled nothing warns rather than invalidating the rollout. --- src/openenv/core/harness/capture/detection.py | 8 +- src/openenv/harbor/models.py | 15 + tests/envs/test_harbor_capture_graph.py | 185 ++++++++++ tests/envs/test_harbor_capture_validate.py | 148 ++++++++ tests/envs/test_harbor_result_rendering.py | 330 ++++++++++++++++++ tests/envs/test_harbor_rollout_contract.py | 131 +++++++ tests/envs/test_harbor_tasks_and_dialects.py | 184 ++++++++++ 7 files changed, 999 insertions(+), 2 deletions(-) create mode 100644 tests/envs/test_harbor_capture_graph.py create mode 100644 tests/envs/test_harbor_capture_validate.py create mode 100644 tests/envs/test_harbor_result_rendering.py create mode 100644 tests/envs/test_harbor_rollout_contract.py create mode 100644 tests/envs/test_harbor_tasks_and_dialects.py diff --git a/src/openenv/core/harness/capture/detection.py b/src/openenv/core/harness/capture/detection.py index f20742ca8..076524541 100644 --- a/src/openenv/core/harness/capture/detection.py +++ b/src/openenv/core/harness/capture/detection.py @@ -39,8 +39,12 @@ def detect(path: str, headers: dict[str, str], body: dict[str, Any]) -> APIType: if "/v1/responses" in path: return APIType.OPENAI_RESPONSES # Google puts the method in the path: `/v1beta/models/{model}:generateContent`, and its - # streaming variant `:streamGenerateContent`. Both contain this substring. - if "generateContent" in path: + # streaming variant `:streamGenerateContent`. The comparison must be case-insensitive: the + # streaming form capitalises the G, so a literal `"generateContent" in path` matches the + # non-streaming route and misses every streaming one. A missed Google request is then handed to + # the chat-completions transformer, which finds no `messages` and produces a valid-looking + # response in the wrong envelope, and gemini-cli reports that as nothing at all. + if "generatecontent" in path.lower(): return APIType.GOOGLE if "anthropic-version" in {k.lower() for k in headers}: diff --git a/src/openenv/harbor/models.py b/src/openenv/harbor/models.py index b66645589..bdf795329 100644 --- a/src/openenv/harbor/models.py +++ b/src/openenv/harbor/models.py @@ -168,6 +168,21 @@ def _tool_calls(response: dict[str, Any]) -> list[dict[str, Any]]: "arguments": function.get("arguments", call.get("arguments", "")), } ) + # Anthropic does not use `tool_calls`: it puts tool use in the content block list. Reading only + # the chat-completions shape leaves claude-code's actions out of the result entirely, so + # `contract.json` and the rendered conversation both show it as a stream of text that did + # nothing. + content = response.get("content") + if isinstance(content, list): + for block in content: + if ( + isinstance(block, dict) + and block.get("type") == "tool_use" + and block.get("name") + ): + out.append( + {"name": str(block["name"]), "arguments": block.get("input", "")} + ) return out diff --git a/tests/envs/test_harbor_capture_graph.py b/tests/envs/test_harbor_capture_graph.py new file mode 100644 index 000000000..dd434cdeb --- /dev/null +++ b/tests/envs/test_harbor_capture_graph.py @@ -0,0 +1,185 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""The rollout graph: how captured calls become training sequences. + +This is the load-bearing piece of the capture layer. Turns are linked by exact token prefix and +nothing else, so every structural claim a trainer relies on (this is one conversation, this branch +was abandoned, these tokens are the model's own output) is a consequence of the linking rule. A bug +here does not crash: it silently produces training data that misattributes tokens. +""" + +from __future__ import annotations + +import pytest + +graph_mod = pytest.importorskip("openenv.core.harness.capture.graph") + +RolloutGraph = graph_mod.RolloutGraph +TurnNode = graph_mod.TurnNode +common_prefix_len = graph_mod.common_prefix_len + + +def node(node_id: str, prompt: list[int], sampled: list[int], **kwargs) -> TurnNode: + return TurnNode( + node_id=node_id, + prompt_ids=prompt, + sampled_ids=sampled, + sampled_logprobs=[-0.1] * len(sampled), + **kwargs, + ) + + +def chain( + graph: RolloutGraph, *lengths: int, base: int = 0, context: int = 1 +) -> list[TurnNode]: + """Add a linear conversation. + + `context` is how many tokens the harness inserts between turns: a tool result plus the chat + template's scaffolding. Real rollouts always have some, and turn boundaries are derived from + those mask-0 runs, so a chain built without them is not representative. + """ + added, prompt = [], [base] + for i, length in enumerate(lengths): + if i: + prompt = prompt + [base + 900 + i] * context + sampled = list(range(base + 100 + i * 10, base + 100 + i * 10 + length)) + current = graph.add_turn(node(f"n{base}_{i}", list(prompt), sampled)) + added.append(current) + prompt = current.end_ids + return added + + +# --- the linking rule ------------------------------------------------------- +def test_common_prefix_len(): + assert common_prefix_len([1, 2, 3], [1, 2, 9]) == 2 + assert common_prefix_len([1, 2], [1, 2, 3]) == 2 + assert common_prefix_len([], [1]) == 0 + assert common_prefix_len([1], [2]) == 0 + + +def test_a_turn_whose_prompt_extends_another_becomes_its_child(): + g = RolloutGraph() + first, second = chain(g, 3, 4) + assert second.parent_id == first.node_id + assert g.children(first.node_id) == [second] + + +def test_an_unrelated_prompt_starts_a_new_root(): + """A different system prompt breaks the prefix, which is what makes it a separate conversation.""" + g = RolloutGraph() + chain(g, 3) + chain(g, 3, base=5000) + assert len(g.roots()) == 2 + + +def test_linking_ignores_arrival_order(): + """Order of arrival must not decide structure; only the token prefix may.""" + g = RolloutGraph() + parent = g.add_turn(node("p", [1, 2], [3, 4])) + other = g.add_turn(node("other", [9, 9], [8])) + child = g.add_turn(node("c", [1, 2, 3, 4], [5])) + assert child.parent_id == parent.node_id + assert other.parent_id is None + + +# --- forks and discards ----------------------------------------------------- +def test_two_children_of_one_node_are_a_fork(): + g = RolloutGraph() + parent = g.add_turn(node("p", [1], [2, 3])) + g.add_turn(node("a", [1, 2, 3], [4])) + g.add_turn(node("b", [1, 2, 3], [5])) + forks = g.forks() + assert len(forks) == 1 + assert forks[0][0] == parent.node_id + assert len(forks[0][1]) == 2 + + +def test_an_abandoned_branch_is_discarded_and_the_continued_one_is_not(): + """A retry the agent walked away from must not be trained with the task's reward.""" + g = RolloutGraph() + g.add_turn(node("p", [1], [2, 3])) + g.add_turn(node("dead", [1, 2, 3], [4])) # never extended + g.add_turn(node("live", [1, 2, 3], [5])) + g.add_turn(node("live2", [1, 2, 3, 5], [6])) # extends `live` + + discarded = {n.node_id for n in g.discarded_nodes()} + assert "dead" in discarded + assert "live" not in discarded and "live2" not in discarded + + +# --- sequences, the actual training rows ------------------------------------- +def test_sequence_masks_prompt_and_marks_only_sampled_tokens(): + g = RolloutGraph() + first, second = chain(g, 2, 3) # noqa: F841 + seq = g.sequence_for(second.node_id) + + assert len(seq.input_ids) == len(seq.loss_mask) == len(seq.logprobs) + # Exactly the sampled tokens are trainable: 2 from the first turn, 3 from the second. + assert sum(seq.loss_mask) == 5 + assert seq.n_trainable == 5 + trainable = [i for i, m in zip(seq.input_ids, seq.loss_mask) if m] + assert trainable == first.sampled_ids + second.sampled_ids + + +def test_turn_lengths_match_what_each_turn_sampled(): + g = RolloutGraph() + turns = chain(g, 2, 3, 4) + seq = g.sequence_for(turns[-1].node_id) + assert seq.turn_lengths() == [2, 3, 4] + + +def test_turn_lengths_merge_when_a_turn_adds_no_context(): + """A documented edge, not an endorsement. + + Turn boundaries are runs of mask-0 tokens, so two turns with nothing between them read as one. + Every real harness inserts at least the chat template's turn scaffolding, which is why this has + never been observed, but `turns_from_document` zips `node_ids` against `turn_lengths` and a + shorter list would silently drop turns rather than fail. Pinned so a future change to the + linking rule surfaces here instead of in training data. + """ + g = RolloutGraph() + turns = chain(g, 2, 3, context=0) + seq = g.sequence_for(turns[-1].node_id) + assert seq.turn_lengths() == [5] + assert len(seq.node_ids) == 2 + + +def test_context_tokens_are_conditioned_on_but_never_trained(): + """Tool results are real tokens the model saw and did not produce: mask 0, not absent.""" + g = RolloutGraph() + parent = g.add_turn(node("p", [1, 2], [3])) + # The harness inserted a tool result (99) between the turns. + child = g.add_turn(node("c", [1, 2, 3, 99], [4])) + + assert child.context_ids(parent) == [99] + seq = g.sequence_for(child.node_id) + assert 99 in seq.input_ids + assert seq.loss_mask[seq.input_ids.index(99)] == 0 + + +def test_one_sequence_per_leaf(): + g = RolloutGraph() + g.add_turn(node("p", [1], [2])) + g.add_turn(node("a", [1, 2], [3])) + g.add_turn(node("b", [1, 2], [4])) + assert len(g.sequences()) == len(g.leaves()) == 2 + + +def test_stats_report_the_shape(): + g = RolloutGraph() + chain(g, 2, 2) + chain(g, 2, base=5000) + stats = g.stats() + assert stats["n_turns"] == 3 + assert stats["n_roots"] == 2 + assert stats["n_leaves"] == 2 + + +def test_empty_graph_is_not_an_error(): + g = RolloutGraph() + assert g.nodes() == [] and g.roots() == [] and g.sequences() == [] + assert g.stats()["n_turns"] == 0 diff --git a/tests/envs/test_harbor_capture_validate.py b/tests/envs/test_harbor_capture_validate.py new file mode 100644 index 000000000..60a9e691d --- /dev/null +++ b/tests/envs/test_harbor_capture_validate.py @@ -0,0 +1,148 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Validation on ingest, and the capability checks that run before a rollout starts. + +Both exist to convert a silent wrong answer into a loud one. A turn whose logprobs are misaligned +has to be caught while we still know which turn it was; a sandbox that cannot be constructed has to +be caught before it is offered rather than 90 seconds into a run. +""" + +from __future__ import annotations + +import sys +import types + +import pytest + +validate = pytest.importorskip("openenv.core.harness.capture.validate") +capabilities = pytest.importorskip("openenv.harbor.capabilities") + +check_turn = validate.check_turn + + +def codes(report) -> set[str]: + return {f.code for f in report.findings} + + +# --- per-turn validation ---------------------------------------------------- +def test_a_well_formed_turn_passes(): + report = check_turn([1, 2, 3], [4, 5], [-0.1, -0.2], finish_reason="stop") + assert report.ok + + +def test_missing_prompt_ids_is_fatal(): + """The endpoint was started without token-id capture: every rebuilt row would be empty.""" + report = check_turn([], [4], [-0.1]) + assert not report.ok + assert "no_prompt_ids" in codes(report) + + +def test_logprob_count_must_match_sampled_count(): + """Off-by-one here silently trains on the wrong token's probability.""" + report = check_turn([1], [4, 5, 6], [-0.1, -0.2]) + assert not report.ok + + +def test_a_turn_that_sampled_nothing_is_reported_but_not_fatal(): + """Reported, not rejected: a model can legitimately stop without emitting a token. + + `ok` means usable, so an empty completion warns rather than invalidating the rollout. It still + has to be visible, because a run of these means the agent is looping without producing anything. + """ + report = check_turn([1, 2], [], []) + assert report.ok + assert "no_sampled_ids" in codes(report) + + +def test_findings_name_the_turn(): + """A misalignment is only actionable if you know which call produced it.""" + report = check_turn([], [1], [-0.1], index=7) + assert any("turn 7" in str(f) for f in report.findings) + + +# --- sandbox capability ----------------------------------------------------- +def _module_with(**flags): + module = types.ModuleType("fake_backend_module") + for key, value in flags.items(): + setattr(module, key, value) + sys.modules[module.__name__] = module + return module + + +def test_missing_sdk_is_detected_from_the_backends_own_flag(): + """Harbor guards each SDK with a module-level `_HAS_X` and raises from `__init__`. + + So the module imports, the class loads, and the check passes, with the failure arriving only + once a rollout tries to build a sandbox, where it reads as a broken rollout rather than a + missing dependency. This is the case that shipped a Space offering `e2b` it could never run. + """ + module = _module_with(_HAS_E2B=False) + cls = type("E2BEnvironment", (), {"__module__": module.__name__}) + detail = capabilities._missing_sdk(cls) + assert detail and "e2b" in detail + assert "harbor[cloud]" in detail or "openenv[harbor]" in detail + + +def test_present_sdk_reports_nothing(): + module = _module_with(_HAS_E2B=True) + cls = type("E2BEnvironment", (), {"__module__": module.__name__}) + assert capabilities._missing_sdk(cls) == "" + + +def test_a_backend_without_flags_is_not_assumed_broken(): + module = _module_with(SOMETHING_ELSE=1) + cls = type("Plain", (), {"__module__": module.__name__}) + assert capabilities._missing_sdk(cls) == "" + + +def test_several_missing_extras_are_all_named(): + module = _module_with(_HAS_MODAL=False, _HAS_DOCKERFILE_PARSE=False) + cls = type("ModalEnvironment", (), {"__module__": module.__name__}) + detail = capabilities._missing_sdk(cls) + assert "modal" in detail and "dockerfile_parse" in detail + + +def test_an_unimportable_module_is_not_a_crash(): + cls = type("Ghost", (), {"__module__": "module.that.does.not.exist"}) + assert capabilities._missing_sdk(cls) == "" + + +def test_unknown_sandbox_names_are_rejected_by_name(): + status = capabilities.check_sandbox("not-a-real-backend") + assert status.available is False + assert "unknown" in status.detail.lower() or "harbor" in status.detail.lower() + + +# --- capability reporting --------------------------------------------------- +def test_render_says_why_a_sandbox_is_unavailable(): + """The commonest cause of a rollout dying 90s in is a missing key, so it belongs at startup.""" + caps = capabilities.Capabilities( + sandboxes=[ + capabilities.SandboxStatus("e2b", True), + capabilities.SandboxStatus("daytona", False, "DAYTONA_API_KEY is not set"), + ] + ) + out = caps.render() + assert "DAYTONA_API_KEY is not set" in out + assert caps.available_sandboxes == ["e2b"] + + +def test_render_warns_when_nothing_is_usable(): + caps = capabilities.Capabilities( + sandboxes=[capabilities.SandboxStatus("e2b", False, "no key")] + ) + assert "WARNING" in caps.render() + + +def test_capabilities_serialise_for_the_wire(): + caps = capabilities.Capabilities( + sandboxes=[capabilities.SandboxStatus("e2b", True)], + llm={"model": "m", "ok": True}, + ) + payload = caps.to_dict() + assert set(payload) == {"harnesses", "sandboxes", "datasets", "llm"} + assert payload["llm"]["ok"] is True diff --git a/tests/envs/test_harbor_result_rendering.py b/tests/envs/test_harbor_result_rendering.py new file mode 100644 index 000000000..28b1428e7 --- /dev/null +++ b/tests/envs/test_harbor_result_rendering.py @@ -0,0 +1,330 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Turning a capture document into a result, and a result into something readable. + +Two dialect families put tool calls in different places (`tool_calls` vs `tool_use` content blocks), +so a reader that knows only one shows an agent as a stream of text with no visible actions. And a +forked conversation appears once per path, which rendered as several near-identical transcripts all +claiming to be the main one. +""" + +from __future__ import annotations + +import json + +import pytest + +models = pytest.importorskip("openenv.harbor.models") +ui = pytest.importorskip("openenv.harbor.ui") + +conversations_from_document = models.conversations_from_document +turns_from_document = models.turns_from_document + + +def document(sequences, turns): + return {"sequences": sequences, "turns": turns} + + +# --- conversations ---------------------------------------------------------- +def test_a_forked_root_yields_one_conversation_not_one_per_path(): + """The regression: two paths through one root rendered as two 'main conversation' blocks.""" + doc = document( + sequences=[ + {"root_id": "r1", "role": "agent", "node_ids": ["a"], "n_turns": 1}, + {"root_id": "r1", "role": "agent", "node_ids": ["a", "b"], "n_turns": 2}, + ], + turns=[ + { + "node_id": "a", + "request_messages": [{"role": "user", "content": "hi"}], + "response_message": {"content": "one"}, + }, + { + "node_id": "b", + "request_messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "one"}, + ], + "response_message": {"content": "two"}, + }, + ], + ) + convos = conversations_from_document(doc) + assert len(convos) == 1 + assert convos[0].n_turns == 2, "the longest path is the complete one" + + +def test_separate_roots_stay_separate(): + doc = document( + sequences=[ + {"root_id": "r1", "role": "agent", "node_ids": ["a"], "n_turns": 1}, + {"root_id": "r2", "role": "auxiliary", "node_ids": ["b"], "n_turns": 1}, + ], + turns=[ + { + "node_id": "a", + "request_messages": [{"role": "user", "content": "task"}], + "response_message": {"content": "working"}, + }, + { + "node_id": "b", + "request_messages": [{"role": "user", "content": "who next?"}], + "response_message": {"content": "agent"}, + }, + ], + ) + convos = conversations_from_document(doc) + assert {c.role for c in convos} == {"agent", "auxiliary"} + + +def test_a_conversation_keeps_the_system_prompt_and_tool_results(): + doc = document( + sequences=[{"root_id": "r", "role": "agent", "node_ids": ["a"], "n_turns": 1}], + turns=[ + { + "node_id": "a", + "request_messages": [ + {"role": "system", "content": "You are an assistant."}, + {"role": "user", "content": "count rows"}, + {"role": "tool", "content": "42"}, + ], + "response_message": {"content": "42 rows"}, + } + ], + ) + roles = [m["role"] for m in conversations_from_document(doc)[0].messages] + assert roles == ["system", "user", "tool", "assistant"] + + +def test_sequences_without_nodes_or_messages_are_skipped(): + doc = document( + sequences=[{"root_id": "r", "role": "agent", "node_ids": [], "n_turns": 0}], + turns=[], + ) + assert conversations_from_document(doc) == [] + + +# --- turns ------------------------------------------------------------------ +def _agent_doc(response): + return document( + sequences=[ + { + "root_id": "r", + "role": "agent", + "node_ids": ["a"], + "n_turns": 1, + "input_ids": [1, 2, 3], + "loss_mask": [0, 1, 1], + "logprobs": [0.0, -0.1, -0.2], + "prompt_len": 1, + "turn_lengths": [2], + } + ], + turns=[ + { + "node_id": "a", + "finish_reason": "stop", + "n_tools": 3, + "response_message": response, + } + ], + ) + + +def test_turn_text_and_tool_calls_from_chat_completions(): + turns = turns_from_document( + _agent_doc( + { + "content": "Reading the file.", + "tool_calls": [ + {"function": {"name": "bash", "arguments": '{"cmd":"ls"}'}} + ], + } + ) + ) + assert turns[0].text == "Reading the file." + assert turns[0].tool_calls == [{"name": "bash", "arguments": '{"cmd":"ls"}'}] + + +def test_turn_text_and_tool_calls_from_anthropic_blocks(): + """claude-code puts tool use in content blocks; reading only `tool_calls` shows no actions.""" + turns = turns_from_document( + _agent_doc( + { + "content": [ + {"type": "text", "text": "Checking."}, + {"type": "tool_use", "name": "Bash", "input": {"command": "ls"}}, + ], + } + ) + ) + assert turns[0].text == "Checking." + assert turns[0].tool_calls[0]["name"] == "Bash" + + +def test_a_turn_with_no_response_is_still_a_turn(): + turns = turns_from_document(_agent_doc({})) + assert len(turns) == 1 and turns[0].text == "" and turns[0].tool_calls == [] + + +def test_only_agent_sequences_become_turns(): + """An auxiliary call must never be credited with the reward for solving the task.""" + doc = _agent_doc({"content": "x"}) + doc["sequences"][0]["role"] = "auxiliary" + assert turns_from_document(doc) == [] + + +# --- rendering -------------------------------------------------------------- +@pytest.mark.parametrize( + "result,marker", + [ + ({"ok": True, "reward": 1.0}, "Solved"), + ({"ok": True, "reward": 0.0}, "Not solved"), + ({"ok": True, "reward": None}, "Not graded"), + ({"ok": False, "reward": None, "exception_type": "Boom"}, "Failed"), + ], +) +def test_every_verdict_state_renders(result, marker): + assert marker in ui._result_html({**result, "turns": []}) + + +def test_ungraded_shows_a_dash_rather_than_a_zero(): + """A dead sandbox rendered as 0.00 reads as the model getting the answer wrong.""" + out = ui._result_html({"ok": True, "reward": None, "turns": []}) + assert "0.00" not in out + + +def test_findings_are_grouped_by_severity(): + out = ui._findings_html(["[FATAL] gone", "[WARN] odd", "[INFO] fyi"]) + assert "FATAL" in out and "WARN" in out and "INFO" in out + assert out.index("FATAL") < out.index("WARN"), "worst first" + + +def test_no_findings_renders_nothing(): + assert ui._findings_html([]) == "" + + +def test_conversation_labels_are_unambiguous_when_several_exist(): + convos = [ + {"role": "agent", "n_turns": 1, "messages": [{"role": "user", "content": "a"}]}, + {"role": "agent", "n_turns": 1, "messages": [{"role": "user", "content": "b"}]}, + ] + out = ui._conversation_html({"conversations": convos}) + assert out.count("main conversation") == 0 + assert "conversation 1 of 2" in out and "conversation 2 of 2" in out + + +def test_turns_html_does_not_show_the_tools_offered_count(): + """That number is a property of the harness, identical on every row, and told nobody anything.""" + out = ui._turns_html( + { + "turns": [ + { + "turn": 0, + "completion_token_ids": [1, 2], + "per_token_logps": [-0.1, -0.2], + "tool_calls": [{"name": "bash", "arguments": "ls"}], + "n_tools": 24, + "finish_reason": "tool_calls", + } + ] + } + ) + assert "24 tools" not in out + assert "bash" in out and "confidence" in out + + +def test_escaping_prevents_markup_injection_from_a_model_reply(): + out = ui._conversation_html( + { + "conversations": [ + { + "role": "agent", + "n_turns": 1, + "messages": [ + {"role": "assistant", "content": ""} + ], + } + ] + } + ) + assert "