Spaces:
Sleeping
Sleeping
| # 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. | |
| """Protocol One OpenEnv Environment. | |
| One instance == one training rollout. On reset(), we: | |
| 1. Ask the Designer for a (possibly mutated) copy of SPEC. | |
| 2. Build a fresh MockProtocolServer off that spec, wrap in a TestClient. | |
| 3. Reset belief_graph, probe_history, probe count. | |
| 4. Return the initial observation with tool docs and a starter token. | |
| Each step() receives a ProtocolOneAction with a tool discriminator and | |
| dispatches to probe / update_model / finalize. Reward is only set at | |
| terminal step — we deliberately avoid per-probe rewards to prevent the | |
| model from reward-hacking by spamming update_model. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import os | |
| from typing import Any | |
| from uuid import uuid4 | |
| from fastapi.testclient import TestClient | |
| from openenv.core.env_server.interfaces import Environment | |
| from openenv.core.env_server.types import State | |
| try: | |
| from ..models import ProtocolOneAction, ProtocolOneObservation | |
| except ImportError: | |
| from models import ProtocolOneAction, ProtocolOneObservation | |
| from .designer import Designer | |
| from .matcher import score | |
| from .protocol_server import MockProtocolServer, create_server | |
| from .spec import INITIAL_TOKEN, SPEC | |
| MAX_PROBES_PER_EPISODE = int(os.environ.get("MAX_PROBES_PER_EPISODE", "12")) | |
| RESPONSE_TRUNCATE_CHARS = 2000 | |
| INSTRUCTIONS_TEMPLATE = """You are reverse-engineering an undocumented HTTP API. | |
| Task: discover the API's structure and produce a complete belief graph of its \ | |
| endpoints, resources, authentication, and state transitions. | |
| Base URL: (implicit; the probe tool handles it — paths start with /) | |
| Starting auth token: Bearer {token} | |
| Additional tokens with other scopes may exist. Discovering them is part of the \ | |
| task. You have {max_probes} probes; unauthenticated probes are cheap ways to \ | |
| discover endpoints (401 vs 404 is informative). | |
| Tools: | |
| 1. probe(method, path, headers, body): send an HTTP request, get back status + body | |
| 2. update_model(delta): merge findings into your belief graph (incremental) | |
| 3. finalize(final_belief_graph=None): commit and end the episode (reward is computed here) | |
| Belief-graph shape (emit via update_model): | |
| {{ | |
| "endpoints": [ | |
| {{"method": "GET", "path": "/users", | |
| "auth_required": true, "auth_scope": "users:read", | |
| "params": [{{"name": "limit", "type": "int", "location": "query"}}], | |
| "responses": {{"200": {{"shape": "list<User>"}}, "401": {{"shape": "error"}}}}}} | |
| ], | |
| "resources": [ | |
| {{"name": "User", | |
| "fields": [{{"name": "id", "type": "string"}}, {{"name": "email", "type": "string"}}], | |
| "state_machine": {{"states": ["active", "suspended"], | |
| "transitions": [{{"from": "active", "to": "suspended", | |
| "via": "POST /users/{{id}}/suspend"}}]}}}} | |
| ], | |
| "auth": {{"type": "bearer", "scopes_observed": ["users:read"]}} | |
| }} | |
| You are scored on: | |
| - Endpoints correctly identified (method + path) | |
| - Correct auth requirements and scopes | |
| - Correct params and response codes | |
| - Correct resource fields and state-machine transitions | |
| False claims reduce your score. When uncertain, omit rather than guess. Call \ | |
| finalize() before the probe budget runs out; you get zero reward if you never finalize. | |
| Begin probing.""" | |
| class ProtocolOneEnvironment(Environment): | |
| """OpenEnv environment for the Protocol One task.""" | |
| # Each WebSocket session gets its own Environment instance; state never leaks | |
| # between rollouts, so we can safely run many concurrent sessions for GRPO. | |
| SUPPORTS_CONCURRENT_SESSIONS: bool = True | |
| def __init__(self) -> None: | |
| super().__init__() | |
| self._state = State(episode_id=str(uuid4()), step_count=0) | |
| # Persistent across episodes (mutation counter lives in the designer) | |
| self.designer = Designer(SPEC) | |
| # Refreshed per-episode in reset() | |
| self.current_spec: dict = SPEC | |
| self.server: MockProtocolServer | None = None | |
| self.client: TestClient | None = None | |
| self.belief_graph: dict[str, Any] = {"endpoints": [], "resources": [], "auth": {}} | |
| self.probe_history: list[dict[str, Any]] = [] | |
| self.probes_used: int = 0 | |
| self.done: bool = False | |
| self.last_reward: float = 0.0 | |
| # ------------------------------------------------------------------ | |
| # OpenEnv contract | |
| # ------------------------------------------------------------------ | |
| def reset(self, **kwargs: Any) -> ProtocolOneObservation: | |
| self._state = State(episode_id=str(uuid4()), step_count=0) | |
| self.current_spec = self.designer.maybe_mutate() | |
| self.server = create_server(self.current_spec) | |
| self.client = TestClient(self.server.app) | |
| self.belief_graph = {"endpoints": [], "resources": [], "auth": {}} | |
| self.probe_history = [] | |
| self.probes_used = 0 | |
| self.done = False | |
| self.last_reward = 0.0 | |
| return ProtocolOneObservation( | |
| text=INSTRUCTIONS_TEMPLATE.format( | |
| token=INITIAL_TOKEN, | |
| max_probes=MAX_PROBES_PER_EPISODE, | |
| ), | |
| probes_used=0, | |
| probes_remaining=MAX_PROBES_PER_EPISODE, | |
| belief_graph_stats=self._stats(), | |
| done=False, | |
| reward=None, | |
| metadata={ | |
| "episode_id": self._state.episode_id, | |
| "mutation_log": self.designer.last_mutation_log, | |
| }, | |
| ) | |
| def step(self, action: ProtocolOneAction, **kwargs: Any) -> ProtocolOneObservation: # type: ignore[override] | |
| # Defensive: if a client steps without having called reset() first | |
| # (e.g., in a stateless HTTP smoke test), initialise episode state | |
| # on demand rather than crashing. | |
| if self.client is None: | |
| self.reset() | |
| if self.done: | |
| return self._obs("Episode already ended.", done=True) | |
| self._state.step_count += 1 | |
| tool = action.tool | |
| args = action.args or {} | |
| if tool == "probe": | |
| return self._handle_probe(args) | |
| if tool == "update_model": | |
| return self._handle_update(args) | |
| if tool == "finalize": | |
| return self._handle_finalize(args) | |
| return self._obs( | |
| f"Invalid tool: {tool!r}. Must be 'probe', 'update_model', or 'finalize'.", | |
| done=False, | |
| ) | |
| def state(self) -> State: | |
| return self._state | |
| # ------------------------------------------------------------------ | |
| # Tool handlers | |
| # ------------------------------------------------------------------ | |
| def _handle_probe(self, args: dict[str, Any]) -> ProtocolOneObservation: | |
| if self.probes_used >= MAX_PROBES_PER_EPISODE: | |
| return self._end_episode("Probe budget exhausted.") | |
| method = str(args.get("method", "GET")).upper() | |
| path = str(args.get("path", "/")) | |
| headers = args.get("headers") or {} | |
| body = args.get("body") | |
| if not isinstance(headers, dict): | |
| headers = {} | |
| self.probes_used += 1 | |
| assert self.client is not None | |
| try: | |
| if body is not None and method in ("GET", "DELETE"): | |
| # Silently drop body for methods that don't use it; saves the | |
| # agent from confusing 400s. | |
| body = None | |
| if body is not None and not isinstance(body, (dict, list, str)): | |
| body = None | |
| resp = self.client.request( | |
| method=method, | |
| url=path, | |
| headers={str(k): str(v) for k, v in headers.items()}, | |
| json=body if body is not None else None, | |
| ) | |
| resp_body = resp.text | |
| if len(resp_body) > RESPONSE_TRUNCATE_CHARS: | |
| resp_body = resp_body[:RESPONSE_TRUNCATE_CHARS] + "… [truncated]" | |
| text = ( | |
| f"[Probe {self.probes_used}/{MAX_PROBES_PER_EPISODE}] " | |
| f"HTTP {resp.status_code}\nBody: {resp_body}" | |
| ) | |
| except Exception as e: | |
| text = f"[Probe {self.probes_used}/{MAX_PROBES_PER_EPISODE}] Request error: {e}" | |
| self.probe_history.append({"method": method, "path": path, "response_text": text}) | |
| return self._obs(text, done=False) | |
| def _handle_update(self, args: dict[str, Any]) -> ProtocolOneObservation: | |
| delta = args.get("delta") | |
| if delta is None: | |
| # Be forgiving: sometimes the agent emits the belief graph at | |
| # the top level instead of wrapping in {"delta": ...}. | |
| delta = args | |
| if isinstance(delta, str): | |
| # Forgive stringified JSON | |
| try: | |
| delta = json.loads(delta) | |
| except Exception: | |
| return self._obs( | |
| f"update_model: could not parse delta as JSON. " | |
| f"Expected {{'endpoints': [...], 'resources': [...], 'auth': {{...}}}}", | |
| done=False, | |
| ) | |
| if not isinstance(delta, dict): | |
| return self._obs( | |
| "update_model: delta must be a JSON object with endpoints/resources/auth keys.", | |
| done=False, | |
| ) | |
| self._merge_delta(delta) | |
| stats = self._stats() | |
| return self._obs( | |
| f"Belief graph updated. Stats: {stats}. Call finalize() when done.", | |
| done=False, | |
| ) | |
| def _handle_finalize(self, args: dict[str, Any]) -> ProtocolOneObservation: | |
| final = args.get("final_belief_graph") | |
| if isinstance(final, dict): | |
| # Replace, not merge — agent is committing a final answer. | |
| self.belief_graph = { | |
| "endpoints": final.get("endpoints", []), | |
| "resources": final.get("resources", []), | |
| "auth": final.get("auth", {}), | |
| } | |
| return self._end_episode("Finalized.") | |
| # ------------------------------------------------------------------ | |
| # Internal | |
| # ------------------------------------------------------------------ | |
| def _merge_delta(self, delta: dict) -> None: | |
| # endpoints: upsert by (method, path) key | |
| if isinstance(delta.get("endpoints"), list): | |
| by_key: dict[tuple[str, str], dict] = {} | |
| for e in self.belief_graph["endpoints"]: | |
| by_key[(str(e.get("method", "")).upper(), str(e.get("path", "")))] = e | |
| for new_e in delta["endpoints"]: | |
| if not isinstance(new_e, dict): | |
| continue | |
| key = (str(new_e.get("method", "")).upper(), str(new_e.get("path", ""))) | |
| if key in by_key: | |
| by_key[key] = {**by_key[key], **new_e} | |
| else: | |
| by_key[key] = new_e | |
| self.belief_graph["endpoints"] = list(by_key.values()) | |
| # resources: upsert by name | |
| if isinstance(delta.get("resources"), list): | |
| by_name: dict[str, dict] = {} | |
| for r in self.belief_graph["resources"]: | |
| by_name[str(r.get("name", ""))] = r | |
| for new_r in delta["resources"]: | |
| if not isinstance(new_r, dict): | |
| continue | |
| name = str(new_r.get("name", "")) | |
| if name in by_name: | |
| by_name[name] = {**by_name[name], **new_r} | |
| else: | |
| by_name[name] = new_r | |
| self.belief_graph["resources"] = list(by_name.values()) | |
| # auth: shallow merge, dedup scopes_observed | |
| if isinstance(delta.get("auth"), dict): | |
| new_auth = {**self.belief_graph.get("auth", {}), **delta["auth"]} | |
| if isinstance(new_auth.get("scopes_observed"), list): | |
| new_auth["scopes_observed"] = sorted({ | |
| s for s in new_auth["scopes_observed"] if isinstance(s, str) | |
| }) | |
| self.belief_graph["auth"] = new_auth | |
| def _stats(self) -> dict[str, int]: | |
| auth = self.belief_graph.get("auth", {}) or {} | |
| scopes = auth.get("scopes_observed", []) or [] | |
| return { | |
| "endpoints": len(self.belief_graph.get("endpoints", [])), | |
| "resources": len(self.belief_graph.get("resources", [])), | |
| "auth_scopes_observed": len(scopes) if isinstance(scopes, list) else 0, | |
| } | |
| def _obs( | |
| self, | |
| text: str, | |
| done: bool, | |
| reward: float | None = None, | |
| extra_metadata: dict | None = None, | |
| ) -> ProtocolOneObservation: | |
| metadata: dict[str, Any] = {"episode_id": self._state.episode_id} | |
| if extra_metadata: | |
| metadata.update(extra_metadata) | |
| return ProtocolOneObservation( | |
| text=text, | |
| probes_used=self.probes_used, | |
| probes_remaining=max(0, MAX_PROBES_PER_EPISODE - self.probes_used), | |
| belief_graph_stats=self._stats(), | |
| done=done, | |
| reward=reward, | |
| metadata=metadata, | |
| ) | |
| def _end_episode(self, reason: str) -> ProtocolOneObservation: | |
| result = score(self.belief_graph, self.current_spec) | |
| self.last_reward = result.total | |
| self.done = True | |
| breakdown_str = ", ".join(f"{k}={v:.2f}" for k, v in sorted(result.breakdown.items())) | |
| text = ( | |
| f"{reason}\n" | |
| f"Final reward: {result.total:.3f} (endpoints_found={result.endpoints_found}" | |
| f"/{result.endpoints_total}, false_claims={result.false_claims})\n" | |
| f"Breakdown: {breakdown_str}" | |
| ) | |
| return self._obs( | |
| text, | |
| done=True, | |
| reward=result.total, | |
| extra_metadata={ | |
| "breakdown": result.breakdown, | |
| "endpoints_found": result.endpoints_found, | |
| "endpoints_total": result.endpoints_total, | |
| "false_claims": result.false_claims, | |
| "mutation_log": self.designer.last_mutation_log, | |
| }, | |
| ) | |