fix: bound model stream aggregation and close failed responses

artemis/llm/argument_fingerprint.py

@@ -16,12 +16,38 @@ class ArgumentFingerprint:
         self.encoding_valid = True
         self.digest = hashlib.sha256()
         self.prefixes: list[dict] = []
+        self.whitespace_chars = 0
+        self.whitespace_run = 0
+        self.max_whitespace_run = 0
+        self.min_delta_chars: int | None = None
+        self.max_delta_chars = 0
+        self.same_delta_run = 0
+        self.max_same_delta_run = 0
+        self.previous_delta_digest: bytes | None = None
 
     def observe(self, arguments: str) -> None:
         self.characters += len(arguments)
         if not arguments:
             return
         self.deltas += 1
+        length = len(arguments)
+        self.min_delta_chars = min(self.min_delta_chars or length, length)
+        self.max_delta_chars = max(self.max_delta_chars, length)
+        for character in arguments:
+            if character in " \t\r\n":
+                self.whitespace_chars += 1
+                self.whitespace_run += 1
+                self.max_whitespace_run = max(self.max_whitespace_run, self.whitespace_run)
+            else:
+                self.whitespace_run = 0
+        # One digest, never the previous argument body. Surrogate handling here
+        # compares SDK strings only; wire UTF-8 fingerprints remain invalid below.
+        delta_digest = hashlib.sha256(arguments.encode("utf-8", "surrogatepass")).digest()
+        self.same_delta_run = (
+            self.same_delta_run + 1 if delta_digest == self.previous_delta_digest else 1
+        )
+        self.max_same_delta_run = max(self.max_same_delta_run, self.same_delta_run)
+        self.previous_delta_digest = delta_digest
         if not self.encoding_valid:
             return
         try:
@@ -50,4 +76,11 @@ class ArgumentFingerprint:
             "argument_deltas": self.deltas,
             "argument_prefixes": [dict(prefix) for prefix in self.prefixes],
             "argument_encoding_valid": self.encoding_valid,
+            "argument_json_whitespace_chars": self.whitespace_chars,
+            "argument_json_whitespace_run": self.whitespace_run,
+            "argument_json_whitespace_run_max": self.max_whitespace_run,
+            "argument_delta_chars_min": self.min_delta_chars,
+            "argument_delta_chars_max": self.max_delta_chars,
+            "argument_same_delta_run": self.same_delta_run,
+            "argument_same_delta_run_max": self.max_same_delta_run,
         }

artemis/llm/chunk_accumulator.py

@@ -0,0 +1,37 @@
+"""Amortize LangChain's partial-JSON parsing without buffering live UI output."""
+
+from langchain_core.messages import AIMessageChunk
+from langchain_core.messages.ai import add_ai_message_chunks
+
+CHUNK_BATCH_SIZE = 128
+
+
+class ChunkAccumulator:
+    def __init__(self):
+        self.response = None
+        self.pending = []
+
+    @property
+    def has_output(self):
+        return self.response is not None or bool(self.pending)
+
+    @property
+    def ready(self):
+        return len(self.pending) >= CHUNK_BATCH_SIZE
+
+    def append(self, chunk):
+        self.pending.append(chunk)
+
+    def flush(self):
+        chunks = ([self.response] if self.response is not None else []) + self.pending
+        if not chunks:
+            return None
+        if all(isinstance(chunk, AIMessageChunk) for chunk in chunks):
+            self.response = add_ai_message_chunks(chunks[0], *chunks[1:])
+        else:
+            result = chunks[0]
+            for chunk in chunks[1:]:
+                result = result + chunk
+            self.response = result
+        self.pending.clear()
+        return self.response

artemis/llm/diagnostics.py

@@ -1,6 +1,7 @@
 """Small, credential-free request snapshots in the existing task trace store."""
 
 import base64
+import asyncio
 from contextvars import ContextVar
 from copy import deepcopy
 import io
@@ -11,6 +12,8 @@ import time
 from uuid import uuid4
 
 from PIL import Image
+import httpx
+import httpx2
 
 from artemis.data_engine.context_vars import CURRENT_NODE_NAME, CURRENT_TRACE_ID
 from artemis.llm.argument_fingerprint import ArgumentFingerprint, MAX_TRACKED_TOOLS
