diff --git a/pyproject.toml b/pyproject.toml index 056ebad80..427c18c1d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,7 +46,7 @@ dependencies = [ "aiolimiter>=1.2.1", "setproctitle>=1.3.0", "httpx>=0.27.0", - "aiohttp>=3.9.0", + "aiohttp>=3.14.1", "prime-pydantic-config[toml]>=0.4.3", "uvloop>=0.21.0; sys_platform != 'win32' and sys_platform != 'cygwin' and platform_python_implementation != 'PyPy'", "loguru>=0.7.0", diff --git a/tests/v1/test_e2e.py b/tests/v1/test_e2e.py index cef9bf7e1..77588af12 100644 --- a/tests/v1/test_e2e.py +++ b/tests/v1/test_e2e.py @@ -273,6 +273,12 @@ async def test_acp_resume_with_tool(run_v1, harness, harness_runtime, tmp_path): assert segments[1]["tool_outputs"] if harness.id == "rlm": assert "turns_since_last_compaction" in trace.metrics + assert all(call.acp is not None for call in trace.calls) + assert all( + trace.nodes[parent.node].sampled and node.sampled + for node in trace.nodes + for parent in node.semantic_parents + ) if harness.id == "prime-agent": lifecycle = trace.info["acp_lifecycle"]["ai.primeintellect.prime-agent"] assert len(lifecycle) == 2 diff --git a/tests/v1/test_trace.py b/tests/v1/test_trace.py index 08bca54f0..85a720356 100644 --- a/tests/v1/test_trace.py +++ b/tests/v1/test_trace.py @@ -11,7 +11,17 @@ import verifiers.v1 as vf from verifiers.v1.agent import Interaction from verifiers.v1.graph import MessageNode +from verifiers.v1.harnesses.rlm.harness import ( + RLM_SESSION_METADATA_KEY, + RLMHarness, + RLMHarnessConfig, +) from verifiers.v1.rollout import Rollout, RolloutTimeouts +from verifiers.v1.semantic import ( + ACP_EXTENSION_HEADERS, + ACP_SEMANTIC_EDGES_METADATA_KEY, + extract_acp_info, +) from verifiers.v1.types import AssistantMessage, UserMessage @@ -170,3 +180,301 @@ def test_wire_trace_round_trip(): # the env-server wire form (a plain model_dump) loads too assert vf.WireTrace.model_validate(tr.model_dump()).num_branches == 2 + + +def _semantic_edge_set() -> vf.SemanticEdgeSet: + return vf.SemanticEdgeSet( + edges=[ + vf.SemanticEdge( + source_request_id="root-turn", + target_request_id="root-compact", + type="continuation", + ), + vf.SemanticEdge( + source_request_id="root-turn", + target_request_id="child-turn", + type="subagent_call", + ), + vf.SemanticEdge( + source_request_id="child-turn", + target_request_id="root-after", + type="subagent_return", + ), + vf.SemanticEdge( + source_request_id="root-compact", + target_request_id="root-after", + type="compaction", + ), + vf.SemanticEdge( + source_request_id="root-turn", + target_request_id="root-after", + type="critic_review", + ), + ], + ) + + +def test_semantic_edges_resolve_to_message_nodes_and_round_trip(): + """Request edges resolve by exact IDs, not call adjacency or graph shape.""" + tr = vf.Trace( + agent=vf.AgentInfo(config=vf.AgentConfig()), + task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="q")), + nodes=[ + MessageNode(parent=None, message=UserMessage(content="root")), + MessageNode( + parent=0, message=AssistantMessage(content="root turn"), sampled=True + ), + MessageNode(parent=None, message=UserMessage(content="child")), + MessageNode( + parent=2, message=AssistantMessage(content="child turn"), sampled=True + ), + MessageNode(parent=None, message=UserMessage(content="summarize")), + MessageNode( + parent=4, message=AssistantMessage(content="summary"), sampled=True + ), + MessageNode(parent=None, message=UserMessage(content="resume")), + MessageNode( + parent=6, message=AssistantMessage(content="done"), sampled=True + ), + ], + ) + tr.calls = [ + vf.ModelCall( + node=1, + acp=vf.ACPInfo(request_id="root-turn"), + ), + vf.ModelCall( + node=3, + acp=vf.ACPInfo(request_id="child-turn"), + ), + vf.ModelCall( + node=5, + acp=vf.ACPInfo(request_id="root-compact"), + ), + vf.ModelCall( + node=7, + acp=vf.ACPInfo(request_id="root-after"), + ), + ] + + edge_set = _semantic_edge_set() + tr.add_semantic_edges(vf.SemanticEdgeSet(edges=edge_set.edges[:2])) + first_semantic_parents = tr.nodes[3].semantic_parents + tr.add_semantic_edges(vf.SemanticEdgeSet.model_validate(edge_set.model_dump())) + expected_parents = [ + [], + [], + [], + [vf.ParentLink(node=1, type="subagent_call")], + [], + [vf.ParentLink(node=1, type="continuation")], + [], + [ + vf.ParentLink(node=3, type="subagent_return"), + vf.ParentLink(node=5, type="compaction"), + vf.ParentLink(node=1, type="critic_review"), + ], + ] + assert [node.semantic_parents for node in tr.nodes] == expected_parents + assert tr.nodes[3].semantic_parents is first_semantic_parents + + restored = vf.WireTrace.model_validate_json(tr.model_dump_json()) + assert [node.semantic_parents for node in restored.nodes] == expected_parents + assert [call.acp for call in restored.calls] == [call.acp for call in tr.calls] + + # The base ACP layer resolves the generic edge set before harness-owned metadata. + harness = RLMHarness(RLMHarnessConfig(id="rlm")) + turn_metadata = { + ACP_SEMANTIC_EDGES_METADATA_KEY: _semantic_edge_set().model_dump(mode="json"), + RLM_SESSION_METADATA_KEY: { + "session_id": restored.id, + "metrics": {"turns": 4}, + }, + } + harness._consume_protocol_metadata(restored, turn_metadata) + harness.acp_turn_result( + restored, vf.ACPTurn(reply="done", response_metadata=turn_metadata) + ) + assert restored.metrics["turns"] == 4 + assert [node.semantic_parents for node in restored.nodes] == expected_parents + + # session/close may publish the same cumulative edge set again. + close_metadata = { + ACP_SEMANTIC_EDGES_METADATA_KEY: _semantic_edge_set().model_dump(mode="json"), + RLM_SESSION_METADATA_KEY: { + "session_id": restored.id, + "metrics": {"turns": 4}, + }, + } + harness._consume_protocol_metadata(restored, close_metadata) + harness.acp_close_result(restored, close_metadata) + assert restored.metrics["turns"] == 4 + assert [node.semantic_parents for node in restored.nodes] == expected_parents + + # A failed provider exchange and its SDK retry share one logical request ID. + restored.calls.append( + vf.ModelCall( + acp=restored.calls[0].acp, + error=vf.Error(type="E", message="x"), + ) + ) + restored.add_semantic_edges(_semantic_edge_set()) + assert [node.semantic_parents for node in restored.nodes] == expected_parents + + +def test_semantic_edge_uses_last_committed_retry_node(): + tr = vf.Trace( + agent=vf.AgentInfo(config=vf.AgentConfig()), + task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="q")), + nodes=[ + MessageNode(parent=None, message=UserMessage(content="root")), + MessageNode( + parent=0, message=AssistantMessage(content="attempt 1"), sampled=True + ), + MessageNode( + parent=0, message=AssistantMessage(content="attempt 2"), sampled=True + ), + MessageNode( + parent=None, message=AssistantMessage(content="next"), sampled=True + ), + ], + calls=[ + vf.ModelCall(node=1, acp=vf.ACPInfo(request_id="retried")), + vf.ModelCall(node=2, acp=vf.ACPInfo(request_id="retried")), + vf.ModelCall(node=3, acp=vf.ACPInfo(request_id="next")), + ], + ) + + tr.add_semantic_edges( + vf.SemanticEdgeSet( + edges=[ + vf.SemanticEdge( + source_request_id="retried", + target_request_id="next", + type="continuation", + ) + ] + ) + ) + + assert tr.nodes[3].semantic_parents == [vf.ParentLink(node=2, type="continuation")] + + +def test_semantic_edge_cycle_is_rejected_without_partial_mutation(): + tr = vf.Trace( + agent=vf.AgentInfo(config=vf.AgentConfig()), + task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="q")), + nodes=[ + MessageNode(parent=None, message=UserMessage(content="start")), + MessageNode( + parent=0, message=AssistantMessage(content="first"), sampled=True + ), + MessageNode(parent=1, message=UserMessage(content="continue")), + MessageNode( + parent=2, message=AssistantMessage(content="second"), sampled=True + ), + ], + calls=[ + vf.ModelCall(node=1, acp=vf.ACPInfo(request_id="first")), + vf.ModelCall(node=3, acp=vf.ACPInfo(request_id="second")), + ], + ) + + with pytest.raises(ValueError, match="cycle in the message graph"): + tr.add_semantic_edges( + vf.SemanticEdgeSet( + edges=[ + vf.SemanticEdge( + source_request_id="second", + target_request_id="first", + type="custom", + ) + ] + ) + ) + + assert all(not node.semantic_parents for node in tr.nodes) + + +def test_acp_info_is_validated_and_stripped(): + headers = { + "Authorization": "Bearer local", + "Idempotency-Key": "provider-key", + "X-ACP-Model-Request-ID": "request-1", + "OpenAI-Beta": "feature", + } + acp, forwarded = extract_acp_info(headers) + assert acp == vf.ACPInfo(request_id="request-1") + assert not ACP_EXTENSION_HEADERS.intersection(map(str.lower, forwarded)) + assert forwarded["Idempotency-Key"] == "provider-key" + assert forwarded["OpenAI-Beta"] == "feature" + + absent, unchanged = extract_acp_info({"OpenAI-Beta": "feature"}) + assert absent is None and unchanged == {"OpenAI-Beta": "feature"} + + with pytest.raises(ValueError, match="not a valid ACP request ID"): + extract_acp_info({"X-ACP-Model-Request-ID": "not/a/valid/id"}) + + +def test_acp_semantic_edge_metadata_is_optional(): + trace = vf.Trace( + agent=vf.AgentInfo(config=vf.AgentConfig()), + task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="q")), + ) + harness = RLMHarness(RLMHarnessConfig(id="rlm")) + + harness._consume_protocol_metadata(trace, {}) + + assert all(not node.semantic_parents for node in trace.nodes) + + harness._consume_protocol_metadata( + trace, {ACP_SEMANTIC_EDGES_METADATA_KEY: {"edges": []}} + ) + + assert all(not node.semantic_parents for node in trace.nodes) + + +def test_semantic_edge_set_rejects_duplicate_self_and_cyclic_edges(): + edge_set = _semantic_edge_set().model_dump(mode="json") + edge_set["edges"].append(edge_set["edges"][0]) + with pytest.raises(ValueError, match="duplicate semantic edge"): + vf.SemanticEdgeSet.model_validate(edge_set) + + with pytest.raises(ValueError, match="cannot link a request to itself"): + vf.SemanticEdgeSet.model_validate( + { + "edges": [ + { + "source_request_id": "request-1", + "target_request_id": "request-1", + "type": "custom", + } + ] + } + ) + + edge_set = _semantic_edge_set().model_dump(mode="json") + edge_set["edges"].append( + { + "source_request_id": "root-after", + "target_request_id": "root-turn", + "type": "custom", + } + ) + with pytest.raises(ValueError, match="semantic edge cycle"): + vf.SemanticEdgeSet.model_validate(edge_set) + + +def test_semantic_edge_set_accepts_deep_acyclic_chain(): + edge_set = vf.SemanticEdgeSet( + edges=[ + vf.SemanticEdge( + source_request_id=f"request-{index}", + target_request_id=f"request-{index + 1}", + type="continuation", + ) + for index in range(2_000) + ] + ) + + assert len(edge_set.edges) == 2_000 diff --git a/verifiers/v1/__init__.py b/verifiers/v1/__init__.py index 438a45639..d3012c8c5 100644 --- a/verifiers/v1/__init__.py +++ b/verifiers/v1/__init__.py @@ -82,6 +82,13 @@ RuntimeProcess, SubprocessConfig, ) +from verifiers.v1.semantic import ( + ACP_SEMANTIC_EDGES_METADATA_KEY, + ACPInfo, + ParentLink, + SemanticEdge, + SemanticEdgeSet, +) from verifiers.v1.state import State, StateT from verifiers.v1.task import Task, TaskData, TaskResources, TaskTimeout, WireTaskData from verifiers.v1.taskset import Taskset @@ -226,6 +233,11 @@ "EvalRunInfo", "EvalWorkInfo", "ModelCall", + "ACPInfo", + "ACP_SEMANTIC_EDGES_METADATA_KEY", + "ParentLink", + "SemanticEdge", + "SemanticEdgeSet", "PolicyEvent", "TrainRunInfo", "TrainWorkInfo", diff --git a/verifiers/v1/acp/__init__.py b/verifiers/v1/acp/__init__.py index 02701dfaf..a08efb59e 100644 --- a/verifiers/v1/acp/__init__.py +++ b/verifiers/v1/acp/__init__.py @@ -16,6 +16,10 @@ from verifiers.v1.errors import HarnessError from verifiers.v1.harness import Harness, HarnessSession from verifiers.v1.runtimes import ProgramResult, Runtime, RuntimeProcess +from verifiers.v1.semantic import ( + ACP_SEMANTIC_EDGES_METADATA_KEY, + SemanticEdgeSet, +) from verifiers.v1.task import TaskData from verifiers.v1.trace import Trace from verifiers.v1.types import Messages @@ -67,6 +71,20 @@ async def setup(self, runtime: Runtime) -> None: def acp_turn_result(self, trace: Trace, result: ACPTurn) -> None: """Consume the typed result of one ACP prompt.""" + def acp_close_result(self, trace: Trace, response_metadata: dict[str, Any]) -> None: + """Consume extension metadata returned by `session/close`, when supported.""" + + def _consume_protocol_metadata( + self, trace: Trace, response_metadata: dict[str, Any] + ) -> None: + """Attach optional protocol extensions understood by every ACP harness.""" + if ACP_SEMANTIC_EDGES_METADATA_KEY not in response_metadata: + return + edge_set = SemanticEdgeSet.model_validate( + response_metadata[ACP_SEMANTIC_EDGES_METADATA_KEY] + ) + trace.add_semantic_edges(edge_set) + @abstractmethod async def prepare_acp( self, @@ -274,22 +292,30 @@ async def _run(self, messages: Messages | None) -> ProgramResult: if stderr := self._stderr(): detail = f"{detail}\n\nACP process stderr:\n{stderr}" raise RuntimeError(detail) - cast(ACPHarness, self.harness).acp_turn_result(self.trace, turn) + harness = cast(ACPHarness, self.harness) + harness._consume_protocol_metadata(self.trace, turn.response_metadata) + harness.acp_turn_result(self.trace, turn) result = ProgramResult(exit_code=0, stdout=turn.reply, stderr="") _require_model_turn(self.trace, calls_before, result) return result - async def _stop(self, *, graceful: bool) -> None: + async def _stop(self, *, graceful: bool) -> dict[str, Any]: process, self._process = self._process, None reader, self._reader = self._reader, None stderr_task, self._stderr_task = self._stderr_task, None if process is None: - return + return {} + response_metadata: dict[str, Any] = {} try: if graceful and reader is not None: with contextlib.suppress(BaseException): await process.write(_packet({"operation": "shutdown"})) - await asyncio.wait_for(reader.read(), timeout=10) + response = await asyncio.wait_for(reader.read(), timeout=10) + result = response.get("result") + if response.get("ok") and isinstance(result, dict): + metadata = result.get("response_metadata") + if isinstance(metadata, dict): + response_metadata = metadata for timeout, stop in ( (10 if graceful else 0.1, None), (5, process.terminate), @@ -309,6 +335,7 @@ async def _stop(self, *, graceful: bool) -> None: stderr_task.cancel() with contextlib.suppress(BaseException): await stderr_task + return response_metadata async def close(self) -> None: if self._closed: @@ -320,6 +347,10 @@ async def close(self) -> None: async def close_process() -> None: async with self._lock: - await self._stop(graceful=True) + response_metadata = await self._stop(graceful=True) + if response_metadata: + harness = cast(ACPHarness, self.harness) + harness._consume_protocol_metadata(self.trace, response_metadata) + harness.acp_close_result(self.trace, response_metadata) await run_shielded(close_process()) diff --git a/verifiers/v1/acp/runner.py b/verifiers/v1/acp/runner.py index 16a4a84bb..499ead313 100644 --- a/verifiers/v1/acp/runner.py +++ b/verifiers/v1/acp/runner.py @@ -225,7 +225,8 @@ async def run(self, config: dict) -> ACPTurn: self.is_new = False return result - async def close(self) -> None: + async def close(self) -> dict[str, Any]: + response_metadata: dict[str, Any] = {} try: if self.connection is not None and self.session_id is not None: session_capabilities = ( @@ -233,12 +234,16 @@ async def close(self) -> None: ) if session_capabilities and session_capabilities.close is not None: with suppress(Exception): - await self.connection.close_session(session_id=self.session_id) + response = await self.connection.close_session( + session_id=self.session_id + ) + response_metadata = dict(response.field_meta or {}) finally: try: await self.stack.aclose() finally: self._reset() + return response_metadata async def read_packet(stream: asyncio.StreamReader) -> dict | None: @@ -285,8 +290,10 @@ async def serve_stream() -> None: } elif operation == "shutdown": stop = True - await session.close() - response = {"ok": True} + response = { + "ok": True, + "result": {"response_metadata": await session.close()}, + } else: raise ValueError(f"unknown ACP session operation: {operation!r}") except Exception as error: # noqa: BLE001 - serialize protocol failures diff --git a/verifiers/v1/clients/eval.py b/verifiers/v1/clients/eval.py index 74ac1c8d5..d57af5fbf 100644 --- a/verifiers/v1/clients/eval.py +++ b/verifiers/v1/clients/eval.py @@ -13,41 +13,45 @@ from verifiers.v1.dialects import Dialect from verifiers.v1.errors import model_error from verifiers.v1.graph import PendingTurn +from verifiers.v1.semantic import ACP_EXTENSION_HEADERS from verifiers.v1.types import Response, SamplingConfig # These fields describe the localhost request, its original bytes, or its connection. HTTPX # rebuilds the provider request from JSON; endpoint configuration and provider auth apply last. -_BLOCKED_REQUEST_HEADERS = frozenset( - { - # The harness uses this rollout secret to authenticate with the localhost server. - # The dialect adds the actual provider authorization after filtering. - "authorization", - # HTTPX recalculates these for the provider URL, JSON bytes, and supported decoders. - "accept-encoding", - "content-encoding", - "content-length", - "content-type", - "host", - "transfer-encoding", - # These control only the localhost HTTP exchange. - "expect", - "keep-alive", - "proxy-authorization", - "proxy-connection", - "te", - "trailer", - "upgrade", - # Provider affinity must not compete with the rollout-wide session header below. - "session_id", - # The eval owns the model and sampling settings, so it changes those JSON fields before - # sending upstream. Hashes and signatures calculated from the intercepted body are stale. - "content-digest", - "content-md5", - "digest", - "repr-digest", - "signature", - "signature-input", - } +_BLOCKED_REQUEST_HEADERS = ( + frozenset( + { + # The harness uses this rollout secret to authenticate with the localhost server. + # The dialect adds the actual provider authorization after filtering. + "authorization", + # HTTPX recalculates these for the provider URL, JSON bytes, and supported decoders. + "accept-encoding", + "content-encoding", + "content-length", + "content-type", + "host", + "transfer-encoding", + # These control only the localhost HTTP exchange. + "expect", + "keep-alive", + "proxy-authorization", + "proxy-connection", + "te", + "trailer", + "upgrade", + # Provider affinity must not compete with the rollout-wide session header below. + "session_id", + # The eval owns the model and sampling settings, so it changes those JSON fields before + # sending upstream. Hashes and signatures calculated from the intercepted body are stale. + "content-digest", + "content-md5", + "digest", + "repr-digest", + "signature", + "signature-input", + } + ) + | ACP_EXTENSION_HEADERS ) diff --git a/verifiers/v1/graph.py b/verifiers/v1/graph.py index 165504bcd..9c26a4522 100644 --- a/verifiers/v1/graph.py +++ b/verifiers/v1/graph.py @@ -1,11 +1,14 @@ """Message-graph trajectory: store each message once, recover branches by walking. A rollout is a graph of `MessageNode`s — one per distinct message, each linked to its -predecessor. The conversation is a path from a root to a leaf; branches (compaction, -subagents) are simply multiple leaves, so branching falls out of the walk. Each node stores -only the tokens it *adds* to the cumulative sequence, keeping size linear in turns and -making a branch's training sample a cheap concat of node `token_ids`/`mask`/`logprobs` along -its path. +predecessor. A conversation is a path from a root to a leaf; prompt divergence from +compaction, subagents, or other history rewrites produces multiple leaves, so branching +falls out of the walk. A branch often corresponds to one harness context window, but it is +a physical, exact-prefix training view rather than a semantic context identity: a prefix +break can split one context and prefix reuse can preserve an ancestral path. Each node +stores only the tokens it *adds* to the cumulative sequence, keeping size linear in turns +and making a branch's training sample a cheap concat of node +`token_ids`/`mask`/`logprobs` along its path. Token attribution (renderer client): the renderer reports, per prompt, each message's token span (`RenderedTokens.message_token_spans()`, carried on `TurnTokens.message_spans`). A new @@ -29,6 +32,7 @@ from pydantic.json_schema import SkipJsonSchema from renderers.base import MultiModalData, PlaceholderRange, RenderedTokens +from verifiers.v1.semantic import ParentLink from verifiers.v1.types import ( AssistantMessage, KeptTokens, @@ -68,6 +72,14 @@ class MessageNode(BaseModel): parent: int | None = None """Index into `Trace.nodes` of the predecessor message; None for a root.""" + semantic_parents: list[ParentLink] = Field(default_factory=list) + """Additional harness-declared parents in the semantic execution graph. + + Unlike ``parent``, these links do not imply an exact token prefix and therefore do + not affect physical branch construction. A list permits multiple parents of the same + type, supports incremental appends, and preserves their advertised wire order; edge + application prevents duplicate ``(node, type)`` links. + """ message: Message """The message this node carries (system / user / assistant / tool).""" sampled: bool = False diff --git a/verifiers/v1/harnesses/rlm/harness.py b/verifiers/v1/harnesses/rlm/harness.py index 639d556c2..94d7dfea4 100644 --- a/verifiers/v1/harnesses/rlm/harness.py +++ b/verifiers/v1/harnesses/rlm/harness.py @@ -1,9 +1,10 @@ """RLM over ACP, with MCP tools exposed as pre-imported IPython skills.""" +import hashlib import logging import random import shlex -from typing import Literal +from typing import Any, Literal from pydantic import BaseModel, ConfigDict, Field, PositiveInt, model_validator @@ -19,8 +20,7 @@ BuiltinSkill = Literal["edit", "search"] RLM_REPO = "github.com/PrimeIntellect-ai/nano-rlm.git" -RLM_DIR = "/tmp/vf-rlm" -RLM_BIN = f"{RLM_DIR}/bin/rlm" +RLM_CACHE_DIR = "/tmp/vf-rlm" SKILLS_DIR = "/task/rlm-skills" RLM_STATE_DIR = ".vf-rlm" RLM_RUNTIME_METADATA_KEY = "ai.prime.rlm/runtime-v1" @@ -36,7 +36,7 @@ class _SessionSnapshot(BaseModel): class RLMHarnessConfig(HarnessConfig): version: str = Field( - default="d4ce3e10e63b359f4f3d432d58a77471e9e21fe7", min_length=1 + default="4a6369611c06d3943ac40681f374a464feb706b9", min_length=1 ) """Git ref (branch, tag, or commit) of nano-rlm to install.""" max_depth: int = 0 @@ -77,18 +77,26 @@ class RLMHarness(ACPHarness[RLMHarnessConfig]): async def setup(self, runtime: Runtime) -> None: # Before the installer: install.sh packages the skills it finds. await self.install_skills(runtime, SKILLS_DIR) + directory = self._install_dir() + binary = f"{directory}/bin/rlm" + checkout = f"{directory}/checkout" + ready = f"{directory}/.ready" # install.sh fetches curl/uv itself; add git only when the image lacks it. install = ( - "command -v git >/dev/null 2>&1 || " + f"rm -f {ready} && " + "(command -v git >/dev/null 2>&1 || " "{ apt-get update -qq && apt-get install -y -qq git; } && " - f"rm -rf /tmp/rlm && git clone https://{RLM_REPO} /tmp/rlm && " - f"git -C /tmp/rlm checkout {shlex.quote(self.config.version)} && " - f"UV_INSTALL_DIR={RLM_DIR}/bin UV_TOOL_BIN_DIR={RLM_DIR}/bin " - f"RLM_CHECKOUT_PATH=/tmp/rlm bash /tmp/rlm/install.sh" + f"rm -rf {checkout} && git clone https://{RLM_REPO} {checkout} && " + f"git -C {checkout} checkout {shlex.quote(self.config.version)} && " + f"UV_INSTALL_DIR={directory}/bin UV_TOOL_BIN_DIR={directory}/bin " + f"RLM_CHECKOUT_PATH={checkout} bash {checkout}/install.sh && " + f"touch {ready})" ) logger.info("rlm: ensuring rlm is installed (version=%s)", self.config.version) - ensure = shlex.quote(f"[ -x {RLM_BIN} ] || ({install})") - guarded = f"mkdir -p {RLM_DIR} && flock {RLM_DIR}/install.lock sh -c {ensure}" + ensure = shlex.quote(f"[ -f {ready} ] && [ -x {binary} ] || ({install})") + guarded = ( + f"mkdir -p {directory} && flock {directory}/install.lock sh -c {ensure}" + ) env = self.config.resolved_env.copy() extra_uv_args = env.get("RLM_EXTRA_UV_ARGS", "") env["RLM_EXTRA_UV_ARGS"] = f"{extra_uv_args} --with mcp~=1.28".strip() @@ -150,24 +158,35 @@ async def prepare_acp( system_prompt, prompt = self.resolve_prompt(data) return ACPConfig( env={**self.config.resolved_env, "RLM_HOME": self._home(trace)}, - command=[RLM_BIN, "--acp"], + command=[f"{self._install_dir()}/bin/rlm", "--acp"], prompt=prompt, session_meta=self._runtime_metadata( ctx, trace, runtime, endpoint, secret, data, system_prompt ), ) - def acp_turn_result(self, trace: Trace, result: ACPTurn) -> None: + def _consume_snapshot(self, trace: Trace, metadata: dict[str, Any]) -> None: snapshot = _SessionSnapshot.model_validate( - result.response_metadata.get(RLM_SESSION_METADATA_KEY) + metadata.get(RLM_SESSION_METADATA_KEY) ) if snapshot.session_id != trace.id: raise ValueError("RLM session snapshot does not match the rollout") trace.record_metrics(snapshot.metrics) + def acp_turn_result(self, trace: Trace, result: ACPTurn) -> None: + self._consume_snapshot(trace, result.response_metadata) + + def acp_close_result(self, trace: Trace, response_metadata: dict[str, Any]) -> None: + if RLM_SESSION_METADATA_KEY in response_metadata: + self._consume_snapshot(trace, response_metadata) + async def cleanup(self, trace: Trace, runtime: Runtime) -> None: await runtime.run(["rm", "-rf", f"{RLM_STATE_DIR}/{trace.id}"], {}) @staticmethod def _home(trace: Trace) -> str: return f"{RLM_STATE_DIR}/{trace.id}/home" + + def _install_dir(self) -> str: + cache_key = hashlib.sha256(self.config.version.encode()).hexdigest() + return f"{RLM_CACHE_DIR}-{cache_key}" diff --git a/verifiers/v1/interception/server.py b/verifiers/v1/interception/server.py index fcb5cc9b4..a89c7fc93 100644 --- a/verifiers/v1/interception/server.py +++ b/verifiers/v1/interception/server.py @@ -58,6 +58,7 @@ TunnelConfig, make_tunnel, ) +from verifiers.v1.semantic import ACPInfo, extract_acp_info from verifiers.v1.session import IdempotentRequest, ReplayResponse, RolloutSession from verifiers.v1.trace import Error, ModelCall, PolicyEvent, TimeSpan from verifiers.v1.types import FinishReason, Request, Response, Usage @@ -421,6 +422,7 @@ def record_call( usage: "Usage | None" = None, error: BaseException | None = None, policy_paths: list[str] | None = None, + acp: ACPInfo | None = None, ) -> None: """Append one provider exchange to the trace's per-call records (`Trace.calls`): the model + effective settings that went upstream, timing, and — when the call @@ -468,6 +470,7 @@ def record_call( ) if policy_paths else None, + acp=acp, ) ) @@ -491,7 +494,10 @@ async def handle_request( del raw body = dialect.apply_overrides(body, session.ctx.model, session.ctx.sampling) streaming = dialect.streaming(body) - upstream_headers = dict(request.headers) + try: + acp, upstream_headers = extract_acp_info(request.headers) + except ValueError as error: + return web.json_response(dialect.error_body(str(error)), status=400) logger.debug( "intercept %s: id=%s stream=%s", request.path, @@ -508,19 +514,15 @@ async def handle_request( replay_key: str | None = None binding = (request.path, req_hash) if idempotency_key: - if streaming: + if streaming and acp is None: return web.json_response( dialect.error_body( "Idempotency-Key is not supported for streaming requests" ), status=400, ) - replay_key = f"explicit:{idempotency_key}" - upstream_headers = { - name: value - for name, value in upstream_headers.items() - if name.lower() != IDEMPOTENCY_KEY_HEADER.lower() - } + if not streaming: + replay_key = f"explicit:{idempotency_key}" elif not streaming: replay_key = f"retry:{request.path}:{req_hash.hex()}" @@ -643,6 +645,8 @@ async def coalesced( turn=turn, inspect_response=inspect_response, policy_paths=policy_paths, + acp=acp, + upstream_headers=upstream_headers, ) def serve(response: Response) -> web.Response: @@ -757,6 +761,7 @@ async def sample() -> web.Response: usage=call_response.usage if call_response else None, error=error, policy_paths=policy_paths, + acp=acp, ) return serve(call_response) @@ -773,6 +778,8 @@ async def _stream( turn: graph.PendingTurn, inspect_response: bool, policy_paths: list[str] | None = None, + acp: ACPInfo | None = None, + upstream_headers: Mapping[str, str] | None = None, ) -> web.StreamResponse: """A streamed (SSE) model turn: relay the provider's stream through to the program, incrementally assembling the response to record on the trace (the only client that @@ -788,7 +795,7 @@ async def _stream( reply = await session.client.relay( dialect, body, - headers=request.headers, + headers=upstream_headers, session_id=session.trace.id, ) except RolloutError as e: @@ -1017,6 +1024,7 @@ async def _stream( usage=response.usage if response is not None else None, error=error, policy_paths=policy_paths, + acp=acp, ) async def handle_aux( diff --git a/verifiers/v1/semantic.py b/verifiers/v1/semantic.py new file mode 100644 index 000000000..d489089fa --- /dev/null +++ b/verifiers/v1/semantic.py @@ -0,0 +1,132 @@ +"""ACP semantic edges carried beside the physical training-message graph.""" + +from __future__ import annotations + +import re +from collections.abc import Mapping + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +ACP_REQUEST_ID_PATTERN = r"^[A-Za-z0-9._:-]{1,128}$" +EDGE_TYPE_PATTERN = r"^[A-Za-z][A-Za-z0-9._:-]{0,127}$" +"""Semantic label syntax, for example ``subagent_return`` or ``vendor:review``.""" + +ACP_SEMANTIC_EDGES_METADATA_KEY = "ai.prime.acp/semantic-edges-v1" +ACP_MODEL_REQUEST_ID_HEADER = "X-ACP-Model-Request-ID" + +ACP_EXTENSION_HEADERS = frozenset({ACP_MODEL_REQUEST_ID_HEADER.lower()}) +"""Private ACP extension headers consumed at interception and never sent upstream.""" + + +class ACPInfo(BaseModel): + """Metadata advertised by an ACP harness for one intercepted model request.""" + + model_config = ConfigDict(extra="forbid", strict=True) + + request_id: str = Field(pattern=ACP_REQUEST_ID_PATTERN) + + +class SemanticEdge(BaseModel): + """A harness-declared relationship between two logical model requests. + + Request IDs are wire-level correlation handles. They let a harness describe the + relationship before it can know Verifiers-local message-node indexes. + """ + + model_config = ConfigDict(extra="forbid", strict=True) + + source_request_id: str = Field(pattern=ACP_REQUEST_ID_PATTERN) + target_request_id: str = Field(pattern=ACP_REQUEST_ID_PATTERN) + type: str = Field(pattern=EDGE_TYPE_PATTERN) + + @model_validator(mode="after") + def reject_self_edge(self) -> SemanticEdge: + if self.source_request_id == self.target_request_id: + raise ValueError("semantic edge cannot link a request to itself") + return self + + +class ParentLink(BaseModel): + """One semantic parent of a ``MessageNode``. + + ``node`` is an index into the containing ``Trace.nodes``. The child is the + ``MessageNode`` carrying this link. + """ + + model_config = ConfigDict(extra="forbid", strict=True) + + node: int = Field(ge=0) + type: str = Field(pattern=EDGE_TYPE_PATTERN) + + +class SemanticEdgeSet(BaseModel): + """A harness-published set of semantic edges over logical request IDs. + + Edge labels are intentionally extensible. Initial harnesses use ``continuation``, + ``compaction``, ``subagent_call``, and ``subagent_return``; consumers must preserve + unknown labels. + """ + + model_config = ConfigDict(extra="forbid", strict=True) + + edges: list[SemanticEdge] + + @model_validator(mode="after") + def validate_edges(self) -> SemanticEdgeSet: + identities: set[tuple[str, str, str]] = set() + children: dict[str, list[str]] = {} + nodes: set[str] = set() + for edge in self.edges: + identity = ( + edge.source_request_id, + edge.target_request_id, + edge.type, + ) + if identity in identities: + raise ValueError(f"duplicate semantic edge: {identity!r}") + identities.add(identity) + children.setdefault(edge.source_request_id, []).append( + edge.target_request_id + ) + nodes.update((edge.source_request_id, edge.target_request_id)) + + indegree = dict.fromkeys(nodes, 0) + for targets in children.values(): + for target in targets: + indegree[target] += 1 + + stack = [node for node, degree in indegree.items() if degree == 0] + visited = 0 + while stack: + node = stack.pop() + visited += 1 + for child in children.get(node, ()): + indegree[child] -= 1 + if indegree[child] == 0: + stack.append(child) + if visited != len(nodes): + raise ValueError("semantic edge cycle detected") + return self + + +def extract_acp_info( + headers: Mapping[str, str], +) -> tuple[ACPInfo | None, dict[str, str]]: + """Parse ACP request metadata and remove its private transport header. + + The semantic edge set arrives independently in ACP metadata. Ordinary requests + without the private correlation header remain unchanged. + """ + + forwarded = { + name: value + for name, value in headers.items() + if name.lower() not in ACP_EXTENSION_HEADERS + } + normalized = {name.lower(): value for name, value in headers.items()} + request_id = normalized.get(ACP_MODEL_REQUEST_ID_HEADER.lower()) + if request_id is None: + return None, forwarded + if re.fullmatch(ACP_REQUEST_ID_PATTERN, request_id) is None: + raise ValueError(f"{ACP_MODEL_REQUEST_ID_HEADER} is not a valid ACP request ID") + return ACPInfo(request_id=request_id), forwarded diff --git a/verifiers/v1/trace.py b/verifiers/v1/trace.py index 8e1bafdbc..560becbb7 100644 --- a/verifiers/v1/trace.py +++ b/verifiers/v1/trace.py @@ -19,6 +19,11 @@ from verifiers.v1.errors import ProviderError from verifiers.v1.graph import MessageNode from verifiers.v1.runtimes import RuntimeInfo +from verifiers.v1.semantic import ( + ACPInfo, + ParentLink, + SemanticEdgeSet, +) from verifiers.v1.state import State, StateT from verifiers.v1.task import DataT, WireTaskData from verifiers.v1.types import ( @@ -171,6 +176,8 @@ class ModelCall(BaseModel): """The failure that ended this call, coupled to the exchange that caused it.""" policy: PolicyEvent | None = None """Policy mediation applied to the request before this call.""" + acp: ACPInfo | None = None + """Metadata advertised by the ACP harness for this model request.""" def min_new_input_tokens(calls: Iterable[ModelCall]) -> Iterator[tuple[ModelCall, int]]: @@ -188,7 +195,7 @@ def min_new_input_tokens(calls: Iterable[ModelCall]) -> Iterator[tuple[ModelCall class Branch(BaseModel): - """A root-to-leaf graph path; each branch becomes one training sample.""" + """A root-to-leaf message-graph path; each branch becomes one training sample.""" index: int nodes: list[MessageNode] @@ -387,7 +394,7 @@ class Trace(BaseModel, Generic[DataT, StateT, AgentConfigT]): """The tools advertised to the agent, automatically recorded from last intercepted turn.""" nodes: list[MessageNode] = Field(default_factory=list) - """The message graph; branches are derived views and storage stays linear in turns.""" + """The message graph, including physical and semantic parent links.""" calls: list[ModelCall] = Field(default_factory=list) """Every model call; automatically recorded at intercept time + linked into `nodes`.""" mm_token_type_id_map: dict[int, int] = Field(default_factory=dict) @@ -493,6 +500,72 @@ def branches(self) -> list[Branch]: ) return branches + def add_semantic_edges(self, edge_set: SemanticEdgeSet) -> None: + """Add newly advertised ACP edges; cumulative replays are idempotent.""" + node_by_request: dict[str, int] = {} + for call in self.calls: + if call.acp is None or call.node is None: + continue + if not 0 <= call.node < len(self.nodes): + raise ValueError(f"model call has invalid message node {call.node}") + if not self.nodes[call.node].sampled: + raise ValueError(f"model call node {call.node} is not sampled") + # SDK retries are sequential. If more than one attempt commits, the last + # sampled response is the logical request result consumed by the harness. + node_by_request[call.acp.request_id] = call.node + + resolved_identities: set[tuple[int, int, str]] = set() + additions: list[tuple[int, ParentLink]] = [] + pending_parents: dict[int, list[ParentLink]] = {} + for edge in edge_set.edges: + endpoints: list[int] = [] + for request_id in ( + edge.source_request_id, + edge.target_request_id, + ): + node = node_by_request.get(request_id) + if node is None: + raise ValueError( + f"semantic edge request {request_id!r} has no committed message node" + ) + endpoints.append(node) + source, target = endpoints + if source == target: + raise ValueError("semantic edge resolves to the same message node") + identity = (source, target, edge.type) + if identity in resolved_identities: + raise ValueError(f"duplicate resolved semantic edge: {identity!r}") + resolved_identities.add(identity) + link = ParentLink(node=source, type=edge.type) + if link in self.nodes[target].semantic_parents: + continue + + # Adding source -> target creates a cycle exactly when target is already an + # ancestor of source. Walk parent links directly so existing nodes and links + # are never rebuilt as cumulative ACP edge sets arrive. + stack = [source] + visited: set[int] = set() + while stack: + node_id = stack.pop() + if node_id == target: + raise ValueError( + "semantic edges create a cycle in the message graph" + ) + if node_id in visited: + continue + visited.add(node_id) + node = self.nodes[node_id] + if node.parent is not None: + stack.append(node.parent) + stack.extend(parent.node for parent in node.semantic_parents) + stack.extend(parent.node for parent in pending_parents.get(node_id, ())) + + additions.append((target, link)) + pending_parents.setdefault(target, []).append(link) + + for target, link in additions: + self.nodes[target].semantic_parents.append(link) + @property def messages(self) -> Messages: """Messages on the final branch."""