feat(session): serve Anthropic Messages through SessionCore

miles/rollout/session/anthropic_adapter.py

@@ -0,0 +1,130 @@
+"""Anthropic protocol helpers for the session HTTP adapter."""
+
+import json
+
+from pydantic import ValidationError
+from sglang.srt.entrypoints.anthropic import utils as anthropic_utils
+from sglang.srt.entrypoints.anthropic.protocol import AnthropicMessagesRequest, is_server_tool
+from starlette.responses import Response
+
+from miles.rollout.session.core import JSON_MEDIA_TYPE, _render_json
+
+# Preserve end-to-end error metadata; drop headers tied to the replaced body.
+_ANTHROPIC_ERROR_HEADER_ALLOWLIST = ("www-authenticate", "retry-after", "x-request-id")
+_ANTHROPIC_ERROR_HEADER_PREFIXES = ("x-ratelimit-", "anthropic-ratelimit-")
+
+
+def _anthropic_wire_json(model) -> bytes:
+    return _render_json(model.model_dump(mode="json", exclude_none=True, by_alias=True))
+
+
+def _anthropic_error_response(status_code: int, body: bytes, headers: dict | None = None) -> Response:
+    envelope = anthropic_utils.to_anthropic_error(status_code, body)
+    kept_headers = {
+        k: v
+        for k, v in (headers or {}).items()
+        if k.lower() in _ANTHROPIC_ERROR_HEADER_ALLOWLIST or k.lower().startswith(_ANTHROPIC_ERROR_HEADER_PREFIXES)
+    }
+    return Response(
+        content=_anthropic_wire_json(envelope),
+        status_code=status_code,
+        headers=kept_headers,
+        media_type=JSON_MEDIA_TYPE,
+    )
+
+
+def _anthropic_sse_body(events) -> bytes:
+    return b"".join(
+        f"event: {event.type}\ndata: ".encode() + _anthropic_wire_json(event) + b"\n\n" for event in events
+    )
+
+
+def _parse_anthropic_request(body: bytes) -> AnthropicMessagesRequest:
+    try:
+        payload = json.loads(body)
+    except (json.JSONDecodeError, UnicodeDecodeError, ValueError) as exc:
+        raise ValueError(f"invalid JSON body: {exc}") from exc
+    try:
+        return AnthropicMessagesRequest.model_validate(payload)
+    except ValidationError as exc:
+        raise ValueError(str(exc)) from exc
+
+
+def _validate_anthropic_content_block(block, *, allow_thinking: bool = False) -> None:
+    if block.type == "thinking":
+        if allow_thinking:
+            return
+        raise ValueError("thinking content blocks are only supported in assistant history")
+    if block.type == "redacted_thinking":
+        raise ValueError("redacted_thinking content blocks are not supported by this endpoint")
+    if block.type == "image":
+        raise ValueError("image content blocks are not enabled for this deployment")
+    if block.type == "tool_reference":
+        raise ValueError("tool_reference content blocks are not enabled for this deployment")
+    if block.type == "search_result":
+        raise ValueError("search_result content blocks are not enabled for this deployment")
+    if block.type == "tool_result" and block.is_error is True:
+        raise ValueError("tool_result is_error=true is not supported by this endpoint")
+
+
+def _validate_anthropic_features(request: AnthropicMessagesRequest) -> None:
+    if request.thinking is not None:
+        raise ValueError("thinking is not supported by this endpoint")
+    if request.output_config is not None:
+        raise ValueError("output_config is not enabled for this deployment")
+    if request.betas:
+        raise ValueError("betas is not enabled for this deployment")
+    if request.tools:
+        for tool in request.tools:
+            if is_server_tool(tool):
+                raise ValueError(f"server tool {tool.name!r} (type={tool.type!r}) is not enabled for this deployment")
+
+    if request.system is not None and not isinstance(request.system, str):
+        for block in request.system:
+            _validate_anthropic_content_block(block)
+    for message in request.messages:
+        if isinstance(message.content, str):
+            continue
+        for block in message.content:
+            _validate_anthropic_content_block(block, allow_thinking=message.role == "assistant")
+            if block.type == "tool_result" and isinstance(block.content, list):
+                for nested_block in block.content:
+                    _validate_anthropic_content_block(nested_block)
+
+
+def _strip_anthropic_reasoning_history(
+    anthropic_request: AnthropicMessagesRequest,
+) -> tuple[AnthropicMessagesRequest, list[str | None]]:
+    """Return a conversion copy plus canonical assistant reasoning history."""
+    conversion_messages = []
+    reasoning_history: list[str | None] = []
+    for message in anthropic_request.messages:
+        if message.role != "assistant":
+            conversion_messages.append(message)
+            continue
+        if isinstance(message.content, str):
+            conversion_messages.append(message)
+            reasoning_history.append(None)
+            continue
+        thinking_blocks = [block for block in message.content if block.type == "thinking"]
+        thinking_parts = [block.thinking for block in thinking_blocks if block.thinking]
+        reasoning_history.append("\n".join(thinking_parts) or None)
+        if thinking_blocks:
+            message = message.model_copy(
+                update={"content": [block for block in message.content if block.type != "thinking"]}
+            )
+        conversion_messages.append(message)
+    return anthropic_request.model_copy(update={"messages": conversion_messages}), reasoning_history
+
+
+def _restore_anthropic_reasoning_history(openai_body: dict, reasoning_history: list[str | None]) -> None:
+    """Map replayed assistant thinking blocks back to canonical reasoning content."""
+    assistants = [message for message in openai_body["messages"] if message["role"] == "assistant"]
+    if len(assistants) != len(reasoning_history):
+        raise ValueError(
+            f"assistant history count changed during Anthropic conversion: "
+            f"{len(reasoning_history)} before, {len(assistants)} after"
+        )
+    for message, reasoning_content in zip(assistants, reasoning_history, strict=True):
+        if reasoning_content is not None:
+            message["reasoning_content"] = reasoning_content