@@ -20,6 +23,8 @@ from artemis.llm.usage import usage_from_message
 from artemis.llm.stream_guard import ToolArgumentOverflow
 from artemis.llm.response_guard import ResponseIntegrityError
 from artemis.runtime.cancellation import CancellationReason, cancellation_reason
+from artemis.llm.responses_diagnostics import ObservedResponseStream
+from artemis.llm.tool_constraints import request_constraints
 
 logger = logging.getLogger(__name__)
 ACTIVE_ATTEMPT: ContextVar["RequestAttempt | None"] = ContextVar("llm_attempt", default=None)
@@ -109,6 +114,7 @@ class RequestAttempt:
         self.last_publish = self.clock
         self.tools: dict[int | str, dict] = {}
         self.tool_fingerprints: dict[int | str, ArgumentFingerprint] = {}
+        self.wire_streams: list[ObservedResponseStream] = []
         self.review = ACTIVE_FINAL_REVIEW.get()
         if self.review:
             self.review.request_started(request_id, attempt)
@@ -138,6 +144,8 @@ class RequestAttempt:
         self.publish()
 
     def publish(self):
+        if self.wire_streams:
+            self.data["wire_responses"] = [s.observation.snapshot() for s in self.wire_streams]
         self.data["sequence"] += 1
         self.data["updated_at"] = time.time()
         self.data["elapsed_seconds"] = round(time.monotonic() - self.clock, 3)
@@ -168,6 +176,35 @@ class RequestAttempt:
         except (OSError, RuntimeError, ValueError, TypeError, AttributeError):
             logger.warning("Could not persist model request diagnostics", exc_info=True)
 
