From d6bcb9ea33e38968a5ee801a435eedf8cd324b79 Mon Sep 17 00:00:00 2001 From: Lingrui Mei Date: Fri, 24 Jul 2026 05:25:29 +0800 Subject: [PATCH 1/5] trace: lock in byte-fidelity of gen_ai span content capture MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Spike verdict: agentix's span-population and capture paths do NOT truncate attribute values — populate_*_span stores full strings, Span.set_attribute keeps raw values, JsonlProcessor writes Span.export() verbatim, and the OTel exporter applies no value-length limit (SDK default is unlimited). A downstream consumer's captured-span fixture showing gen_ai.prompt.*.content cut at exactly 120 chars was hand-trimmed fixture data, not a capture defect. Two deterministic tests pin the property: a >2000-char prompt survives populate_anthropic_span -> in-memory processor -> JsonlProcessor round-trip byte-for-byte, and the same length survives the OTelTraceProcessor export leg untruncated. Co-Authored-By: Claude Fable 5 --- plugins/abridge/tests/test_tracing.py | 59 +++++++++++++++++++ .../trace-otel/tests/test_otel_processor.py | 21 +++++++ 2 files changed, 80 insertions(+) diff --git a/plugins/abridge/tests/test_tracing.py b/plugins/abridge/tests/test_tracing.py index 82987c7..4ccea17 100644 --- a/plugins/abridge/tests/test_tracing.py +++ b/plugins/abridge/tests/test_tracing.py @@ -95,6 +95,65 @@ async def test_openai_client_emits_genai_span(wired, capture_spans) -> None: assert sp.attrs["gen_ai.completion.0.content"] == "hello from upstream" +def test_genai_span_content_is_byte_faithful_for_long_prompts(capture_spans, tmp_path) -> None: + """Byte-fidelity spike: the span-population path must NOT truncate + gen_ai content attributes. + + A downstream consumer's captured-span fixture showed every + `gen_ai.prompt.*.content` cut at exactly 120 chars, raising the + question whether agentix's capture truncates attribute values. + Verdict from source: it does not — `populate_*_span` stores the full + string, `Span.set_attribute` keeps raw values, and `JsonlProcessor` + writes `Span.export()` verbatim (no OTel attribute-length limit is + involved on this path). That fixture was hand-trimmed. This test + locks the property in: a >2000-char prompt survives population, + in-memory capture, AND the JSONL sink byte-for-byte. + """ + from agentix.bridge.clients import populate_anthropic_span + + # Deterministic, marker-rich content: mid-string and end markers would + # be destroyed by any length-limited capture. + prompt = ("0123456789" * 250) + "" + completion = ("abcdefghij" * 300) + "" + assert len(prompt) > 2000 and len(completion) > 2000 + + sink_path = tmp_path / "spans.jsonl" + sink = trace.JsonlProcessor(sink_path) + trace.add_processor(sink) + try: + with trace.span("anthropic messages fidelity"): + populate_anthropic_span( + request={ + "model": "claude-3-haiku", + "system": prompt, + "messages": [{"role": "user", "content": prompt}], + }, + response={ + "model": "m", + "content": [{"type": "text", "text": completion}], + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + finally: + trace.remove_processor(sink) + sink.shutdown() + + # In-memory capture: full fidelity. + (sp,) = capture_spans.spans + assert sp.attrs["gen_ai.prompt.0.content"] == prompt + assert sp.attrs["gen_ai.prompt.1.content"] == prompt + assert sp.attrs["gen_ai.completion.0.content"] == completion + + # JSONL capture channel (`trace.collect` / `JsonlProcessor`): what is + # read back from disk is byte-identical to what was populated. + import json as _json + + (record,) = [_json.loads(line) for line in sink_path.read_text().splitlines()] + assert record["attrs"]["gen_ai.prompt.0.content"] == prompt + assert record["attrs"]["gen_ai.prompt.1.content"] == prompt + assert record["attrs"]["gen_ai.completion.0.content"] == completion + + @pytest.mark.asyncio async def test_silent_client_emits_no_spans(capture_spans) -> None: """A custom client that opens no spans and calls no populate helper diff --git a/plugins/trace-otel/tests/test_otel_processor.py b/plugins/trace-otel/tests/test_otel_processor.py index bf15485..caf655e 100644 --- a/plugins/trace-otel/tests/test_otel_processor.py +++ b/plugins/trace-otel/tests/test_otel_processor.py @@ -125,3 +125,24 @@ def test_resource_carries_service_name() -> None: resource = spans[0].resource assert resource.attributes["service.name"] == "rollout-eval" assert resource.attributes["deployment.environment"] == "test" + + +def test_long_attribute_values_export_untruncated(otel_recorder) -> None: + """Byte-fidelity guard for the OTel leg: the exporter must not + truncate long attribute values. The OTel Python SDK's attribute + value-length limit defaults to None (unlimited) and + `OTelTraceProcessor` sets no explicit limit; gen_ai content + attributes (full prompts/completions) rely on that. If a future SDK + or processor change introduces a limit, this fails loudly instead of + silently clipping capture data.""" + exporter, _ = otel_recorder + + long_value = ("0123456789" * 250) + "" + assert len(long_value) > 2000 + + with trace.trace("workflow"): + with trace.span("llm.request") as s: + s.set_attribute("gen_ai.prompt.0.content", long_value) + + span = _by_name(list(exporter.get_finished_spans()), "llm.request") + assert span.attributes["gen_ai.prompt.0.content"] == long_value From a0dc9410ffe3e7c788662c4b285926a0d3ef09a7 Mon Sep 17 00:00:00 2001 From: Lingrui Mei Date: Fri, 24 Jul 2026 05:42:46 +0800 Subject: [PATCH 2/5] tito: per-turn token record persistence + session lifecycle (production capture) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The gateway's token truth used to live only in an unbounded in-process dict — no durability, no TTL, and an interleaved turn was silently dropped from the trajectory with just a warning. This makes the capture production-grade: - --record-dir (env TITO_RECORD_DIR): one flushed tito.record.v1 JSON line per committed turn in /.jsonl — session_id, turn_index, request_id (x-request-id echo), model, backend_kind, prompt/completion token ids, completion logprobs (1:1, retained from the same forward pass on BOTH backends — the vLLM path previously validated then discarded them), prompt_segments (render/prefix/ per-role/generation_prompt spans: loss-mask material), assistant message, finish_reason, tokenizer fingerprint (checkpoint + chat_template_sha256), prefix_stable, render_skew (vLLM per-turn render-vs-accumulated probe), ts. A final tito.session.v1 line lands when the session closes. No record dir -> behavior unchanged. - prefix_stable is computed against the last COMMITTED checkpoint: a retry rollback or history rewrite is served and recorded but flagged, so a trainer can split or reject instead of splicing a lie. - Interleaved turns on one session are now an explicit 409 (ConcurrentSessionUpdateError) — silent capture loss is unacceptable. - Optional --session-ttl-seconds / --max-sessions eviction for long-running gateways: record files are flushed+finalized before a session becomes unreachable, in-flight sessions are never evicted, and eviction drops the pool's sticky pin like DELETE does. The documented flow stays harvest-then-DELETE. - Engine refactor to expose segment boundaries without changing render behavior: TITOTokenizer.appended_segments + fix_prefix (Qwen3's newline fixup moved there); LinearTrajectory.prepare_prompt returns ids + segments + stability, with prepare_pretokenized kept as a thin compatible wrapper. Tests: record shape + crash-safety (flush per line), prefix_stable true/false, 409 interleave, TTL/capacity eviction flush + in-flight immunity, shutdown finalize, vLLM logprobs/skew/finish_reason records, and a real-Qwen3-tokenizer golden (network-marked, cache-or-skip offline) proving incremental == from-scratch token-exactly across a multi-turn tool-calling session with segment boundaries decoding to the expected role markers. Co-Authored-By: Claude Fable 5 --- plugins/tito/README.md | 48 +++ plugins/tito/agentix/tito/cli.py | 28 ++ plugins/tito/agentix/tito/config.py | 24 ++ plugins/tito/agentix/tito/engine/errors.py | 9 + .../tito/agentix/tito/engine/pretokenize.py | 60 ++- plugins/tito/agentix/tito/engine/record.py | 201 +++++++++ .../tito/agentix/tito/engine/session_app.py | 61 ++- .../tito/agentix/tito/engine/trajectory.py | 172 +++++++- plugins/tito/agentix/tito/engine/upstream.py | 48 ++- plugins/tito/agentix/tito/server.py | 6 + plugins/tito/pyproject.toml | 3 + plugins/tito/tests/test_gateway_vllm_http.py | 69 ++++ plugins/tito/tests/test_qwen3_golden.py | 224 ++++++++++ plugins/tito/tests/test_record.py | 381 ++++++++++++++++++ pyproject.toml | 5 +- 15 files changed, 1296 insertions(+), 43 deletions(-) create mode 100644 plugins/tito/agentix/tito/engine/record.py create mode 100644 plugins/tito/tests/test_qwen3_golden.py create mode 100644 plugins/tito/tests/test_record.py diff --git a/plugins/tito/README.md b/plugins/tito/README.md index 5d63628..0cb2540 100644 --- a/plugins/tito/README.md +++ b/plugins/tito/README.md @@ -78,6 +78,54 @@ tokenizer's own template). `--backend-kind` selects the backend token dialect a local backend (see `agentix.tito.discovery`). Run `agentix-tito serve -h` for the full list. +## Per-turn record persistence + +With `--record-dir DIR` (env `TITO_RECORD_DIR`), every committed turn appends +one `tito.record.v1` JSON line to `DIR/.jsonl` and flushes it — +the file is complete up to the last committed turn even if the process dies +mid-rollout. Each line carries the exact token truth of one turn: + +- `prompt_token_ids` / `completion_token_ids` / `completion_logprobs` + (sampled-token logprobs from the same forward pass, 1:1 with the ids); +- `prompt_segments` — `{start, end, source}` spans over the prompt ids + (`render`, `prefix`, per appended role, `generation_prompt`): the material + a trainer needs to build a loss mask without re-tokenizing anything; +- `prefix_stable` — whether this prompt extends the previous committed + checkpoint; `false` (retry rollback / history rewrite) means the turn must + not be spliced into one linear token stream; +- `request_id` (echoed from the caller's `x-request-id` header), `model`, + `backend_kind`, `finish_reason`, `assistant_message`, and a + `tokenizer_fingerprint` (`checkpoint` + `chat_template_sha256`) pinning the + render rules; +- `render_skew` (vLLM only) — a cheap per-turn probe comparing the render + endpoint's from-scratch ids against the gateway's accumulated prompt ids + (recorded, never enforced). + +Closing a session appends one final `tito.session.v1` metadata line +(`reason`: `deleted` / `ttl_evicted` / `capacity_evicted` / `shutdown`) and +closes the file. Without `--record-dir` nothing is written and the in-memory +behavior is unchanged. + +## Session lifecycle + +Sessions live in memory until deleted. The intended rollout flow is +**harvest, then delete**: run the rollout, `GET /sessions/{id}` to read the +records + accumulated ids (and the mismatch audit), then +`DELETE /sessions/{id}` — the delete finalizes the record file and frees the +in-memory trajectory and its pool pin. + +For long-running gateways two optional guards bound memory: +`--session-ttl-seconds` evicts sessions idle beyond the TTL, and +`--max-sessions` LRU-evicts beyond a count. Eviction always finalizes the +session's record file first and **never** touches a session with an in-flight +request; capacity may transiently overflow rather than kill a live rollout. + +Interleaved turns on one session (a second request racing the first) are +rejected with an explicit **409**: the losing turn's completion cannot be +committed to a trajectory that changed under it, and silently serving an +unrecorded response would be capture data loss. Callers retry on the current +session state. + ## HTTP surface - `POST /sessions` → `{session_id}` diff --git a/plugins/tito/agentix/tito/cli.py b/plugins/tito/agentix/tito/cli.py index b3986dd..340c44a 100644 --- a/plugins/tito/agentix/tito/cli.py +++ b/plugins/tito/agentix/tito/cli.py @@ -3,6 +3,7 @@ from __future__ import annotations import argparse +import os import sys from .config import TITOGatewayConfig @@ -68,6 +69,30 @@ def _add_serve_arguments(parser: argparse.ArgumentParser) -> None: parser.add_argument("--session-server-ip", default="127.0.0.1", help="Gateway bind host.") parser.add_argument("--session-server-port", type=int, default=30000, help="Gateway bind port.") parser.add_argument("--router-timeout", type=float, default=600.0, help="Proxy timeout in seconds.") + parser.add_argument( + "--record-dir", + default=os.environ.get("TITO_RECORD_DIR") or None, + help=( + "Persist one tito.record.v1 JSON line per committed turn to " + "/.jsonl, flushed per turn (env: TITO_RECORD_DIR). " + "Unset = in-memory trajectories only." + ), + ) + parser.add_argument( + "--session-ttl-seconds", + type=float, + default=None, + help=( + "Evict sessions idle longer than this many seconds (record files are " + "flushed+finalized first; in-flight sessions are never evicted). Unset = no TTL." + ), + ) + parser.add_argument( + "--max-sessions", + type=int, + default=None, + help="LRU-evict sessions beyond this count (same flush-first, never-in-flight rules). Unset = unbounded.", + ) parser.add_argument( "--backend-probe-candidate", action="append", @@ -100,6 +125,9 @@ def _serve(args: argparse.Namespace) -> int: router_timeout=args.router_timeout, backend_probe_candidates=args.backend_probe_candidate, backend_probe_timeout=args.backend_probe_timeout, + record_dir=args.record_dir, + session_ttl_seconds=args.session_ttl_seconds, + max_sessions=args.max_sessions, ) TITOGateway(config).run() return 0 diff --git a/plugins/tito/agentix/tito/config.py b/plugins/tito/agentix/tito/config.py index d1d3959..2a72e74 100644 --- a/plugins/tito/agentix/tito/config.py +++ b/plugins/tito/agentix/tito/config.py @@ -34,6 +34,16 @@ class TITOGatewayConfig: router_timeout: float = 600.0 backend_probe_candidates: tuple[str, ...] = field(default_factory=lambda: DEFAULT_BACKEND_PROBE_CANDIDATES) backend_probe_timeout: float = 0.25 + # Durable capture: when set, every committed turn appends one + # `tito.record.v1` line to `/.jsonl` (flushed per + # turn). Unset = in-memory trajectories only (unchanged default). + record_dir: str | None = None + # Long-running-gateway lifecycle (both optional, default off): evict + # sessions idle beyond the TTL, and LRU-evict beyond max_sessions. + # Eviction flushes+finalizes the session's record file first and never + # touches a session with in-flight requests. + session_ttl_seconds: float | None = None + max_sessions: int | None = None def __post_init__(self) -> None: if not self.hf_checkpoint: @@ -53,6 +63,11 @@ def __post_init__(self) -> None: if self.backend_kind not in ("sglang", "vllm"): raise ValueError(f"backend_kind must be 'sglang' or 'vllm'; got {self.backend_kind!r}") + if self.session_ttl_seconds is not None and self.session_ttl_seconds <= 0: + raise ValueError(f"session_ttl_seconds must be > 0; got {self.session_ttl_seconds!r}") + if self.max_sessions is not None and self.max_sessions < 1: + raise ValueError(f"max_sessions must be >= 1; got {self.max_sessions!r}") + @classmethod def from_cli_values( cls, @@ -71,6 +86,9 @@ def from_cli_values( trust_remote_code: bool = False, backend_probe_candidates: list[str] | None = None, backend_probe_timeout: float = 0.25, + record_dir: str | None = None, + session_ttl_seconds: float | None = None, + max_sessions: int | None = None, ) -> TITOGatewayConfig: return cls( hf_checkpoint=hf_checkpoint, @@ -87,6 +105,9 @@ def from_cli_values( router_timeout=router_timeout, backend_probe_candidates=tuple(backend_probe_candidates or DEFAULT_BACKEND_PROBE_CANDIDATES), backend_probe_timeout=backend_probe_timeout, + record_dir=record_dir, + session_ttl_seconds=session_ttl_seconds, + max_sessions=max_sessions, ) def as_session_args(self): @@ -103,4 +124,7 @@ def as_session_args(self): session_server_ip=self.session_server_ip, session_server_port=self.session_server_port, router_timeout=self.router_timeout, + record_dir=self.record_dir, + session_ttl_seconds=self.session_ttl_seconds, + max_sessions=self.max_sessions, ) diff --git a/plugins/tito/agentix/tito/engine/errors.py b/plugins/tito/agentix/tito/engine/errors.py index 1a8673e..01e882d 100644 --- a/plugins/tito/agentix/tito/engine/errors.py +++ b/plugins/tito/agentix/tito/engine/errors.py @@ -21,6 +21,15 @@ class MessageValidationError(SessionError): status_code: int = 400 +class ConcurrentSessionUpdateError(SessionError): + """The session changed (or was deleted) while the turn was in flight, so + the completed turn cannot be committed to the trajectory or recorded. + Silent data loss is unacceptable for production capture — the caller gets + an explicit 409 and must retry on the current session state.""" + + status_code: int = 409 + + class TokenizationError(SessionError): """A TITO tokenization invariant was violated (e.g. pretokenized prefix mismatch).""" diff --git a/plugins/tito/agentix/tito/engine/pretokenize.py b/plugins/tito/agentix/tito/engine/pretokenize.py index 4d22aa0..8717e4d 100644 --- a/plugins/tito/agentix/tito/engine/pretokenize.py +++ b/plugins/tito/agentix/tito/engine/pretokenize.py @@ -138,28 +138,52 @@ def _tokenize_user_and_system_segment( ) -> list[int]: return self._tokenize_rendered_suffix([_DUMMY_SYSTEM], [appended_message], tools=tools) - def tokenize_additional_non_assistant( + def appended_segments( self, old_messages: list[dict[str, Any]], new_messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None = None, - ) -> list[int]: - """Incremental token IDs (incl. the next generation prompt) for the - non-assistant messages appended after the pretokenized prefix.""" + ) -> list[tuple[str, list[int]]]: + """Per-segment incremental token IDs for the non-assistant messages + appended after the pretokenized prefix, as ``(source, token_ids)`` + pairs — one per appended role segment plus a final + ``("generation_prompt", ...)`` entry. The flat concatenation equals + `tokenize_additional_non_assistant`; the segment boundaries feed the + per-turn record's ``prompt_segments``.""" assert_messages_append_only_with_allowed_role(old_messages, new_messages, self.allowed_append_roles) appended_messages = new_messages[len(old_messages):] - incremental: list[int] = [] + segments: list[tuple[str, list[int]]] = [] for segment in self._split_appended_segments(appended_messages): role = segment[0]["role"] if role == "tool": - incremental.extend(self._tokenize_tool_segment(segment, tools)) + segments.append((role, self._tokenize_tool_segment(segment, tools))) elif role in ("user", "system"): - incremental.extend(self._tokenize_user_and_system_segment(segment[0], tools)) + segments.append((role, self._tokenize_user_and_system_segment(segment[0], tools))) else: raise ValueError(f"unsupported appended role for TITO tokenization: {role}") - return incremental + self._tokenize_rendered_suffix( - new_messages, [], tools=tools, add_generation_prompt=True + segments.append( + ( + "generation_prompt", + self._tokenize_rendered_suffix(new_messages, [], tools=tools, add_generation_prompt=True), + ) ) + return segments + + def tokenize_additional_non_assistant( + self, + old_messages: list[dict[str, Any]], + new_messages: list[dict[str, Any]], + tools: list[dict[str, Any]] | None = None, + ) -> list[int]: + """Incremental token IDs (incl. the next generation prompt) for the + non-assistant messages appended after the pretokenized prefix.""" + return [tid for _, ids in self.appended_segments(old_messages, new_messages, tools) for tid in ids] + + def fix_prefix(self, pretokenized_token_ids: list[int]) -> list[int]: + """Canonicalize a stored prefix before merging (returns a new list). + The base engine is a no-op; a model subclass re-inserts boundary + tokens the model omits when it stops (see `Qwen3TITOTokenizer`).""" + return list(pretokenized_token_ids) def merge_tokens( self, @@ -168,14 +192,15 @@ def merge_tokens( pretokenized_token_ids: list[int], tools: list[dict[str, Any]] | None = None, ) -> list[int]: - """Default: concatenate the stored prefix with the incremental tokens.""" + """Default: concatenate the (canonicalized) stored prefix with the + incremental tokens.""" incremental = self.tokenize_additional_non_assistant(old_messages, new_messages, tools) - return list(pretokenized_token_ids) + incremental + return self.fix_prefix(pretokenized_token_ids) + incremental class Qwen3TITOTokenizer(TITOTokenizer): """Qwen3: the model stops at `<|im_end|>` without the trailing `\\n` the template - emits, so `merge_tokens` re-inserts it so the stored prefix stays canonical.""" + emits, so `fix_prefix` re-inserts it so the stored prefix stays canonical.""" reasoning_parser = "qwen3" tool_call_parser = "qwen25" @@ -201,18 +226,11 @@ def __init__( self._im_end_id: int = tokenizer.convert_tokens_to_ids("<|im_end|>") self.trailing_token_ids = frozenset({self._newline_id}) - def merge_tokens( - self, - old_messages: list[dict[str, Any]], - new_messages: list[dict[str, Any]], - pretokenized_token_ids: list[int], - tools: list[dict[str, Any]] | None = None, - ) -> list[int]: - incremental = self.tokenize_additional_non_assistant(old_messages, new_messages, tools) + def fix_prefix(self, pretokenized_token_ids: list[int]) -> list[int]: prefix = list(pretokenized_token_ids) if prefix and prefix[-1] == self._im_end_id: prefix.append(self._newline_id) - return prefix + incremental + return prefix _QWEN3_FIXED = "qwen3_fixed.jinja" diff --git a/plugins/tito/agentix/tito/engine/record.py b/plugins/tito/agentix/tito/engine/record.py new file mode 100644 index 0000000..f04a8ad --- /dev/null +++ b/plugins/tito/agentix/tito/engine/record.py @@ -0,0 +1,201 @@ +"""Per-turn token record persistence — the gateway's durable capture. + +The in-process trajectory (`LinearTrajectory`) is the live state machine; +this module is the crash-safe, append-only export of it. When the gateway is +started with a record directory, every committed chat turn appends exactly +one JSON line to ``/.jsonl`` and flushes it, so the +file is complete up to the last committed turn even if the process dies +mid-rollout. Closing a session (DELETE, TTL/capacity eviction, shutdown) +appends one final ``tito.session.v1`` metadata line and closes the file. + +Two line schemas: + +``tito.record.v1`` — one per committed turn:: + + {"schema_version": "tito.record.v1", "session_id": ..., "turn_index": 0, + "request_id": ..., "model": ..., "backend_kind": "sglang" | "vllm", + "prompt_token_ids": [...], "completion_token_ids": [...], + "completion_logprobs": [...], # len == len(completion_token_ids) + "prompt_segments": [{"start": 0, "end": N, "source": ...}, ...], + "assistant_message": {...}, "finish_reason": ..., + "tokenizer_fingerprint": {"checkpoint": ..., "chat_template_sha256": ...}, + "prefix_stable": true, "render_skew": null | {"equal": ..., "first_divergence": ...}, + "ts": } + +``prompt_segments`` sources: ``render`` (a from-scratch chat-template render +— the first turn, or the fallback when the token prefix window was outrun), +``prefix`` (the reused accumulated checkpoint), one segment per appended +role (``tool`` / ``user`` / ``system``), and ``generation_prompt``. + +``prefix_stable`` is true iff this turn's prompt token ids extend the last +*committed* checkpoint (previous prompt + completion). A retry rollback or a +history rewrite records ``false`` — the turn is still served and recorded, +but a trainer must not splice it into one linear token stream. + +``render_skew`` (vLLM backend only) is the cheap per-turn probe comparing +the render endpoint's from-scratch token ids against the gateway's +accumulated prompt ids; non-equal skew is expected mid-conversation (render +re-renders the echoed history) and is recorded, never enforced. + +``tito.session.v1`` — one final line when the session closes:: + + {"schema_version": "tito.session.v1", "session_id": ..., "turns": N, + "reason": "deleted" | "ttl_evicted" | "capacity_evicted" | "shutdown", + "ts": } +""" + +from __future__ import annotations + +import json +import logging +import time +from pathlib import Path +from typing import IO, Any + +logger = logging.getLogger(__name__) + +RECORD_SCHEMA_VERSION = "tito.record.v1" +SESSION_META_SCHEMA_VERSION = "tito.session.v1" + + +def compute_render_skew( + render_token_ids: list[int] | None, prompt_token_ids: list[int] +) -> dict[str, Any] | None: + """Compare the backend render's from-scratch prompt ids against the + gateway's accumulated prompt ids. ``None`` when the backend kind exposes + no render ids (sglang).""" + if render_token_ids is None: + return None + if render_token_ids == prompt_token_ids: + return {"equal": True, "first_divergence": None} + first = next( + (i for i, (a, b) in enumerate(zip(render_token_ids, prompt_token_ids, strict=False)) if a != b), + min(len(render_token_ids), len(prompt_token_ids)), + ) + return {"equal": False, "first_divergence": first} + + +def build_turn_record( + *, + session_id: str, + request_id: str | None, + model: str | None, + backend_kind: str, + prompt_token_ids: list[int], + prompt_segments: list[dict[str, Any]], + prefix_stable: bool, + completion_token_ids: list[int], + completion_logprobs: list[float], + assistant_message: dict[str, Any], + finish_reason: str | None, + tokenizer_fingerprint: dict[str, str], + render_token_ids: list[int] | None, +) -> dict[str, Any]: + """One JSON-serializable ``tito.record.v1`` line (without ``turn_index``, + which the sink assigns monotonically per session file).""" + if len(completion_logprobs) != len(completion_token_ids): + raise ValueError( + f"len(completion_logprobs)={len(completion_logprobs)} != " + f"len(completion_token_ids)={len(completion_token_ids)}" + ) + return { + "schema_version": RECORD_SCHEMA_VERSION, + "session_id": session_id, + "request_id": request_id, + "model": model, + "backend_kind": backend_kind, + "prompt_token_ids": list(prompt_token_ids), + "completion_token_ids": list(completion_token_ids), + "completion_logprobs": list(completion_logprobs), + "prompt_segments": list(prompt_segments), + "assistant_message": assistant_message, + "finish_reason": finish_reason, + "tokenizer_fingerprint": dict(tokenizer_fingerprint), + "prefix_stable": bool(prefix_stable), + "render_skew": compute_render_skew(render_token_ids, prompt_token_ids), + "ts": time.time(), + } + + +class TurnRecordSink: + """Per-session JSONL files under ``record_dir``, appended and flushed one + line per committed turn. + + Files open lazily on the first turn and are closed by ``finalize`` (which + also appends the ``tito.session.v1`` metadata line). ``finalize`` is + idempotent; ``close_all`` finalizes every open file (process shutdown). + A failed disk write is logged and never fails the live request — the + in-memory trajectory still holds the turn for a read-time harvest. + """ + + def __init__(self, record_dir: str | Path) -> None: + self.record_dir = Path(record_dir) + self.record_dir.mkdir(parents=True, exist_ok=True) + self._files: dict[str, IO[str]] = {} + self._turns: dict[str, int] = {} + + def path_for(self, session_id: str) -> Path: + return self.record_dir / f"{session_id}.jsonl" + + def append_turn(self, session_id: str, record: dict[str, Any]) -> int: + """Append one turn record (assigning ``turn_index``), flush, and + return the assigned index.""" + turn_index = self._turns.get(session_id, 0) + record = {**record, "turn_index": turn_index} + self._write(session_id, record) + self._turns[session_id] = turn_index + 1 + return turn_index + + def finalize(self, session_id: str, *, reason: str) -> None: + """Append the final session metadata line and close the file. + No-op for sessions that never recorded a turn; idempotent.""" + if session_id not in self._files and session_id not in self._turns: + return + meta = { + "schema_version": SESSION_META_SCHEMA_VERSION, + "session_id": session_id, + "turns": self._turns.get(session_id, 0), + "reason": reason, + "ts": time.time(), + } + try: + self._write(session_id, meta) + finally: + self._turns.pop(session_id, None) + handle = self._files.pop(session_id, None) + if handle is not None and not handle.closed: + try: + handle.close() + except OSError: + logger.exception("tito record: failed to close %s", self.path_for(session_id)) + + def close_all(self, *, reason: str = "shutdown") -> None: + for session_id in list(self._files): + self.finalize(session_id, reason=reason) + + def _write(self, session_id: str, record: dict[str, Any]) -> None: + try: + handle = self._files.get(session_id) + if handle is None or handle.closed: + handle = self.path_for(session_id).open("a", encoding="utf-8") + self._files[session_id] = handle + # `default=repr` keeps a stray non-JSON value (and `allow_nan` + # keeps a NaN logprob the backend passed through) from sinking + # the line — best-effort capture beats failing the live turn. + handle.write(json.dumps(record, ensure_ascii=False, default=repr) + "\n") + handle.flush() + except OSError: + # Capture must never take down the serving path: the turn is + # already committed in memory and remains harvestable via GET. + logger.exception( + "tito record: failed to append to %s — turn NOT persisted", self.path_for(session_id) + ) + + +__all__ = [ + "RECORD_SCHEMA_VERSION", + "SESSION_META_SCHEMA_VERSION", + "TurnRecordSink", + "build_turn_record", + "compute_render_skew", +] diff --git a/plugins/tito/agentix/tito/engine/session_app.py b/plugins/tito/agentix/tito/engine/session_app.py index 141b700..3fababd 100644 --- a/plugins/tito/agentix/tito/engine/session_app.py +++ b/plugins/tito/agentix/tito/engine/session_app.py @@ -23,6 +23,7 @@ from starlette.responses import Response from .errors import ( + ConcurrentSessionUpdateError, MessageValidationError, SessionError, SessionNotFoundError, @@ -30,7 +31,8 @@ ) from .pretokenize import get_tito_tokenizer from .processing import load_tokenizer -from .trajectory import GetSessionResponse, SessionRecord, SessionRegistry +from .record import build_turn_record +from .trajectory import GetSessionResponse, LinearTrajectory, SessionRecord, SessionRegistry from .upstream import Backend, get_upstream logger = logging.getLogger(__name__) @@ -64,8 +66,12 @@ def setup_session_routes(app: FastAPI, backend: Backend, args: Any) -> None: return # Exposed for operational introspection (session counts, tests). app.state.tito_registry = registry + # Finalize every open per-session record file on shutdown — the durable + # capture must be complete and closed however the process exits cleanly. + app.router.on_shutdown.append(registry.close) adapter = get_upstream(getattr(args, "backend_kind", "sglang")) + backend_kind = str(getattr(args, "backend_kind", "sglang") or "sglang") instance_id = getattr(args, "session_server_instance_id", None) @@ -117,10 +123,22 @@ async def delete_session(session_id: str) -> Response: @app.post("/sessions/{session_id}/v1/chat/completions") async def chat_completions(request: Request, session_id: str) -> Response: + # Opportunistic lifecycle sweep: a long-running gateway must not need + # new sessions to arrive before idle ones expire. + registry.sweep() session = registry.get_session(session_id) if session.closing: raise SessionNotFoundError(f"session not found: session_id={session_id}") + # In-flight guard for the WHOLE turn (including the un-locked + # upstream exchange): the registry never evicts while this is held. + session.inflight += 1 + try: + return await _chat_turn(request, session_id, session) + finally: + session.inflight -= 1 + + async def _chat_turn(request: Request, session_id: str, session: LinearTrajectory) -> Response: # Read + parse the body BEFORE taking the lock: the read lasts as long # as the client's upload — under the lock, one dribbling client would # wedge DELETE and every other operation on the session. @@ -140,9 +158,10 @@ async def chat_completions(request: Request, session_id: str) -> Response: if session.closing: raise SessionNotFoundError(f"session not found: session_id={session_id}") request_messages = request_body.get("messages", []) - prompt_token_ids = session.prepare_pretokenized( + prepared = session.prepare_prompt( request_messages, tools=request_body.get("tools"), tito_tokenizer=registry.tito_tokenizer ) + prompt_token_ids = prepared.token_ids # `version` advances on BOTH update and rollback; `num_assistant` # would be ABA-prone (a concurrent rollback+update restores it). expected_version = session.version @@ -156,12 +175,23 @@ async def chat_completions(request: Request, session_id: str) -> Response: harvest = turn.harvest # Phase 3: append the trajectory checkpoint (lock held briefly). + # A session that changed (or was deleted) while the turn was in + # flight can NOT absorb this completion — the trajectory would lie + # and the turn would silently vanish from the capture. Explicit 409: + # the tokens the agent would have received are discarded, the caller + # retries against the current session state. Interleave turns on one + # session was never supported; now it is loudly unsupported. async with session.lock: if session.closing: - return backend.build_proxy_response(turn.proxy_result) + raise ConcurrentSessionUpdateError( + f"session was deleted while the turn was in flight: session_id={session_id}; " + "the completed turn was not recorded" + ) if session.version != expected_version: - logger.warning("session %s changed during proxy; skipping state update", session_id) - return backend.build_proxy_response(turn.proxy_result) + raise ConcurrentSessionUpdateError( + f"session changed while the turn was in flight: session_id={session_id}; " + "interleaved turns on one session are rejected — retry on the current state" + ) session.update_pretokenized_state( request_messages, harvest.assistant_message, @@ -179,6 +209,27 @@ async def chat_completions(request: Request, session_id: str) -> Response: response=harvest.response, ) ) + if registry.record_sink is not None: + # Inside the lock so file order == trajectory order. Sink + # failures are logged, never fail the served turn. + registry.record_sink.append_turn( + session_id, + build_turn_record( + session_id=session_id, + request_id=request.headers.get("x-request-id"), + model=request_body.get("model"), + backend_kind=backend_kind, + prompt_token_ids=prompt_token_ids, + prompt_segments=prepared.segments, + prefix_stable=prepared.prefix_stable, + completion_token_ids=harvest.completion_token_ids, + completion_logprobs=harvest.completion_logprobs, + assistant_message=harvest.assistant_message, + finish_reason=harvest.finish_reason, + tokenizer_fingerprint=registry.tokenizer_fingerprint, + render_token_ids=harvest.render_token_ids, + ), + ) return backend.build_proxy_response(turn.proxy_result) @app.api_route("/sessions/{session_id}/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"]) diff --git a/plugins/tito/agentix/tito/engine/trajectory.py b/plugins/tito/agentix/tito/engine/trajectory.py index 1e1ac43..62e86ab 100644 --- a/plugins/tito/agentix/tito/engine/trajectory.py +++ b/plugins/tito/agentix/tito/engine/trajectory.py @@ -11,8 +11,12 @@ from __future__ import annotations import asyncio +import hashlib import logging +import time import uuid +from collections import OrderedDict +from collections.abc import Callable from dataclasses import dataclass, field from typing import Any @@ -22,6 +26,7 @@ from .errors import MessageValidationError, SessionNotFoundError, TokenizationError from .messages import assert_messages_append_only_with_allowed_role, message_matches from .pretokenize import TITOTokenizer +from .record import TurnRecordSink logger = logging.getLogger(__name__) @@ -54,12 +59,51 @@ class _RollbackPlan: discard_count: int +@dataclass(frozen=True, slots=True) +class PreparedPrompt: + """The pretokenized prompt for one turn, plus the capture metadata the + per-turn record needs. + + `segments` are half-open ``{"start", "end", "source"}`` spans over + `token_ids`; sources are ``render`` (from-scratch chat-template render), + ``prefix`` (the reused accumulated checkpoint), one per appended role + (``tool``/``user``/``system``), and ``generation_prompt``. + + `prefix_stable` is True iff `token_ids` extends the last *committed* + checkpoint (previous prompt + completion) — False after a retry rollback + or when the checkpoint window was outrun and the prompt was re-rendered. + """ + + token_ids: list[int] + segments: list[dict[str, Any]] + prefix_stable: bool + + +def _spans(parts: list[tuple[str, list[int]]]) -> list[dict[str, Any]]: + """Turn ``(source, ids)`` parts into cumulative-offset segment spans, + dropping empty parts (a template may render an empty suffix).""" + segments: list[dict[str, Any]] = [] + offset = 0 + for source, ids in parts: + if not ids: + continue + segments.append({"start": offset, "end": offset + len(ids), "source": source}) + offset += len(ids) + return segments + + @dataclass class LinearTrajectory: """Message history + accumulated token-ID checkpoints for one session.""" lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False, compare=False) closing: bool = field(default=False, repr=False, compare=False) + # Requests currently being served on this session (any phase, including + # the un-locked upstream exchange). The registry never evicts a session + # whose count is non-zero. + inflight: int = field(default=0, repr=False, compare=False) + # Idle clock for TTL eviction; the registry touches it on every lookup. + last_used: float = field(default_factory=time.monotonic, repr=False, compare=False) messages: list[dict[str, Any]] = field(default_factory=list) records: list[SessionRecord] = field(default_factory=list) trajectory_token_ids: list[list[int]] = field(default_factory=list) @@ -87,10 +131,27 @@ def prepare_pretokenized( """Build the full prompt token IDs for *request_messages*. First turn renders from scratch; later turns reuse the stored token prefix (rolling back at most one assistant step on a retry). Must be called under ``self.lock``.""" + return self.prepare_prompt(request_messages, tools, tito_tokenizer=tito_tokenizer).token_ids + + def prepare_prompt( + self, + request_messages: list[dict[str, Any]], + tools: list[dict[str, Any]] | None = None, + *, + tito_tokenizer: TITOTokenizer, + ) -> PreparedPrompt: + """`prepare_pretokenized` plus the capture metadata: segment spans over + the prompt ids and the prefix-stability verdict against the last + committed checkpoint. Must be called under ``self.lock``.""" + # The stability baseline is the checkpoint as COMMITTED — captured + # before any rollback below mutates it. + committed = list(self.token_ids) + if not self.messages: - return tito_tokenizer.render_messages( + ids = tito_tokenizer.render_messages( request_messages, tools=tools, add_generation_prompt=True, tokenize=True ) + return PreparedPrompt(token_ids=ids, segments=_spans([("render", ids)]), prefix_stable=True) # Plan first, mutate last: a request that fails validation must be a # pure 4xx with NO committed rollback side effects — otherwise a @@ -113,15 +174,20 @@ def prepare_pretokenized( # its update. Re-render from scratch: incremental == from-scratch # is the engine invariant, so this is lossless (never merge onto # an empty prefix, which would drop the whole stored history). - return tito_tokenizer.render_messages( + ids = tito_tokenizer.render_messages( request_messages, tools=tools, add_generation_prompt=True, tokenize=True ) - return tito_tokenizer.merge_tokens( - old_messages=self.messages, - new_messages=request_messages, - pretokenized_token_ids=self.token_ids, - tools=tools, - ) + parts: list[tuple[str, list[int]]] = [("render", ids)] + else: + # Same construction as `TITOTokenizer.merge_tokens`, decomposed so + # the segment boundaries survive into the per-turn record. + prefix = tito_tokenizer.fix_prefix(self.token_ids) + appended = tito_tokenizer.appended_segments(self.messages, request_messages, tools) + ids = prefix + [tid for _, seg_ids in appended for tid in seg_ids] + parts = [("prefix", prefix), *appended] + + prefix_stable = not committed or ids[: len(committed)] == committed + return PreparedPrompt(token_ids=ids, segments=_spans(parts), prefix_stable=prefix_stable) def update_pretokenized_state( self, @@ -235,17 +301,40 @@ def _apply_rollback(self, plan: _RollbackPlan) -> None: class SessionRegistry: - """Session ID -> trajectory map + shared tokenizer/comparator. Pure CRUD plus the - read-only mismatch computation; never mutates trajectory state itself.""" + """Session ID -> trajectory map + shared tokenizer/comparator, plus the + optional durable capture and lifecycle policy for long-running gateways. + + CRUD plus the read-only mismatch computation; never mutates trajectory + *token* state itself. From ``args`` (all optional, default off): + + - ``record_dir`` — per-session JSONL turn records (`TurnRecordSink`). + - ``session_ttl_seconds`` — evict sessions idle longer than this. + - ``max_sessions`` — LRU-evict beyond this many sessions. + + Eviction always flushes + finalizes the session's record file first and + NEVER touches a session with in-flight requests (``inflight`` > 0, a held + lock, or a pending close). ``on_evict`` lets the server layer drop + per-session routing state (e.g. the pool's sticky pin). + """ def __init__(self, args: Any, tokenizer: Any, *, tito_tokenizer: TITOTokenizer) -> None: - self.sessions: dict[str, LinearTrajectory] = {} + self.sessions: OrderedDict[str, LinearTrajectory] = OrderedDict() self.args = args self.tokenizer = tokenizer self.tito_tokenizer = tito_tokenizer self.comparator: TokenSeqComparator = tito_tokenizer.create_comparator() + record_dir = getattr(args, "record_dir", None) + self.record_sink: TurnRecordSink | None = TurnRecordSink(record_dir) if record_dir else None + self.session_ttl_seconds: float | None = getattr(args, "session_ttl_seconds", None) or None + self.max_sessions: int | None = getattr(args, "max_sessions", None) or None + self.on_evict: Callable[[str], None] | None = None + self.tokenizer_fingerprint: dict[str, str] = { + "checkpoint": str(getattr(args, "hf_checkpoint", "") or ""), + "chat_template_sha256": _chat_template_sha256(tokenizer, tito_tokenizer), + } def create_session(self) -> str: + self.sweep() session_id = uuid.uuid4().hex self.sessions[session_id] = LinearTrajectory() return session_id @@ -254,11 +343,61 @@ def get_session(self, session_id: str) -> LinearTrajectory: session = self.sessions.get(session_id) if session is None: raise SessionNotFoundError(f"session not found: session_id={session_id}") + session.last_used = time.monotonic() + self.sessions.move_to_end(session_id) return session def remove_session(self, session_id: str) -> None: if self.sessions.pop(session_id, None) is None: raise SessionNotFoundError(f"session not found: session_id={session_id}") + if self.record_sink is not None: + self.record_sink.finalize(session_id, reason="deleted") + + def sweep(self) -> list[str]: + """Evict expired-idle sessions (TTL) and least-recently-used overflow + (max-sessions). Returns the evicted ids. Skips any session that is + in-flight, locked, or closing — capacity may transiently overflow + rather than ever killing a live rollout.""" + evicted: list[str] = [] + ttl = self.session_ttl_seconds + if ttl is not None: + now = time.monotonic() + for session_id, session in list(self.sessions.items()): + if now - session.last_used < ttl: + break # LRU order: everything later was used more recently + if self._evict(session_id, session, reason="ttl_evicted"): + evicted.append(session_id) + if self.max_sessions is not None: + overflow = len(self.sessions) - self.max_sessions + if overflow > 0: + for session_id, session in list(self.sessions.items()): + if overflow <= 0: + break + if self._evict(session_id, session, reason="capacity_evicted"): + evicted.append(session_id) + overflow -= 1 + return evicted + + def _evict(self, session_id: str, session: LinearTrajectory, *, reason: str) -> bool: + if session.inflight > 0 or session.lock.locked() or session.closing: + return False + # Flush + finalize the durable record BEFORE the session becomes + # unreachable — eviction must never lose committed capture. + if self.record_sink is not None: + self.record_sink.finalize(session_id, reason=reason) + self.sessions.pop(session_id, None) + logger.info("evicted session %s (%s)", session_id, reason) + if self.on_evict is not None: + try: + self.on_evict(session_id) + except Exception: + logger.exception("on_evict callback failed for session %s", session_id) + return True + + def close(self) -> None: + """Finalize every open record file (process shutdown).""" + if self.record_sink is not None: + self.record_sink.close_all(reason="shutdown") def compute_session_mismatch(self, session: LinearTrajectory) -> list[dict] | None: """Compare accumulated token IDs against a from-scratch render. Read-only.""" @@ -273,3 +412,14 @@ def compute_session_mismatch(self, session: LinearTrajectory) -> list[dict] | No return [m.to_dict() for m in mismatches] except Exception as e: raise TokenizationError(f"failed to compute tito_session_mismatch: {e}") from e + + +def _chat_template_sha256(tokenizer: Any, tito_tokenizer: TITOTokenizer) -> str: + """SHA-256 of the chat template actually in effect for this gateway — + the fixed-template override when set (e.g. qwen3_fixed.jinja), otherwise + the tokenizer's own template. Pins the exact render rules a record's + token ids were produced under.""" + template = tito_tokenizer.chat_template_kwargs.get("chat_template") or getattr( + tokenizer, "chat_template", None + ) + return hashlib.sha256(str(template or "").encode()).hexdigest() diff --git a/plugins/tito/agentix/tito/engine/upstream.py b/plugins/tito/agentix/tito/engine/upstream.py index 40e0f7b..dc9777c 100644 --- a/plugins/tito/agentix/tito/engine/upstream.py +++ b/plugins/tito/agentix/tito/engine/upstream.py @@ -39,11 +39,22 @@ def build_proxy_response(self, result: dict) -> Response: ... @dataclass(frozen=True) class TurnHarvest: - """The token-exact outcome of a successful chat turn.""" + """The token-exact outcome of a successful chat turn. + + ``completion_logprobs`` pairs 1:1 with ``completion_token_ids`` — the + sampled-token logprobs both backends are forced to return (they are the + per-token cross-check AND the recorded rollout logprobs, so they are + retained here, never validated-then-discarded). ``render_token_ids`` is + vLLM-only: the render endpoint's from-scratch prompt ids, kept for the + cheap per-turn skew probe against the gateway's accumulated prompt ids. + """ response: dict assistant_message: dict completion_token_ids: list[int] + completion_logprobs: list[float] + finish_reason: str | None = None + render_token_ids: list[int] | None = None @dataclass(frozen=True) @@ -146,8 +157,10 @@ async def chat_turn( f"len(output_token_logprobs)={len(output_token_logprobs)} != completion_tokens={completion_tokens}" ) try: + # sglang entries are [logprob, token_id, text]. completion_token_ids = [t[1] for t in output_token_logprobs] - except (TypeError, IndexError, KeyError) as e: + completion_logprobs = [float(t[0]) for t in output_token_logprobs] + except (TypeError, ValueError, IndexError, KeyError) as e: raise UpstreamResponseError(f"malformed output_token_logprobs entry: {e}") from e return ChatTurn( @@ -156,6 +169,8 @@ async def chat_turn( response=response, assistant_message=assistant_message, completion_token_ids=completion_token_ids, + completion_logprobs=completion_logprobs, + finish_reason=_finish_reason(choice), ), ) @@ -219,6 +234,12 @@ async def chat_turn( if render_result["status_code"] != 200: return ChatTurn(proxy_result=render_result) generate_request = _parse_response_object(render_result["response_body"]) + # Keep render's from-scratch ids for the recorded skew probe before + # they are discarded from the generate request below. + rendered = generate_request.get("token_ids") + render_token_ids = ( + list(rendered) if isinstance(rendered, list) and all(isinstance(t, int) for t in rendered) else None + ) # The whole point: generate from the session's accumulated prompt ids, # not render's from-scratch re-render of the message history. generate_request["token_ids"] = prompt_token_ids @@ -232,7 +253,7 @@ async def chat_turn( if generate_result["status_code"] != 200: return ChatTurn(proxy_result=generate_result) generate_response = _parse_response_object(generate_result["response_body"]) - completion_token_ids = _harvest_generate_token_ids(generate_response) + completion_token_ids, completion_logprobs = _harvest_generate_tokens(generate_response) derender_request = { "model": model, @@ -267,11 +288,24 @@ async def chat_turn( response=response, assistant_message=assistant_message, completion_token_ids=completion_token_ids, + completion_logprobs=completion_logprobs, + # After the tool-call rewrite: the recorded finish_reason is + # what the agent actually received. + finish_reason=_finish_reason(_first_choice(response)), + render_token_ids=render_token_ids, ), ) -def _harvest_generate_token_ids(generate_response: dict) -> list[int]: +def _finish_reason(choice: dict) -> str | None: + reason = choice.get("finish_reason") + return reason if isinstance(reason, str) else None + + +def _harvest_generate_tokens(generate_response: dict) -> tuple[list[int], list[float]]: + """Completion (token_ids, logprobs) from a generate response — logprobs + are the sampled-token values from the SAME forward pass, kept for the + turn record, not just validated.""" choice = _first_choice(generate_response) token_ids = choice.get("token_ids") if not isinstance(token_ids, list) or not token_ids or not all(isinstance(t, int) for t in token_ids): @@ -287,7 +321,11 @@ def _harvest_generate_token_ids(generate_response: dict) -> list[int]: raise UpstreamResponseError( f"len(logprobs.content)={len(content)} != len(token_ids)={len(token_ids)}" ) - return list(token_ids) + try: + completion_logprobs = [float(entry["logprob"]) for entry in content] + except (TypeError, ValueError, KeyError) as e: + raise UpstreamResponseError(f"malformed logprobs.content entry: {e}") from e + return list(token_ids), completion_logprobs def get_upstream(kind: str) -> UpstreamAdapter: diff --git a/plugins/tito/agentix/tito/server.py b/plugins/tito/agentix/tito/server.py index cceb63a..22e1676 100644 --- a/plugins/tito/agentix/tito/server.py +++ b/plugins/tito/agentix/tito/server.py @@ -104,6 +104,12 @@ def __init__(self, args: Any, pool: BackendPool) -> None: self._backend = _PooledBackend(args, pool) self.app.router.on_shutdown.append(self._backend.aclose) setup_session_routes(self.app, self._backend, args) + # Registry eviction (TTL / max-sessions) ends a session exactly like + # DELETE does — drop its sticky pool pin too, so a reused id can't + # inherit a stale replica assignment. + registry = getattr(self.app.state, "tito_registry", None) + if registry is not None: + registry.on_evict = pool.forget self.app.middleware("http")(self._forget_on_delete) async def _forget_on_delete(self, request: Request, call_next: Any) -> Response: diff --git a/plugins/tito/pyproject.toml b/plugins/tito/pyproject.toml index 374e95d..e38d4ae 100644 --- a/plugins/tito/pyproject.toml +++ b/plugins/tito/pyproject.toml @@ -67,3 +67,6 @@ namespace = true testpaths = ["tests"] pythonpath = ["."] addopts = "-q" +markers = [ + "network: downloads a real tokenizer from the Hugging Face Hub — skips cleanly offline", +] diff --git a/plugins/tito/tests/test_gateway_vllm_http.py b/plugins/tito/tests/test_gateway_vllm_http.py index 9fff5a7..3ac1109 100644 --- a/plugins/tito/tests/test_gateway_vllm_http.py +++ b/plugins/tito/tests/test_gateway_vllm_http.py @@ -528,3 +528,72 @@ def test_unknown_backend_kind_fails_at_startup(tok, monkeypatch): args.backend_kind = "tgi" with pytest.raises(ValueError, match="backend_kind"): SessionServer(args, BackendPool([A])) + + +@pytest.mark.asyncio +async def test_vllm_turn_record_retains_logprobs_and_render_skew(tok, monkeypatch, tmp_path): + """The vLLM path used to validate-then-discard the generate logprobs; + with a record dir they must land in the tito.record.v1 line 1:1 with the + completion ids, alongside the per-turn render-skew probe (render's + from-scratch ids vs the gateway's accumulated prompt ids).""" + monkeypatch.setattr("agentix.tito.engine.session_app.load_tokenizer", lambda *a, **k: tok) + args = _args() + args.record_dir = str(tmp_path) + srv = SessionServer(args, BackendPool([A])) + replica = _VllmReplica() + srv._backend.client = httpx.AsyncClient(transport=httpx.MockTransport(replica.handler), timeout=5.0) + client = httpx.AsyncClient(transport=httpx.ASGITransport(app=srv.app), base_url="http://gw", timeout=5.0) + + sid = (await client.post("/sessions")).json()["session_id"] + r = await client.post( + f"/sessions/{sid}/v1/chat/completions", json=_CHAT, headers={"x-request-id": "req-v1"} + ) + assert r.status_code == 200 + + lines = [json.loads(line) for line in (tmp_path / f"{sid}.jsonl").read_text().splitlines()] + (rec,) = [line for line in lines if line["schema_version"] == "tito.record.v1"] + prompt_ids = replica.calls["generate"][0]["token_ids"] + + assert rec["backend_kind"] == "vllm" + assert rec["request_id"] == "req-v1" + assert rec["prompt_token_ids"] == prompt_ids + assert rec["completion_token_ids"] == [7, 8] + assert rec["completion_logprobs"] == [-0.1, -0.1] + assert len(rec["completion_logprobs"]) == len(rec["completion_token_ids"]) + assert rec["finish_reason"] == "stop" + # render returned [99, 98] (a from-scratch render); the gateway generated + # from its own pretokenized ids -> skew observed, recorded, non-blocking. + assert rec["render_skew"] == {"equal": False, "first_divergence": 0} + assert rec["prefix_stable"] is True + + +@pytest.mark.asyncio +async def test_vllm_tool_call_record_carries_rewritten_finish_reason(tok, monkeypatch, tmp_path): + """The recorded finish_reason is what the agent actually received — + i.e. AFTER the gateway's tool_calls rewrite of derender's verbatim + "stop".""" + monkeypatch.setattr("agentix.tito.engine.session_app.load_tokenizer", lambda *a, **k: tok) + args = _args() + args.record_dir = str(tmp_path) + srv = SessionServer(args, BackendPool([A])) + replica = _VllmReplica() + replica.message = { + "role": "assistant", + "content": None, + "tool_calls": [{ + "id": "call_1", "type": "function", + "function": {"name": "compute", "arguments": "{}"}, + }], + } + srv._backend.client = httpx.AsyncClient(transport=httpx.MockTransport(replica.handler), timeout=5.0) + client = httpx.AsyncClient(transport=httpx.ASGITransport(app=srv.app), base_url="http://gw", timeout=5.0) + + sid = (await client.post("/sessions")).json()["session_id"] + r = await client.post(f"/sessions/{sid}/v1/chat/completions", json={**_CHAT, "tools": _TOOLS}) + assert r.status_code == 200 + assert r.json()["choices"][0]["finish_reason"] == "tool_calls" + + lines = [json.loads(line) for line in (tmp_path / f"{sid}.jsonl").read_text().splitlines()] + (rec,) = [line for line in lines if line["schema_version"] == "tito.record.v1"] + assert rec["finish_reason"] == "tool_calls" + assert rec["assistant_message"]["tool_calls"][0]["id"] == "call_1" diff --git a/plugins/tito/tests/test_qwen3_golden.py b/plugins/tito/tests/test_qwen3_golden.py new file mode 100644 index 0000000..5c4b162 --- /dev/null +++ b/plugins/tito/tests/test_qwen3_golden.py @@ -0,0 +1,224 @@ +"""Golden test: the incremental==from-scratch invariant on the REAL Qwen3 +tokenizer + the bundled fixed chat template. + +Every other engine test runs on a tiny in-memory WordLevel tokenizer, which +cannot catch real-tokenizer failure modes (BPE merge boundaries at segment +junctions, ``/tool-tag added-token splitting, the missing trailing +newline after `<|im_end|>`). This module downloads the tokenizer-only files +for Qwen/Qwen3-0.6B at test time (a few hundred KB; the HF cache is reused on +later runs) and drives a full multi-turn tool-calling session through +`LinearTrajectory.prepare_prompt`, asserting: + +- each turn's incrementally merged prompt ids equal a from-scratch render of + the same request (token-exact, not just segment-equivalent); +- the accumulated trajectory equals the from-scratch render of the whole + conversation (modulo the trailing newline Qwen3 omits at stop); +- the read-time mismatch audit reports clean; +- the `prompt_segments` spans the per-turn record persists decode to the + expected role boundaries. + +Offline behavior: if the tokenizer is neither cached nor downloadable the +module SKIPS (marker: `network`) — it never fails a disconnected run. +""" + +from __future__ import annotations + +import json + +import pytest +from agentix.tito.engine.pretokenize import get_tito_tokenizer +from agentix.tito.engine.trajectory import LinearTrajectory, SessionRecord, SessionRegistry + +pytestmark = pytest.mark.network + +_REPO = "Qwen/Qwen3-0.6B" + + +@pytest.fixture(scope="module") +def qwen3_tok(): + from transformers import AutoTokenizer + + try: # cache first: offline runs with a warm cache still exercise the golden + return AutoTokenizer.from_pretrained(_REPO, local_files_only=True) + except Exception: + pass + try: # one download attempt; a blocked network is a skip, not a failure + return AutoTokenizer.from_pretrained(_REPO) + except Exception as exc: # noqa: BLE001 - hub errors vary by transport + pytest.skip(f"Qwen3 tokenizer unavailable (offline?): {type(exc).__name__}: {exc}") + + +_TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Look up the current weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } +] + + +def _simulate_completion(tt, request_messages, assistant_message, prompt_ids, tools): + """The completion token ids a template-canonical model would emit for + *assistant_message*: the from-scratch render of request+assistant minus + the prompt prefix, with the trailing newline dropped (Qwen3 stops at + `<|im_end|>` and never emits the newline the template writes).""" + full = tt.render_messages( + request_messages + [assistant_message], tools=tools, add_generation_prompt=False, tokenize=True + ) + assert full[: len(prompt_ids)] == prompt_ids, "assistant render must extend the generation prompt" + completion = full[len(prompt_ids):] + newline_id = tt.tokenizer.encode("\n", add_special_tokens=False)[0] + assert completion and completion[-1] == newline_id + return completion[:-1] + + +def test_qwen3_incremental_equals_from_scratch_multi_turn_tool_calls(qwen3_tok): + tt = get_tito_tokenizer(qwen3_tok, "qwen3", allowed_append_roles=("tool", "user")) + registry = SessionRegistry(None, qwen3_tok, tito_tokenizer=tt) + tr = LinearTrajectory() + + system = {"role": "system", "content": "You are a terse weather assistant."} + user1 = {"role": "user", "content": "What's the weather in Paris right now?"} + assistant1 = { + "role": "assistant", + "content": "", + "reasoning_content": "The user wants current weather; call the tool.", + "tool_calls": [ + { + "id": "call_0001", + "type": "function", + "function": {"name": "get_weather", "arguments": json.dumps({"city": "Paris"})}, + } + ], + } + tool1 = {"role": "tool", "content": '{"temp_c": 21, "sky": "clear"}', "tool_call_id": "call_0001"} + assistant2 = { + "role": "assistant", + "content": "Paris is 21°C and clear.", + "reasoning_content": "Tool says 21C, clear. Answer briefly.", + } + user2 = {"role": "user", "content": "And in London?"} + assistant3 = { + "role": "assistant", + "content": "", + "reasoning_content": "Same tool, city London.", + "tool_calls": [ + { + "id": "call_0002", + "type": "function", + "function": {"name": "get_weather", "arguments": json.dumps({"city": "London"})}, + } + ], + } + + turns = [ + ([system, user1], assistant1), + ([system, user1, assistant1, tool1], assistant2), + ([system, user1, assistant1, tool1, assistant2, user2], assistant3), + ] + + all_segment_sources: list[list[str]] = [] + for request_messages, assistant in turns: + prepared = tr.prepare_prompt(request_messages, _TOOLS, tito_tokenizer=tt) + + # THE invariant, token-exact on the real tokenizer: the incrementally + # merged prompt equals a from-scratch render of the same request. + from_scratch = tt.render_messages( + request_messages, tools=_TOOLS, add_generation_prompt=True, tokenize=True + ) + assert prepared.token_ids == from_scratch + assert prepared.prefix_stable is True + + # Segment spans tile the prompt exactly. + assert prepared.segments[0]["start"] == 0 + assert prepared.segments[-1]["end"] == len(prepared.token_ids) + for left, right in zip(prepared.segments, prepared.segments[1:], strict=False): + assert left["end"] == right["start"] + all_segment_sources.append([s["source"] for s in prepared.segments]) + + completion = _simulate_completion(tt, request_messages, assistant, prepared.token_ids, _TOOLS) + tr.update_pretokenized_state( + request_messages, + assistant, + prompt_token_ids=prepared.token_ids, + completion_token_ids=completion, + max_trim_tokens=tt.max_trim_tokens, + ) + # As the gateway does — the audit reads tools off the last record. + tr.append_record(SessionRecord( + timestamp=0.0, method="POST", path="/v1/chat/completions", status_code=200, + request={"model": "m", "messages": request_messages, "tools": _TOOLS}, response={}, + )) + + assert all_segment_sources == [ + ["render"], + ["prefix", "tool", "generation_prompt"], + ["prefix", "user", "generation_prompt"], + ] + + # Accumulated trajectory == from-scratch render of the full conversation + # (fix_prefix restores the trailing newline Qwen3 omitted at stop). + final_messages = turns[-1][0] + [assistant3] + assert tt.fix_prefix(tr.token_ids) == tt.render_messages( + final_messages, tools=_TOOLS, add_generation_prompt=False, tokenize=True + ) + + # The read-time audit agrees: no structural or content mismatch. + assert registry.compute_session_mismatch(tr) == [] + + +def test_qwen3_prompt_segment_boundaries_decode_to_role_markers(qwen3_tok): + """The spans persisted as record.prompt_segments decode to the expected + template boundaries — the loss-mask construction material downstream.""" + tt = get_tito_tokenizer(qwen3_tok, "qwen3", allowed_append_roles=("tool", "user")) + tr = LinearTrajectory() + + request1 = [ + {"role": "system", "content": "You are a terse weather assistant."}, + {"role": "user", "content": "What's the weather in Paris right now?"}, + ] + assistant1 = { + "role": "assistant", + "content": "", + "reasoning_content": "Call the tool.", + "tool_calls": [ + { + "id": "call_0001", + "type": "function", + "function": {"name": "get_weather", "arguments": json.dumps({"city": "Paris"})}, + } + ], + } + prepared1 = tr.prepare_prompt(request1, _TOOLS, tito_tokenizer=tt) + completion1 = _simulate_completion(tt, request1, assistant1, prepared1.token_ids, _TOOLS) + tr.update_pretokenized_state( + request1, assistant1, + prompt_token_ids=prepared1.token_ids, + completion_token_ids=completion1, + max_trim_tokens=tt.max_trim_tokens, + ) + + request2 = request1 + [assistant1, {"role": "tool", "content": "21C clear", "tool_call_id": "call_0001"}] + prepared2 = tr.prepare_prompt(request2, _TOOLS, tito_tokenizer=tt) + + def decode(segment): + return qwen3_tok.decode( + prepared2.token_ids[segment["start"]:segment["end"]], skip_special_tokens=False + ) + + prefix_seg, tool_seg, gen_seg = prepared2.segments + assert prefix_seg["source"] == "prefix" + # fix_prefix restored the newline after the completion's final <|im_end|>. + assert decode(prefix_seg).endswith("<|im_end|>\n") + assert "" in decode(prefix_seg) # the assistant turn lives in the prefix + assert tool_seg["source"] == "tool" + assert "" in decode(tool_seg) and "21C clear" in decode(tool_seg) + assert gen_seg["source"] == "generation_prompt" + assert decode(gen_seg) == "<|im_start|>assistant\n" diff --git a/plugins/tito/tests/test_record.py b/plugins/tito/tests/test_record.py new file mode 100644 index 0000000..ee03597 --- /dev/null +++ b/plugins/tito/tests/test_record.py @@ -0,0 +1,381 @@ +"""Contract tests for per-turn record persistence and session lifecycle. + +Drive the REAL gateway app over ASGI (same harness as test_gateway_http / +test_gateway_vllm_http: tiny in-memory WordLevel tokenizer, fake replica via +httpx.MockTransport) and assert the durable-capture contract: + +- one flushed `tito.record.v1` JSON line per committed turn (crash-safe: + the file is complete after every turn, before any close); +- the record shape: ids/logprobs 1:1, prompt segment spans, tokenizer + fingerprint, request_id from the x-request-id header; +- `prefix_stable` true on the linear path, false (but still served + + recorded) after a history rewrite that rolls back a checkpoint; +- interleaved turns on one session are an explicit 409, never a silently + dropped update; +- TTL / capacity eviction finalizes the record file first and never touches + an in-flight session; DELETE appends the final `tito.session.v1` line; +- no --record-dir -> no files, behavior unchanged. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import json +import types +from pathlib import Path + +import httpx +import pytest +from agentix.tito.engine.record import compute_render_skew +from agentix.tito.pool import BackendPool +from agentix.tito.server import SessionServer +from tokenizers import Tokenizer, models, pre_tokenizers +from transformers import PreTrainedTokenizerFast + +A = "http://replica-a:8000" + +_CHAT = {"model": "m", "messages": [{"role": "user", "content": "Hello"}]} + + +@pytest.fixture(scope="module") +def tok(): + specials = ["", "", "", "<|im_start|>", "<|im_end|>"] + words = ["system", "user", "assistant", "tool", "You", "are", "ok", "done", "Hello"] + vocab = {t: i for i, t in enumerate(specials + words)} + tk = Tokenizer(models.WordLevel(vocab=vocab, unk_token="")) + tk.pre_tokenizer = pre_tokenizers.Whitespace() + t = PreTrainedTokenizerFast( + tokenizer_object=tk, unk_token="", bos_token="", eos_token="", + additional_special_tokens=["<|im_start|>", "<|im_end|>"], + ) + t.chat_template = ( + "{%- for m in messages -%}<|im_start|>{{ m['role'] }} {{ m['content'] or '' }}<|im_end|>{%- endfor -%}" + "{%- if add_generation_prompt -%}<|im_start|>assistant {%- endif -%}" + ) + return t + + +def _args(**overrides): + ns = types.SimpleNamespace( + hf_checkpoint="tiny-in-memory", + chat_template_path=None, + tito_allowed_append_roles=None, + tito_model="default", + session_server_instance_id=None, + router_timeout=5.0, + ) + for key, value in overrides.items(): + setattr(ns, key, value) + return ns + + +class _Replica: + """sglang-shaped fake replica; optionally parks a request on an event + (for the interleave test).""" + + def __init__(self) -> None: + self.calls: list[dict] = [] + self.completion_ids = [7, 8] # "ok done" + self.logprobs = [-0.25, -0.5] + self.message: dict = {"role": "assistant", "content": "ok done"} + self.hold: asyncio.Event | None = None + self.hold_marker: str | None = None + + async def handler(self, request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + self.calls.append(body) + if self.hold is not None and self.hold_marker in json.dumps(body): + await self.hold.wait() + ids = list(self.completion_ids) + return httpx.Response(200, json={ + "id": "c1", "object": "chat.completion", "model": "m", + "choices": [{ + "index": 0, + "finish_reason": "stop", + "message": dict(self.message), + "meta_info": { + "output_token_logprobs": [[lp, t, ""] for lp, t in zip(self.logprobs, ids, strict=True)], + "completion_tokens": len(ids), + }, + }], + "usage": {"prompt_tokens": 3, "completion_tokens": len(ids), "total_tokens": 5}, + }) + + +def _make_gateway(tok, monkeypatch, **arg_overrides): + monkeypatch.setattr("agentix.tito.engine.session_app.load_tokenizer", lambda *a, **k: tok) + pool = BackendPool([A]) + srv = SessionServer(_args(**arg_overrides), pool) + replica = _Replica() + srv._backend.client = httpx.AsyncClient(transport=httpx.MockTransport(replica.handler), timeout=5.0) + client = httpx.AsyncClient(transport=httpx.ASGITransport(app=srv.app), base_url="http://gw", timeout=5.0) + return client, replica, srv, pool + + +def _lines(path: Path) -> list[dict]: + return [json.loads(line) for line in path.read_text().splitlines()] + + +def _records(path: Path) -> list[dict]: + return [r for r in _lines(path) if r["schema_version"] == "tito.record.v1"] + + +def _meta(path: Path) -> list[dict]: + return [r for r in _lines(path) if r["schema_version"] == "tito.session.v1"] + + +@pytest.mark.asyncio +async def test_record_line_shape_and_crash_safety(tok, monkeypatch, tmp_path): + """One complete, flushed record line per committed turn — readable after + EVERY turn with the file still open (a crash mid-rollout loses nothing + already committed) — carrying the full tito.record.v1 shape.""" + client, replica, srv, _ = _make_gateway(tok, monkeypatch, record_dir=str(tmp_path)) + sid = (await client.post("/sessions")).json()["session_id"] + path = tmp_path / f"{sid}.jsonl" + + r = await client.post( + f"/sessions/{sid}/v1/chat/completions", json=_CHAT, headers={"x-request-id": "req-001"} + ) + assert r.status_code == 200 + + # Crash safety: the line is on disk NOW, without any close/flush call. + (rec,) = _records(path) + accumulated = (await client.get(f"/sessions/{sid}")).json()["metadata"]["accumulated_token_ids"] + + assert rec["schema_version"] == "tito.record.v1" + assert rec["session_id"] == sid + assert rec["turn_index"] == 0 + assert rec["request_id"] == "req-001" + assert rec["model"] == "m" + assert rec["backend_kind"] == "sglang" + assert rec["prompt_token_ids"] + rec["completion_token_ids"] == accumulated + assert rec["completion_token_ids"] == [7, 8] + assert rec["completion_logprobs"] == [-0.25, -0.5] + assert len(rec["completion_logprobs"]) == len(rec["completion_token_ids"]) + assert rec["assistant_message"]["content"] == "ok done" + assert rec["finish_reason"] == "stop" + assert rec["prefix_stable"] is True + assert rec["render_skew"] is None # sglang exposes no render ids + assert rec["tokenizer_fingerprint"] == { + "checkpoint": "tiny-in-memory", + "chat_template_sha256": hashlib.sha256(tok.chat_template.encode()).hexdigest(), + } + # First turn: one from-scratch render segment covering the whole prompt. + assert rec["prompt_segments"] == [ + {"start": 0, "end": len(rec["prompt_token_ids"]), "source": "render"} + ] + assert "ts" in rec + + # Turn 2 (tool append): line 2 is on disk immediately; segments show the + # reused prefix boundary and the per-role suffixes. + followup = { + "model": "m", + "messages": [ + *_CHAT["messages"], + {"role": "assistant", "content": "ok done"}, + {"role": "tool", "content": "done"}, + ], + } + r = await client.post(f"/sessions/{sid}/v1/chat/completions", json=followup) + assert r.status_code == 200 + rec1, rec2 = _records(path) + assert rec2["turn_index"] == 1 + assert rec2["request_id"] is None # no x-request-id header sent + assert rec2["prefix_stable"] is True + assert [s["source"] for s in rec2["prompt_segments"]] == ["prefix", "tool", "generation_prompt"] + # Segment spans tile the prompt exactly, and the prefix span IS the + # previous checkpoint (prompt + completion of turn 1). + assert rec2["prompt_segments"][0]["start"] == 0 + prefix_end = rec2["prompt_segments"][0]["end"] + assert rec2["prompt_token_ids"][:prefix_end] == rec1["prompt_token_ids"] + rec1["completion_token_ids"] + for left, right in zip(rec2["prompt_segments"], rec2["prompt_segments"][1:], strict=False): + assert left["end"] == right["start"] + assert rec2["prompt_segments"][-1]["end"] == len(rec2["prompt_token_ids"]) + + +@pytest.mark.asyncio +async def test_history_rewrite_records_prefix_stable_false_and_still_serves(tok, monkeypatch, tmp_path): + """A compaction-style rewrite of the last tool turn rolls back one + checkpoint: the turn is SERVED and RECORDED, flagged prefix_stable=false + so a trainer can split or reject it instead of splicing a lie.""" + client, replica, srv, _ = _make_gateway(tok, monkeypatch, record_dir=str(tmp_path)) + sid = (await client.post("/sessions")).json()["session_id"] + path = tmp_path / f"{sid}.jsonl" + + await client.post(f"/sessions/{sid}/v1/chat/completions", json=_CHAT) + base = [*_CHAT["messages"], {"role": "assistant", "content": "ok done"}] + r = await client.post( + f"/sessions/{sid}/v1/chat/completions", + json={"model": "m", "messages": [*base, {"role": "tool", "content": "done"}]}, + ) + assert r.status_code == 200 + + # Rewrite the tool turn (divergent content) — single-step rollback. + r = await client.post( + f"/sessions/{sid}/v1/chat/completions", + json={"model": "m", "messages": [*base, {"role": "tool", "content": "You"}]}, + ) + assert r.status_code == 200 + + recs = _records(path) + assert [rec["prefix_stable"] for rec in recs] == [True, True, False] + assert [rec["turn_index"] for rec in recs] == [0, 1, 2] + # The rewritten turn still tiles cleanly over its own prompt. + assert [s["source"] for s in recs[2]["prompt_segments"]] == ["prefix", "tool", "generation_prompt"] + + +@pytest.mark.asyncio +async def test_interleaved_turn_is_explicit_409(tok, monkeypatch, tmp_path): + """Two in-flight turns on one session: the loser's completion cannot be + committed to the rewound/advanced trajectory. It must be an explicit 409 + — the previous behavior served the response and silently dropped it from + the capture, which is unacceptable data loss in production.""" + client, replica, srv, _ = _make_gateway(tok, monkeypatch, record_dir=str(tmp_path)) + sid = (await client.post("/sessions")).json()["session_id"] + + replica.hold = asyncio.Event() + replica.hold_marker = "HOLD-ME" + + slow = asyncio.create_task(client.post( + f"/sessions/{sid}/v1/chat/completions", + json={"model": "m", "messages": [{"role": "user", "content": "Hello HOLD-ME"}]}, + )) + while not replica.calls: # let the slow turn reach the backend + await asyncio.sleep(0.01) + + fast = await client.post(f"/sessions/{sid}/v1/chat/completions", json=_CHAT) + assert fast.status_code == 200 # the interleaver wins and is recorded + + replica.hold.set() + lost = await slow + assert lost.status_code == 409 + assert "changed while the turn was in flight" in lost.json()["error"] + + # Exactly the winner is recorded — no phantom line for the 409. + recs = _records(tmp_path / f"{sid}.jsonl") + assert len(recs) == 1 + assert recs[0]["prompt_token_ids"] # the fast turn's record + + +@pytest.mark.asyncio +async def test_delete_finalizes_record_file(tok, monkeypatch, tmp_path): + """DELETE-after-harvest: the documented rollout end appends the final + tito.session.v1 line and closes the file.""" + client, _, srv, _ = _make_gateway(tok, monkeypatch, record_dir=str(tmp_path)) + sid = (await client.post("/sessions")).json()["session_id"] + await client.post(f"/sessions/{sid}/v1/chat/completions", json=_CHAT) + assert (await client.delete(f"/sessions/{sid}")).status_code == 204 + + path = tmp_path / f"{sid}.jsonl" + (meta,) = _meta(path) + assert meta == { + "schema_version": "tito.session.v1", + "session_id": sid, + "turns": 1, + "reason": "deleted", + "ts": meta["ts"], + } + assert _lines(path)[-1] == meta # the meta line is the LAST line + + +@pytest.mark.asyncio +async def test_ttl_eviction_flushes_record_file_and_forgets_pin(tok, monkeypatch, tmp_path): + client, _, srv, pool = _make_gateway( + tok, monkeypatch, record_dir=str(tmp_path), session_ttl_seconds=0.05 + ) + sid = (await client.post("/sessions")).json()["session_id"] + await client.post(f"/sessions/{sid}/v1/chat/completions", json=_CHAT) + assert sid in pool._assigned # noqa: SLF001 + + await asyncio.sleep(0.06) + registry = srv.app.state.tito_registry + assert registry.sweep() == [sid] + + path = tmp_path / f"{sid}.jsonl" + (meta,) = _meta(path) + assert meta["reason"] == "ttl_evicted" + assert meta["turns"] == 1 + assert len(_records(path)) == 1 # committed turns survived the eviction + assert sid not in pool._assigned # noqa: SLF001 - eviction drops the sticky pin + assert (await client.get(f"/sessions/{sid}")).status_code == 404 + + +@pytest.mark.asyncio +async def test_eviction_never_touches_inflight_sessions(tok, monkeypatch, tmp_path): + """A session with a turn parked on the backend (lock NOT held — phase 2) + must survive TTL and capacity sweeps until the turn finishes.""" + client, replica, srv, _ = _make_gateway( + tok, monkeypatch, record_dir=str(tmp_path), session_ttl_seconds=0.01, max_sessions=1 + ) + sid = (await client.post("/sessions")).json()["session_id"] + + replica.hold = asyncio.Event() + replica.hold_marker = "HOLD-ME" + slow = asyncio.create_task(client.post( + f"/sessions/{sid}/v1/chat/completions", + json={"model": "m", "messages": [{"role": "user", "content": "Hello HOLD-ME"}]}, + )) + while not replica.calls: + await asyncio.sleep(0.01) + + await asyncio.sleep(0.02) # idle past the TTL while in flight + registry = srv.app.state.tito_registry + assert registry.sweep() == [] # in-flight: skipped by TTL AND capacity + assert sid in registry.sessions + + replica.hold.set() + assert (await slow).status_code == 200 + await asyncio.sleep(0.02) + assert registry.sweep() == [sid] # idle again -> now evictable + + +@pytest.mark.asyncio +async def test_capacity_eviction_is_lru_and_flushes(tok, monkeypatch, tmp_path): + client, _, srv, _ = _make_gateway(tok, monkeypatch, record_dir=str(tmp_path), max_sessions=1) + sid_old = (await client.post("/sessions")).json()["session_id"] + await client.post(f"/sessions/{sid_old}/v1/chat/completions", json=_CHAT) + sid_new = (await client.post("/sessions")).json()["session_id"] + + # The next request's sweep trims the overflow: LRU (sid_old) goes. + r = await client.post(f"/sessions/{sid_new}/v1/chat/completions", json=_CHAT) + assert r.status_code == 200 + + registry = srv.app.state.tito_registry + assert sid_old not in registry.sessions + assert sid_new in registry.sessions + (meta,) = _meta(tmp_path / f"{sid_old}.jsonl") + assert meta["reason"] == "capacity_evicted" + + +@pytest.mark.asyncio +async def test_no_record_dir_writes_nothing_and_behavior_is_unchanged(tok, monkeypatch, tmp_path): + client, _, srv, _ = _make_gateway(tok, monkeypatch) # no record_dir + sid = (await client.post("/sessions")).json()["session_id"] + r = await client.post(f"/sessions/{sid}/v1/chat/completions", json=_CHAT) + assert r.status_code == 200 + got = (await client.get(f"/sessions/{sid}")).json() + assert len(got["records"]) == 1 + assert srv.app.state.tito_registry.record_sink is None + assert list(tmp_path.iterdir()) == [] # nothing written anywhere + + +@pytest.mark.asyncio +async def test_shutdown_finalizes_open_record_files(tok, monkeypatch, tmp_path): + client, _, srv, _ = _make_gateway(tok, monkeypatch, record_dir=str(tmp_path)) + sid = (await client.post("/sessions")).json()["session_id"] + await client.post(f"/sessions/{sid}/v1/chat/completions", json=_CHAT) + + registry = srv.app.state.tito_registry + assert registry.close in srv.app.router.on_shutdown # lifespan wiring + registry.close() + (meta,) = _meta(tmp_path / f"{sid}.jsonl") + assert meta["reason"] == "shutdown" + + +def test_compute_render_skew_contract(): + assert compute_render_skew(None, [1, 2]) is None + assert compute_render_skew([1, 2], [1, 2]) == {"equal": True, "first_divergence": None} + assert compute_render_skew([9, 2], [1, 2]) == {"equal": False, "first_divergence": 0} + assert compute_render_skew([1, 2], [1, 2, 3]) == {"equal": False, "first_divergence": 2} + assert compute_render_skew([1, 2, 3], [1, 2]) == {"equal": False, "first_divergence": 2} diff --git a/pyproject.toml b/pyproject.toml index ada8e75..9c732d3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -168,7 +168,10 @@ python_files = ["test_*.py"] # `e2e` tests build a real bundle image (need docker, take minutes). # `addopts` excludes them by default — the `e2e` CI job runs them with # `-m e2e`, which overrides this. -markers = ["e2e: builds a real bundle image — needs docker, slow"] +markers = [ + "e2e: builds a real bundle image — needs docker, slow", + "network: downloads a real tokenizer from the Hugging Face Hub — skips cleanly offline", +] # Coverage runs against `agentix` and the plugin packages; the # subprocess worker (`agentix/runtime/server/worker/process.py`) is # tracked separately because it executes in a child process and the From e74e95b00a644d0a78431f5e80b266f84173efb0 Mon Sep 17 00:00:00 2001 From: Lingrui Mei Date: Fri, 24 Jul 2026 05:53:26 +0800 Subject: [PATCH 3/5] abridge: tito composition in serve mode + capture-aligned Recorder rows MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wire the direct-mode server onto the token-recording session gateway and make the message-level capture joinable against its token records: - serve --tito-url (env TITO_URL, mutually exclusive with --upstream-base-url): each caller session composes AnthropicToOpenAI(SessionForward(tito_url).handler(), model=...) — the gateway sees the OpenAI chat body, owns render/generate and the token record, one caller key == one gateway session. Default mode unchanged. - serve --record-dir (env ABRIDGE_RECORD_DIR), either mode: wrap each session's client in Recorder(/.jsonl). - Recorder rows gain request_id + optional session_id. The request id is minted at the capture layer and bound on a context var; Forward and AnthropicFromOpenAIClient reuse it for the upstream x-request-id, so a message row and the gateway's per-turn token record share one join key (and orphan rows from tunnel-504 agent retries become deduplicable). The record file now opens lazily — a route-enumeration probe leaves no empty file. - Translation contract pin: TRANSLATION_SPEC_SHA (sha256 of the transform module source) exported and reported on GET /_health so downstream data contracts can pin the exact translation in effect. - Transform gaps closed: Anthropic tool_choice now maps to OpenAI (auto->auto, any->required, none->none, {type:tool,name}->named function); the four tool_choice xfails flip green. disable_parallel_tool_use stays dropped, matching the reference translator. Assistant thinking history is now forwarded as reasoning_content (the vLLM/sglang dialect our downstreams read) — the 35_assistant_thinking_history xfail stays, re-documented as a deliberate dialect divergence from the LiteLLM golden (thinking_blocks + crypto signature), with a positive contract test. Tests: composition e2e over a real-HTTP fake gateway (session-per-key, forced non-streaming, model override, identity headers, SSE replay, record-dir row/record join), Recorder id fields + transport alignment + lazy open, /_health sha pin, mode flag validation. Co-Authored-By: Claude Fable 5 --- plugins/abridge/README.md | 33 ++- plugins/abridge/agentix/bridge/_request_id.py | 37 ++++ .../bridge/clients/_anthropic_transforms.py | 63 +++++- .../bridge/clients/anthropic_from_openai.py | 5 +- plugins/abridge/agentix/bridge/forward.py | 5 +- plugins/abridge/agentix/bridge/recorder.py | 41 +++- plugins/abridge/agentix/bridge/serve.py | 133 ++++++++++-- .../test_anthropic_transform_fixtures.py | 29 ++- plugins/abridge/tests/test_recorder.py | 65 ++++++ plugins/abridge/tests/test_serve.py | 17 +- .../abridge/tests/test_tito_composition.py | 191 ++++++++++++++++++ 11 files changed, 575 insertions(+), 44 deletions(-) create mode 100644 plugins/abridge/agentix/bridge/_request_id.py create mode 100644 plugins/abridge/tests/test_tito_composition.py diff --git a/plugins/abridge/README.md b/plugins/abridge/README.md index a33d561..05cece5 100644 --- a/plugins/abridge/README.md +++ b/plugins/abridge/README.md @@ -132,17 +132,22 @@ Claude-speaking agent in front of an OpenAI-shaped recording gateway. ### Record the tunnel traffic `Recorder` wraps any handler client and appends one JSONL line per served -call — `{ts, path, request, response}` — flushed as it goes, so the file is -complete up to the last call even if the host dies mid-rollout. It exposes -the wrapped client's routes and closes it on teardown, so it drops in -transparently: +call — `{ts, path, request_id, session_id?, request, response}` — flushed +as it goes, so the file is complete up to the last call even if the host +dies mid-rollout. It exposes the wrapped client's routes and closes it on +teardown, so it drops in transparently: ```python from agentix.bridge import Proxy, Recorder -proxy = Proxy(Recorder(client, "runs/rollout-42.jsonl")) +proxy = Proxy(Recorder(client, "runs/rollout-42.jsonl", session_id="rollout-42")) ``` +The `request_id` in each row is the same id the transport stamps as +`x-request-id` on the upstream hop (bound through a context var), so a +message-level row joins a downstream token recorder's per-turn record; +`session_id`, when given, tags every row with the rollout identity. + ## Writing your own handler Any class with `@on(path)`-decorated methods works. No base class to @@ -258,6 +263,24 @@ code, so when you expose it to them (`--host`), also set `build_session_app`) so only keys your harness minted are served; everything else gets a 401. +Two more serve options: + +* `--tito-url http://tito:30000` (mutually exclusive with + `--upstream-base-url`) — put the Anthropic shell in front of a + token-recording session gateway instead of a plain engine: each caller + session composes `AnthropicToOpenAI(SessionForward(tito_url).handler())`, + so the gateway sees the OpenAI chat body, owns render/generate and the + token record, and one caller key maps to one gateway session. +* `--record-dir DIR` (either mode) — wrap each session's client in a + `Recorder` writing `DIR/.jsonl`; rows carry `session_id` + + `request_id`, and the same `request_id` reaches the upstream as + `x-request-id`, so message rows join the gateway's token records. + +`GET /_health` reports `translation_spec_sha` — the SHA-256 of the +Anthropic↔OpenAI transform module source — so downstream data contracts +can pin the exact translation their captured trajectories were produced +under. + Programmatic surface in `agentix.bridge.serve`: `build_app(*clients)` (shared session) and `build_session_app(factory)` (one client per caller key; LRU-bounded with in-flight-safe eviction — an evicted diff --git a/plugins/abridge/agentix/bridge/_request_id.py b/plugins/abridge/agentix/bridge/_request_id.py new file mode 100644 index 0000000..9ac344f --- /dev/null +++ b/plugins/abridge/agentix/bridge/_request_id.py @@ -0,0 +1,37 @@ +"""Per-call request-id propagation between capture and transport layers. + +The tunnel deliberately carries no HTTP metadata (see `proxy.Request`), so a +request id can't ride the `Request` object. But capture and transport must +agree on ONE id per call: the `Recorder` writes a `request_id` into its JSONL +row, and the transport layer (`Forward`, the SDK clients) stamps +`x-request-id` on the upstream hop — downstream token recorders (the TITO +gateway) echo that header into their own per-turn records. If each layer +minted its own id, the message-level row and the token-level record for the +same call could never be joined. + +A `ContextVar` is the seam: the outermost interested layer (the `Recorder`, +when present) mints the id and binds it for the duration of the handler call; +inner layers reuse a bound id and only mint their own when nothing upstream +bound one. Works unchanged across `await` within one handler invocation and +never leaks across concurrent calls. +""" + +from __future__ import annotations + +import uuid +from contextvars import ContextVar + +current_request_id: ContextVar[str | None] = ContextVar("abridge_request_id", default=None) + + +def mint_request_id() -> str: + return uuid.uuid4().hex + + +def get_or_mint_request_id() -> str: + """The id bound by an outer capture layer, or a fresh one.""" + bound = current_request_id.get() + return bound if bound else mint_request_id() + + +__all__ = ["current_request_id", "get_or_mint_request_id", "mint_request_id"] diff --git a/plugins/abridge/agentix/bridge/clients/_anthropic_transforms.py b/plugins/abridge/agentix/bridge/clients/_anthropic_transforms.py index e40dcf4..7607dc4 100644 --- a/plugins/abridge/agentix/bridge/clients/_anthropic_transforms.py +++ b/plugins/abridge/agentix/bridge/clients/_anthropic_transforms.py @@ -1,15 +1,24 @@ """Pure Anthropic ↔ OpenAI shape converters. -Used only by `clients.anthropic_from_openai`. The functions here are -JSON-in, JSON-out — no I/O, no SDK calls, no spans. Anyone writing a -custom Anthropic-on-OpenAI client can import these directly. +Used by `clients.anthropic_from_openai` and `clients.anthropic_to_openai`. +The functions here are JSON-in, JSON-out — no I/O, no SDK calls, no spans. +Anyone writing a custom Anthropic-on-OpenAI client can import these directly. + +This module IS the translation contract for downstream consumers: whatever +the agent "actually said" to a recording backend is defined by these +functions, so a change here changes the byte identity of captured +trajectories. `TRANSLATION_SPEC_SHA` (bottom of the module) hashes this +file's source so consumers can pin the exact translation their data was +produced under; abridge-serve surfaces it on `/_health`. """ from __future__ import annotations import dataclasses +import hashlib import json import uuid +from pathlib import Path from typing import Any @@ -55,6 +64,10 @@ def anthropic_messages_to_openai( if tools: out["tools"] = tools + tool_choice = _tool_choice_anthropic_to_openai(body.get("tool_choice")) + if tool_choice is not None: + out["tool_choice"] = tool_choice + if extra_body: out.update(extra_body) @@ -275,6 +288,7 @@ def _messages_anthropic_to_openai(messages: list[Any]) -> list[dict[str, Any]]: continue text_parts: list[str] = [] + thinking_parts: list[str] = [] tool_calls: list[dict[str, Any]] = [] tool_results: list[dict[str, Any]] = [] for block in content: @@ -286,6 +300,13 @@ def _messages_anthropic_to_openai(messages: list[Any]) -> list[dict[str, Any]]: block_type = block.get("type") if block_type == "text": text_parts.append(str(block.get("text", ""))) + elif block_type == "thinking": + # Assistant thinking history maps to the reasoning_content + # key the vLLM/sglang dialect reads (Anthropic's crypto + # `signature` has no OpenAI-side meaning and is dropped). + # Keeping the reasoning in the echoed history preserves byte + # identity with what a session-recording backend stored. + thinking_parts.append(str(block.get("thinking", ""))) elif block_type == "tool_use": tool_calls.append( { @@ -309,6 +330,8 @@ def _messages_anthropic_to_openai(messages: list[Any]) -> list[dict[str, Any]]: text = "\n".join(part for part in text_parts if part) if role == "assistant": message: dict[str, Any] = {"role": "assistant", "content": text or None} + if thinking_parts: + message["reasoning_content"] = "\n".join(part for part in thinking_parts if part) if tool_calls: message["tool_calls"] = tool_calls if message["content"] is not None or tool_calls: @@ -320,6 +343,29 @@ def _messages_anthropic_to_openai(messages: list[Any]) -> list[dict[str, Any]]: return out +def _tool_choice_anthropic_to_openai(tool_choice: Any) -> Any: + """Map Anthropic `tool_choice` to the OpenAI field. + + `auto` -> "auto", `any` -> "required", `none` -> "none", and + `{type: tool, name}` -> a named function choice. + `disable_parallel_tool_use` is intentionally not mapped (matching the + reference translator): OpenAI's `parallel_tool_calls` is not honored by + every OpenAI-compatible engine, and a silently ignored knob is worse + than a documented drop. Unknown shapes are dropped, not guessed.""" + if not isinstance(tool_choice, dict): + return None + kind = tool_choice.get("type") + if kind == "auto": + return "auto" + if kind == "any": + return "required" + if kind == "none": + return "none" + if kind == "tool" and tool_choice.get("name"): + return {"type": "function", "function": {"name": str(tool_choice["name"])}} + return None + + def _tools_anthropic_to_openai(tools: Any) -> list[dict[str, Any]]: if not isinstance(tools, list): return [] @@ -359,7 +405,18 @@ def _sse(event: str, data: dict[str, Any]) -> bytes: return f"event: {event}\ndata: {payload}\n\n".encode() +# The translation contract version: SHA-256 over this module's source. Any +# edit to the transforms — however small — changes the byte identity of the +# OpenAI bodies a recording backend sees, so downstream data contracts pin +# this value (abridge-serve reports it on `/_health`). Deliberately the +# file's bytes, not a semantic hash: comments and docstrings changing the sha +# is a false positive we accept; a behavior change slipping through unhashed +# is not. +TRANSLATION_SPEC_SHA: str = hashlib.sha256(Path(__file__).read_bytes()).hexdigest() + + __all__ = [ + "TRANSLATION_SPEC_SHA", "AnthropicCountTokens", "anthropic_messages_to_openai", "anthropic_sse", diff --git a/plugins/abridge/agentix/bridge/clients/anthropic_from_openai.py b/plugins/abridge/agentix/bridge/clients/anthropic_from_openai.py index 7356419..3b67508 100644 --- a/plugins/abridge/agentix/bridge/clients/anthropic_from_openai.py +++ b/plugins/abridge/agentix/bridge/clients/anthropic_from_openai.py @@ -19,6 +19,7 @@ from agentix.utils import trace +from .._request_id import get_or_mint_request_id from ..proxy import ( AbridgeError, ClientResponse, @@ -104,7 +105,9 @@ async def messages(self, request: Request) -> ClientResponse: openai_body = anthropic_messages_to_openai(request.body, upstream_model=self._model) openai_body.update(self._upstream_params) openai_body["stream"] = False - record_id = uuid.uuid4().hex + # Reuses the id a wrapping capture layer (Recorder) bound for this + # call, so its JSONL row and the upstream header share one id. + record_id = get_or_mint_request_id() extra_headers = { "x-session-id": self.session_id, "x-request-id": record_id, diff --git a/plugins/abridge/agentix/bridge/forward.py b/plugins/abridge/agentix/bridge/forward.py index e65a979..901ef0d 100644 --- a/plugins/abridge/agentix/bridge/forward.py +++ b/plugins/abridge/agentix/bridge/forward.py @@ -35,6 +35,7 @@ import httpx +from ._request_id import get_or_mint_request_id from .proxy import AbridgeError, ClientResponse, Handler, Request logger = logging.getLogger(__name__) @@ -132,7 +133,9 @@ def handler(self, path: str | None = None) -> Handler: return _OwnedHandler(routes[path], self) async def _forward(self, path: str, request: Request) -> ClientResponse: - record_id = uuid.uuid4().hex + # Reuses the id a wrapping capture layer (Recorder) bound for this + # call, so its JSONL row and the sidecar's token record share one id. + record_id = get_or_mint_request_id() headers = { **self._headers, "x-session-id": self.session_id, diff --git a/plugins/abridge/agentix/bridge/recorder.py b/plugins/abridge/agentix/bridge/recorder.py index e3e3b53..0315f59 100644 --- a/plugins/abridge/agentix/bridge/recorder.py +++ b/plugins/abridge/agentix/bridge/recorder.py @@ -7,16 +7,30 @@ Each served request appends one line:: - {"ts": ..., "path": "/v1/messages", "request": {...}, + {"ts": ..., "path": "/v1/messages", "request_id": "<32 hex>", + "session_id": ..., # only when the Recorder has one + "request": {...}, "response": {"status_code": 200, "media_type": "...", "body": ...}} +`request_id` is minted per call and bound on the `current_request_id` +context var for the duration of the handler, so the transport layer +(`Forward` / the SDK clients) stamps the SAME id as `x-request-id` on the +upstream hop — a downstream token recorder's per-turn record and this row +join on it. `session_id`, when given, identifies the rollout the wrapped +client serves (pass the same value as the client's session identity). +Without these keys, rows from a retried call (e.g. an agent retry after a +tunnel 504 produced an orphan success row) are only deduplicable by +request-body equality. + A handler that raises records `{"error": ...}` instead of `"response"` and re-raises — a failed call is signal, not something to lose. JSON bodies are recorded as objects; anything else (e.g. a pre-rendered SSE blob) as text. Handlers run on the event loop, so appends never interleave; each line is flushed as it is written so the file is complete up to the last call even -if the process dies mid-rollout. +if the process dies mid-rollout. The file opens lazily on the first record, +so a Recorder that never serves (e.g. a route-enumeration probe) leaves no +empty file behind. """ from __future__ import annotations @@ -26,6 +40,7 @@ from pathlib import Path from typing import IO, Any +from ._request_id import current_request_id, mint_request_id from .proxy import ClientResponse, Handler, Request, _collect_handlers @@ -38,24 +53,35 @@ class Recorder: whole stack down once, as usual. """ - def __init__(self, client: Any, path: str | Path) -> None: + def __init__(self, client: Any, path: str | Path, *, session_id: str | None = None) -> None: self._client = client self._path = Path(path) self._path.parent.mkdir(parents=True, exist_ok=True) - self._file: IO[str] = self._path.open("a", encoding="utf-8") + self._session_id = session_id + self._file: IO[str] | None = None def abridge_routes(self) -> dict[str, Handler]: return {path: self._recording(path, handler) for path, handler in _collect_handlers(self._client).items()} def _recording(self, path: str, handler: Handler) -> Handler: async def record(request: Request) -> ClientResponse: - line: dict[str, Any] = {"ts": time.time(), "path": path, "request": request.body} + # Reuse an id bound by an even-outer layer; otherwise mint here. + # Binding it makes the transport's upstream `x-request-id` equal + # this row's `request_id`. + request_id = current_request_id.get() or mint_request_id() + line: dict[str, Any] = {"ts": time.time(), "path": path, "request_id": request_id} + if self._session_id is not None: + line["session_id"] = self._session_id + line["request"] = request.body + token = current_request_id.set(request_id) try: response = await handler(request) except BaseException as exc: line["error"] = f"{type(exc).__name__}: {exc}" self._write(line) raise + finally: + current_request_id.reset(token) line["response"] = { "status_code": response.status_code, "media_type": response.media_type, @@ -67,6 +93,8 @@ async def record(request: Request) -> ClientResponse: return record def _write(self, line: dict[str, Any]) -> None: + if self._file is None or self._file.closed: + self._file = self._path.open("a", encoding="utf-8") self._file.write(json.dumps(line, ensure_ascii=False, default=repr) + "\n") self._file.flush() @@ -79,7 +107,8 @@ async def aclose(self) -> None: if aclose is not None: await aclose() finally: - self._file.close() + if self._file is not None: + self._file.close() def _decode_body(response: ClientResponse) -> Any: diff --git a/plugins/abridge/agentix/bridge/serve.py b/plugins/abridge/agentix/bridge/serve.py index 3c212f2..fd47e25 100644 --- a/plugins/abridge/agentix/bridge/serve.py +++ b/plugins/abridge/agentix/bridge/serve.py @@ -27,9 +27,27 @@ `--require-key-prefix`) so only keys your harness minted are served and everything else gets a 401. -Anything beyond grouping (multi-backend routing, token capture) is the -full gateway's job — this is deliberately just "the tunnel without the -tunnel", and the tunnel remains the mode for sandboxes with no egress. +Two upstream modes, mutually exclusive: + +* `--upstream-base-url` — plain OpenAI-compatible engine; translation + + transport live in `AnthropicFromOpenAIClient` (the openai SDK owns the + HTTP). The original mode, unchanged. +* `--tito-url` — a token-recording session gateway (the TITO gateway): + each caller session composes + `AnthropicToOpenAI(SessionForward(tito_url).handler(), model=...)`, so + the gateway sees the OpenAI chat body, owns render/generate and the + token record, and one abridge caller session maps to one gateway + session. Token capture stays the gateway's job; this server only adds + the Anthropic shell and identity stamping. + +`--record-dir` (either mode) wraps each session's client in a `Recorder` +writing message-level rows to `/.jsonl`; rows +carry `session_id` + `request_id`, and the same `request_id` is stamped +as `x-request-id` on the upstream hop so message rows join token records. + +`GET /_health` reports `translation_spec_sha` — the SHA-256 of the +Anthropic<->OpenAI transform module — so downstream data contracts can +pin the exact translation in effect. OPENAI_API_KEY=EMPTY agentix-bridge-serve \ --upstream-base-url http://vllm:8000/v1 --upstream-model qwen3-32b \ @@ -47,6 +65,7 @@ from collections.abc import AsyncIterator, Callable from contextlib import AbstractAsyncContextManager, asynccontextmanager from dataclasses import dataclass, field +from pathlib import Path import uvicorn from fastapi import FastAPI @@ -253,7 +272,19 @@ def resolve(request: FastAPIRequest) -> AbstractAsyncContextManager[Handler]: async def _health() -> dict[str, str]: - return {"status": "ok"} + body = {"status": "ok"} + # The translation contract pin: SHA-256 of the transform module source. + # Downstream data contracts compare it against the value their captured + # trajectories were produced under. Lazy import: the transforms are + # SDK-free, but apps serving only custom handlers shouldn't fail health + # over a broken clients package. + try: + from .clients._anthropic_transforms import TRANSLATION_SPEC_SHA + + body["translation_spec_sha"] = TRANSLATION_SPEC_SHA + except ImportError: # pragma: no cover - clients package always ships + pass + return body class _UnknownKey(Exception): @@ -313,9 +344,7 @@ async def _read_json(request: FastAPIRequest) -> dict: return parsed if isinstance(parsed, dict) else {} -def main(argv: list[str] | None = None) -> None: - """`agentix-bridge-serve` — an Anthropic-speaking front for an - OpenAI-compatible engine, one session per caller key.""" +def _build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( prog="agentix-bridge-serve", description=( @@ -329,6 +358,17 @@ def main(argv: list[str] | None = None) -> None: default=os.environ.get("OPENAI_BASE_URL"), help="OpenAI-compatible endpoint, e.g. http://vllm:8000/v1 (env: OPENAI_BASE_URL)", ) + parser.add_argument( + "--tito-url", + default=os.environ.get("TITO_URL"), + help=( + "session-scoped token-recording gateway (the TITO gateway), e.g. " + "http://tito:30000 — mutually exclusive with --upstream-base-url: the " + "gateway becomes the upstream, owns render/generate and the token " + "record, and each caller session maps to one gateway session " + "(env: TITO_URL)" + ), + ) parser.add_argument( "--upstream-api-key", default=os.environ.get("OPENAI_API_KEY", "EMPTY"), @@ -358,22 +398,71 @@ def main(argv: list[str] | None = None) -> None: help="openai SDK retries per upstream call; keep total occupancy = timeout x (1+retries) explicit", ) parser.add_argument("--max-sessions", type=int, default=256) + parser.add_argument( + "--record-dir", + default=os.environ.get("ABRIDGE_RECORD_DIR"), + help=( + "record every served (request, response) pair to " + "/.jsonl via Recorder — message-level rows with " + "session_id + request_id, flushed per line (env: ABRIDGE_RECORD_DIR)" + ), + ) + return parser + + +def _client_factory(args: argparse.Namespace) -> Callable[[str], Client]: + """The per-caller-session client for the chosen upstream mode, wrapped in + a `Recorder` when `--record-dir` is set.""" + if args.tito_url: + # Composition seam (transport-blind): the TITO gateway sees the + # OpenAI chat body and owns tokens + recording; abridge adds only the + # Anthropic shell. One SessionForward per caller session == one + # gateway session per rollout key. + from .clients import AnthropicToOpenAI + from .forward import SessionForward + + def build(session_id: str) -> Client: + forward = SessionForward(args.tito_url, paths=["/v1/chat/completions"], timeout=args.upstream_timeout) + return AnthropicToOpenAI(forward.handler(), model=args.upstream_model) + else: + # Lazy: the translation client needs the `openai` extra. + from .clients import AnthropicFromOpenAIClient + + def build(session_id: str) -> Client: + return AnthropicFromOpenAIClient( + base_url=args.upstream_base_url, + api_key=args.upstream_api_key, + model=args.upstream_model, + timeout=args.upstream_timeout, + max_retries=args.upstream_max_retries, + session_id=session_id, + ) + + if not args.record_dir: + return build + + from .recorder import Recorder + + record_dir = Path(args.record_dir) + + def build_recorded(session_id: str) -> Client: + return Recorder(build(session_id), record_dir / f"{session_id}.jsonl", session_id=session_id) + + return build_recorded + + +def main(argv: list[str] | None = None) -> None: + """`agentix-bridge-serve` — an Anthropic-speaking front for an + OpenAI-compatible engine or a token-recording session gateway, one + session per caller key.""" + parser = _build_parser() args = parser.parse_args(argv) - if not args.upstream_base_url: - parser.error("--upstream-base-url (or OPENAI_BASE_URL) is required") - - # Lazy: the translation client needs the `openai` extra. - from .clients import AnthropicFromOpenAIClient - - def factory(session_id: str) -> AnthropicFromOpenAIClient: - return AnthropicFromOpenAIClient( - base_url=args.upstream_base_url, - api_key=args.upstream_api_key, - model=args.upstream_model, - timeout=args.upstream_timeout, - max_retries=args.upstream_max_retries, - session_id=session_id, - ) + if args.upstream_base_url and args.tito_url: + parser.error("--upstream-base-url and --tito-url are mutually exclusive: the gateway IS the upstream") + if not args.upstream_base_url and not args.tito_url: + parser.error("one of --upstream-base-url (env OPENAI_BASE_URL) or --tito-url (env TITO_URL) is required") + + factory = _client_factory(args) verify_key: Callable[[str], bool] | None = None if args.require_key_prefix: diff --git a/plugins/abridge/tests/test_anthropic_transform_fixtures.py b/plugins/abridge/tests/test_anthropic_transform_fixtures.py index f576375..952ce97 100644 --- a/plugins/abridge/tests/test_anthropic_transform_fixtures.py +++ b/plugins/abridge/tests/test_anthropic_transform_fixtures.py @@ -48,17 +48,19 @@ "06_user_image_url": "a user message whose content is only an image block is dropped entirely", "11_user_tool_result_multipart": "image parts inside a multipart tool_result are dropped (only text is forwarded)", "13_long_tool_name": ">64-char tool names are not sanitized (no truncate+hash rename, no tool_map)", - "14_tool_choice_any": "tool_choice {type: any} is dropped (should map to OpenAI tool_choice 'required')", - "15_tool_choice_named": "tool_choice {type: tool, name} is dropped (should map to a named function tool_choice)", "16_metadata_user_id": "metadata.user_id is dropped (should map to the OpenAI 'user' field)", "17_thinking_medium": "thinking {budget_tokens: 5000} is dropped (no reasoning_effort='medium' mapping)", "18_top_k_dropped": "top_k is dropped; the oracle passes it through for OpenAI-compatible upstreams", "19_stream_include_usage": "stream=true is dropped (the adapter always makes a non-streaming upstream call)", - "35_assistant_thinking_history": "assistant thinking blocks are dropped instead of forwarded as thinking_blocks", + "35_assistant_thinking_history": ( + "DELIBERATE dialect divergence, not a drop: assistant thinking history is forwarded as " + "`reasoning_content` (the vLLM/sglang key our downstreams read), while the LiteLLM golden " + "expects `thinking_blocks` (+ the Anthropic crypto `signature`, meaningless to an " + "OpenAI-compatible engine). The reasoning text itself is preserved — see " + "test_assistant_thinking_history_maps_to_reasoning_content for the positive contract." + ), "36_user_mixed_content": "the image part of mixed text+image user content is dropped", "37_empty_string_content": "a message with empty-string content is forwarded; the oracle drops it", - "39_tool_choice_auto_no_parallel": "tool_choice {type: auto} is dropped (should map to OpenAI tool_choice 'auto')", - "40_tool_choice_none": "tool_choice {type: none} is dropped (should map to OpenAI tool_choice 'none')", "41_thinking_high": "thinking with a large budget is dropped (no reasoning_effort='high' mapping)", "42_thinking_low": "thinking {budget_tokens: 2000} is dropped (no reasoning_effort='low' mapping)", "43_stop_sequences": "stop_sequences is dropped (never mapped to an OpenAI stop parameter)", @@ -182,6 +184,23 @@ def test_request_transform_matches_golden(case: str) -> None: assert _normalized_openai_request(actual) == _normalized_openai_request(expected) +def test_assistant_thinking_history_maps_to_reasoning_content() -> None: + """The positive contract behind the 35_assistant_thinking_history xfail: + the reasoning text is NOT dropped — it rides the `reasoning_content` key + our OpenAI-compatible downstreams (vLLM/sglang dialect) read, so an + echoed assistant turn keeps byte identity with what a session-recording + backend stored. Only the layout differs from the LiteLLM golden + (`thinking_blocks` + `signature`).""" + anthropic_body = _load("requests", "anthropic_35_assistant_thinking_history.json") + + actual = anthropic_messages_to_openai(anthropic_body) + + (assistant,) = [m for m in actual["messages"] if m["role"] == "assistant"] + assert assistant["reasoning_content"] == "Let me work this out step by step..." + assert assistant["content"] == "The answer is 42." + assert "signature" not in json.dumps(actual) # the crypto signature is dropped + + # ── responses: OpenAI -> Anthropic ────────────────────────────────────────── diff --git a/plugins/abridge/tests/test_recorder.py b/plugins/abridge/tests/test_recorder.py index 96e77da..3d79322 100644 --- a/plugins/abridge/tests/test_recorder.py +++ b/plugins/abridge/tests/test_recorder.py @@ -117,3 +117,68 @@ def environ(self, handle) -> dict[str, str]: recorder = Recorder(_EnvClient(), tmp_path / "run.jsonl") assert recorder.environ(None) == {"X": "y"} + + +@pytest.mark.asyncio +async def test_recorder_rows_carry_session_and_request_ids(tmp_path) -> None: + """Rows are joinable against downstream token records: `session_id` (the + rollout identity the Recorder was built with) and a per-call + `request_id`, unique across calls.""" + out = tmp_path / "run.jsonl" + recorder = Recorder(_EchoClient(), out, session_id="sess-42") + routes = recorder.abridge_routes() + await routes["/v1/messages"](Request(path="/v1/messages", body={"msg": "a"})) + await routes["/v1/messages"](Request(path="/v1/messages", body={"msg": "b"})) + + first, second = _lines(out) + assert first["session_id"] == second["session_id"] == "sess-42" + assert first["request_id"] and second["request_id"] + assert first["request_id"] != second["request_id"] + + +@pytest.mark.asyncio +async def test_recorder_request_id_matches_upstream_x_request_id(tmp_path, monkeypatch) -> None: + """The alignment contract: the id in the Recorder row IS the + `x-request-id` the transport stamps on the upstream hop (via the + `current_request_id` context var), so a message-level row and a + token-level sidecar record for the same call join on one key.""" + import httpx + from agentix.bridge import Forward + + fwd = Forward("http://side.car", paths=["/v1/messages"], session_id="sess-1") + seen_headers: list[dict] = [] + + async def fake_post(url, *, json, headers): + seen_headers.append(dict(headers)) + return httpx.Response(200, content=b'{"ok": true}', headers={"content-type": "application/json"}) + + monkeypatch.setattr(fwd._client, "post", fake_post) + + out = tmp_path / "run.jsonl" + routes = Recorder(fwd, out, session_id="sess-1").abridge_routes() + await routes["/v1/messages"](Request(path="/v1/messages", body={"x": 1})) + + (row,) = _lines(out) + assert seen_headers[0]["x-request-id"] == row["request_id"] + assert seen_headers[0]["x-session-id"] == row["session_id"] == "sess-1" + + +@pytest.mark.asyncio +async def test_recorder_error_rows_also_carry_ids(tmp_path) -> None: + out = tmp_path / "run.jsonl" + routes = Recorder(_FailingClient(), out, session_id="sess-9").abridge_routes() + with pytest.raises(AbridgeError): + await routes["/v1/messages"](Request(path="/v1/messages", body={})) + (row,) = _lines(out) + assert row["session_id"] == "sess-9" + assert row["request_id"] + assert "error" in row + + +def test_recorder_opens_file_lazily(tmp_path) -> None: + """A Recorder that never serves (e.g. build_session_app's route + enumeration probe) must leave no empty file behind.""" + out = tmp_path / "probe.jsonl" + recorder = Recorder(_EchoClient(), out) + recorder.abridge_routes() + assert not out.exists() diff --git a/plugins/abridge/tests/test_serve.py b/plugins/abridge/tests/test_serve.py index a5e4249..24247cc 100644 --- a/plugins/abridge/tests/test_serve.py +++ b/plugins/abridge/tests/test_serve.py @@ -40,13 +40,28 @@ async def aclose(self) -> None: def test_build_app_serves_handlers_and_health() -> None: tc = TestClient(build_app(EchoClient())) - assert tc.get("/_health").json() == {"status": "ok"} + health = tc.get("/_health").json() + assert health["status"] == "ok" r = tc.post("/v1/echo", json={"x": 1}) assert r.status_code == 200 assert r.json()["echo"] == {"x": 1} assert tc.post("/nope", json={}).status_code == 404 +def test_health_pins_the_translation_spec() -> None: + """`/_health` reports the SHA-256 of the Anthropic<->OpenAI transform + module, so downstream data contracts can pin the exact translation their + captured trajectories were produced under.""" + import hashlib + from pathlib import Path + + import agentix.bridge.clients._anthropic_transforms as transforms + + health = TestClient(build_app(EchoClient())).get("/_health").json() + expected = hashlib.sha256(Path(transforms.__file__).read_bytes()).hexdigest() + assert health["translation_spec_sha"] == expected == transforms.TRANSLATION_SPEC_SHA + + def test_handler_errors_become_wire_errors() -> None: tc = TestClient(build_app(EchoClient())) r = tc.post("/v1/teapot", json={}) diff --git a/plugins/abridge/tests/test_tito_composition.py b/plugins/abridge/tests/test_tito_composition.py new file mode 100644 index 0000000..d996458 --- /dev/null +++ b/plugins/abridge/tests/test_tito_composition.py @@ -0,0 +1,191 @@ +"""End-to-end: serve mode composed onto a token-recording session gateway. + +`agentix-bridge-serve --tito-url ...` builds, per caller session, +`AnthropicToOpenAI(SessionForward(tito_url).handler(), model=...)` — the +agent speaks Anthropic, the gateway receives the OpenAI chat body on its +session-scoped route and owns tokens + recording. A tiny in-process HTTP +server plays the gateway (create-session + session-scoped chat completions, +the same wire shapes as the real TITO gateway), so the whole path is real: +FastAPI serve app -> translation -> SessionForward -> real httpx -> gateway. +""" + +from __future__ import annotations + +import json +import socket +from collections.abc import Iterator +from http.server import BaseHTTPRequestHandler, HTTPServer +from threading import Thread +from typing import Any + +import pytest +from agentix.bridge.serve import _build_parser, _client_factory, build_session_app, session_id_for +from fastapi.testclient import TestClient + + +class _FakeTito(BaseHTTPRequestHandler): + """Session-scoped fake gateway: POST /sessions mints ids; the chat route + records the exact OpenAI body + identity headers it received.""" + + sessions: list[str] = [] + chat_calls: list[dict[str, Any]] = [] # {"session", "body", "headers"} + + @classmethod + def reset(cls) -> None: + cls.sessions = [] + cls.chat_calls = [] + + def do_POST(self) -> None: # noqa: N802 - http.server convention + length = int(self.headers.get("Content-Length") or 0) + raw = self.rfile.read(length) + if self.path == "/sessions": + session_id = f"tito-sess-{len(_FakeTito.sessions)}" + _FakeTito.sessions.append(session_id) + self._json(200, {"session_id": session_id}) + return + parts = self.path.strip("/").split("/") + if len(parts) >= 2 and parts[0] == "sessions" and self.path.endswith("/v1/chat/completions"): + _FakeTito.chat_calls.append({ + "session": parts[1], + "body": json.loads(raw), + "headers": {k.lower(): v for k, v in self.headers.items()}, + }) + self._json(200, { + "id": "chatcmpl-tito", "object": "chat.completion", "model": "qwen3-4b", + "choices": [{ + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "hello from tito"}, + }], + "usage": {"prompt_tokens": 11, "completion_tokens": 3, "total_tokens": 14}, + }) + return + self._json(404, {"error": f"no route {self.path}"}) + + def _json(self, status: int, body: dict) -> None: + blob = json.dumps(body).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(blob))) + self.end_headers() + self.wfile.write(blob) + + def log_message(self, *_: Any) -> None: + return + + +@pytest.fixture +def fake_tito() -> Iterator[str]: + _FakeTito.reset() + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + server = HTTPServer(("127.0.0.1", port), _FakeTito) + thread = Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{port}" + finally: + server.shutdown() + thread.join(timeout=2) + + +_ANTHROPIC_BODY = { + "model": "claude-sonnet-4-5", + "max_tokens": 64, + "system": "be brief", + "messages": [{"role": "user", "content": "hi"}], +} + + +def _tito_app(fake_tito: str, *extra_args: str): + args = _build_parser().parse_args(["--tito-url", fake_tito, "--upstream-model", "qwen3-4b", *extra_args]) + return build_session_app(_client_factory(args)) + + +def test_tito_composition_end_to_end(fake_tito: str) -> None: + tc = TestClient(_tito_app(fake_tito)) + + r = tc.post("/v1/messages", json=_ANTHROPIC_BODY, headers={"x-api-key": "rollout-1"}) + assert r.status_code == 200 + + # Agent side: a faithful Anthropic response. + body = r.json() + assert body["role"] == "assistant" + assert body["content"] == [{"type": "text", "text": "hello from tito"}] + assert body["model"] == "claude-sonnet-4-5" # the agent's model id echoes back + assert body["usage"] == {"input_tokens": 11, "output_tokens": 3} + + # Gateway side: one session created, the chat body is OpenAI-shaped, + # forced non-streaming (the token recorder needs the full completion), + # model overridden to what the engine serves. + (call,) = _FakeTito.chat_calls + assert call["session"] == _FakeTito.sessions[0] + assert call["body"]["stream"] is False + assert call["body"]["model"] == "qwen3-4b" + assert call["body"]["messages"][0] == {"role": "system", "content": "be brief"} + assert call["body"]["messages"][1] == {"role": "user", "content": "hi"} + + # Identity stamping: the gateway groups by x-session-id (its own session + # id — SessionForward adopts the gateway-assigned id) + per-call + # x-request-id. + assert call["headers"]["x-session-id"] == _FakeTito.sessions[0] + assert call["headers"]["x-request-id"] + + +def test_tito_composition_one_gateway_session_per_caller_key(fake_tito: str) -> None: + tc = TestClient(_tito_app(fake_tito)) + assert tc.post("/v1/messages", json=_ANTHROPIC_BODY, headers={"x-api-key": "rollout-a"}).status_code == 200 + assert tc.post("/v1/messages", json=_ANTHROPIC_BODY, headers={"x-api-key": "rollout-a"}).status_code == 200 + assert tc.post("/v1/messages", json=_ANTHROPIC_BODY, headers={"x-api-key": "rollout-b"}).status_code == 200 + + assert len(_FakeTito.sessions) == 2 # one gateway session per caller key + assert [c["session"] for c in _FakeTito.chat_calls] == [ + _FakeTito.sessions[0], _FakeTito.sessions[0], _FakeTito.sessions[1], + ] + + +def test_tito_composition_with_record_dir_joins_rows_to_gateway_calls(fake_tito: str, tmp_path) -> None: + """--record-dir in tito mode: the message-level Recorder row and the + gateway's x-request-id share one id, and the row's session_id is the + caller-derived serve session (the row->record join key set).""" + tc = TestClient(_tito_app(fake_tito, "--record-dir", str(tmp_path))) + r = tc.post("/v1/messages", json=_ANTHROPIC_BODY, headers={"x-api-key": "rollout-1"}) + assert r.status_code == 200 + + serve_session = session_id_for("rollout-1") + (row,) = [json.loads(line) for line in (tmp_path / f"{serve_session}.jsonl").read_text().splitlines()] + (call,) = _FakeTito.chat_calls + assert row["session_id"] == serve_session + assert row["request_id"] == call["headers"]["x-request-id"] + assert row["path"] == "/v1/messages" + assert row["request"] == _ANTHROPIC_BODY # the agent-side (Anthropic) shape + assert row["response"]["body"]["content"] == [{"type": "text", "text": "hello from tito"}] + # No file for the route-enumeration probe session. + assert sorted(p.name for p in tmp_path.iterdir()) == [f"{serve_session}.jsonl"] + + +def test_tito_streaming_agent_gets_replayed_sse(fake_tito: str) -> None: + """stream:true agents get the locally rendered SSE replay while the + gateway still saw a non-streaming call.""" + tc = TestClient(_tito_app(fake_tito)) + r = tc.post( + "/v1/messages", + json={**_ANTHROPIC_BODY, "stream": True}, + headers={"x-api-key": "rollout-1"}, + ) + assert r.status_code == 200 + assert r.headers["content-type"].startswith("text/event-stream") + assert "hello from tito" in r.text + assert _FakeTito.chat_calls[0]["body"]["stream"] is False + + +def test_serve_mode_flags_are_mutually_exclusive_and_required(monkeypatch) -> None: + from agentix.bridge.serve import main + + for var in ("OPENAI_BASE_URL", "TITO_URL"): + monkeypatch.delenv(var, raising=False) + with pytest.raises(SystemExit): + main(["--upstream-base-url", "http://engine:8000/v1", "--tito-url", "http://tito:30000"]) + with pytest.raises(SystemExit): + main([]) # neither mode selected From f309757c564482c38b47b9ec206ee5f9abc2df9f Mon Sep 17 00:00:00 2001 From: Lingrui Mei Date: Fri, 24 Jul 2026 07:24:33 +0800 Subject: [PATCH 4/5] tito: harden the record contract per adversarial review (schema, stability baseline, lifecycle) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit tito.record.v1 is now the normative, documented contract (the README carries the schema document; downstream consumers adapt to it). The structure stays flat; three data additions close the review's contract-alignment findings: - sampling: whitelisted sampling params lifted verbatim from the chat request (temperature/top_p/top_k/min_p/max_tokens/... ), {} = backend defaults. - tokenizer block replaces tokenizer_fingerprint: {checkpoint, tokenizer_sha256, chat_template_sha256}. tokenizer_sha256 is really computed — SHA-256 over the tokenizer DEFINITION bytes (fast tokenizers: the complete serialized tokenizer.json via backend_tokenizer.to_str(); slow: sorted-JSON vocab); method is part of the documented contract. - thread_id: passthrough of the x-thread-id request header, key omitted when absent. - prompt_segments keep the gateway-native source vocabulary {render, prefix, system, user, tool, generation_prompt}, documented as exhaustive. Token-exactness blocker (reviewer repro): prefix_stable was computed against the in-memory checkpoint, so a rollback applied in phase 1 whose turn never produced a record line (upstream non-200/timeout/409) vanished from the JSONL — the next successful line claimed stability while its prompt did not extend the previous LINE. The sink now owns prefix_stable, computing it against the last line it actually wrote (per-session last prompt+completion), so unrecorded discontinuities surface as false. PreparedPrompt.prefix_stable remains as advisory engine metadata with the distinction documented. Lifecycle (reviewer repro): the idle clock now starts when a turn ENDS (last_used refreshed in the chat handler's finally) — previously a single generation slower than the TTL let the agent's next request sweep its own live session away mid-rollout. Sink strictness: append failures catch Exception (not just OSError — e.g. UnicodeEncodeError from a lone surrogate no longer 500s a committed turn), and the turn_index increment moved to finally so ANY failed line leaves a detectable index gap; lines are strict JSON (allow_nan=False) and the sink rejects non-int token ids / non-finite logprobs (log + gap) instead of coercing them into training data. Tests: both reviewer repros regressed (unrecorded-rollback stability break, TTL-vs-generation-time), NaN/strictness gap detection, sampling whitelist + thread_id passthrough, tokenizer block hashes. Co-Authored-By: Claude Fable 5 --- plugins/tito/README.md | 91 +++++-- plugins/tito/agentix/tito/engine/record.py | 224 +++++++++++++----- .../tito/agentix/tito/engine/session_app.py | 53 +++-- .../tito/agentix/tito/engine/trajectory.py | 37 ++- plugins/tito/tests/test_record.py | 136 ++++++++++- 5 files changed, 431 insertions(+), 110 deletions(-) diff --git a/plugins/tito/README.md b/plugins/tito/README.md index 0cb2540..7b261f3 100644 --- a/plugins/tito/README.md +++ b/plugins/tito/README.md @@ -78,31 +78,86 @@ tokenizer's own template). `--backend-kind` selects the backend token dialect a local backend (see `agentix.tito.discovery`). Run `agentix-tito serve -h` for the full list. -## Per-turn record persistence +## Per-turn record persistence — `tito.record.v1` (normative) With `--record-dir DIR` (env `TITO_RECORD_DIR`), every committed turn appends one `tito.record.v1` JSON line to `DIR/.jsonl` and flushes it — the file is complete up to the last committed turn even if the process dies -mid-rollout. Each line carries the exact token truth of one turn: - -- `prompt_token_ids` / `completion_token_ids` / `completion_logprobs` - (sampled-token logprobs from the same forward pass, 1:1 with the ids); -- `prompt_segments` — `{start, end, source}` spans over the prompt ids - (`render`, `prefix`, per appended role, `generation_prompt`): the material - a trainer needs to build a loss mask without re-tokenizing anything; -- `prefix_stable` — whether this prompt extends the previous committed - checkpoint; `false` (retry rollback / history rewrite) means the turn must - not be spliced into one linear token stream; -- `request_id` (echoed from the caller's `x-request-id` header), `model`, - `backend_kind`, `finish_reason`, `assistant_message`, and a - `tokenizer_fingerprint` (`checkpoint` + `chat_template_sha256`) pinning the - render rules; +mid-rollout. **This section is the normative schema document: the gateway is +the producing side of the contract, and downstream consumers adapt to the +shape defined here.** The structure is flat (no nested `prompt`/`completion` +objects). + +```jsonc +{ + "schema_version": "tito.record.v1", + "session_id": "", + "thread_id": "", // OPTIONAL: key omitted when the header is absent + "request_id": "", + "turn_index": 0, // monotonic per session file; see gap semantics below + "ts": 1750000000.0, // unix seconds, request commit time + "model": "", + "backend_kind": "sglang" | "vllm", + "sampling": { "temperature": 0.6, "top_p": 0.95, "max_tokens": 4096 }, + "tokenizer": { + "checkpoint": "", + "tokenizer_sha256": "<64 hex>", + "chat_template_sha256": "<64 hex>" + }, + "prompt_token_ids": [ ... ], // exactly what generation ran on + "prompt_segments": [ {"start": 0, "end": 42, "source": "prefix"}, ... ], + "completion_token_ids": [ ... ], // exactly what was sampled + "completion_logprobs": [ ... ], // finite floats, len == len(completion_token_ids) + "assistant_message": { ... }, // the parsed assistant turn as stored + "finish_reason": "stop" | "tool_calls" | ... | null, + "prefix_stable": true, + "render_skew": null | {"equal": false, "first_divergence": 17} +} +``` + +Field semantics: + +- `sampling` — the whitelisted sampling parameters lifted verbatim from the + chat request body (`temperature`, `top_p`, `top_k`, `min_p`, `max_tokens`, + `max_completion_tokens`, `frequency_penalty`, `presence_penalty`, + `repetition_penalty`, `seed`, `stop`, `n`); keys absent from the request + are absent here, so the empty object means "backend defaults". +- `tokenizer.tokenizer_sha256` — SHA-256 over the tokenizer **definition** + bytes: for fast tokenizers the complete serialized `tokenizer.json` + definition (`backend_tokenizer.to_str()` — vocab, merges, normalizer, + added tokens); for slow tokenizers the sorted-JSON vocabulary. Two + gateways report the same value iff they tokenize identically. + `tokenizer.chat_template_sha256` hashes the chat template actually in + effect (the fixed-template override when set, else the tokenizer's own). +- `prompt_segments.source` vocabulary (gateway-native, exhaustive): + `render` (a from-scratch chat-template render — the first turn, or the + fallback when the prefix window was outrun), `prefix` (the reused + accumulated checkpoint = all previous prompt+completion tokens), `system` + / `user` / `tool` (one segment per appended role), `generation_prompt`. + Segments tile `prompt_token_ids` exactly; this is the loss-mask + construction material — no re-tokenization needed downstream. +- `prefix_stable` — true iff this turn's `prompt_token_ids` extend the + previous **recorded** line's `prompt_token_ids + completion_token_ids`. + The baseline is the record stream itself, not the in-memory checkpoint: a + rollback applied for a turn that never produced a line (upstream error, + timeout, 409) surfaces as `false` on the next recorded turn. `false` + (retry rollback / history rewrite) means the turn must not be spliced + into one linear token stream. +- `turn_index` — assigned by the sink, monotonic per session, and advances + even when a line fails to persist: a write/validation failure is logged + and leaves a detectable index gap instead of an undetectable missing turn. - `render_skew` (vLLM only) — a cheap per-turn probe comparing the render endpoint's from-scratch ids against the gateway's accumulated prompt ids - (recorded, never enforced). + (recorded, never enforced); `null` on sglang. + +Strictness guarantees: every line is strict JSON (`allow_nan=False` — never +`NaN`/`Infinity` literals); token ids are Python ints and logprobs finite +floats, validated at the sink — a violating record is dropped with a logged +error and an index gap, never silently coerced. -Closing a session appends one final `tito.session.v1` metadata line -(`reason`: `deleted` / `ttl_evicted` / `capacity_evicted` / `shutdown`) and +Closing a session appends one final `tito.session.v1` metadata line — +`{"schema_version": "tito.session.v1", "session_id", "turns", "reason": +"deleted" | "ttl_evicted" | "capacity_evicted" | "shutdown", "ts"}` — and closes the file. Without `--record-dir` nothing is written and the in-memory behavior is unchanged. diff --git a/plugins/tito/agentix/tito/engine/record.py b/plugins/tito/agentix/tito/engine/record.py index f04a8ad..cb2fcfb 100644 --- a/plugins/tito/agentix/tito/engine/record.py +++ b/plugins/tito/agentix/tito/engine/record.py @@ -3,40 +3,59 @@ The in-process trajectory (`LinearTrajectory`) is the live state machine; this module is the crash-safe, append-only export of it. When the gateway is started with a record directory, every committed chat turn appends exactly -one JSON line to ``/.jsonl`` and flushes it, so the -file is complete up to the last committed turn even if the process dies -mid-rollout. Closing a session (DELETE, TTL/capacity eviction, shutdown) -appends one final ``tito.session.v1`` metadata line and closes the file. +one strict-JSON line to ``/.jsonl`` and flushes it, +so the file is complete up to the last committed turn even if the process +dies mid-rollout. Closing a session (DELETE, TTL/capacity eviction, +shutdown) appends one final ``tito.session.v1`` metadata line and closes the +file. + +This module is the NORMATIVE definition of the record shape — downstream +consumers adapt to it (see the plugin README for the full schema document). Two line schemas: -``tito.record.v1`` — one per committed turn:: +``tito.record.v1`` — one per committed turn (flat structure):: - {"schema_version": "tito.record.v1", "session_id": ..., "turn_index": 0, - "request_id": ..., "model": ..., "backend_kind": "sglang" | "vllm", - "prompt_token_ids": [...], "completion_token_ids": [...], - "completion_logprobs": [...], # len == len(completion_token_ids) + {"schema_version": "tito.record.v1", "session_id": ..., + "thread_id": ..., # only when x-thread-id was sent + "request_id": ..., "turn_index": 0, "ts": , + "model": ..., "backend_kind": "sglang" | "vllm", + "sampling": {"temperature": ..., ...}, # whitelist from the request body + "tokenizer": {"checkpoint": ..., "tokenizer_sha256": ..., + "chat_template_sha256": ...}, + "prompt_token_ids": [...], "prompt_segments": [{"start": 0, "end": N, "source": ...}, ...], + "completion_token_ids": [...], + "completion_logprobs": [...], # len == len(completion_token_ids) "assistant_message": {...}, "finish_reason": ..., - "tokenizer_fingerprint": {"checkpoint": ..., "chat_template_sha256": ...}, - "prefix_stable": true, "render_skew": null | {"equal": ..., "first_divergence": ...}, - "ts": } + "prefix_stable": true, + "render_skew": null | {"equal": ..., "first_divergence": ...}} ``prompt_segments`` sources: ``render`` (a from-scratch chat-template render — the first turn, or the fallback when the token prefix window was outrun), ``prefix`` (the reused accumulated checkpoint), one segment per appended role (``tool`` / ``user`` / ``system``), and ``generation_prompt``. -``prefix_stable`` is true iff this turn's prompt token ids extend the last -*committed* checkpoint (previous prompt + completion). A retry rollback or a -history rewrite records ``false`` — the turn is still served and recorded, -but a trainer must not splice it into one linear token stream. +``prefix_stable`` is computed by the sink against the last line it actually +wrote for the session: true iff this turn's prompt token ids extend the +previous RECORDED turn's ``prompt_token_ids + completion_token_ids``. This +is deliberately not the in-memory checkpoint — a rollback applied for a turn +that never produced a line (upstream non-200, timeout, 409) must surface as +``false`` on the next recorded turn, or a trainer following the contract +would splice a broken token stream. + +``turn_index`` is monotonic per session file and advances even when a line +fails to persist, so any dropped line leaves a detectable index gap. ``render_skew`` (vLLM backend only) is the cheap per-turn probe comparing the render endpoint's from-scratch token ids against the gateway's accumulated prompt ids; non-equal skew is expected mid-conversation (render re-renders the echoed history) and is recorded, never enforced. +Strictness: lines are strict JSON (``allow_nan=False``); token ids must be +Python ints and logprobs finite floats — a violating record is NOT written +(logged + index gap), never silently coerced. + ``tito.session.v1`` — one final line when the session closes:: {"schema_version": "tito.session.v1", "session_id": ..., "turns": N, @@ -48,6 +67,7 @@ import json import logging +import math import time from pathlib import Path from typing import IO, Any @@ -57,6 +77,30 @@ RECORD_SCHEMA_VERSION = "tito.record.v1" SESSION_META_SCHEMA_VERSION = "tito.session.v1" +# Request-body keys lifted verbatim into the record's `sampling` object. +# Whitelist, not passthrough: the body also carries messages/tools/stream and +# the fields the gateway itself forces (logprobs et al.), none of which are +# sampling parameters. +SAMPLING_KEYS = ( + "temperature", + "top_p", + "top_k", + "min_p", + "max_tokens", + "max_completion_tokens", + "frequency_penalty", + "presence_penalty", + "repetition_penalty", + "seed", + "stop", + "n", +) + + +def sampling_from_request(request_body: dict[str, Any]) -> dict[str, Any]: + """The whitelisted sampling parameters present in the chat request.""" + return {key: request_body[key] for key in SAMPLING_KEYS if request_body.get(key) is not None} + def compute_render_skew( render_token_ids: list[int] | None, prompt_token_ids: list[int] @@ -79,52 +123,78 @@ def build_turn_record( *, session_id: str, request_id: str | None, + thread_id: str | None, model: str | None, backend_kind: str, + sampling: dict[str, Any], prompt_token_ids: list[int], prompt_segments: list[dict[str, Any]], - prefix_stable: bool, completion_token_ids: list[int], completion_logprobs: list[float], assistant_message: dict[str, Any], finish_reason: str | None, - tokenizer_fingerprint: dict[str, str], + tokenizer: dict[str, str], render_token_ids: list[int] | None, ) -> dict[str, Any]: - """One JSON-serializable ``tito.record.v1`` line (without ``turn_index``, - which the sink assigns monotonically per session file).""" - if len(completion_logprobs) != len(completion_token_ids): - raise ValueError( - f"len(completion_logprobs)={len(completion_logprobs)} != " - f"len(completion_token_ids)={len(completion_token_ids)}" - ) - return { + """One ``tito.record.v1`` line, without the fields the sink assigns + (``turn_index``, ``prefix_stable``). Pure assembly — validation happens + in the sink so every rejected record leaves the same detectable gap.""" + record: dict[str, Any] = { "schema_version": RECORD_SCHEMA_VERSION, "session_id": session_id, - "request_id": request_id, - "model": model, - "backend_kind": backend_kind, - "prompt_token_ids": list(prompt_token_ids), - "completion_token_ids": list(completion_token_ids), - "completion_logprobs": list(completion_logprobs), - "prompt_segments": list(prompt_segments), - "assistant_message": assistant_message, - "finish_reason": finish_reason, - "tokenizer_fingerprint": dict(tokenizer_fingerprint), - "prefix_stable": bool(prefix_stable), - "render_skew": compute_render_skew(render_token_ids, prompt_token_ids), - "ts": time.time(), } + if thread_id is not None: + record["thread_id"] = thread_id + record.update( + { + "request_id": request_id, + "ts": time.time(), + "model": model, + "backend_kind": backend_kind, + "sampling": dict(sampling), + "tokenizer": dict(tokenizer), + "prompt_token_ids": list(prompt_token_ids), + "prompt_segments": list(prompt_segments), + "completion_token_ids": list(completion_token_ids), + "completion_logprobs": list(completion_logprobs), + "assistant_message": assistant_message, + "finish_reason": finish_reason, + "render_skew": compute_render_skew(render_token_ids, prompt_token_ids), + } + ) + return record + + +def _validate_turn_record(record: dict[str, Any]) -> None: + """Strictness gate for token truth: ids are Python ints, logprobs are + finite floats paired 1:1 with the completion ids. Anything else (numpy + scalars repr'd to strings, NaN logprobs a backend passed through) must be + rejected, not silently coerced into training data.""" + for field in ("prompt_token_ids", "completion_token_ids"): + ids = record[field] + if not all(type(t) is int for t in ids): + raise ValueError(f"{field} must be Python ints; got {[type(t).__name__ for t in ids[:5]]}") + logprobs = record["completion_logprobs"] + if len(logprobs) != len(record["completion_token_ids"]): + raise ValueError( + f"len(completion_logprobs)={len(logprobs)} != " + f"len(completion_token_ids)={len(record['completion_token_ids'])}" + ) + if not all(isinstance(lp, float) and math.isfinite(lp) for lp in logprobs): + raise ValueError("completion_logprobs must be finite floats") class TurnRecordSink: """Per-session JSONL files under ``record_dir``, appended and flushed one line per committed turn. - Files open lazily on the first turn and are closed by ``finalize`` (which - also appends the ``tito.session.v1`` metadata line). ``finalize`` is - idempotent; ``close_all`` finalizes every open file (process shutdown). - A failed disk write is logged and never fails the live request — the + The sink assigns ``turn_index`` (monotonic per session, advancing even + when a line fails so gaps are detectable) and computes ``prefix_stable`` + against the last line it actually wrote. Files open lazily on the first + turn and are closed by ``finalize`` (which also appends the + ``tito.session.v1`` metadata line). ``finalize`` is idempotent; + ``close_all`` finalizes every open file (process shutdown). A failed or + rejected write is logged and never fails the live request — the in-memory trajectory still holds the turn for a read-time harvest. """ @@ -133,17 +203,41 @@ def __init__(self, record_dir: str | Path) -> None: self.record_dir.mkdir(parents=True, exist_ok=True) self._files: dict[str, IO[str]] = {} self._turns: dict[str, int] = {} + # prompt+completion ids of the last successfully written line — + # the baseline `prefix_stable` is defined against. + self._last_recorded: dict[str, list[int]] = {} def path_for(self, session_id: str) -> Path: return self.record_dir / f"{session_id}.jsonl" def append_turn(self, session_id: str, record: dict[str, Any]) -> int: - """Append one turn record (assigning ``turn_index``), flush, and - return the assigned index.""" + """Validate and append one turn record (assigning ``turn_index`` and + ``prefix_stable``), flush, and return the assigned index. Any failure + is logged and leaves an index gap instead of breaking the turn.""" turn_index = self._turns.get(session_id, 0) - record = {**record, "turn_index": turn_index} - self._write(session_id, record) - self._turns[session_id] = turn_index + 1 + try: + _validate_turn_record(record) + last = self._last_recorded.get(session_id) + prompt_ids = record["prompt_token_ids"] + record = { + **record, + "turn_index": turn_index, + "prefix_stable": not last or prompt_ids[: len(last)] == last, + } + self._write(session_id, record) + self._last_recorded[session_id] = list(prompt_ids) + list(record["completion_token_ids"]) + except Exception: + # Capture must never take down the serving path: the turn is + # already committed in memory and remains harvestable via GET. + # The finally-increment below leaves a detectable turn_index gap. + logger.exception( + "tito record: turn %d for session %s NOT persisted (%s)", + turn_index, + session_id, + self.path_for(session_id), + ) + finally: + self._turns[session_id] = turn_index + 1 return turn_index def finalize(self, session_id: str, *, reason: str) -> None: @@ -160,8 +254,11 @@ def finalize(self, session_id: str, *, reason: str) -> None: } try: self._write(session_id, meta) + except Exception: + logger.exception("tito record: failed to finalize %s", self.path_for(session_id)) finally: self._turns.pop(session_id, None) + self._last_recorded.pop(session_id, None) handle = self._files.pop(session_id, None) if handle is not None and not handle.closed: try: @@ -174,28 +271,27 @@ def close_all(self, *, reason: str = "shutdown") -> None: self.finalize(session_id, reason=reason) def _write(self, session_id: str, record: dict[str, Any]) -> None: - try: - handle = self._files.get(session_id) - if handle is None or handle.closed: - handle = self.path_for(session_id).open("a", encoding="utf-8") - self._files[session_id] = handle - # `default=repr` keeps a stray non-JSON value (and `allow_nan` - # keeps a NaN logprob the backend passed through) from sinking - # the line — best-effort capture beats failing the live turn. - handle.write(json.dumps(record, ensure_ascii=False, default=repr) + "\n") - handle.flush() - except OSError: - # Capture must never take down the serving path: the turn is - # already committed in memory and remains harvestable via GET. - logger.exception( - "tito record: failed to append to %s — turn NOT persisted", self.path_for(session_id) - ) + """Append one strict-JSON line and flush. Raises on any failure — + the callers own the log-and-continue policy and the gap semantics.""" + handle = self._files.get(session_id) + if handle is None or handle.closed: + handle = self.path_for(session_id).open("a", encoding="utf-8") + self._files[session_id] = handle + # Strict JSON: no NaN/Infinity literals (allow_nan=False) — a + # non-finite value slipping past validation fails the line (leaving a + # gap) rather than emitting non-standard JSON. `default=repr` only + # covers non-token metadata (e.g. an exotic value inside a sampling + # field); token ids and logprobs are strictly validated above. + handle.write(json.dumps(record, ensure_ascii=False, allow_nan=False, default=repr) + "\n") + handle.flush() __all__ = [ "RECORD_SCHEMA_VERSION", + "SAMPLING_KEYS", "SESSION_META_SCHEMA_VERSION", "TurnRecordSink", "build_turn_record", "compute_render_skew", + "sampling_from_request", ] diff --git a/plugins/tito/agentix/tito/engine/session_app.py b/plugins/tito/agentix/tito/engine/session_app.py index 3fababd..4a7f287 100644 --- a/plugins/tito/agentix/tito/engine/session_app.py +++ b/plugins/tito/agentix/tito/engine/session_app.py @@ -31,7 +31,7 @@ ) from .pretokenize import get_tito_tokenizer from .processing import load_tokenizer -from .record import build_turn_record +from .record import build_turn_record, sampling_from_request from .trajectory import GetSessionResponse, LinearTrajectory, SessionRecord, SessionRegistry from .upstream import Backend, get_upstream @@ -137,6 +137,11 @@ async def chat_completions(request: Request, session_id: str) -> Response: return await _chat_turn(request, session_id, session) finally: session.inflight -= 1 + # The idle clock starts when the turn ENDS: `get_session` touched + # it at request start, so without this a generation that takes + # longer than the TTL would count as idle time and the agent's + # very next request would sweep its own live session away. + session.last_used = time.monotonic() async def _chat_turn(request: Request, session_id: str, session: LinearTrajectory) -> Response: # Read + parse the body BEFORE taking the lock: the read lasts as long @@ -210,26 +215,32 @@ async def _chat_turn(request: Request, session_id: str, session: LinearTrajector ) ) if registry.record_sink is not None: - # Inside the lock so file order == trajectory order. Sink - # failures are logged, never fail the served turn. - registry.record_sink.append_turn( - session_id, - build_turn_record( - session_id=session_id, - request_id=request.headers.get("x-request-id"), - model=request_body.get("model"), - backend_kind=backend_kind, - prompt_token_ids=prompt_token_ids, - prompt_segments=prepared.segments, - prefix_stable=prepared.prefix_stable, - completion_token_ids=harvest.completion_token_ids, - completion_logprobs=harvest.completion_logprobs, - assistant_message=harvest.assistant_message, - finish_reason=harvest.finish_reason, - tokenizer_fingerprint=registry.tokenizer_fingerprint, - render_token_ids=harvest.render_token_ids, - ), - ) + # Inside the lock so file order == trajectory order. Capture + # failures (build OR write) are logged, never fail the served + # turn; sink-level failures additionally leave a turn_index + # gap so a dropped line is detectable downstream. + try: + registry.record_sink.append_turn( + session_id, + build_turn_record( + session_id=session_id, + request_id=request.headers.get("x-request-id"), + thread_id=request.headers.get("x-thread-id"), + model=request_body.get("model"), + backend_kind=backend_kind, + sampling=sampling_from_request(request_body), + prompt_token_ids=prompt_token_ids, + prompt_segments=prepared.segments, + completion_token_ids=harvest.completion_token_ids, + completion_logprobs=harvest.completion_logprobs, + assistant_message=harvest.assistant_message, + finish_reason=harvest.finish_reason, + tokenizer=registry.tokenizer_info, + render_token_ids=harvest.render_token_ids, + ), + ) + except Exception: + logger.exception("tito record: failed to capture turn for session %s", session_id) return backend.build_proxy_response(turn.proxy_result) @app.api_route("/sessions/{session_id}/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"]) diff --git a/plugins/tito/agentix/tito/engine/trajectory.py b/plugins/tito/agentix/tito/engine/trajectory.py index 62e86ab..42d1120 100644 --- a/plugins/tito/agentix/tito/engine/trajectory.py +++ b/plugins/tito/agentix/tito/engine/trajectory.py @@ -12,6 +12,7 @@ import asyncio import hashlib +import json import logging import time import uuid @@ -70,8 +71,13 @@ class PreparedPrompt: (``tool``/``user``/``system``), and ``generation_prompt``. `prefix_stable` is True iff `token_ids` extends the last *committed* - checkpoint (previous prompt + completion) — False after a retry rollback - or when the checkpoint window was outrun and the prompt was re-rendered. + in-memory checkpoint (previous prompt + completion) — False after a retry + rollback or when the checkpoint window was outrun and the prompt was + re-rendered. NOTE: this is engine-level, advisory metadata. The + `prefix_stable` field in a persisted `tito.record.v1` line is computed by + the record sink against the last RECORDED line instead — an applied + rollback whose turn never produced a line (upstream error) must still + surface as a break in the record stream (see `record.TurnRecordSink`). """ token_ids: list[int] @@ -328,8 +334,11 @@ def __init__(self, args: Any, tokenizer: Any, *, tito_tokenizer: TITOTokenizer) self.session_ttl_seconds: float | None = getattr(args, "session_ttl_seconds", None) or None self.max_sessions: int | None = getattr(args, "max_sessions", None) or None self.on_evict: Callable[[str], None] | None = None - self.tokenizer_fingerprint: dict[str, str] = { + # The record's `tokenizer` block: pins the checkpoint name, the + # tokenizer definition bytes, and the chat template in effect. + self.tokenizer_info: dict[str, str] = { "checkpoint": str(getattr(args, "hf_checkpoint", "") or ""), + "tokenizer_sha256": _tokenizer_sha256(tokenizer), "chat_template_sha256": _chat_template_sha256(tokenizer, tito_tokenizer), } @@ -423,3 +432,25 @@ def _chat_template_sha256(tokenizer: Any, tito_tokenizer: TITOTokenizer) -> str: tokenizer, "chat_template", None ) return hashlib.sha256(str(template or "").encode()).hexdigest() + + +def _tokenizer_sha256(tokenizer: Any) -> str: + """SHA-256 over the tokenizer DEFINITION bytes — the identity a record's + token ids depend on. + + Method (documented as part of the record contract): for fast tokenizers, + hash ``backend_tokenizer.to_str()`` (the complete serialized + ``tokenizer.json`` definition: vocab, merges, normalizer, added tokens — + deterministic for a given tokenizer); for slow tokenizers, hash the + sorted-JSON vocabulary. Two gateways report the same value iff they + tokenize identically.""" + backend = getattr(tokenizer, "backend_tokenizer", None) + if backend is not None: + try: + return hashlib.sha256(backend.to_str().encode()).hexdigest() + except Exception: # pragma: no cover - serialization quirks per version + logger.exception("tokenizer_sha256: backend serialization failed; falling back to vocab") + get_vocab = getattr(tokenizer, "get_vocab", None) + vocab = get_vocab() if callable(get_vocab) else {} + blob = json.dumps(vocab, sort_keys=True, ensure_ascii=False) + return hashlib.sha256(blob.encode()).hexdigest() diff --git a/plugins/tito/tests/test_record.py b/plugins/tito/tests/test_record.py index ee03597..57b5b11 100644 --- a/plugins/tito/tests/test_record.py +++ b/plugins/tito/tests/test_record.py @@ -81,14 +81,23 @@ def __init__(self) -> None: self.message: dict = {"role": "assistant", "content": "ok done"} self.hold: asyncio.Event | None = None self.hold_marker: str | None = None + self.fail_next = False + self.delay = 0.0 async def handler(self, request: httpx.Request) -> httpx.Response: body = json.loads(request.content) self.calls.append(body) if self.hold is not None and self.hold_marker in json.dumps(body): await self.hold.wait() + if self.delay: + await asyncio.sleep(self.delay) + if self.fail_next: + self.fail_next = False + return httpx.Response(503, json={"error": "upstream busy"}) ids = list(self.completion_ids) - return httpx.Response(200, json={ + # Serialize with stdlib json (allows NaN like a lenient real backend + # would) and ship verbatim — httpx's own `json=` encoder is strict. + blob = json.dumps({ "id": "c1", "object": "chat.completion", "model": "m", "choices": [{ "index": 0, @@ -100,7 +109,8 @@ async def handler(self, request: httpx.Request) -> httpx.Response: }, }], "usage": {"prompt_tokens": 3, "completion_tokens": len(ids), "total_tokens": 5}, - }) + }).encode() + return httpx.Response(200, content=blob, headers={"content-type": "application/json"}) def _make_gateway(tok, monkeypatch, **arg_overrides): @@ -135,7 +145,9 @@ async def test_record_line_shape_and_crash_safety(tok, monkeypatch, tmp_path): path = tmp_path / f"{sid}.jsonl" r = await client.post( - f"/sessions/{sid}/v1/chat/completions", json=_CHAT, headers={"x-request-id": "req-001"} + f"/sessions/{sid}/v1/chat/completions", + json={**_CHAT, "temperature": 0.6, "top_p": 0.95, "max_tokens": 128}, + headers={"x-request-id": "req-001", "x-thread-id": "thread-7"}, ) assert r.status_code == 200 @@ -145,10 +157,13 @@ async def test_record_line_shape_and_crash_safety(tok, monkeypatch, tmp_path): assert rec["schema_version"] == "tito.record.v1" assert rec["session_id"] == sid + assert rec["thread_id"] == "thread-7" # x-thread-id passthrough assert rec["turn_index"] == 0 assert rec["request_id"] == "req-001" assert rec["model"] == "m" assert rec["backend_kind"] == "sglang" + # Whitelisted sampling params lifted verbatim from the request body. + assert rec["sampling"] == {"temperature": 0.6, "top_p": 0.95, "max_tokens": 128} assert rec["prompt_token_ids"] + rec["completion_token_ids"] == accumulated assert rec["completion_token_ids"] == [7, 8] assert rec["completion_logprobs"] == [-0.25, -0.5] @@ -157,8 +172,9 @@ async def test_record_line_shape_and_crash_safety(tok, monkeypatch, tmp_path): assert rec["finish_reason"] == "stop" assert rec["prefix_stable"] is True assert rec["render_skew"] is None # sglang exposes no render ids - assert rec["tokenizer_fingerprint"] == { + assert rec["tokenizer"] == { "checkpoint": "tiny-in-memory", + "tokenizer_sha256": hashlib.sha256(tok.backend_tokenizer.to_str().encode()).hexdigest(), "chat_template_sha256": hashlib.sha256(tok.chat_template.encode()).hexdigest(), } # First turn: one from-scratch render segment covering the whole prompt. @@ -182,6 +198,8 @@ async def test_record_line_shape_and_crash_safety(tok, monkeypatch, tmp_path): rec1, rec2 = _records(path) assert rec2["turn_index"] == 1 assert rec2["request_id"] is None # no x-request-id header sent + assert "thread_id" not in rec2 # key omitted when x-thread-id absent + assert rec2["sampling"] == {} # no sampling params in the request assert rec2["prefix_stable"] is True assert [s["source"] for s in rec2["prompt_segments"]] == ["prefix", "tool", "generation_prompt"] # Segment spans tile the prompt exactly, and the prefix span IS the @@ -373,6 +391,116 @@ async def test_shutdown_finalizes_open_record_files(tok, monkeypatch, tmp_path): assert meta["reason"] == "shutdown" +@pytest.mark.asyncio +async def test_unrecorded_rollback_turn_still_breaks_prefix_stability(tok, monkeypatch, tmp_path): + """Adversarial-review regression: a rollback applied in phase 1 whose + turn never produced a record line (upstream 503) must NOT vanish from the + record stream — the next successful line's prefix_stable is computed + against the last RECORDED line, not the in-memory checkpoint, so the + discontinuity is flagged.""" + client, replica, srv, _ = _make_gateway(tok, monkeypatch, record_dir=str(tmp_path)) + sid = (await client.post("/sessions")).json()["session_id"] + path = tmp_path / f"{sid}.jsonl" + + assert (await client.post(f"/sessions/{sid}/v1/chat/completions", json=_CHAT)).status_code == 200 + base = [*_CHAT["messages"], {"role": "assistant", "content": "ok done"}] + assert (await client.post( + f"/sessions/{sid}/v1/chat/completions", + json={"model": "m", "messages": [*base, {"role": "tool", "content": "done"}]}, + )).status_code == 200 + + # History rewrite (rollback applied in phase 1), upstream 503s -> passed + # through, NO record line, but the rollback stays committed in memory. + replica.fail_next = True + r = await client.post( + f"/sessions/{sid}/v1/chat/completions", + json={"model": "m", "messages": [*base, {"role": "tool", "content": "You"}]}, + ) + assert r.status_code == 503 + + # The agent retries the same rewritten history; the in-memory checkpoint + # now extends cleanly — but the record stream does not. + r = await client.post( + f"/sessions/{sid}/v1/chat/completions", + json={"model": "m", "messages": [*base, {"role": "tool", "content": "You"}]}, + ) + assert r.status_code == 200 + + recs = _records(path) + assert len(recs) == 3 # the 503 turn left no line (and no index gap: it never committed) + prev_stream = recs[1]["prompt_token_ids"] + recs[1]["completion_token_ids"] + assert recs[2]["prompt_token_ids"][: len(prev_stream)] != prev_stream + assert [rec["prefix_stable"] for rec in recs] == [True, True, False] + + +@pytest.mark.asyncio +async def test_generation_time_does_not_count_as_idle_for_ttl(tok, monkeypatch, tmp_path): + """Adversarial-review regression: the idle clock starts when a turn ENDS. + A single generation slower than the TTL must not let the agent's + immediately following request sweep its own live session away.""" + client, replica, srv, _ = _make_gateway( + tok, monkeypatch, record_dir=str(tmp_path), session_ttl_seconds=0.5 + ) + sid = (await client.post("/sessions")).json()["session_id"] + + replica.delay = 0.8 # one turn's generation exceeds the TTL + assert (await client.post(f"/sessions/{sid}/v1/chat/completions", json=_CHAT)).status_code == 200 + replica.delay = 0.0 + + followup = { + "model": "m", + "messages": [ + *_CHAT["messages"], + {"role": "assistant", "content": "ok done"}, + {"role": "tool", "content": "done"}, + ], + } + r = await client.post(f"/sessions/{sid}/v1/chat/completions", json=followup) + assert r.status_code == 200 # the session survived: zero actual idle time + assert len(_records(tmp_path / f"{sid}.jsonl")) == 2 + + +@pytest.mark.asyncio +async def test_invalid_token_data_is_rejected_with_detectable_gap(tok, monkeypatch, tmp_path): + """Strict-JSON sink contract: a NaN logprob (JSON-parseable upstream, not + strict JSON) or a non-int token id must never be coerced into the record + file — the line is dropped, the turn still serves, and the next line's + turn_index shows a detectable gap.""" + client, replica, srv, _ = _make_gateway(tok, monkeypatch, record_dir=str(tmp_path)) + sid = (await client.post("/sessions")).json()["session_id"] + path = tmp_path / f"{sid}.jsonl" + + # Turn 0: NaN logprob passed through by the backend (stdlib json accepts it). + replica.logprobs = [float("nan"), -0.5] + r = await client.post(f"/sessions/{sid}/v1/chat/completions", json=_CHAT) + assert r.status_code == 200 # capture failure never fails the served turn + assert not path.exists() or _records(path) == [] + + # Turn 1: clean — recorded with a turn_index gap exposing the dropped line. + replica.logprobs = [-0.25, -0.5] + followup = { + "model": "m", + "messages": [ + *_CHAT["messages"], + {"role": "assistant", "content": "ok done"}, + {"role": "tool", "content": "done"}, + ], + } + assert (await client.post(f"/sessions/{sid}/v1/chat/completions", json=followup)).status_code == 200 + (rec,) = _records(path) + assert rec["turn_index"] == 1 # index 0 is the detectable hole + # The dropped line is also the stability baseline hole: turn 1 extends + # nothing recorded, so it is the stream's first line and reads stable. + assert rec["prefix_stable"] is True + # Every persisted line is strict JSON (no NaN/Infinity literals). + for line in path.read_text().splitlines(): + json.loads(line, parse_constant=lambda name: pytest.fail(f"non-strict JSON literal {name}")) + # The session meta line still counts every turn slot, recorded or not. + assert (await client.delete(f"/sessions/{sid}")).status_code == 204 + (meta,) = _meta(path) + assert meta["turns"] == 2 + + def test_compute_render_skew_contract(): assert compute_render_skew(None, [1, 2]) is None assert compute_render_skew([1, 2], [1, 2]) == {"equal": True, "first_divergence": None} From 13c25f2809c0be682940cf608bbfce2e9386edb2 Mon Sep 17 00:00:00 2001 From: Lingrui Mei Date: Fri, 24 Jul 2026 07:24:50 +0800 Subject: [PATCH 5/5] abridge: conversation demux for tito serve mode + capture hardening (review fixes) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Real Anthropic agents multiplex several logical conversations over one API key (helper calls, Task subagents, reruns), while a TITO gateway session is one append-only linear history — binding a key to a single gateway session bricked the key at the first unrelated history (400, 409 under concurrency). serve --tito-url now DEMUXES per caller session: requests are keyed by the canonicalized (system, first user message) pair and each distinct conversation lazily gets its own AnthropicToOpenAI(SessionForward) — its own gateway session, its own assistant-replay memory. Documented boundaries (loud, not silent): deep compaction that keeps the opening rides the gateway's rollback / from-scratch paths; byte-identical openings collide; a caller key LRU-evicted mid-rollout continues in fresh gateway sessions (capture splits there — size --max-sessions above the concurrent key count). New --tito-delete-on-evict (default off) reaps a caller session's gateway sessions when it closes; default keeps them for harvest. Join-key fidelity: the transport now publishes the session id it stamped upstream (context var), and Recorder rows gain gateway_session_id — the gateway's OWN session id, i.e. the session_id in its token records — restoring the session-level join that the caller-key hash cannot provide (session_id_for docstring carries the tito-mode caveat). Recorder hardening: _write is log-and-serve (any capture failure is logged, the agent's call still succeeds — matching the gateway sink's policy), and a straggler dispatch after aclose() drops its row loudly instead of resurrecting the closed file handle. Translation contract: TRANSLATION_SPEC_SHA now hashes every module that shapes the upstream body — the pure transforms plus both client modules (assistant replay, forced stream=False, model override, upstream_params) — so behavior drift outside the transforms file can no longer hide behind an unmoved pin. Thinking-only assistant history messages (legal Anthropic shape) are no longer dropped: their reasoning rides reasoning_content through the append guard. Tests: multiplexed-conversations regression (subagent/helper openings get separate gateway sessions while each conversation's turns stay in its own), delete-on-evict reap, gateway_session_id row join, recorder log-and-serve + closed-drop, thinking-only preservation, joint spec-sha pin. Co-Authored-By: Claude Fable 5 --- plugins/abridge/README.md | 45 ++-- plugins/abridge/agentix/bridge/_request_id.py | 19 +- .../bridge/clients/_anthropic_transforms.py | 44 +++- plugins/abridge/agentix/bridge/forward.py | 7 +- plugins/abridge/agentix/bridge/recorder.py | 63 +++++- plugins/abridge/agentix/bridge/serve.py | 196 ++++++++++++++++-- .../abridge/tests/test_anthropic_transform.py | 25 +++ plugins/abridge/tests/test_recorder.py | 42 ++++ plugins/abridge/tests/test_serve.py | 18 +- .../abridge/tests/test_tito_composition.py | 71 ++++++- 10 files changed, 471 insertions(+), 59 deletions(-) diff --git a/plugins/abridge/README.md b/plugins/abridge/README.md index 05cece5..fe2b2d5 100644 --- a/plugins/abridge/README.md +++ b/plugins/abridge/README.md @@ -263,23 +263,42 @@ code, so when you expose it to them (`--host`), also set `build_session_app`) so only keys your harness minted are served; everything else gets a 401. -Two more serve options: +More serve options: * `--tito-url http://tito:30000` (mutually exclusive with `--upstream-base-url`) — put the Anthropic shell in front of a - token-recording session gateway instead of a plain engine: each caller - session composes `AnthropicToOpenAI(SessionForward(tito_url).handler())`, - so the gateway sees the OpenAI chat body, owns render/generate and the - token record, and one caller key maps to one gateway session. + token-recording session gateway instead of a plain engine. The gateway + keeps one append-only linear conversation per session, while real + Anthropic agents multiplex several conversations over one key (helper + calls, subagents, reruns), so each caller session **demuxes by + conversation**: requests are keyed by the canonicalized + `(system, first user message)` pair, and each distinct key gets its own + `AnthropicToOpenAI(SessionForward(tito_url).handler())` — its own + gateway session. Known boundaries (by design, fail loudly rather than + silently): a mid-conversation history rewrite that keeps the opening + (deep compaction) still lands in the same gateway session and rides the + gateway's rollback / from-scratch paths (rewrites past its rollback + window are its documented 400); two genuinely different conversations + with a byte-identical opening collide into one session; and a caller + key evicted by the serve LRU (`--max-sessions`) mid-rollout continues + in fresh gateway sessions — the gateway-side capture splits there + (`turn_index` restarts), so size `--max-sessions` above your concurrent + rollout-key count. +* `--tito-delete-on-evict` — reap a caller session's gateway sessions + when it closes (LRU eviction or shutdown). Default off: gateway + sessions stay alive for harvest and the harvester deletes them. * `--record-dir DIR` (either mode) — wrap each session's client in a - `Recorder` writing `DIR/.jsonl`; rows carry `session_id` + - `request_id`, and the same `request_id` reaches the upstream as - `x-request-id`, so message rows join the gateway's token records. - -`GET /_health` reports `translation_spec_sha` — the SHA-256 of the -Anthropic↔OpenAI transform module source — so downstream data contracts -can pin the exact translation their captured trajectories were produced -under. + `Recorder` writing `DIR/.jsonl`; rows carry `session_id`, + `request_id`, and (in tito mode) `gateway_session_id` — the gateway's + own session id, i.e. the `session_id` in its token records. The same + `request_id` reaches the upstream as `x-request-id`, so message rows + join the gateway's token records per call as well as per session. + +`GET /_health` reports `translation_spec_sha` — one SHA-256 over the +source of the Anthropic↔OpenAI transform module and both client modules +that shape the upstream body (assistant replay, forced non-streaming, +model override, operator params) — so downstream data contracts can pin +the exact translation their captured trajectories were produced under. Programmatic surface in `agentix.bridge.serve`: `build_app(*clients)` (shared session) and `build_session_app(factory)` (one client per diff --git a/plugins/abridge/agentix/bridge/_request_id.py b/plugins/abridge/agentix/bridge/_request_id.py index 9ac344f..a44ee9e 100644 --- a/plugins/abridge/agentix/bridge/_request_id.py +++ b/plugins/abridge/agentix/bridge/_request_id.py @@ -14,6 +14,14 @@ inner layers reuse a bound id and only mint their own when nothing upstream bound one. Works unchanged across `await` within one handler invocation and never leaks across concurrent calls. + +`current_upstream_session_id` flows the OTHER way on the same principle: the +transport layer (`Forward`) publishes the session id it stamped upstream as +`x-session-id` — for a `SessionForward` that is the gateway-assigned session +id, which the caller-side capture cannot otherwise know (it exists only +after the lazy session create). The `Recorder` clears it before each handler +call and reads it afterwards into the row's `gateway_session_id`, restoring +the session-level join between caller-side rows and gateway-side records. """ from __future__ import annotations @@ -23,6 +31,10 @@ current_request_id: ContextVar[str | None] = ContextVar("abridge_request_id", default=None) +current_upstream_session_id: ContextVar[str | None] = ContextVar( + "abridge_upstream_session_id", default=None +) + def mint_request_id() -> str: return uuid.uuid4().hex @@ -34,4 +46,9 @@ def get_or_mint_request_id() -> str: return bound if bound else mint_request_id() -__all__ = ["current_request_id", "get_or_mint_request_id", "mint_request_id"] +__all__ = [ + "current_request_id", + "current_upstream_session_id", + "get_or_mint_request_id", + "mint_request_id", +] diff --git a/plugins/abridge/agentix/bridge/clients/_anthropic_transforms.py b/plugins/abridge/agentix/bridge/clients/_anthropic_transforms.py index 7607dc4..d935cd8 100644 --- a/plugins/abridge/agentix/bridge/clients/_anthropic_transforms.py +++ b/plugins/abridge/agentix/bridge/clients/_anthropic_transforms.py @@ -334,7 +334,11 @@ def _messages_anthropic_to_openai(messages: list[Any]) -> list[dict[str, Any]]: message["reasoning_content"] = "\n".join(part for part in thinking_parts if part) if tool_calls: message["tool_calls"] = tool_calls - if message["content"] is not None or tool_calls: + # A thinking-only assistant turn (legal Anthropic shape — e.g. + # extended thinking cut off at max_tokens) maps to reasoning_content + # and must survive: dropping the whole message would leave two + # adjacent user turns and silently lose the forwarded reasoning. + if message["content"] is not None or tool_calls or thinking_parts: out.append(message) else: out.extend(tool_results) @@ -405,14 +409,36 @@ def _sse(event: str, data: dict[str, Any]) -> bytes: return f"event: {event}\ndata: {payload}\n\n".encode() -# The translation contract version: SHA-256 over this module's source. Any -# edit to the transforms — however small — changes the byte identity of the -# OpenAI bodies a recording backend sees, so downstream data contracts pin -# this value (abridge-serve reports it on `/_health`). Deliberately the -# file's bytes, not a semantic hash: comments and docstrings changing the sha -# is a false positive we accept; a behavior change slipping through unhashed -# is not. -TRANSLATION_SPEC_SHA: str = hashlib.sha256(Path(__file__).read_bytes()).hexdigest() +# The translation contract version: SHA-256 over the source of EVERY module +# that shapes what an upstream/recording backend receives — the pure +# transforms here, plus the two client modules that rewrite the body around +# them (assistant-replay memory, forced stream=False, model override, +# operator upstream_params). Hashing only this file would let those rewrites +# drift without moving the pin. Any edit — however small — changes the byte +# identity of the OpenAI bodies a recording backend sees, so downstream data +# contracts pin this value (abridge-serve reports it on `/_health`). +# Deliberately file bytes, not a semantic hash: comments changing the sha is +# a false positive we accept; a behavior change slipping through unhashed is +# not. Sibling sources are read directly (no import) to avoid a cycle with +# the client modules, which import this one. +_TRANSLATION_SPEC_FILES = ( + "_anthropic_transforms.py", + "anthropic_to_openai.py", + "anthropic_from_openai.py", +) + + +def _translation_spec_sha() -> str: + digest = hashlib.sha256() + for name in _TRANSLATION_SPEC_FILES: + digest.update(name.encode()) + digest.update(b"\x00") + digest.update((Path(__file__).parent / name).read_bytes()) + digest.update(b"\x00") + return digest.hexdigest() + + +TRANSLATION_SPEC_SHA: str = _translation_spec_sha() __all__ = [ diff --git a/plugins/abridge/agentix/bridge/forward.py b/plugins/abridge/agentix/bridge/forward.py index 901ef0d..95add78 100644 --- a/plugins/abridge/agentix/bridge/forward.py +++ b/plugins/abridge/agentix/bridge/forward.py @@ -35,7 +35,7 @@ import httpx -from ._request_id import get_or_mint_request_id +from ._request_id import current_upstream_session_id, get_or_mint_request_id from .proxy import AbridgeError, ClientResponse, Handler, Request logger = logging.getLogger(__name__) @@ -136,6 +136,11 @@ async def _forward(self, path: str, request: Request) -> ClientResponse: # Reuses the id a wrapping capture layer (Recorder) bound for this # call, so its JSONL row and the sidecar's token record share one id. record_id = get_or_mint_request_id() + # Publish the upstream session identity for the capture layer: for a + # SessionForward this is the gateway-assigned session id (known only + # here, after the lazy create) — the Recorder reads it back into the + # row's `gateway_session_id` join key. + current_upstream_session_id.set(self.session_id) headers = { **self._headers, "x-session-id": self.session_id, diff --git a/plugins/abridge/agentix/bridge/recorder.py b/plugins/abridge/agentix/bridge/recorder.py index 0315f59..b34e054 100644 --- a/plugins/abridge/agentix/bridge/recorder.py +++ b/plugins/abridge/agentix/bridge/recorder.py @@ -9,6 +9,7 @@ {"ts": ..., "path": "/v1/messages", "request_id": "<32 hex>", "session_id": ..., # only when the Recorder has one + "gateway_session_id": ..., # only when the transport published one "request": {...}, "response": {"status_code": 200, "media_type": "...", "body": ...}} @@ -18,9 +19,13 @@ upstream hop — a downstream token recorder's per-turn record and this row join on it. `session_id`, when given, identifies the rollout the wrapped client serves (pass the same value as the client's session identity). -Without these keys, rows from a retried call (e.g. an agent retry after a -tunnel 504 produced an orphan success row) are only deduplicable by -request-body equality. +`gateway_session_id` is read back from the transport after the call (via +`current_upstream_session_id`): when the downstream is a session-scoped +gateway (`SessionForward`), it is the gateway's OWN session id — i.e. the +`session_id` in the gateway's token records — restoring the session-level +join that the caller-side hash alone cannot provide. Without these keys, +rows from a retried call (e.g. an agent retry after a tunnel 504 produced an +orphan success row) are only deduplicable by request-body equality. A handler that raises records `{"error": ...}` instead of `"response"` and re-raises — a failed call is signal, not something to lose. JSON bodies are @@ -30,19 +35,27 @@ flushed as it is written so the file is complete up to the last call even if the process dies mid-rollout. The file opens lazily on the first record, so a Recorder that never serves (e.g. a route-enumeration probe) leaves no -empty file behind. +empty file behind. Capture is log-and-serve: a failed row write (disk full, +unencodable text) is logged and the agent's call still succeeds — matching +the token-recording gateway's policy, so the two capture layers never +disagree about whether a turn happened. After `aclose()` a straggler +in-flight call's row is dropped (logged), never written to a resurrected +file handle. """ from __future__ import annotations import json +import logging import time from pathlib import Path from typing import IO, Any -from ._request_id import current_request_id, mint_request_id +from ._request_id import current_request_id, current_upstream_session_id, mint_request_id from .proxy import ClientResponse, Handler, Request, _collect_handlers +logger = logging.getLogger(__name__) + class Recorder: """Wrap a handler client; record every (request, response) pair it serves. @@ -59,6 +72,7 @@ def __init__(self, client: Any, path: str | Path, *, session_id: str | None = No self._path.parent.mkdir(parents=True, exist_ok=True) self._session_id = session_id self._file: IO[str] | None = None + self._closed = False def abridge_routes(self) -> dict[str, Handler]: return {path: self._recording(path, handler) for path, handler in _collect_handlers(self._client).items()} @@ -73,15 +87,23 @@ async def record(request: Request) -> ClientResponse: if self._session_id is not None: line["session_id"] = self._session_id line["request"] = request.body - token = current_request_id.set(request_id) + rid_token = current_request_id.set(request_id) + # Cleared per call so a value published by a PREVIOUS call on + # this task never leaks into an unrelated row. + upstream_token = current_upstream_session_id.set(None) try: response = await handler(request) except BaseException as exc: line["error"] = f"{type(exc).__name__}: {exc}" + self._stamp_gateway_session(line) self._write(line) raise finally: - current_request_id.reset(token) + current_request_id.reset(rid_token) + gateway_session_id = current_upstream_session_id.get() + current_upstream_session_id.reset(upstream_token) + if gateway_session_id is not None: + line["gateway_session_id"] = gateway_session_id line["response"] = { "status_code": response.status_code, "media_type": response.media_type, @@ -92,11 +114,29 @@ async def record(request: Request) -> ClientResponse: return record + @staticmethod + def _stamp_gateway_session(line: dict[str, Any]) -> None: + gateway_session_id = current_upstream_session_id.get() + if gateway_session_id is not None: + line["gateway_session_id"] = gateway_session_id + def _write(self, line: dict[str, Any]) -> None: - if self._file is None or self._file.closed: - self._file = self._path.open("a", encoding="utf-8") - self._file.write(json.dumps(line, ensure_ascii=False, default=repr) + "\n") - self._file.flush() + # Log-and-serve, mirroring the token-recording gateway's policy: the + # upstream call already succeeded (or its error is being re-raised), + # so a capture failure must not turn it into a wire error. + try: + if self._closed: + # A straggler dispatch outlived aclose(): the file is closed + # for good — dropping the row (loudly) beats resurrecting a + # file handle nobody will ever close. + logger.warning("abridge recorder: dropping row for %s — recorder is closed", self._path) + return + if self._file is None or self._file.closed: + self._file = self._path.open("a", encoding="utf-8") + self._file.write(json.dumps(line, ensure_ascii=False, default=repr) + "\n") + self._file.flush() + except Exception: # noqa: BLE001 - capture must never fail the served call + logger.exception("abridge recorder: failed to append to %s — row NOT persisted", self._path) def environ(self, handle: Any) -> dict[str, str]: return self._client.environ(handle) @@ -107,6 +147,7 @@ async def aclose(self) -> None: if aclose is not None: await aclose() finally: + self._closed = True if self._file is not None: self._file.close() diff --git a/plugins/abridge/agentix/bridge/serve.py b/plugins/abridge/agentix/bridge/serve.py index fd47e25..dd44a90 100644 --- a/plugins/abridge/agentix/bridge/serve.py +++ b/plugins/abridge/agentix/bridge/serve.py @@ -32,13 +32,30 @@ * `--upstream-base-url` — plain OpenAI-compatible engine; translation + transport live in `AnthropicFromOpenAIClient` (the openai SDK owns the HTTP). The original mode, unchanged. -* `--tito-url` — a token-recording session gateway (the TITO gateway): - each caller session composes - `AnthropicToOpenAI(SessionForward(tito_url).handler(), model=...)`, so - the gateway sees the OpenAI chat body, owns render/generate and the - token record, and one abridge caller session maps to one gateway - session. Token capture stays the gateway's job; this server only adds - the Anthropic shell and identity stamping. +* `--tito-url` — a token-recording session gateway (the TITO gateway). + The gateway keeps ONE append-only linear conversation per session, + while real Anthropic agents multiplex several logical conversations + over a single API key (helper calls, subagents, a rerun). So each + caller session DEMUXES by conversation: requests are keyed by + `sha256(canonical system + first user message)` and each distinct key + gets its own `AnthropicToOpenAI(SessionForward(tito_url).handler())` — + its own gateway session. The gateway sees the OpenAI chat body, owns + render/generate and the token record; this server only adds the + Anthropic shell, identity stamping, and the demux. + + Demux boundaries (documented, not silently papered over): a history + rewrite that keeps the first user message (e.g. mid-conversation + compaction) still lands in the SAME gateway session and is handled by + the gateway's rollback / from-scratch paths — rewrites deeper than the + gateway's rollback window are its documented 400. Two genuinely + different conversations with byte-identical (system, first user + message) collide into one session. And when the caller-session LRU + evicts a key (max_sessions), a still-active rollout on that key + continues in FRESH gateway sessions — the gateway-side capture splits + at that boundary (turn_index restarts); size `--max-sessions` above + the number of concurrent rollout keys. `--tito-delete-on-evict` opts + into reaping the gateway sessions when a caller session closes; the + default keeps them for harvest. `--record-dir` (either mode) wraps each session's client in a `Recorder` writing message-level rows to `/.jsonl`; rows @@ -59,6 +76,7 @@ import argparse import asyncio import hashlib +import json import logging import os from collections import OrderedDict @@ -66,13 +84,23 @@ from contextlib import AbstractAsyncContextManager, asynccontextmanager from dataclasses import dataclass, field from pathlib import Path +from typing import Any import uvicorn from fastapi import FastAPI from fastapi import Request as FastAPIRequest from fastapi.responses import JSONResponse, Response -from .proxy import AbridgeError, Client, Handler, Request, _AsyncCloseable, _collect_handlers +from .proxy import ( + AbridgeError, + Client, + ClientResponse, + Handler, + Request, + _AsyncCloseable, + _collect_handlers, + on, +) logger = logging.getLogger(__name__) @@ -91,6 +119,13 @@ def session_id_for(caller_key: str) -> str: `"anonymous"`. Public so the side minting per-rollout keys can compute the same id and correlate upstream `x-session-id` values (the raw key is never echoed into logs or upstream headers). + + Correlation caveat for `--tito-url` mode: there the upstream + `x-session-id` is the GATEWAY's own session id (assigned per + conversation), not this hash — the caller-side↔gateway-side join + lives in the Recorder rows instead (`session_id` = this hash, + `gateway_session_id` = the gateway's id, `request_id` = the per-call + `x-request-id` echoed into the gateway's token records). """ if not caller_key: return "anonymous" @@ -404,26 +439,153 @@ def _build_parser() -> argparse.ArgumentParser: help=( "record every served (request, response) pair to " "/.jsonl via Recorder — message-level rows with " - "session_id + request_id, flushed per line (env: ABRIDGE_RECORD_DIR)" + "session_id + request_id (+ gateway_session_id in tito mode), flushed " + "per line (env: ABRIDGE_RECORD_DIR)" + ), + ) + parser.add_argument( + "--tito-delete-on-evict", + action="store_true", + help=( + "tito mode only: DELETE the gateway sessions when a caller session " + "closes (LRU eviction or shutdown). Default off — gateway sessions " + "stay alive for harvest, and the harvester deletes them" ), ) return parser +def _canonical_text(value: Any) -> str: + """Flatten Anthropic content (str or block list) to conversation-identity + text: text blocks contribute their text (cache_control and other + tokenization-irrelevant decorations are ignored), other blocks their + sorted-JSON form.""" + if value is None: + return "" + if isinstance(value, str): + return value + if isinstance(value, list): + parts: list[str] = [] + for block in value: + if isinstance(block, str): + parts.append(block) + elif isinstance(block, dict) and block.get("type") == "text": + parts.append(str(block.get("text", ""))) + else: + parts.append(json.dumps(block, sort_keys=True, ensure_ascii=False, default=repr)) + return "\n".join(parts) + return json.dumps(value, sort_keys=True, ensure_ascii=False, default=repr) + + +def _conversation_key(body: dict[str, Any]) -> str: + """A logical-conversation key for an Anthropic Messages request: the + canonicalized system prompt + first user message. Turns of one + conversation share it (the head of the history is append-only in normal + operation); a helper call, subagent, or rerun with a different opening + gets a different key.""" + first_user = next( + (m.get("content") for m in body.get("messages") or [] if isinstance(m, dict) and m.get("role") == "user"), + None, + ) + blob = _canonical_text(body.get("system")) + "\x00" + _canonical_text(first_user) + return hashlib.sha256(blob.encode()).hexdigest()[:16] + + +class _TitoConversationDemux: + """One caller key, many logical conversations → one gateway session each. + + The TITO gateway's session is a single append-only linear history, but an + Anthropic agent multiplexes conversations over one API key (helper calls, + subagents, reruns). Routing everything into one gateway session bricks + the key at the first unrelated history (the gateway rightly 400s a + request sharing no prefix). This client keys each request by + `_conversation_key` and lazily builds one + `AnthropicToOpenAI(SessionForward(...).handler())` per conversation — + scoped like the conversation itself, including the converter's + assistant-replay memory. The conversation map is bounded by the caller + session's own lifetime (LRU eviction / shutdown closes all of them). + + `delete_on_close=True` reaps the gateway sessions on `aclose()` + (`--tito-delete-on-evict`); the default leaves them alive for harvest. + """ + + def __init__( + self, + tito_url: str, + *, + model: str | None = None, + timeout: float = 540.0, + delete_on_close: bool = False, + ) -> None: + self._tito_url = tito_url + self._model = model + self._timeout = timeout + self._delete_on_close = delete_on_close + self._conversations: dict[str, Any] = {} # key -> AnthropicToOpenAI + self._forwards: dict[str, Any] = {} # key -> SessionForward + + def _converter_for(self, body: dict[str, Any]) -> Any: + key = _conversation_key(body) + converter = self._conversations.get(key) + if converter is None: + from .clients import AnthropicToOpenAI + from .forward import SessionForward + + forward = SessionForward(self._tito_url, paths=["/v1/chat/completions"], timeout=self._timeout) + converter = AnthropicToOpenAI(forward.handler(), model=self._model) + self._conversations[key] = converter + self._forwards[key] = forward + logger.info("tito demux: new conversation %s -> fresh gateway session", key) + return converter + + @on("/v1/messages") + async def messages(self, request: Request) -> ClientResponse: + return await self._converter_for(request.body).messages(request) + + @on("/v1/messages/count_tokens") + async def count_tokens(self, request: Request) -> ClientResponse: + # Answered locally (same estimate as the converters) — counting must + # not create a gateway session for a conversation that never runs. + from .clients._anthropic_transforms import count_anthropic_tokens + + return ClientResponse.json({"input_tokens": count_anthropic_tokens(request.body).input_tokens}) + + def environ(self, handle: Any) -> dict[str, str]: + from .clients.anthropic import PLACEHOLDER_API_KEY + + return {"ANTHROPIC_BASE_URL": handle.url, "ANTHROPIC_API_KEY": PLACEHOLDER_API_KEY} + + async def aclose(self) -> None: + for key, forward in self._forwards.items(): + if self._delete_on_close: + try: + await forward.delete_session() + except Exception: # noqa: BLE001 - best-effort reap + logger.exception("tito demux: failed to delete gateway session for conversation %s", key) + for key, converter in self._conversations.items(): + try: + await converter.aclose() + except Exception: # noqa: BLE001 - close every conversation + logger.exception("tito demux: failed to close conversation %s", key) + self._conversations.clear() + self._forwards.clear() + + def _client_factory(args: argparse.Namespace) -> Callable[[str], Client]: """The per-caller-session client for the chosen upstream mode, wrapped in a `Recorder` when `--record-dir` is set.""" if args.tito_url: # Composition seam (transport-blind): the TITO gateway sees the - # OpenAI chat body and owns tokens + recording; abridge adds only the - # Anthropic shell. One SessionForward per caller session == one - # gateway session per rollout key. - from .clients import AnthropicToOpenAI - from .forward import SessionForward - + # OpenAI chat body and owns tokens + recording; abridge adds the + # Anthropic shell and the per-conversation demux (one gateway + # session per logical conversation on the caller key). def build(session_id: str) -> Client: - forward = SessionForward(args.tito_url, paths=["/v1/chat/completions"], timeout=args.upstream_timeout) - return AnthropicToOpenAI(forward.handler(), model=args.upstream_model) + return _TitoConversationDemux( + args.tito_url, + model=args.upstream_model, + timeout=args.upstream_timeout, + delete_on_close=bool(getattr(args, "tito_delete_on_evict", False)), + ) else: # Lazy: the translation client needs the `openai` extra. from .clients import AnthropicFromOpenAIClient diff --git a/plugins/abridge/tests/test_anthropic_transform.py b/plugins/abridge/tests/test_anthropic_transform.py index 1b7f05a..a8166a4 100644 --- a/plugins/abridge/tests/test_anthropic_transform.py +++ b/plugins/abridge/tests/test_anthropic_transform.py @@ -226,3 +226,28 @@ def test_tool_loop_turn_ordering_matches_openai_protocol(): assert roles[ai + 1] == "tool", f"tool result must follow tool_calls: {roles}" assert out[ai + 1]["tool_call_id"] == "tu_1" assert "file.txt" in out[ai + 1]["content"] + + +def test_thinking_only_assistant_history_is_preserved(): + """A thinking-only assistant turn (legal Anthropic shape — e.g. extended + thinking cut off at max_tokens) must not be dropped from the history: its + reasoning rides `reasoning_content`, and the surrounding user turns stay + non-adjacent.""" + from agentix.bridge.clients._anthropic_transforms import anthropic_messages_to_openai + + body = { + "model": "m", + "max_tokens": 16, + "messages": [ + {"role": "user", "content": "q"}, + {"role": "assistant", "content": [ + {"type": "thinking", "thinking": "hmm, hard", "signature": "sig"}, + ]}, + {"role": "user", "content": "continue"}, + ], + } + out = anthropic_messages_to_openai(body) + assert [m["role"] for m in out["messages"]] == ["user", "assistant", "user"] + assistant = out["messages"][1] + assert assistant["reasoning_content"] == "hmm, hard" + assert assistant["content"] is None diff --git a/plugins/abridge/tests/test_recorder.py b/plugins/abridge/tests/test_recorder.py index 3d79322..395c475 100644 --- a/plugins/abridge/tests/test_recorder.py +++ b/plugins/abridge/tests/test_recorder.py @@ -182,3 +182,45 @@ def test_recorder_opens_file_lazily(tmp_path) -> None: recorder = Recorder(_EchoClient(), out) recorder.abridge_routes() assert not out.exists() + + +@pytest.mark.asyncio +async def test_recorder_write_failure_is_log_and_serve(tmp_path, monkeypatch, caplog) -> None: + """Capture failures never fail the served call (matching the token + gateway's policy): with the record path unwritable the agent still gets + its response and the drop is logged.""" + import logging + + out = tmp_path / "run.jsonl" + recorder = Recorder(_EchoClient(), out) + routes = recorder.abridge_routes() + + def broken_open(*args, **kwargs): + raise OSError("disk full") + + monkeypatch.setattr(type(out), "open", broken_open) + with caplog.at_level(logging.ERROR, logger="agentix.bridge.recorder"): + resp = await routes["/v1/messages"](Request(path="/v1/messages", body={"msg": "hi"})) + assert json.loads(resp.body)["echo"] == "hi" # the call succeeded + assert any("row NOT persisted" in r.message for r in caplog.records) + + +@pytest.mark.asyncio +async def test_recorder_drops_rows_after_aclose(tmp_path, caplog) -> None: + """A straggler dispatch that outlives aclose() must not resurrect the + record file: the row is dropped (logged), the file stays closed, and no + orphan handle is created.""" + import logging + + out = tmp_path / "run.jsonl" + recorder = Recorder(_EchoClient(), out) + routes = recorder.abridge_routes() + await routes["/v1/messages"](Request(path="/v1/messages", body={"msg": "before"})) + await recorder.aclose() + + with caplog.at_level(logging.WARNING, logger="agentix.bridge.recorder"): + resp = await routes["/v1/messages"](Request(path="/v1/messages", body={"msg": "late"})) + assert json.loads(resp.body)["echo"] == "late" # still served + assert any("recorder is closed" in r.message for r in caplog.records) + assert [r["request"]["msg"] for r in _lines(out)] == ["before"] # no late row + assert recorder._file is not None and recorder._file.closed # noqa: SLF001 - not reopened diff --git a/plugins/abridge/tests/test_serve.py b/plugins/abridge/tests/test_serve.py index 24247cc..abfd955 100644 --- a/plugins/abridge/tests/test_serve.py +++ b/plugins/abridge/tests/test_serve.py @@ -49,17 +49,25 @@ def test_build_app_serves_handlers_and_health() -> None: def test_health_pins_the_translation_spec() -> None: - """`/_health` reports the SHA-256 of the Anthropic<->OpenAI transform - module, so downstream data contracts can pin the exact translation their - captured trajectories were produced under.""" + """`/_health` reports one SHA-256 over the source of every module that + shapes what an upstream/recording backend receives — the pure transforms + AND the client modules that rewrite the body around them (assistant + replay, forced stream=False, model override, upstream_params). A change + to any of them must move the pin.""" import hashlib from pathlib import Path import agentix.bridge.clients._anthropic_transforms as transforms health = TestClient(build_app(EchoClient())).get("/_health").json() - expected = hashlib.sha256(Path(transforms.__file__).read_bytes()).hexdigest() - assert health["translation_spec_sha"] == expected == transforms.TRANSLATION_SPEC_SHA + clients_dir = Path(transforms.__file__).parent + digest = hashlib.sha256() + for name in ("_anthropic_transforms.py", "anthropic_to_openai.py", "anthropic_from_openai.py"): + digest.update(name.encode()) + digest.update(b"\x00") + digest.update((clients_dir / name).read_bytes()) + digest.update(b"\x00") + assert health["translation_spec_sha"] == digest.hexdigest() == transforms.TRANSLATION_SPEC_SHA def test_handler_errors_become_wire_errors() -> None: diff --git a/plugins/abridge/tests/test_tito_composition.py b/plugins/abridge/tests/test_tito_composition.py index d996458..ed1a971 100644 --- a/plugins/abridge/tests/test_tito_composition.py +++ b/plugins/abridge/tests/test_tito_composition.py @@ -29,11 +29,22 @@ class _FakeTito(BaseHTTPRequestHandler): sessions: list[str] = [] chat_calls: list[dict[str, Any]] = [] # {"session", "body", "headers"} + deletes: list[str] = [] @classmethod def reset(cls) -> None: cls.sessions = [] cls.chat_calls = [] + cls.deletes = [] + + def do_DELETE(self) -> None: # noqa: N802 - http.server convention + parts = self.path.strip("/").split("/") + if len(parts) == 2 and parts[0] == "sessions": + _FakeTito.deletes.append(parts[1]) + self.send_response(204) + self.end_headers() + return + self._json(404, {"error": f"no route {self.path}"}) def do_POST(self) -> None: # noqa: N802 - http.server convention length = int(self.headers.get("Content-Length") or 0) @@ -133,18 +144,70 @@ def test_tito_composition_end_to_end(fake_tito: str) -> None: assert call["headers"]["x-request-id"] -def test_tito_composition_one_gateway_session_per_caller_key(fake_tito: str) -> None: +def test_tito_composition_one_gateway_session_per_conversation_per_caller(fake_tito: str) -> None: + """The demux unit is (caller key, logical conversation): repeated turns + of one conversation share a gateway session; another caller's identical + conversation gets its own.""" tc = TestClient(_tito_app(fake_tito)) assert tc.post("/v1/messages", json=_ANTHROPIC_BODY, headers={"x-api-key": "rollout-a"}).status_code == 200 assert tc.post("/v1/messages", json=_ANTHROPIC_BODY, headers={"x-api-key": "rollout-a"}).status_code == 200 assert tc.post("/v1/messages", json=_ANTHROPIC_BODY, headers={"x-api-key": "rollout-b"}).status_code == 200 - assert len(_FakeTito.sessions) == 2 # one gateway session per caller key + assert len(_FakeTito.sessions) == 2 assert [c["session"] for c in _FakeTito.chat_calls] == [ _FakeTito.sessions[0], _FakeTito.sessions[0], _FakeTito.sessions[1], ] +def test_multiplexed_conversations_on_one_key_get_separate_gateway_sessions(fake_tito: str) -> None: + """Adversarial-review regression: real Anthropic agents multiplex + unrelated conversations over ONE key (helper calls, Task subagents, + reruns). Each distinct (system, first user message) must land in its own + gateway session — not 400 against the first conversation's linear + history — while later turns of each conversation stay in theirs.""" + tc = TestClient(_tito_app(fake_tito)) + key = {"x-api-key": "rollout-1"} + + # Conversation A, turn 1. + assert tc.post("/v1/messages", json=_ANTHROPIC_BODY, headers=key).status_code == 200 + # Conversation B: same key, different opening (a subagent-style call). + conv_b = {**_ANTHROPIC_BODY, "messages": [{"role": "user", "content": "summarize the repo"}]} + assert tc.post("/v1/messages", json=conv_b, headers=key).status_code == 200 + # Conversation C: different system prompt (a helper-call profile). + conv_c = {**_ANTHROPIC_BODY, "system": "you detect topics"} + assert tc.post("/v1/messages", json=conv_c, headers=key).status_code == 200 + # Conversation A, turn 2 (history extended) routes back to A's session. + turn2 = { + **_ANTHROPIC_BODY, + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello from tito"}, + {"role": "user", "content": "and?"}, + ], + } + assert tc.post("/v1/messages", json=turn2, headers=key).status_code == 200 + + assert len(_FakeTito.sessions) == 3 + by_session = [c["session"] for c in _FakeTito.chat_calls] + assert by_session == [ + _FakeTito.sessions[0], _FakeTito.sessions[1], _FakeTito.sessions[2], _FakeTito.sessions[0], + ] + + +def test_tito_delete_on_evict_reaps_gateway_sessions(fake_tito: str) -> None: + """--tito-delete-on-evict: closing a caller session (here: app shutdown + via the TestClient context manager) DELETEs its gateway sessions. The + default keeps them alive for harvest.""" + with TestClient(_tito_app(fake_tito, "--tito-delete-on-evict")) as tc: + assert tc.post("/v1/messages", json=_ANTHROPIC_BODY, headers={"x-api-key": "r1"}).status_code == 200 + assert _FakeTito.deletes == [_FakeTito.sessions[0]] + + _FakeTito.reset() + with TestClient(_tito_app(fake_tito)) as tc: # default: no reap + assert tc.post("/v1/messages", json=_ANTHROPIC_BODY, headers={"x-api-key": "r1"}).status_code == 200 + assert _FakeTito.deletes == [] + + def test_tito_composition_with_record_dir_joins_rows_to_gateway_calls(fake_tito: str, tmp_path) -> None: """--record-dir in tito mode: the message-level Recorder row and the gateway's x-request-id share one id, and the row's session_id is the @@ -158,6 +221,10 @@ def test_tito_composition_with_record_dir_joins_rows_to_gateway_calls(fake_tito: (call,) = _FakeTito.chat_calls assert row["session_id"] == serve_session assert row["request_id"] == call["headers"]["x-request-id"] + # The gateway's OWN session id (== session_id in its token records) is + # read back from the transport, restoring the session-level join the + # caller-key hash cannot provide. + assert row["gateway_session_id"] == call["session"] == _FakeTito.sessions[0] assert row["path"] == "/v1/messages" assert row["request"] == _ANTHROPIC_BODY # the agent-side (Anthropic) shape assert row["response"]["body"]["content"] == [{"type": "text", "text": "hello from tito"}]