miles/rollout/session/sessions.py

@@ -9,8 +9,31 @@
 
 from fastapi import Request
 from fastapi.responses import JSONResponse
-
-from miles.rollout.session.core import SessionCore
+from sglang.srt.entrypoints.anthropic import utils as anthropic_utils
+from sglang.srt.entrypoints.anthropic.serving import convert_response, convert_to_chat_completion_request
+from sglang.srt.entrypoints.openai.protocol import ChatCompletionResponse
+from sglang.srt.parser.template_detection import detect_inline_system_support
+from starlette.responses import Response
+
+from miles.rollout.session.anthropic_adapter import (
+    _ANTHROPIC_ERROR_HEADER_ALLOWLIST as _ANTHROPIC_ERROR_HEADER_ALLOWLIST,
+)
+from miles.rollout.session.anthropic_adapter import (
+    _ANTHROPIC_ERROR_HEADER_PREFIXES as _ANTHROPIC_ERROR_HEADER_PREFIXES,
+)
+from miles.rollout.session.anthropic_adapter import (
+    _anthropic_error_response,
+    _anthropic_sse_body,
+    _anthropic_wire_json,
+    _parse_anthropic_request,
+    _restore_anthropic_reasoning_history,
+    _strip_anthropic_reasoning_history,
+)
+from miles.rollout.session.anthropic_adapter import (
+    _validate_anthropic_content_block as _validate_anthropic_content_block,
+)
+from miles.rollout.session.anthropic_adapter import _validate_anthropic_features
+from miles.rollout.session.core import JSON_MEDIA_TYPE, SessionCore, _render_json
 from miles.rollout.session.errors import SessionError
 from miles.rollout.session.linear_trajectory import SessionRegistry
 from miles.utils.chat_template_utils import get_tito_tokenizer
@@ -44,6 +67,7 @@ def setup_session_routes(app, backend, args, *, use_addition_r3: bool = False):
         tokenizer_type=getattr(args, "tito_model", "default"),
         chat_template_kwargs=getattr(args, "apply_chat_template_kwargs", None),
     )
+    merge_inline_system = not detect_inline_system_support(getattr(tokenizer, "chat_template", None))
 
     use_v2 = getattr(args, "use_session_server", None) == "v2"
     if use_v2:
@@ -66,7 +90,10 @@ async def session_message_matcher_error_handler(request: Request, exc: SessionMe
 
     @app.get("/health")
     async def health():
-        return await core.health()
+        response = await core.health()
+        body = json.loads(response.body)
+        body["anthropic_intermediate_system_supported"] = not merge_inline_system
+        return Response(content=_render_json(body), status_code=response.status_code, media_type=JSON_MEDIA_TYPE)
 
     @app.post("/sessions")
     async def create_session():
@@ -91,6 +118,78 @@ async def chat_completions(request: Request, session_id: str):
             body=body,
         )
 