+    def publish_if_due(self):
+        if time.monotonic() - self.last_publish >= 5:
+            self.publish()
+
+    def stream_cancelled(self):
+        self.data.setdefault("cancel_observed_at", time.time())
+        self.data.setdefault("cancel_observed_phase", self.data["phase"])
+        self.data.setdefault("cancel_phase", self.data["phase"])
+
+    def stream_work(self, stage: str, elapsed: float):
+        stats = self.data.setdefault("stream_work", {}).setdefault(
+            stage, {"calls": 0, "total_seconds": 0.0, "max_seconds": 0.0}
+        )
+        stats["calls"] += 1
+        stats["total_seconds"] += elapsed
+        stats["max_seconds"] = max(stats["max_seconds"], elapsed)
+
+    async def close_streams(self):
+        try:
+            async with asyncio.timeout(1):
+                for stream in self.wire_streams:
+                    await stream.aclose()
+        except (TimeoutError, OSError, RuntimeError, httpx.HTTPError, httpx2.HTTPError) as error:
+            self.data["stream_close_error_type"] = type(error).__name__
+            logger.warning("Model stream cleanup failed: %s", type(error).__name__)
+        finally:
+            if self.wire_streams:
+                self.publish()
+
     def headers(self, response):
         self.data.update(
             headers_at=time.time(),
@@ -207,6 +244,9 @@ class RequestAttempt:
             arguments = tool.get("args")
             if isinstance(arguments, str):
                 self.data["tool_argument_chars"] += len(arguments)
+                if arguments:
+                    self.data.setdefault("first_tool_arguments_at", time.time())
+                    self.data["last_tool_arguments_at"] = time.time()
             index = tool.get("index")
             call_id = _identifier(tool.get("id"))
             key = index if index is not None else call_id or "unknown"
@@ -297,6 +337,7 @@ class RequestAttempt:
         self.publish()
 
     def cancel(self, error: BaseException | None = None):
+        self.stream_cancelled()
         reason = cancellation_reason(error)
         if self.review and self.review.expired and reason is CancellationReason.UNKNOWN:
             self.data["timeout_scope"] = "final_review"
@@ -321,6 +362,10 @@ async def observe_request_headers(request):
     attempt = ACTIVE_ATTEMPT.get()
     if attempt is None:
         return
+    try:
+        attempt.data["wire_request_constraints"] = request_constraints(request.content)
+    except (httpx.RequestNotRead, httpx2.RequestNotRead):
+        attempt.data["wire_request_constraints"] = {"observation": "streaming_body_unavailable"}
     request.headers.setdefault("X-Client-Request-Id", attempt.data["request_id"])
     request.headers.setdefault("X-Client-Attempt-Id", attempt.data["attempt_id"])
     routing = {}
@@ -342,7 +387,25 @@ async def observe_request_headers(request):
 
 
 async def observe_response_headers(response):
-    """httpx hook: only allowlisted headers, never body/auth/request headers."""
+    """Correlate allowlisted headers and passively observe bounded SSE frames."""
     attempt = ACTIVE_ATTEMPT.get()
     if attempt is not None:
         attempt.headers(response)
+        content_type = response.headers.get("content-type", "").split(";", 1)[0].lower()
+        if (
+            response.status_code == 200
+            and content_type == "text/event-stream"
+            and isinstance(response.stream, (httpx.AsyncByteStream, httpx2.AsyncByteStream))
+            and not isinstance(response.stream, ObservedResponseStream)
+        ):
+            if len(attempt.wire_streams) < 4:
+                stream = ObservedResponseStream(
+                    response.stream,
+                    attempt,
+                    encoded=response.headers.get("content-encoding", "identity").lower().strip()
+                    not in {"", "identity"},
+                )
+                response.stream = stream
+                attempt.wire_streams.append(stream)
+            else:
+                attempt.data["wire_streams_truncated"] = True

artemis/llm/responses_diagnostics.py

@@ -0,0 +1,301 @@
+"""Bounded observations at the client's raw SSE read boundary; never prefetch."""
+
+import asyncio
+import json
+import logging
+import re
+import time
+
+import httpx
+import httpx2
+
+from artemis.llm.argument_fingerprint import ArgumentFingerprint, MAX_TRACKED_TOOLS
+from artemis.llm.tool_constraints import safe_label, tool_constraints
+
+MAX_FRAME_BYTES = 256 * 1024
+MAX_RESPONSES = 16
+logger = logging.getLogger(__name__)
+
+
+def identifier(value):
+    return (
+        value
+        if isinstance(value, str) and re.fullmatch(r"[A-Za-z0-9_.:/-]{1,256}", value)
+        else None
+    )
+
+
+class ResponsesObservation:
+    def __init__(self, *, observe_sse=True):
+        self.observe_sse = observe_sse
+        self.data = {
+            "observation_boundary": "httpx_response_body_read",
+            "body_bytes": 0,
+            "sse_data_events": 0,
+            "parse_errors": 0,
+            "oversized_frames": 0,
+            "unknown_response_events": 0,
+            "statistics_truncated": False,
+            "body_eof_observed": False,
+        }
+        if not observe_sse:
+            self.data["sse_observation"] = "skipped_content_encoding"
+        self.responses: dict[str, dict] = {}
+        self.tools: dict[tuple, dict] = {}
+        self.fingerprints: dict[tuple, ArgumentFingerprint] = {}
+        self.aliases: dict[tuple, tuple] = {}
+        self.current_response: str | None = None
+        self.line = bytearray()
+        self.lines: list[bytes] = []
+        self.frame_bytes = 0
+        self.dropping = False
+        self.discarded_line = False
+
+    def feed(self, data: bytes):
+        self.data["body_bytes"] += len(data)
+        self.data["last_read_at"] = time.time()
+        self.data.setdefault("first_read_at", self.data["last_read_at"])
+        if not self.observe_sse:
+            return
+        offset = 0
+        while offset < len(data):
+            newline = data.find(b"\n", offset)
+            end = len(data) if newline < 0 else newline
+            size = end - offset
+            if self.dropping:
+                if size:
+                    if not self.discarded_line and not self.line and data[offset:end] == b"\r":
+                        self.line.extend(b"\r")
+                    else:
+                        self.discarded_line = True
+                        self.line.clear()
+            elif self.frame_bytes + len(self.line) + size + (newline >= 0) > MAX_FRAME_BYTES:
+                self.data["oversized_frames"] += 1
+                self.dropping = True
+                self.line.clear()
+                self.lines.clear()
+                self.discarded_line = True
+            elif not self.dropping:
+                self.line.extend(data[offset:end])
+            if newline < 0:
+                break
+            # When dropping, a blank physical line still ends the oversized frame.
+            if self.dropping:
+                if not self.discarded_line:
+                    self.dropping = False
+                    self.frame_bytes = 0
+                self.discarded_line = False
+                self.line.clear()
+            else:
+                line = bytes(self.line).removesuffix(b"\r")
+                self.frame_bytes += len(self.line) + 1
+                self.line.clear()
+                if not line:
+                    self._frame()
+                    self.lines.clear()
+                    self.frame_bytes = 0
+                elif line.startswith(b"data:"):
+                    self.lines.append(line[5:].removeprefix(b" "))
+            offset = newline + 1
+
+    def _frame(self):
+        if not self.lines:
+            return
+        self.data["sse_data_events"] += 1
+        body = b"\n".join(self.lines)
+        if body == b"[DONE]":
+            self.data["sse_done_events"] = self.data.get("sse_done_events", 0) + 1
+            return
+        try:
+            event = json.loads(body)
+        except (ValueError, UnicodeError, RecursionError):
+            self.data["parse_errors"] += 1
+            return
+        if not isinstance(event, dict):
+            self.data["parse_errors"] += 1
+            return
+        self.event(event)
+
+    def event(self, event: dict):
+        kind = event.get("type")
+        kind = kind if isinstance(kind, str) else None
+        response = event.get("response")
+        response = response if isinstance(response, dict) else {}
+        response_id = identifier(event.get("response_id")) or identifier(response.get("id"))
+        if response_id:
+            self.current_response = response_id
+        response_id = self.current_response
+        if response_id is None:
+            self.data["unknown_response_events"] += 1
+            return
+        if response_id not in self.responses and len(self.responses) >= MAX_RESPONSES:
+            self.data["statistics_truncated"] = True
+            return
+        stats = self.responses.setdefault(
+            response_id,
+            {
+                "response_id": response_id,
+                "events": 0,
+                "sequence_missing_events": 0,
+                "sequence_forward_gaps": 0,
+                "sequence_forward_missing_total": 0,
+                "sequence_non_increasing": 0,
+                "sequence_equal_previous": 0,
+            },
+        )
+        stats["events"] += 1
+        sequence = event.get("sequence_number")
+        if type(sequence) is int and sequence >= 0:
+            previous = stats.get("sequence_last")
+            if previous is not None:
+                if sequence <= previous:
+                    stats["sequence_non_increasing"] += 1
+                    stats["sequence_equal_previous"] += int(sequence == previous)
+                elif sequence > previous + 1:
+                    stats["sequence_forward_gaps"] += 1
+                    stats["sequence_forward_missing_total"] += sequence - previous - 1
+            stats.setdefault("sequence_first", sequence)
+            stats["sequence_last"] = sequence
+        else:
+            stats["sequence_missing_events"] += 1
+        if kind in {
+            "response.completed",
+            "response.failed",
+            "response.incomplete",
+            "response.cancelled",
+        }:
+            stats["terminal_event"] = kind
+        if "tools" in response:
+            stats["returned_tool_constraints"] = tool_constraints(response["tools"])
+        if "max_output_tokens" in response:
+            limit = response["max_output_tokens"]
+            stats["returned_max_output_tokens"] = limit if type(limit) is int else None
+        self._tool(event, response_id)
+
+    def _tool(self, event, response_id):
+        kind = event.get("type")
+        if not isinstance(kind, str):
+            return
+        item = event.get("item")
+        item = item if isinstance(item, dict) else {}
+        if kind in {"response.output_item.added", "response.output_item.done"}:
+            if item.get("type") != "function_call":
+                return
+        elif kind not in {
+            "response.function_call_arguments.delta",
+            "response.function_call_arguments.done",
+        }:
+            return
+        item_id = identifier(event.get("item_id")) or identifier(item.get("id"))
+        index = event.get("output_index")
+        index = index if type(index) is int and index >= 0 else None
+        alias = (response_id, index)
+        key = self.aliases.get(alias) if index is not None else None
+        key = key or (response_id, item_id if item_id else index)
+        if key not in self.tools and len(self.tools) >= MAX_TRACKED_TOOLS:
+            self.data["statistics_truncated"] = True
+            return
+        if index is not None:
+            if alias in self.aliases or len(self.aliases) < 2 * MAX_TRACKED_TOOLS:
+                self.aliases[alias] = key
+            else:
+                self.data["statistics_truncated"] = True
+        stats = self.tools.setdefault(key, {"response_id": response_id, "delta_events": 0})
+        if not item_id and index is None:
+            stats["identity_unknown"] = True
+        if item_id and stats.get("item_id") not in (None, item_id):
+            stats["identity_conflicts"] = stats.get("identity_conflicts", 0) + 1
+        if item_id:
+            stats["item_id"] = item_id
+        if (call_id := identifier(item.get("call_id"))) is not None:
+            stats["call_id"] = call_id
+        if (name := safe_label(item.get("name"))) is not None:
+            stats["name"] = name
+        if index is not None:
+            stats["output_index"] = index
+        fingerprint = self.fingerprints.get(key)
+        if fingerprint is None:
+            fingerprint = self.fingerprints[key] = ArgumentFingerprint()
+        arguments = None
+        if kind == "response.function_call_arguments.delta":
+            stats["delta_events"] += 1
+            arguments = event.get("delta")
+        elif kind == "response.output_item.added":
+            arguments = item.get("arguments")
+        elif kind == "response.function_call_arguments.done":
+            stats["arguments_done_events"] = stats.get("arguments_done_events", 0) + 1
+        elif kind == "response.output_item.done":
+            stats["output_item_done_events"] = stats.get("output_item_done_events", 0) + 1
+        if isinstance(arguments, str):
+            fingerprint.observe(arguments)
+            if arguments:
+                stats.setdefault("first_argument_at", time.time())
+                stats["last_argument_at"] = time.time()
+
+    def end(self, *, eof=False):
+        self.data["body_eof_observed"] = self.data["body_eof_observed"] or eof
+        self.data["incomplete_sse_frame"] = self.data.get("incomplete_sse_frame", False) or bool(
+            self.line or self.lines or self.dropping
+        )
+        self.line.clear()
+        self.lines.clear()
+
+    def snapshot(self):
+        return {
+            **self.data,
+            "responses": [dict(r) for r in self.responses.values()],
+            "tools": [{**r, **self.fingerprints[k].snapshot()} for k, r in self.tools.items()],
+        }
+
+
+class ObservedResponseStream(httpx.AsyncByteStream, httpx2.AsyncByteStream):
+    def __init__(
+        self, stream: httpx.AsyncByteStream | httpx2.AsyncByteStream, attempt, *, encoded=False
+    ):
+        self.stream = stream
+        self.attempt = attempt
+        self.observation = ResponsesObservation(observe_sse=not encoded)
+        self.closed = False
+        self.close_attempted = False
+
+    async def __aiter__(self):
+        try:
+            async for data in self.stream:
+                started = time.monotonic()
+                self.observation.feed(data)
+                self.attempt.stream_work("wire_observe", time.monotonic() - started)
+                self.attempt.publish_if_due()
+                yield data
+            self.observation.end(eof=True)
+        except asyncio.CancelledError:
+            self.attempt.stream_cancelled()
+            raise
+
+    async def aclose(self):
+        if self.close_attempted:
+            return
+        self.close_attempted = True
+        stats = self.observation.data
+        stats.setdefault("close_started_at", time.time())
+        try:
+            async with asyncio.timeout(1):
+                await self.stream.aclose()
+        except TimeoutError:
+            # A cleanup deadline must not replace an in-flight user stop with
+            # a retryable model timeout. Failure remains explicit in diagnostics.
+            stats["close_error_type"] = "TimeoutError"
+            logger.warning("Model HTTP stream close exceeded its cleanup deadline")
+        except (
+            asyncio.CancelledError,
+            OSError,
+            RuntimeError,
+            httpx.HTTPError,
+            httpx2.HTTPError,
+        ) as error:
+            stats["close_error_type"] = type(error).__name__
+            raise
+        else:
+            self.closed = True
+            stats["closed_at"] = time.time()
+        finally:
+            self.observation.end(eof=stats["body_eof_observed"])

artemis/llm/tool_constraints.py

@@ -0,0 +1,109 @@
+"""Summarize actual wire tool constraints without persisting schemas or prompts."""
+
+import hashlib
+import json
+import re
+
+from artemis.llm.argument_fingerprint import MAX_TRACKED_TOOLS
+
+MAX_REQUEST_BYTES = 4 * 1024 * 1024
+
+
+def safe_label(value):
+    return value if isinstance(value, str) and re.fullmatch(r"[\w.:-]{1,128}", value) else None
+
+
+def strict_state(tool: dict) -> str:
+    if "strict" not in tool:
+        return "omitted"
+    if tool["strict"] is True:
+        return "true"
+    if tool["strict"] is False:
+        return "false"
+    return "invalid"
+
+
+def schema_summary(schema: dict) -> dict:
+    """Count syntactic schema nodes, not expanded $refs or arbitrary examples."""
+    result = {"objects": 0, "optional_properties": 0, "open_objects": 0, "nodes": 0}
+    pending = [schema]
+    while pending and result["nodes"] < 1024:
+        node = pending.pop()
+        result["nodes"] += 1
+        properties = node.get("properties")
+        if node.get("type") == "object" or isinstance(properties, dict):
+            result["objects"] += 1
+            required = node.get("required")
+            required = required if isinstance(required, list) else []
+            if isinstance(properties, dict):
+                result["optional_properties"] += sum(k not in required for k in properties)
+            if node.get("additionalProperties") is not False:
+                result["open_objects"] += 1
+        children = []
+        for key in ("properties", "$defs", "definitions"):
+            if isinstance(node.get(key), dict):
+                children.extend(node[key].values())
+        for key in ("anyOf", "oneOf", "allOf", "prefixItems"):
+            if isinstance(node.get(key), list):
+                children.extend(node[key])
+        children.extend(node.get(k) for k in ("items", "additionalProperties", "not"))
+        remaining = max(0, 1024 - result["nodes"] - len(pending))
+        eligible = [v for v in children if isinstance(v, dict)]
+        if len(eligible) > remaining:
+            result["truncated"] = True
+        pending.extend(eligible[:remaining])
+    result["truncated"] = bool(pending) or result.get("truncated", False)
+    try:
+        result["sha256"] = hashlib.sha256(
+            json.dumps(schema, sort_keys=True, separators=(",", ":"), ensure_ascii=True).encode()
+        ).hexdigest()
+    except (RecursionError, TypeError, ValueError):
+        result["sha256"] = None
+        result["truncated"] = True
+    return result
+
+
+def tool_constraints(tools) -> dict:
+    records = []
+    pending = tools[:MAX_TRACKED_TOOLS] if isinstance(tools, list) else []
+    truncated = isinstance(tools, list) and len(tools) > MAX_TRACKED_TOOLS
+    visited = 0
+    while pending and visited < MAX_TRACKED_TOOLS:
+        tool = pending.pop(0)
+        visited += 1
+        if not isinstance(tool, dict):
+            continue
+        if tool.get("type") == "namespace":
+            nested = tool.get("tools")
+            if isinstance(nested, list):
+                remaining = MAX_TRACKED_TOOLS - len(pending) - visited
+                truncated = truncated or len(nested) > remaining
+                pending.extend(nested[:remaining])
+            continue
+        function = tool.get("function", tool)
+        if not isinstance(function, dict):
+            continue
+        record = {
+            "type": safe_label(tool.get("type")),
+            "name": safe_label(function.get("name")),
+            "strict": strict_state(function),
+        }
+        if isinstance(function.get("parameters"), dict):
+            record["schema"] = schema_summary(function["parameters"])
+        records.append(record)
+    return {"tools": records, "truncated": truncated or bool(pending), "observation": "available"}
+
+
+def request_constraints(body: bytes) -> dict:
+    if len(body) > MAX_REQUEST_BYTES:
+        return {"observation": "body_too_large"}
+    try:
+        payload = json.loads(body)
+    except (ValueError, UnicodeError, RecursionError):
+        return {"observation": "invalid_json"}
+    if not isinstance(payload, dict):
+        return {"observation": "not_an_object"}
+    result = tool_constraints(payload.get("tools"))
+    limit = payload.get("max_output_tokens")
+    result["max_output_tokens"] = limit if type(limit) is int and limit > 0 else None
+    return result