devcopilot 0.2.0__py3-none-any.whl

core/anthropic/stream_recovery.py

@@ -0,0 +1,346 @@
+"""Always-on recovery helpers for truncated provider streams."""
+
+from __future__ import annotations
+
+import json
+import time
+from collections.abc import Callable
+from copy import deepcopy
+from dataclasses import dataclass
+from typing import Any
+
+import httpx
+import jsonschema
+import openai
+from loguru import logger
+
+EARLY_TRANSPARENT_TOTAL_ATTEMPTS = 5
+EARLY_TRANSPARENT_MAX_RETRIES = EARLY_TRANSPARENT_TOTAL_ATTEMPTS - 1
+MIDSTREAM_RECOVERY_ATTEMPTS = 5
+EARLY_HOLDBACK_SECONDS = 0.75
+RECOVERY_BUFFER_MAX_BYTES = 65_536
+
+_RECOVERY_USER_PREFIX = (
+    "The previous provider stream was interrupted. Continue the assistant response "
+    "exactly where it stopped. Do not repeat text already written."
+)
+
+
+class TruncatedProviderStreamError(RuntimeError):
+    """Raised internally when an upstream stream ends without a terminal marker."""
+
+
+@dataclass(frozen=True, slots=True)
+class ToolSchema:
+    """Tool schema resolved from the original Anthropic request."""
+
+    name: str
+    input_schema: dict[str, Any]
+
+
+@dataclass(frozen=True, slots=True)
+class ToolRepair:
+    """Accepted append-only tool JSON repair."""
+
+    suffix: str
+    parsed_input: dict[str, Any]
+
+
+class RecoveryHoldbackBuffer:
+    """Briefly hold downstream SSE so early stream cutoffs can be retried invisibly."""
+
+    def __init__(
+        self,
+        *,
+        holdback_seconds: float = EARLY_HOLDBACK_SECONDS,
+        max_bytes: int = RECOVERY_BUFFER_MAX_BYTES,
+        now: Callable[[], float] | None = None,
+    ) -> None:
+        self._holdback_seconds = holdback_seconds
+        self._max_bytes = max_bytes
+        self._now = now or time.monotonic
+        self._events: list[str] = []
+        self._bytes = 0
+        self._started_at: float | None = None
+        self.committed = False
+
+    def push(self, event: str) -> list[str]:
+        """Buffer ``event`` until holdback expires or cap is reached."""
+        if self.committed:
+            return [event]
+        if self._started_at is None:
+            self._started_at = self._now()
+        self._events.append(event)
+        self._bytes += len(event.encode("utf-8", errors="replace"))
+        if (
+            self._bytes >= self._max_bytes
+            or self._now() - self._started_at >= self._holdback_seconds
+        ):
+            return self.flush()
+        return []
+
+    def flush(self) -> list[str]:
+        """Commit and return all held events."""
+        if self.committed:
+            return []
+        self.committed = True
+        events = self._events
+        self._events = []
+        self._bytes = 0
+        self._started_at = None
+        return events
+
+    def discard(self) -> None:
+        """Drop held events without committing them downstream."""
+        self._events = []
+        self._bytes = 0
+        self._started_at = None
+
+    @property
+    def has_buffered(self) -> bool:
+        return bool(self._events)
+
+
+def is_retryable_stream_error(exc: BaseException) -> bool:
+    """Return whether a provider stream error can be retried/recovered."""
+    if isinstance(exc, TruncatedProviderStreamError):
+        return True
+    if isinstance(exc, openai.AuthenticationError | openai.BadRequestError):
+        return False
+    if isinstance(exc, httpx.HTTPStatusError):
+        status = exc.response.status_code
+        return status == 429 or 500 <= status <= 599
+    if isinstance(exc, openai.RateLimitError):
+        return True
+    if isinstance(exc, openai.APIStatusError):
+        status = getattr(exc, "status_code", None)
+        return isinstance(status, int) and (status == 429 or 500 <= status <= 599)
+    return isinstance(
+        exc,
+        (
+            TimeoutError,
+            httpx.ReadTimeout,
+            httpx.ReadError,
+            httpx.RemoteProtocolError,
+            httpx.ConnectError,
+            httpx.NetworkError,
+            openai.APITimeoutError,
+            openai.APIConnectionError,
+        ),
+    )
+
+
+def tool_schemas_by_name(request: Any) -> dict[str, ToolSchema]:
+    """Return Anthropic tool input schemas keyed by tool name."""
+    schemas: dict[str, ToolSchema] = {}
+    tools = getattr(request, "tools", None)
+    if not tools:
+        return schemas
+
+    for tool in tools:
+        name = _tool_attr(tool, "name")
+        if not isinstance(name, str) or not name:
+            continue
+        schema = _tool_attr(tool, "input_schema")
+        if not isinstance(schema, dict):
+            schema = {"type": "object"}
+        schemas[name] = ToolSchema(name=name, input_schema=deepcopy(schema))
+    return schemas
+
+
+def validate_tool_input(
+    tool_name: str, parsed_input: dict[str, Any], schemas: dict[str, ToolSchema]
+) -> bool:
+    """Validate tool input against its JSON schema; unknown tools accept any object."""
+    tool_schema = schemas.get(tool_name)
+    if tool_schema is None:
+        return True
+    try:
+        validator_cls = jsonschema.validators.validator_for(tool_schema.input_schema)
+        validator_cls.check_schema(tool_schema.input_schema)
+        validator_cls(tool_schema.input_schema).validate(parsed_input)
+    except jsonschema.exceptions.SchemaError as exc:
+        logger.warning("Skipping invalid tool schema for {}: {}", tool_name, exc)
+        return True
+    except jsonschema.exceptions.ValidationError:
+        return False
+    return True
+
+
+def parse_complete_tool_input(
+    raw_json: str, tool_name: str, schemas: dict[str, ToolSchema]
+) -> dict[str, Any] | None:
+    """Return parsed input when raw JSON is complete and schema-valid."""
+    try:
+        parsed = json.loads(raw_json)
+    except json.JSONDecodeError:
+        return None
+    if not isinstance(parsed, dict):
+        return None
+    if not validate_tool_input(tool_name, parsed, schemas):
+        return None
+    return parsed
+
+
+def accept_tool_json_repair(
+    prefix: str,
+    candidate: str,
+    *,
+    tool_name: str,
+    schemas: dict[str, ToolSchema],
+) -> ToolRepair | None:
+    """Accept only append-only JSON repairs that make ``prefix`` valid."""
+    suffix_candidates = _repair_suffix_candidates(prefix, candidate)
+    for suffix in suffix_candidates:
+        combined = prefix + suffix
+        parsed = parse_complete_tool_input(combined, tool_name, schemas)
+        if parsed is not None:
+            return ToolRepair(suffix=suffix, parsed_input=parsed)
+    return None
+
+
+def continuation_suffix(existing: str, candidate: str) -> str | None:
+    """Return only the new suffix from a text/thinking continuation candidate."""
+    existing = existing or ""
+    candidate = candidate or ""
+    if not candidate:
+        return ""
+    if not existing:
+        return candidate
+    if candidate.startswith(existing):
+        return candidate[len(existing) :]
+
+    max_overlap = min(len(existing), len(candidate))
+    for size in range(max_overlap, 0, -1):
+        if existing.endswith(candidate[:size]):
+            return candidate[size:]
+
+    # Accept short standalone continuations, but reject full unrelated rewrites.
+    if len(candidate) < max(200, len(existing) // 2):
+        return candidate
+    return None
+
+
+def make_openai_text_recovery_body(
+    body: dict[str, Any], partial: str
+) -> dict[str, Any]:
+    """Build a text-only OpenAI-chat continuation request."""
+    recovery = deepcopy(body)
+    recovery.pop("tools", None)
+    recovery.pop("tool_choice", None)
+    recovery["stream"] = True
+    messages = _copied_messages(recovery)
+    if partial:
+        messages.append({"role": "assistant", "content": partial})
+    messages.append({"role": "user", "content": _RECOVERY_USER_PREFIX})
+    recovery["messages"] = messages
+    return recovery
+
+
+def make_openai_tool_repair_body(
+    body: dict[str, Any],
+    *,
+    tool_name: str,
+    prefix: str,
+    input_schema: dict[str, Any] | None,
+) -> dict[str, Any]:
+    """Build a text-only OpenAI-chat request asking for a JSON suffix."""
+    recovery = deepcopy(body)
+    recovery.pop("tools", None)
+    recovery.pop("tool_choice", None)
+    recovery["stream"] = True
+    messages = _copied_messages(recovery)
+    messages.append(
+        {
+            "role": "user",
+            "content": _tool_repair_prompt(
+                tool_name=tool_name, prefix=prefix, input_schema=input_schema
+            ),
+        }
+    )
+    recovery["messages"] = messages
+    return recovery
+
+
+def make_native_text_recovery_body(
+    body: dict[str, Any], partial: str
+) -> dict[str, Any]:
+    """Build a text-only native Anthropic continuation request."""
+    recovery = deepcopy(body)
+    recovery.pop("tools", None)
+    recovery.pop("tool_choice", None)
+    recovery["stream"] = True
+    messages = _copied_messages(recovery)
+    if partial:
+        messages.append({"role": "assistant", "content": partial})
+    messages.append({"role": "user", "content": _RECOVERY_USER_PREFIX})
+    recovery["messages"] = messages
+    return recovery
+
+
+def make_native_tool_repair_body(
+    body: dict[str, Any],
+    *,
+    tool_name: str,
+    prefix: str,
+    input_schema: dict[str, Any] | None,
+) -> dict[str, Any]:
+    """Build a text-only native Anthropic request asking for a JSON suffix."""
+    recovery = deepcopy(body)
+    recovery.pop("tools", None)
+    recovery.pop("tool_choice", None)
+    recovery["stream"] = True
+    messages = _copied_messages(recovery)
+    messages.append(
+        {
+            "role": "user",
+            "content": _tool_repair_prompt(
+                tool_name=tool_name, prefix=prefix, input_schema=input_schema
+            ),
+        }
+    )
+    recovery["messages"] = messages
+    return recovery
+
+
+def _tool_attr(tool: Any, attr: str) -> Any:
+    if isinstance(tool, dict):
+        return tool.get(attr)
+    return getattr(tool, attr, None)
+
+
+def _copied_messages(body: dict[str, Any]) -> list[Any]:
+    messages = body.get("messages")
+    return deepcopy(messages) if isinstance(messages, list) else []
+
+
+def _repair_suffix_candidates(prefix: str, candidate: str) -> list[str]:
+    raw = candidate.strip()
+    if not raw:
+        return []
+    candidates: list[str] = []
+    if raw.startswith("```"):
+        lines = raw.splitlines()
+        if lines and lines[0].startswith("```"):
+            lines = lines[1:]
+        if lines and lines[-1].strip() == "```":
+            lines = lines[:-1]
+        raw = "\n".join(lines).strip()
+    candidates.append(raw)
+    if raw.startswith(prefix):
+        candidates.append(raw[len(prefix) :])
+    return list(dict.fromkeys(candidates))
+
+
+def _tool_repair_prompt(
+    *, tool_name: str, prefix: str, input_schema: dict[str, Any] | None
+) -> str:
+    schema_text = json.dumps(input_schema or {"type": "object"}, separators=(",", ":"))
+    return (
+        "A streamed tool call was interrupted while writing JSON arguments.\n"
+        f"Tool name: {tool_name}\n"
+        f"JSON schema: {schema_text}\n"
+        f"Already emitted JSON prefix: {prefix}\n\n"
+        "Return only the exact missing JSON suffix needed to complete the same object. "
+        "Do not repeat the prefix. Do not include markdown or explanation."
+    )

core/anthropic/stream_recovery_session.py

@@ -0,0 +1,133 @@
+"""Shared stream recovery policy for provider transports."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import StrEnum
+
+from core.anthropic.stream_recovery import (
+    EARLY_TRANSPARENT_MAX_RETRIES,
+    EARLY_TRANSPARENT_TOTAL_ATTEMPTS,
+    RecoveryHoldbackBuffer,
+    is_retryable_stream_error,
+)
+from core.trace import trace_event
+
+
+class StreamFailureAction(StrEnum):
+    """Transport action selected after a provider stream failure."""
+
+    EARLY_RETRY = "early_retry"
+    MIDSTREAM_RECOVERY = "midstream_recovery"
+    FINAL_ERROR = "final_error"
+
+
+@dataclass(frozen=True, slots=True)
+class StreamFailureDecision:
+    """Failure-state snapshot for the current provider stream transition."""
+
+    action: StreamFailureAction
+    retryable: bool
+    committed: bool
+    has_buffered: bool
+    early_retry_attempt: int | None = None
+
+
+class StreamRecoverySession:
+    """Own holdback and retry policy shared by provider stream transports."""
+
+    def __init__(self, *, provider_name: str, request_id: str | None) -> None:
+        self._provider_name = provider_name
+        self._request_id = request_id
+        self._holdback = RecoveryHoldbackBuffer()
+        self._early_retries = 0
+
+    @property
+    def committed(self) -> bool:
+        return self._holdback.committed
+
+    @property
+    def has_buffered(self) -> bool:
+        return self._holdback.has_buffered
+
+    @property
+    def early_retries(self) -> int:
+        return self._early_retries
+
+    def push(self, event: str) -> list[str]:
+        """Buffer one downstream event through the early retry holdback."""
+        return self._holdback.push(event)
+
+    def flush(self) -> list[str]:
+        """Commit and return held events."""
+        return self._holdback.flush()
+
+    def flush_uncommitted(self, decision: StreamFailureDecision) -> list[str]:
+        """Commit held events when the decision snapshot is still uncommitted."""
+        if decision.committed:
+            return []
+        return self.flush()
+
+    def discard(self) -> None:
+        """Drop held events without committing them."""
+        self._holdback.discard()
+
+    def advance_failure(
+        self,
+        error: BaseException,
+        *,
+        stream_opened: bool,
+        generated_output: bool,
+        complete_tool_salvageable: bool,
+    ) -> StreamFailureDecision:
+        """Consume a stream failure and apply shared recovery state changes."""
+        committed = self.committed
+        has_buffered = self.has_buffered
+        retryable = is_retryable_stream_error(error)
+
+        if (
+            not committed
+            and stream_opened
+            and retryable
+            and not complete_tool_salvageable
+            and self._early_retries < EARLY_TRANSPARENT_MAX_RETRIES
+        ):
+            self._early_retries += 1
+            attempt = self._early_retries
+            self._reset_holdback()
+            trace_event(
+                stage="provider",
+                event="provider.recovery.early_retry",
+                source="provider",
+                provider=self._provider_name,
+                request_id=self._request_id,
+                attempt=attempt,
+                max_attempts=EARLY_TRANSPARENT_TOTAL_ATTEMPTS,
+                exc_type=type(error).__name__,
+            )
+            return StreamFailureDecision(
+                action=StreamFailureAction.EARLY_RETRY,
+                retryable=retryable,
+                committed=committed,
+                has_buffered=has_buffered,
+                early_retry_attempt=attempt,
+            )
+
+        if generated_output and retryable:
+            return StreamFailureDecision(
+                action=StreamFailureAction.MIDSTREAM_RECOVERY,
+                retryable=retryable,
+                committed=committed,
+                has_buffered=has_buffered,
+            )
+
+        return StreamFailureDecision(
+            action=StreamFailureAction.FINAL_ERROR,
+            retryable=retryable,
+            committed=committed,
+            has_buffered=has_buffered,
+        )
+
+    def _reset_holdback(self) -> None:
+        self._holdback.discard()
+        self._holdback = RecoveryHoldbackBuffer()

core/anthropic/thinking.py

@@ -0,0 +1,140 @@
+"""Streaming parser for provider-emitted thinking tags."""
+
+from collections.abc import Iterator
+from dataclasses import dataclass
+from enum import Enum
+
+
+class ContentType(Enum):
+    """Type of content chunk."""
+
+    TEXT = "text"
+    THINKING = "thinking"
+
+
+@dataclass
+class ContentChunk:
+    """A chunk of parsed content."""
+
+    type: ContentType
+    content: str
+
+
+class ThinkTagParser:
+    """
+    Streaming parser for ``<think>...</think>`` tags.
+
+    Handles partial tags at chunk boundaries by buffering.
+    """
+
+    OPEN_TAG = "<think>"
+    CLOSE_TAG = "</think>"
+
+    def __init__(self):
+        self._buffer: str = ""
+        self._in_think_tag: bool = False
+
+    @property
+    def in_think_mode(self) -> bool:
+        """Whether currently inside a think tag."""
+        return self._in_think_tag
+
+    def feed(self, content: str) -> Iterator[ContentChunk]:
+        """Feed content and yield parsed chunks."""
+        self._buffer += content
+
+        while self._buffer:
+            prev_len = len(self._buffer)
+            if not self._in_think_tag:
+                chunk = self._parse_outside_think()
+            else:
+                chunk = self._parse_inside_think()
+
+            if chunk:
+                yield chunk
+            elif len(self._buffer) == prev_len:
+                break
+
+    def _parse_outside_think(self) -> ContentChunk | None:
+        """Parse content outside think tags."""
+        think_start = self._buffer.find(self.OPEN_TAG)
+        orphan_close = self._buffer.find(self.CLOSE_TAG)
+
+        if orphan_close != -1 and (think_start == -1 or orphan_close < think_start):
+            pre_orphan = self._buffer[:orphan_close]
+            self._buffer = self._buffer[orphan_close + len(self.CLOSE_TAG) :]
+            if pre_orphan:
+                return ContentChunk(ContentType.TEXT, pre_orphan)
+            return None
+
+        if think_start == -1:
+            last_bracket = self._buffer.rfind("<")
+            if last_bracket != -1:
+                potential_tag = self._buffer[last_bracket:]
+                tag_len = len(potential_tag)
+                if (
+                    tag_len < len(self.OPEN_TAG)
+                    and self.OPEN_TAG.startswith(potential_tag)
+                ) or (
+                    tag_len < len(self.CLOSE_TAG)
+                    and self.CLOSE_TAG.startswith(potential_tag)
+                ):
+                    emit = self._buffer[:last_bracket]
+                    self._buffer = self._buffer[last_bracket:]
+                    if emit:
+                        return ContentChunk(ContentType.TEXT, emit)
+                    return None
+
+            emit = self._buffer
+            self._buffer = ""
+            if emit:
+                return ContentChunk(ContentType.TEXT, emit)
+            return None
+
+        pre_think = self._buffer[:think_start]
+        self._buffer = self._buffer[think_start + len(self.OPEN_TAG) :]
+        self._in_think_tag = True
+        if pre_think:
+            return ContentChunk(ContentType.TEXT, pre_think)
+        return None
+
+    def _parse_inside_think(self) -> ContentChunk | None:
+        """Parse content inside think tags."""
+        think_end = self._buffer.find(self.CLOSE_TAG)
+
+        if think_end == -1:
+            last_bracket = self._buffer.rfind("<")
+            if last_bracket != -1 and len(self._buffer) - last_bracket < len(
+                self.CLOSE_TAG
+            ):
+                potential_tag = self._buffer[last_bracket:]
+                if self.CLOSE_TAG.startswith(potential_tag):
+                    emit = self._buffer[:last_bracket]
+                    self._buffer = self._buffer[last_bracket:]
+                    if emit:
+                        return ContentChunk(ContentType.THINKING, emit)
+                    return None
+
+            emit = self._buffer
+            self._buffer = ""
+            if emit:
+                return ContentChunk(ContentType.THINKING, emit)
+            return None
+
+        thinking_content = self._buffer[:think_end]
+        self._buffer = self._buffer[think_end + len(self.CLOSE_TAG) :]
+        self._in_think_tag = False
+        if thinking_content:
+            return ContentChunk(ContentType.THINKING, thinking_content)
+        return None
+
+    def flush(self) -> ContentChunk | None:
+        """Flush any remaining buffered content."""
+        if self._buffer:
+            chunk_type = (
+                ContentType.THINKING if self._in_think_tag else ContentType.TEXT
+            )
+            content = self._buffer
+            self._buffer = ""
+            return ContentChunk(chunk_type, content)
+        return None