+    # Keep before session_proxy: Starlette's first match must not bypass session/TITO.
+    @app.post("/sessions/{session_id}/v1/messages")
+    async def anthropic_messages(request: Request, session_id: str):
+        """Serve Anthropic Messages through the OpenAI session path."""
+        body = await request.body()
+        try:
+            anthropic_request = _parse_anthropic_request(body)
+            _validate_anthropic_features(anthropic_request)
+            try:
+                conversion_request, reasoning_history = _strip_anthropic_reasoning_history(anthropic_request)
+                openai_request = convert_to_chat_completion_request(
+                    conversion_request, merge_inline_system=merge_inline_system
+                )
+            except Exception as exc:
+                logger.exception("Error converting Anthropic request: %s", exc)
+                raise ValueError(str(exc)) from exc
+            # Core is non-streaming; build fake SSE from its complete response below.
+            openai_request.stream = False
+            openai_request.stream_options = None
+            # Omit defaults so equivalent Anthropic and OpenAI inputs produce the same canonical record.
+            openai_body_dict = openai_request.model_dump(
+                mode="json", exclude_none=True, exclude_unset=True, by_alias=True
+            )
+            _restore_anthropic_reasoning_history(openai_body_dict, reasoning_history)
+            openai_body = _render_json(openai_body_dict)
+        except ValueError as exc:
+            # Parsing and JSON encoding failures are invalid Anthropic requests.
+            return _anthropic_error_response(400, _render_json({"error": str(exc)}))
+
+        anthropic_stream = bool(anthropic_request.stream)
+
+        try:
+            core_response = await core.chat_completions(
+                session_id,
+                method=request.method,
+                query=request.url.query,
+                headers=dict(request.headers),
+                body=openai_body,
+            )
+        except SessionError as exc:
+            return _anthropic_error_response(exc.status_code, _render_json({"error": str(exc)}))
+        except Exception:
+            # Preserve Anthropic error framing; cancellation still propagates.
+            logger.exception("Anthropic chat processing failed for session %s", session_id)
+            return _anthropic_error_response(500, b"")
+
+        if core_response.status_code != 200:
+            return _anthropic_error_response(
+                core_response.status_code, core_response.body, dict(core_response.headers)
+            )
+
+        try:
+            openai_response = ChatCompletionResponse.model_validate_json(core_response.body)
+            if anthropic_stream:
+                events = anthropic_utils.to_anthropic_fake_sse_events(
+                    openai_response,
+                    model=anthropic_request.model,
+                    id_factory=lambda: openai_response.id,
+                )
+                return Response(
+                    content=_anthropic_sse_body(events),
+                    status_code=200,
+                    headers={"cache-control": "no-cache", "x-accel-buffering": "no"},
+                    media_type="text/event-stream",
+                )
+            envelope = convert_response(openai_response).model_copy(update={"id": openai_response.id})
+            return Response(content=_anthropic_wire_json(envelope), status_code=200, media_type=JSON_MEDIA_TYPE)
+        except Exception:
+            # Post-commit failures keep the record and return JSON 500, never partial SSE.
+            logger.exception("Anthropic response conversion failed for session %s", session_id)
+            return _anthropic_error_response(500, b"")
+
     @app.post("/sessions/{session_id}/samples")
     async def collect_samples(request: Request, session_id: str):
         # Starlette matches routes in registration order; keep this before session_proxy.

miles/utils/test_utils/mock_sglang_server.py

@@ -211,12 +211,23 @@ def _compute_chat_completions_response(self, payload: dict) -> dict:
         if prompt_ids is not None:
             choice["prompt_token_ids"] = prompt_ids
 
+        # Real SGLang chat responses always carry ``usage``; clients that
+        # validate the full ChatCompletionResponse schema require it.
+        if payload.get("input_ids") is not None:
+            prompt_token_count = len(payload["input_ids"])
+        else:
+            prompt_token_count = len(self.tokenizer.encode(prompt_str, add_special_tokens=False))
         return {
             "id": f"chatcmpl-{uuid.uuid4().hex[:8]}",
             "object": "chat.completion",
             "created": int(time.time()),
             "model": "mock-model",
             "choices": [choice],
+            "usage": {
+                "prompt_tokens": prompt_token_count,
+                "completion_tokens": len(output_ids),
+                "total_tokens": prompt_token_count + len(output_ids),
+            },
         }