From 7c0c2b59ca8adf4619dd4acaece1d8c2cf8677e6 Mon Sep 17 00:00:00 2001 From: chenmingxuan Date: Mon, 2 Nov 2026 16:48:51 +0800 Subject: [PATCH] fix: bound model stream aggregation and close failed responses --- .../services/request_diagnostics.py | 3 + artemis/llm/argument_fingerprint.py | 33 ++ artemis/llm/chunk_accumulator.py | 37 ++ artemis/llm/diagnostics.py | 65 +++- artemis/llm/responses_diagnostics.py | 301 ++++++++++++++++ artemis/llm/tool_constraints.py | 109 ++++++ artemis/services/llm.py | 44 ++- docs/llm-stream-diagnostics.md | 43 ++- tests/unit/test_argument_fingerprint.py | 20 +- tests/unit/test_chunk_accumulator.py | 83 +++++ tests/unit/test_request_diagnostics.py | 3 +- tests/unit/test_responses_diagnostics.py | 263 ++++++++++++++ tests/unit/test_responses_wire_lifecycle.py | 323 ++++++++++++++++++ 13 files changed, 1315 insertions(+), 12 deletions(-) create mode 100644 artemis/llm/chunk_accumulator.py create mode 100644 artemis/llm/responses_diagnostics.py create mode 100644 artemis/llm/tool_constraints.py create mode 100644 tests/unit/test_chunk_accumulator.py create mode 100644 tests/unit/test_responses_diagnostics.py create mode 100644 tests/unit/test_responses_wire_lifecycle.py diff --git a/apps/admin_console/services/request_diagnostics.py b/apps/admin_console/services/request_diagnostics.py index 8532281..aa30e37 100644 --- a/apps/admin_console/services/request_diagnostics.py +++ b/apps/admin_console/services/request_diagnostics.py @@ -66,6 +66,9 @@ def request_diagnostics( "headers_at", "first_output_at", "last_output_at", + "first_tool_arguments_at", + "last_tool_arguments_at", + "cancel_observed_at", "ended_at", "retry_at", ) diff --git a/artemis/llm/argument_fingerprint.py b/artemis/llm/argument_fingerprint.py index 2630a5d..e6bb51f 100644 --- a/artemis/llm/argument_fingerprint.py +++ b/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, } diff --git a/artemis/llm/chunk_accumulator.py b/artemis/llm/chunk_accumulator.py new file mode 100644 index 0000000..dd98c81 --- /dev/null +++ b/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 diff --git a/artemis/llm/diagnostics.py b/artemis/llm/diagnostics.py index fc4716c..c11e627 100644 --- a/artemis/llm/diagnostics.py +++ b/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 diff --git a/artemis/llm/responses_diagnostics.py b/artemis/llm/responses_diagnostics.py new file mode 100644 index 0000000..5606780 --- /dev/null +++ b/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"]) diff --git a/artemis/llm/tool_constraints.py b/artemis/llm/tool_constraints.py new file mode 100644 index 0000000..3722b29 --- /dev/null +++ b/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 diff --git a/artemis/services/llm.py b/artemis/services/llm.py index f06c25c..c47865e 100644 --- a/artemis/services/llm.py +++ b/artemis/services/llm.py @@ -47,6 +47,7 @@ from artemis.config import ( from artemis.context import ArtemisContext from artemis.data_engine.trace import CURRENT_TRACE_ID, DataEngineCallbackHandler from artemis.llm.diagnostics import ACTIVE_ATTEMPT, RequestAttempt, error_text, input_images +from artemis.llm.chunk_accumulator import ChunkAccumulator from artemis.llm.final_review import ( ACTIVE_FINAL_REVIEW, FinalReviewBudget, @@ -425,8 +426,13 @@ async def _run_with_recovery[T]( observation_token = ACTIVE_ATTEMPT.set(observation) request_timer = asyncio.timeout(timeout) try: - async with request_timer: - result = await call_once() + try: + async with request_timer: + result = await call_once() + finally: + # Release a failed stream before any retry backoff, even + # when the SDK consumer stopped between yielded chunks. + await observation.close_streams() except (KeyboardInterrupt, asyncio.CancelledError) as e: observation.cancel(e) raise @@ -818,24 +824,46 @@ class RobustChatModelWrapper: stream_exec_id = uuid4() full_response = None + accumulator = ChunkAccumulator() argument_guard = ToolArgumentGuard() try: async with aclosing(self.base_model.astream(*args, **kwargs)) as stream: + wait_started = time.monotonic() async for chunk in stream: observation = ACTIVE_ATTEMPT.get() if observation is not None: + observation.stream_work("await_sdk_chunk", time.monotonic() - wait_started) + observe_started = time.monotonic() observation.chunk(chunk) + observation.stream_work("sdk_observe", time.monotonic() - observe_started) argument_guard.observe(chunk) - if full_response is None: - full_response = chunk - else: - full_response = full_response + chunk + accumulator.append(chunk) + if accumulator.ready: + # Buffered transport input may yield without suspending. + # Give stop/deadline signals a turn before a batch merge. + await asyncio.sleep(0) + merge_started = time.monotonic() + full_response = accumulator.flush() + if observation is not None: + observation.stream_work("merge", time.monotonic() - merge_started) if emit_deltas: + emit_started = time.monotonic() self._emit_stream_delta(stream_exec_id, chunk) + if observation is not None: + observation.stream_work("emit", time.monotonic() - emit_started) + wait_started = time.monotonic() + merge_started = time.monotonic() + full_response = accumulator.flush() + observation = ACTIVE_ATTEMPT.get() + if observation is not None: + observation.stream_work("merge", time.monotonic() - merge_started) if full_response is not None: ensure_response_complete(full_response) except (KeyboardInterrupt, asyncio.CancelledError): - if full_response is not None: + observation = ACTIVE_ATTEMPT.get() + if observation is not None: + observation.stream_cancelled() + if accumulator.has_output: _record_llm_event( "llm_stream_reset", { @@ -847,7 +875,7 @@ class RobustChatModelWrapper: ) raise except Exception as stream_error: - if full_response is not None or isinstance( + if accumulator.has_output or isinstance( stream_error, (ToolArgumentOverflow, ResponseIntegrityError) ): if isinstance(stream_error, ResponseIntegrityError): diff --git a/docs/llm-stream-diagnostics.md b/docs/llm-stream-diagnostics.md index d6c388b..bfd4f27 100644 --- a/docs/llm-stream-diagnostics.md +++ b/docs/llm-stream-diagnostics.md @@ -2,6 +2,8 @@ 小型内置动作的参数流仍受 16384 字符保护。超限时关闭输出流、丢弃不完整参数，不执行该次动作；一次新响应重试和可用备用端点耗尽后，只结束当前任务，不进入共享的模型暂停。批次仍遵循“失败即停止”、人工暂停和取消设置。网络故障、认证错误和最终验收保留各自的恢复规则。 +长度保护不是生成进度判断：小于上限但一直没有闭合或终止的参数，仍可能持续到单次请求或外层任务超时。超时、取消和完整性校验共同约束这条路径，未完成的参数不能作为动作执行。 + ## 请求关联 每个模型请求尝试沿用现有 `llm_diagnostic` 记录，不另建日志数据库。 @@ -27,6 +29,9 @@ - `argument_deltas`：非空参数片段数，不能直接与网关原始事件数比较。 - `argument_prefixes`：每 4096 字节的累计 SHA256，最多保留前 16 个。可以在分片大小不同、上游预读或中途取消时比较相同字节位置。 - `argument_encoding_valid`：遇到不能编码为 UTF-8 的片段时为 false，当前总字节数和总哈希留空，避免伪造可比结果；诊断本身不改变执行流程。 +- `argument_json_whitespace_chars` / `argument_json_whitespace_run` / `argument_json_whitespace_run_max`：实际空格、制表、回车和换行字符的总数、当前尾部连续数和最长连续数。统计的是这四种 JSON 空白码点，也包含字符串内部的实际字符；转义文本 `\\n`、不换行空格不计入。不据此判断 JSON 是否有效。 +- `argument_delta_chars_min` / `argument_delta_chars_max`：非空参数增量的最小、最大字符数；尚无非空增量时最小值为空。 +- `argument_same_delta_run` / `argument_same_delta_run_max`：连续相同非空增量的当前、最长出现次数，首个增量算一次，空增量不打断。只保留上一增量的 SHA256，不保留正文。相同增量内容不等于重复的协议事件，例如连续空格也可能来自不同的合法事件。 最多记录 64 个工具；超过时设置 `tool_stats_truncated=true`，总参数字符计数仍继续。诊断不保留参数正文。哈希仅证明已收到的字节前缀，不能证明输出完整或动作成功；完整性仍由响应终止状态和解析校验决定。 @@ -34,6 +39,38 @@ 历史日志没有保存的原始帧无法由新增诊断恢复。HTTP 200 只说明响应头已到达；客户端保护后出现的 499 或 Broken pipe 也不能证明最初原因。网关写入本机下游 writer 成功不等于客户端已消费或执行。 +## 客户端原始 Responses 观测 + +`wire_responses` 被动观察 HTTP 客户端实际读取的 SSE 字节，在 SDK 转换之前统计。同一网络读取中可能包含多个事件，因此原始统计可以领先 SDK；SDK 回调阻塞也会扩大两者差距。这一边界不是服务器写出时刻，也不是工具执行时刻。诊断不增加预读任务、队列或重传，不过滤、改写交给 SDK 的字节。 + +每个响应记录 `response_id`、事件数及 `sequence_number` 的首末值。`sequence_forward_gaps` / `sequence_forward_missing_total` 统计相邻观测间的向前跳变次数及跳过的序号数；乱序事件可能随后补上，不能直接称为永久丢包。`sequence_non_increasing` 统计非递增，`sequence_equal_previous` 统计与前一个序号相等。`sequence_missing_events` 表示事件本身没有有效序号，不表示传输丢失。没有响应编号时不猜测，记录 `unknown_response_events`；标准流中没有编号的后续事件沿用当前响应编号。 + +每个函数调用记录与 SDK 相同的参数指纹和形态统计，另记原始 `delta_events`、首末参数到达时间，以及两类 done 事件的次数。只累计非空初始 arguments 和 arguments.delta，不重复累计 done 的完整快照。关联优先使用响应编号与实际 `call_id`，`item_id` / `output_index` 用于同一原始流内关联；`identity_unknown`、`identity_conflicts` 表明身份不足或冲突，此时不能据此归责。 + +每次尝试最多观察 4 条 HTTP 流；每条流最多 16 个响应、64 个工具、128 个索引别名，每工具最多 16 个固定前缀。单个待解析 SSE 帧最多 256KiB，超出后仅跳过该帧的诊断，继续原样交给 SDK。`oversized_frames`、`parse_errors`、`statistics_truncated`、`wire_streams_truncated` 说明覆盖缺口；`incomplete_sse_frame` 与 `body_eof_observed` 区分未闭合帧和正常字节流结束。待解析字节仅临时保留，关闭时释放，日志不保存帧或参数正文。 + +原始分帧支持 LF / CRLF。HTTP 体声明 gzip 等压缩编码时，标记 `sse_observation=skipped_content_encoding`，只统计读取字节与时间，不把压缩字节误当成 SSE；SDK 仍按原路径解压。这种记录缺少原始事件统计，不能把零计数当成没有收到参数。 + +## 实际工具约束与耗时 + +`wire_request_constraints` 检查 HTTP 请求已经序列化的正文，不记录提示、图片或 Schema 正文。逐工具记录 `strict` 为 `omitted`、`false`、`true`，异常值为 `invalid`；不能把省略当成 false，更不能把 Pydantic 本地严格校验当成上游约束。Schema 只保存 SHA256、对象数、非 required 的属性数及未关闭对象数。这是语法统计，不展开引用，也不是上游 Schema 兼容性验证。 + +当前 `click`、`action_sequence`、`drag` 的真实 SDK 离线请求均省略 strict，保留可选字段。不能只把 strict 开关改为 true：应先按目标接口支持的 Schema 子集处理对象和 required，保留可选值的 nullable 语义，再单独验证兼容性。 + +请求正文超过 4MiB 时跳过检查并记录 `body_too_large`；流式上传正文不可直接读取时记录 `streaming_body_unavailable`，不会为了诊断消费上传流。最多检查 64 个工具或命名空间节点，每个 Schema 最多统计 1024 个节点，达到上限标记截断。响应事件如果回传 tools，则保存同样的 `returned_tool_constraints`；回显不证明服务端实际执行了约束。 + +`max_output_tokens` 与 `returned_max_output_tokens` 只表示本层实际观察到的字段。本次核对的 CodexManager 通道会过滤输出 token 上限；短合成探针成功也不能证明上限被执行。不能把设置 max_tokens / max_output_tokens 当成这条链路可靠的止损办法，客户端长度、超时和取消保护仍然必要。 + +`stream_work` 记录每阶段的调用次数、总耗时和最长耗时：`wire_observe` 为原始帧诊断，`sdk_observe` 为 SDK 片段诊断，`merge` 为聚合，`emit` 为界面回调，`await_sdk_chunk` 为等待 SDK 交出下一个片段。最后一项包含网络读取、原始观测、SDK 转换和 SDK 回调等工作，与其他阶段可能重叠，不能相加或直接当成网络延迟。 + +聚合每 128 个片段合并一次，降低 SDK 反复解析半截 JSON 的开销；长度保护、诊断和界面输出仍逐片段执行，异步压缩继续使用原异步调用路径。分批合并不修复参数 JSON，最终仍须经过完整性和严格解析校验。 + +## 取消与关闭 + +`cancel_observed_at` / `cancel_observed_phase` 记录本次尝试首次在客户端观测到取消的时间与阶段，不等于用户点击停止的时间。`cancel_phase` / `cancel_reason` 沿用恢复层最终记录的取消阶段和来源，可能反映后续重试等待中的停止。首末 SDK 参数时间为 `first_tool_arguments_at` / `last_tool_arguments_at`。原始流内的读取、参数和关闭时间采用 Unix 秒。 + +SDK 读取中、回调中或两次 yield 之间触发长度保护后，尝试结束路径都会主动关闭已观察的 HTTP 流，不等待垃圾回收。`close_started_at` 记录关闭开始，只有实际成功才写 `closed_at`；失败只记录异常类型，关闭等待最多 1 秒。关闭超时不能覆盖用户停止而变成可重试的请求超时。未能关闭时保留错误证据，不声称连接已成功释放。 + ## 已核对的上游异常样本（2026-11-02） 一次真实 `action_sequence` 调用在任务超时前持续接收参数，没有完成该工具调用。该样本发生在 `action_sequence` 纳入参数保护之前。按同一 `response_id + call_id` 核对生产诊断，得到以下字符数；请求编号和设备信息不在文档中保留： @@ -48,10 +85,14 @@ 五个边界在 4096 到 36864 字节的九个累计 SHA256 上一致。NewAPI 使用透传模式，收到和写出的完整参数哈希也一致。上游没有终止事件，也没有完成参数快照。该样本证明参数在官方上游返回时就已持续增长，不能归因于 NewAPI 转换、下游重复累计完成快照或 ARTEMIS 最终消息拼接。取消、预读和缓冲会使各层停止计数的位置不同；尾部总长度不同本身不能证明重复发送或字符丢失。 -诊断没有参数正文，不能进一步断言是空白循环、重复描述或官方服务内部的重复事件。这个定位也不能追溯证明没有原始诊断的历史事故。另一次 `click` 异常同样在 NewAPI 接收时已超过四万字符，说明这种现象并非新动作序列独有。 +诊断没有参数正文，不能进一步断言是模型生成循环、空白循环、重复描述或官方服务内部的重复事件，也不能排除请求结构触发异常。这个定位只排除已核对样本的下游放大作为最初来源，不能追溯证明没有原始诊断的历史事故。另一次 `click` 异常同样在 NewAPI 接收时已超过四万字符，说明这种现象并非新动作序列独有。 + +CodexManager 另行修正了普通函数放入 functions 命名空间的兼容性问题。其短上下文、低思考强度合成测试通过，仅能作为兼容性对照，尚不能代替原事故的长上下文、多图、高思考强度复现，也不能称为参数膨胀的根因修复。 本地保护覆盖 `action_sequence`、`drag` 及其他小型内置动作：超限后关闭流、丢弃参数，并按既有策略有限重试；它限制异常对当前任务的影响，不修复上游生成行为，也不保证业务任务一定成功。缩放面板自动隐藏导致的定位失败由动作序列的 `ensure_visible` 准备步骤处理，与参数流异常分别诊断。 ## 离线验证 `tests/unit/test_llm_stream_guard.py` 验证保护、有限重试、相同/不同备用端点，以及不会创建或清除其他请求的全局暂停。`tests/unit/test_request_diagnostics.py` 验证请求编号、超时重试、响应头及单次参数计数。`tests/unit/test_argument_fingerprint.py` 验证不同分片、中文和 emoji、固定字节前缀、有界统计与异常编码。这些测试不调用真实模型或设备。 + +`tests/unit/test_responses_diagnostics.py` 验证原始分帧、序号、done 去重、有界统计和 strict 三态。`tests/unit/test_responses_wire_lifecycle.py` 使用已安装的真实 OpenAI/LangChain SDK 与内存 HTTP 传输，验证小于长度阈值的未闭合流、SDK 回调背压、请求/任务超时、用户停止、超限后主动关闭和关闭失败，以及实际工具 Schema；不发网络请求。`tests/unit/test_chunk_accumulator.py` 对比逐片段聚合结果，验证交错工具、元数据和使用量保持一致，减少半截 JSON 重复解析，异常参数仍由最终校验拒绝。离线通过不代表原业务任务的稳定性已验收。 diff --git a/tests/unit/test_argument_fingerprint.py b/tests/unit/test_argument_fingerprint.py index cfb1407..9e3bc05 100644 --- a/tests/unit/test_argument_fingerprint.py +++ b/tests/unit/test_argument_fingerprint.py @@ -33,7 +33,7 @@ def test_invalid_utf8_is_marked_without_failing_or_forging_a_digest(): fingerprint = ArgumentFingerprint() for value in ["safe", "\ud800", "later"]: fingerprint.observe(value) - assert fingerprint.snapshot() == { + expected = { "argument_chars": 10, "argument_bytes": None, "argument_sha256": None, @@ -41,6 +41,24 @@ def test_invalid_utf8_is_marked_without_failing_or_forging_a_digest(): "argument_prefixes": [], "argument_encoding_valid": False, } + assert {key: fingerprint.snapshot()[key] for key in expected} == expected + + +def test_shape_counts_cross_delta_whitespace_and_repetition_without_retaining_text(): + fingerprint = ArgumentFingerprint() + for fragment in ['{"text":"private', " ", " \t", "", " \t", " \t", '\u00a0"}']: + fingerprint.observe(fragment) + stats = fingerprint.snapshot() + assert stats["argument_json_whitespace_chars"] == 7 + assert stats["argument_json_whitespace_run_max"] == 7 + assert stats["argument_json_whitespace_run"] == 0 + assert stats["argument_delta_chars_min"] == 1 + assert stats["argument_delta_chars_max"] == len('{"text":"private') + assert stats["argument_same_delta_run_max"] == 3 + assert stats["argument_same_delta_run"] == 1 + assert stats["argument_deltas"] == 6 + assert "private" not in repr(vars(fingerprint)) + assert "private" not in json.dumps(stats) def test_tool_tracking_is_bounded_but_total_count_keeps_all_received_arguments(): diff --git a/tests/unit/test_chunk_accumulator.py b/tests/unit/test_chunk_accumulator.py new file mode 100644 index 0000000..c7f80f1 --- /dev/null +++ b/tests/unit/test_chunk_accumulator.py @@ -0,0 +1,83 @@ +import itertools +import json + +from langchain_core.messages import AIMessageChunk +import langchain_core.messages.ai as ai +import pytest + +from artemis.llm.chunk_accumulator import ChunkAccumulator +from artemis.llm.response_guard import ResponseIntegrityError, ensure_response_complete + + +def test_batching_preserves_interleaved_tools_unicode_metadata_and_usage(): + arguments = [json.dumps({"target": "中文🙂" * 80}), json.dumps({"target": 3})] + chunks = [AIMessageChunk(content="Live text", id="resp_one")] + for index in range(2): + chunks.append( + AIMessageChunk( + content=[], + tool_call_chunks=[ + {"name": "click", "args": "", "id": f"call_{index}", "index": index} + ], + ) + ) + for pair in itertools.zip_longest(*arguments): + for index, value in enumerate(pair): + if value is not None: + chunks.append( + AIMessageChunk(content=[], tool_call_chunks=[{"args": value, "index": index}]) + ) + chunks.append( + AIMessageChunk( + content=[], + response_metadata={"id": "resp_one", "status": "completed"}, + usage_metadata={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + chunk_position="last", + ) + ) + expected = chunks[0] + accumulator = ChunkAccumulator() + max_pending = 0 + for index, chunk in enumerate(chunks): + if index: + expected = expected + chunk + accumulator.append(chunk) + if accumulator.ready: + accumulator.flush() + max_pending = max(max_pending, len(accumulator.pending)) + result = accumulator.flush() + assert result.model_dump() == expected.model_dump() + assert [t["args"] for t in result.tool_calls] == [json.loads(a) for a in arguments] + assert result.usage_metadata["total_tokens"] == 15 + assert max_pending < 256 + assert not accumulator.pending + + +def test_partial_json_is_not_reparsed_for_every_single_character(monkeypatch): + chunks = [ + AIMessageChunk( + content=[], tool_call_chunks=[{"name": "click", "args": '{"target":x', "index": 0}] + ) + ] + chunks.extend( + AIMessageChunk(content=[], tool_call_chunks=[{"args": "x", "index": 0}]) for _ in range(800) + ) + original = ai.parse_partial_json + parsed_lengths = [] + + def counted(arguments, *args, **kwargs): + parsed_lengths.append(len(arguments)) + return original(arguments, *args, **kwargs) + + monkeypatch.setattr(ai, "parse_partial_json", counted) + accumulator = ChunkAccumulator() + for chunk in chunks: + accumulator.append(chunk) + if accumulator.ready: + accumulator.flush() + result = accumulator.flush() + assert result.tool_call_chunks[0]["args"] == '{"target":x' + "x" * 800 + assert len(parsed_lengths) < 20 + result.response_metadata["status"] = "completed" + with pytest.raises(ResponseIntegrityError): + ensure_response_complete(result) diff --git a/tests/unit/test_request_diagnostics.py b/tests/unit/test_request_diagnostics.py index 19f4ba6..d29a6dc 100644 --- a/tests/unit/test_request_diagnostics.py +++ b/tests/unit/test_request_diagnostics.py @@ -226,7 +226,7 @@ def test_tool_deltas_are_counted_once_without_retaining_arguments(): [2, 1], strict=True, ): - assert tool == { + expected = { "name": name, "call_id": call_id, "chunk_index": index, @@ -237,6 +237,7 @@ def test_tool_deltas_are_counted_once_without_retaining_arguments(): "argument_prefixes": [], "argument_encoding_valid": True, } + assert {key: tool[key] for key in expected} == expected assert attempt.data["response_status"] == "completed" assert "private" not in json.dumps(engine.records, default=str) assert "secret" not in json.dumps(engine.records, default=str) diff --git a/tests/unit/test_responses_diagnostics.py b/tests/unit/test_responses_diagnostics.py new file mode 100644 index 0000000..0009edf --- /dev/null +++ b/tests/unit/test_responses_diagnostics.py @@ -0,0 +1,263 @@ +import json +import gzip +from types import SimpleNamespace + +import httpx2 +import pytest + +from artemis.llm.responses_diagnostics import MAX_FRAME_BYTES, ResponsesObservation +from artemis.llm.diagnostics import ACTIVE_ATTEMPT, observe_response_headers +from artemis.llm.tool_constraints import MAX_REQUEST_BYTES, request_constraints, tool_constraints + + +def frame(kind, sequence, **fields): + return ( + "data: " + + json.dumps({"type": kind, "sequence_number": sequence, **fields}, ensure_ascii=False) + + "\r\n\r\n" + ).encode() + + +@pytest.mark.parametrize("size", [1, 37, 100000]) +def test_raw_frames_split_unicode_and_done_snapshots_are_not_double_counted(size): + arguments = '{"text":"私人🙂"}' + data = b"".join( + [ + frame("response.created", 0, response={"id": "resp_one"}), + frame( + "response.output_item.added", + 1, + output_index=2, + item={ + "type": "function_call", + "id": "fc_one", + "call_id": "call_one", + "name": "click", + "arguments": "", + }, + ), + frame( + "response.function_call_arguments.delta", + 2, + output_index=2, + item_id="fc_one", + delta=arguments, + ), + frame( + "response.function_call_arguments.done", + 3, + output_index=2, + item_id="fc_one", + arguments=arguments, + ), + frame( + "response.output_item.done", + 4, + output_index=2, + item={"type": "function_call", "id": "fc_one", "arguments": arguments}, + ), + frame( + "response.completed", + 5, + response={ + "id": "resp_one", + "tools": [{"type": "function", "name": "click", "strict": True}], + }, + ), + ] + ) + observer = ResponsesObservation() + for offset in range(0, len(data), size): + observer.feed(data[offset : offset + size]) + observer.end(eof=True) + stats = observer.snapshot() + assert stats["body_bytes"] == len(data) + assert stats["parse_errors"] == 0 + assert stats["incomplete_sse_frame"] is False + (tool,) = stats["tools"] + assert tool["argument_chars"] == len(arguments) + assert ( + tool["delta_events"] + == tool["arguments_done_events"] + == tool["output_item_done_events"] + == 1 + ) + assert tool["call_id"] == "call_one" + (response,) = stats["responses"] + assert response["sequence_first"] == 0 and response["sequence_last"] == 5 + assert response["sequence_forward_gaps"] == response["sequence_non_increasing"] == 0 + assert response["returned_tool_constraints"]["tools"][0]["strict"] == "true" + assert "私人" not in json.dumps(stats, ensure_ascii=False) + assert not observer.line and not observer.lines + + +def test_sequence_statistics_are_per_response_and_describe_observed_jumps(): + observer = ResponsesObservation() + observer.event({"type": "response.in_progress", "sequence_number": 1}) + observer.event({"type": "response.created", "sequence_number": 0, "response": {"id": "resp_a"}}) + for seq in [2, 2, 1, 3, None]: + observer.event({"type": "response.in_progress", "sequence_number": seq}) + observer.event({"type": "response.created", "sequence_number": 0, "response": {"id": "resp_b"}}) + a, b = observer.snapshot()["responses"] + assert a["sequence_forward_gaps"] == a["sequence_forward_missing_total"] == 2 + assert a["sequence_non_increasing"] == 2 + assert a["sequence_equal_previous"] == 1 + assert a["sequence_missing_events"] == 1 + assert b["sequence_non_increasing"] == b["sequence_forward_gaps"] == 0 + assert observer.data["unknown_response_events"] == 1 + + +def test_oversized_frame_recovery_and_partial_eof_keep_memory_bounded(): + observer = ResponsesObservation() + observer.feed(b"data: " + b"x" * MAX_FRAME_BYTES) + observer.feed(b"\r\n\r\n") + observer.feed(b"data: invalid\n\n") + observer.feed(frame("response.created", 0, response={"id": "resp_valid"})) + observer.feed(b'data: {"private":') + observer.end(eof=True) + stats = observer.snapshot() + assert stats["oversized_frames"] == 1 + assert stats["parse_errors"] == 1 + assert len(stats["responses"]) == 1 + assert stats["incomplete_sse_frame"] is True + assert "private" not in repr(vars(observer)) + assert not observer.line and not observer.lines + observer.end() # HTTP close can follow EOF; retain the original coverage gap. + assert observer.data["incomplete_sse_frame"] is True + + +def test_raw_tool_and_response_statistics_have_fixed_caps(): + observer = ResponsesObservation() + observer.event({"type": "response.created", "response": {"id": "resp_one"}}) + for index in range(100): + observer.event( + { + "type": "response.function_call_arguments.delta", + "item_id": f"fc_{index}", + "output_index": index, + "delta": " ", + } + ) + for index in range(20): + observer.event({"type": "response.created", "response": {"id": f"resp_{index}"}}) + assert len(observer.tools) == len(observer.fingerprints) == len(observer.aliases) == 64 + assert len(observer.responses) == 16 + assert observer.data["statistics_truncated"] is True + + +@pytest.mark.parametrize("size", [1, 37, MAX_FRAME_BYTES * 2]) +def test_oversized_crlf_frame_recovers_across_arbitrary_read_boundaries(size): + observer = ResponsesObservation() + data = b"data: " + b"x" * MAX_FRAME_BYTES + b"\r\n\r\n" + data += frame("response.created", 0, response={"id": "resp_after_large"}) + for offset in range(0, len(data), size): + observer.feed(data[offset : offset + size]) + assert observer.data["oversized_frames"] == 1 + assert observer.data["parse_errors"] == 0 + assert list(observer.responses) == ["resp_after_large"] + assert not observer.line and not observer.lines and not observer.dropping + + +def test_alias_statistics_are_bounded_even_when_item_keeps_changing_index(): + observer = ResponsesObservation() + observer.event({"response": {"id": "resp_one"}}) + for index in range(1000): + observer.event( + { + "type": "response.function_call_arguments.delta", + "item_id": "fc_same", + "output_index": index, + "delta": " ", + } + ) + assert len(observer.tools) == 1 + assert len(observer.aliases) == 128 + assert observer.data["statistics_truncated"] is True + + +def test_malformed_event_type_cannot_break_passive_observation(): + observer = ResponsesObservation() + observer.feed(frame([], 0, response={"id": "resp_one"})) + assert observer.snapshot()["responses"][0]["events"] == 1 + + +@pytest.mark.asyncio +async def test_compressed_body_is_explicitly_skipped_and_sdk_decoding_is_unchanged(): + body = frame("response.created", 0, response={"id": "resp_one"}) + compressed = gzip.compress(body) + + class Stream(httpx2.AsyncByteStream): + async def __aiter__(self): + yield compressed + + attempt = SimpleNamespace( + wire_streams=[], + data={}, + headers=lambda _: None, + stream_work=lambda *_: None, + publish_if_due=lambda: None, + ) + response = httpx2.Response( + 200, + headers={"content-type": "text/event-stream", "content-encoding": "gzip"}, + stream=Stream(), + ) + token = ACTIVE_ATTEMPT.set(attempt) + try: + await observe_response_headers(response) + assert b"".join([chunk async for chunk in response.aiter_bytes()]) == body + stats = attempt.wire_streams[0].observation.snapshot() + assert stats["sse_observation"] == "skipped_content_encoding" + assert stats["body_bytes"] == len(compressed) + assert not stats["responses"] and not stats["parse_errors"] + finally: + await response.aclose() + ACTIVE_ATTEMPT.reset(token) + + +@pytest.mark.parametrize("strict,expected", [(None, "omitted"), (False, "false"), (True, "true")]) +def test_actual_strict_state_and_optional_schema_are_observed_without_mutation(strict, expected): + tool = { + "type": "function", + "name": "click", + "parameters": { + "type": "object", + "properties": {"target": {"anyOf": [{"type": "string"}, {"type": "null"}]}}, + "required": [], + }, + } + if strict is not None: + tool["strict"] = strict + body = json.dumps( + {"tools": [tool], "input": "private prompt", "max_output_tokens": 2048} + ).encode() + result = request_constraints(body) + assert result["max_output_tokens"] == 2048 + (observed,) = result["tools"] + assert observed["strict"] == expected + assert observed["schema"]["optional_properties"] == observed["schema"]["open_objects"] == 1 + assert "strict" not in tool if strict is None else tool["strict"] is strict + assert "private" not in json.dumps(result) + assert json.loads(body)["tools"][0] == tool + + +def test_constraints_bound_large_requests_namespaces_and_deep_schemas(): + assert request_constraints(b"x" * (MAX_REQUEST_BYTES + 1)) == {"observation": "body_too_large"} + assert request_constraints(b"no JSON") == {"observation": "invalid_json"} + tools = [ + { + "type": "namespace", + "name": "functions", + "tools": [{"type": "function", "name": "click"}], + }, + {"type": "function", "name": "drag", "strict": False}, + ] + result = tool_constraints(tools) + assert {t["name"]: t["strict"] for t in result["tools"]} == { + "click": "omitted", + "drag": "false", + } + schema = {"type": "object", "properties": {str(i): {"type": "string"} for i in range(2000)}} + result = tool_constraints([{"type": "function", "name": "wide", "parameters": schema}]) + assert result["tools"][0]["schema"]["truncated"] is True + assert result["tools"][0]["schema"]["nodes"] <= 1024 diff --git a/tests/unit/test_responses_wire_lifecycle.py b/tests/unit/test_responses_wire_lifecycle.py new file mode 100644 index 0000000..3096893 --- /dev/null +++ b/tests/unit/test_responses_wire_lifecycle.py @@ -0,0 +1,323 @@ +import asyncio +from copy import deepcopy +import json + +import httpx2 as httpx +from langchain_core.messages import AIMessageChunk +from langchain_core.callbacks import AsyncCallbackHandler +from pydantic import SecretStr +import pytest +import pytest_asyncio + +from artemis.llm.diagnostics import ACTIVE_ATTEMPT +from artemis.llm.reliability import LLMExhaustedError +from artemis.llm.router import ModelEndpoint, ModelFactory +from artemis.mcp.action_specs import tool_declaration +from artemis.runtime.cancellation import CancellationReason +from artemis.services import llm + + +def event(kind, sequence, **fields): + return ( + "data: " + json.dumps({"type": kind, "sequence_number": sequence, **fields}) + "\n\n" + ).encode() + + +def initial_events(): + return event( + "response.created", + 0, + response={ + "id": "resp_test", + "object": "response", + "created_at": 1, + "status": "in_progress", + "model": "test-model", + "output": [], + }, + ) + event( + "response.output_item.added", + 1, + output_index=0, + item={ + "type": "function_call", + "id": "fc_test", + "call_id": "call_test", + "name": "action_sequence", + "arguments": "", + "status": "in_progress", + }, + ) + + +class Stream(httpx.AsyncByteStream): + def __init__(self): + self.fragments = [ + event( + "response.function_call_arguments.delta", + 2, + output_index=0, + item_id="fc_test", + delta='{"steps":[', + ) + ] + self.stalled = asyncio.Event() + self.close_calls = 0 + + async def __aiter__(self): + yield initial_events() + for fragment in self.fragments: + yield fragment + self.stalled.set() + await asyncio.Event().wait() + + async def aclose(self): + self.close_calls += 1 + + +class Engine: + current_session_id = "wire-lifecycle" + current_step_id = None + + def __init__(self): + self.records = [] + self.retrying = asyncio.Event() + + def record_trace(self, **record): + self.records.append(deepcopy(record)) + if record.get("payload", {}).get("phase") == "retry_wait": + self.retrying.set() + + def _publish(self, *_): + pass + + @property + def data(self): + return next(r["payload"] for r in reversed(self.records) if r["type"] == "llm_diagnostic") + + +@pytest_asyncio.fixture +async def case(monkeypatch, tmp_path): + import openai + from langchain_openai import ChatOpenAI + + stream, engine, clients, requests = Stream(), Engine(), [], [] + + def handle(request): + requests.append(json.loads(request.content)) + return httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=stream) + + def client(**kwargs): + value = httpx.AsyncClient(transport=httpx.MockTransport(handle), **kwargs) + clients.append(value) + return value + + monkeypatch.setattr(openai, "DefaultAsyncHttpxClient", client) + monkeypatch.setattr(llm.settings, "OPENAI_API_KEY", SecretStr("test-secret")) + monkeypatch.setattr(llm.settings, "OPENAI_BASE_URL", "https://responses.example/v1") + monkeypatch.delenv("ARTEMIS_FAKE_LLM", raising=False) + monkeypatch.setattr(llm, "_get_current_data_engine", lambda: engine) + monkeypatch.setattr(llm, "PAUSE_FILE", tmp_path / "pause") + llm._ENDPOINT_BREAKER._states.clear() + llm._NON_STREAMING_ENDPOINTS.clear() + fallback = llm._FALLBACK_AVAILABLE.set(True) + + def wrapper(timeout=30, max_retries=0): + endpoint = ModelEndpoint( + model_name="test-model", timeout_seconds=timeout, max_retries=max_retries + ) + model = ModelFactory.create_model(endpoint) + assert isinstance(model, ChatOpenAI) + return llm.RobustChatModelWrapper(model, endpoint=endpoint).bind_tools( + [tool_declaration(n) for n in ("click", "action_sequence", "drag")] + ) + + yield stream, engine, requests, wrapper + llm._FALLBACK_AVAILABLE.reset(fallback) + for value in clients: + await value.aclose() + llm._ENDPOINT_BREAKER._states.clear() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reason", ["outer_timeout", "request_timeout", "user_stop"]) +async def test_unfinished_arguments_below_length_limit_close_on_all_deadlines(case, reason): + stream, engine, requests, wrapper = case + if reason == "outer_timeout": + running = asyncio.create_task( + llm.invoke_llm_with_timeout_message( + wrapper().complete([]), timeout_seconds=0.1, hard_timeout=0.1 + ) + ) + elif reason == "request_timeout": + running = asyncio.create_task(wrapper(0.1).complete([])) + else: + running = asyncio.create_task(wrapper().complete([])) + await asyncio.wait_for(stream.stalled.wait(), 2) + running.cancel(CancellationReason.USER_STOP) + with pytest.raises((TimeoutError, LLMExhaustedError, asyncio.CancelledError)): + await asyncio.wait_for(running, 3) + data = engine.data + assert data["tool_argument_chars"] == len('{"steps":[') < 16384 + assert ( + data["failure_phase"] == "generating_tool_arguments" + if reason == "request_timeout" + else data["cancel_phase"] == "generating_tool_arguments" + ) + assert ( + data["timeout_scope"] == ("outer" if reason == "outer_timeout" else "request") + if reason != "user_stop" + else data["cancel_reason"] == "user_stop" + ) + assert data["cancel_observed_at"] <= data["ended_at"] + (wire,) = data["wire_responses"] + assert wire["closed_at"] >= wire["close_started_at"] + assert wire["body_eof_observed"] is False + assert wire["tools"][0].get("arguments_done_events", 0) == 0 + assert stream.close_calls == len(requests) == 1 + assert ACTIVE_ATTEMPT.get() is None + assert "test-secret" not in repr(data) + assert '{"steps":[' not in repr(data) + assert not llm.PAUSE_FILE.exists() + + +@pytest.mark.asyncio +async def test_overflow_closes_transport_even_while_consumer_is_between_sdk_yields(case): + stream, engine, requests, wrapper = case + stream.fragments = [ + event( + "response.function_call_arguments.delta", + 2, + output_index=0, + item_id="fc_test", + delta=" " * 17000, + ) + ] + with pytest.raises(LLMExhaustedError): + await wrapper().complete([]) + assert engine.data["stream_guard"]["tool"] == "action_sequence" + assert engine.data["wire_responses"][0]["closed_at"] + assert stream.close_calls == len(requests) == 1 + assert not stream.stalled.is_set() + + +@pytest.mark.asyncio +async def test_overflow_closes_transport_before_waiting_for_retry(case): + stream, engine, requests, wrapper = case + stream.fragments = [ + event( + "response.function_call_arguments.delta", + 2, + output_index=0, + item_id="fc_test", + delta=" " * 17000, + ) + ] + running = asyncio.create_task(wrapper(max_retries=1).complete([])) + await asyncio.wait_for(engine.retrying.wait(), 2) + assert stream.close_calls == len(requests) == 1 + assert engine.data["wire_responses"][0]["closed_at"] + running.cancel(CancellationReason.USER_STOP) + with pytest.raises(asyncio.CancelledError): + await running + assert engine.data["cancel_phase"] == "retry_wait" + + +@pytest.mark.asyncio +async def test_actual_native_schemas_do_not_silently_force_upstream_strict(case): + stream, engine, requests, wrapper = case + running = asyncio.create_task(wrapper().complete([])) + await asyncio.wait_for(stream.stalled.wait(), 2) + running.cancel(CancellationReason.USER_STOP) + with pytest.raises(asyncio.CancelledError): + await running + constraints = engine.data["wire_request_constraints"] + assert constraints["max_output_tokens"] is None + assert {t["name"]: t["strict"] for t in constraints["tools"]} == { + "click": "omitted", + "action_sequence": "omitted", + "drag": "omitted", + } + assert any(t["schema"]["optional_properties"] for t in constraints["tools"]) + assert all("strict" not in tool for tool in requests[0]["tools"]) + + +@pytest.mark.asyncio +async def test_live_deltas_are_emitted_before_a_batch_is_flushed(monkeypatch): + emitted, resume = asyncio.Event(), asyncio.Event() + events = [] + + class Model: + async def astream(self, *_args, **_kwargs): + yield AIMessageChunk(content="first live text") + await resume.wait() + + wrapper = llm.RobustChatModelWrapper(Model(), endpoint=ModelEndpoint(model_name="test")) + monkeypatch.setattr(wrapper, "_emit_stream_delta", lambda *_: emitted.set()) + monkeypatch.setattr(llm, "_record_llm_event", lambda *args: events.append(args)) + running = asyncio.create_task(wrapper._complete_attempt(True, [])) + await asyncio.wait_for(emitted.wait(), 2) + assert not running.done() + running.cancel() + with pytest.raises(asyncio.CancelledError): + await running + assert any(name == "llm_stream_reset" for name, *_ in events) + + +@pytest.mark.asyncio +async def test_callback_backpressure_can_leave_wire_ahead_of_sdk_without_prefetch(case): + stream, engine, requests, wrapper = case + blocked = asyncio.Event() + stream.fragments = [ + b"".join( + event( + "response.function_call_arguments.delta", + n + 2, + output_index=0, + item_id="fc_test", + delta=" ", + ) + for n in range(200) + ) + ] + + class Callback(AsyncCallbackHandler): + async def on_llm_new_token(self, token, *, chunk=None, **kwargs): + if chunk is not None and any(t.get("args") for t in chunk.message.tool_call_chunks): + attempt = ACTIVE_ATTEMPT.get() + attempt.publish() + blocked.set() + await asyncio.Event().wait() + + running = asyncio.create_task(wrapper().complete([], config={"callbacks": [Callback()]})) + await asyncio.wait_for(blocked.wait(), 2) + wire = engine.data["wire_responses"][0] + assert wire["tools"][0]["argument_chars"] == 200 + assert engine.data["tool_argument_chars"] == 0 + assert not stream.stalled.is_set() # no second transport read or background reader + running.cancel(CancellationReason.USER_STOP) + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(running, 2) + assert stream.close_calls == len(requests) == 1 + assert engine.data["wire_responses"][0]["closed_at"] + + +@pytest.mark.asyncio +async def test_unresponsive_close_keeps_user_stop_and_records_cleanup_failure(case, monkeypatch): + stream, engine, requests, wrapper = case + + async def stuck_close(): + stream.close_calls += 1 + await asyncio.Event().wait() + + monkeypatch.setattr(stream, "aclose", stuck_close) + running = asyncio.create_task(wrapper().complete([])) + await asyncio.wait_for(stream.stalled.wait(), 2) + running.cancel(CancellationReason.USER_STOP) + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(running, 3) + assert engine.data["cancel_reason"] == "user_stop" + wire = engine.data["wire_responses"][0] + assert wire["close_error_type"] == "TimeoutError" + assert "closed_at" not in wire + assert stream.close_calls == len(requests) == 1