diff --git a/.github/workflows/windows-nsis-upgrade-regression.yml b/.github/workflows/windows-nsis-upgrade-regression.yml index fdc7570a65..0468a62d9d 100644 --- a/.github/workflows/windows-nsis-upgrade-regression.yml +++ b/.github/workflows/windows-nsis-upgrade-regression.yml @@ -395,6 +395,7 @@ jobs: ${{ runner.temp }}/nsis-packaged-proof/first-interaction/first-send.log ${{ runner.temp }}/security-first-send-trace.zip ${{ runner.temp }}/security-first-send/logs/desktop.log* + ${{ runner.temp }}/security-first-send/logs/gateway.log* ${{ runner.temp }}/security-first-send/first-send-main-console.jsonl candidate/audit-candidate.json if-no-files-found: warn diff --git a/docs/configuration.md b/docs/configuration.md index 0c39382b76..bc6bd749c4 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -271,6 +271,25 @@ OpenSquilla also carries provider registry entries for additional OpenAI-compatible or self-hosted backends. Use `opensquilla providers list` on your install to see the current catalog. +Outgoing model requests have a separate JSON body limit of 24 MiB +(25,165,824 bytes) by default. Set +`OPENSQUILLA_PROVIDER_REQUEST_MAX_BODY_BYTES` in the Gateway environment to +change the limit, or set it to `0` to disable this guard, then restart the +Gateway. This counts the UTF-8 bytes actually sent to the provider, including +base64 image data. The model's token budget measures a different limit: an +image can fit its estimated token allowance while making the HTTP body too +large. When the byte limit is exceeded, OpenSquilla attempts safe request-text +compaction and otherwise reports a budget error before sending the request. +An upstream service may enforce a lower body limit. + +Tool-produced images are prepared before model delivery, with each image +limited to 2048 pixels on its longest edge and 1.5 MiB of base64 data. If an +tool-image batch is larger than four images or 6 MiB, the model request +includes at most four recent usable images and explicitly states how many were +omitted. This applies to the first request after a tool returns images and to +later history replay. The request-only reduction leaves the stored session +unchanged; omitted images must be loaded again when their details are needed. + Read: [`providers-and-models.md`](providers-and-models.md) ## Video Generation diff --git a/opensquilla-webui/src/locales/de.json b/opensquilla-webui/src/locales/de.json index 094e9af6f8..5df16ffc8f 100644 --- a/opensquilla-webui/src/locales/de.json +++ b/opensquilla-webui/src/locales/de.json @@ -4077,6 +4077,8 @@ "credits": "Das Guthaben des Modelldienstkontos reicht nicht aus.", "modelUnavailable": "Das gewählte Modell ist nicht verfügbar.", "contextLimit": "Der Inhalt überschreitet die Modellkapazität. Kürzen Sie die Eingabe oder wählen Sie ein Modell mit größerem Kontextfenster.", + "requestTooLarge": "Der Modelldienst hat diese Anfrage wegen ihrer Größe abgelehnt. Entfernen Sie große Anhänge oder beginnen Sie einen neuen Chat mit einer kurzen Zusammenfassung.", + "outboundRequestTooLarge": "OpenSquilla hat die Anfrage gestoppt, weil der zu sendende Inhalt die Größenbegrenzung überschreitet. Entfernen Sie große Anhänge oder beginnen Sie einen neuen Chat mit einer kurzen Zusammenfassung.", "unsupported": "Dieses Modell unterstützt die Anfrage nicht.", "imageUnsupported": "Dieses Modell unterstützt keine Bilder.", "ensembleImageUnsupported": "Dieser Modus unterstützt keine Bilder.", diff --git a/opensquilla-webui/src/locales/en.json b/opensquilla-webui/src/locales/en.json index 4fadb3a417..f72df0e03f 100644 --- a/opensquilla-webui/src/locales/en.json +++ b/opensquilla-webui/src/locales/en.json @@ -4180,6 +4180,8 @@ "credits": "The model service account has insufficient credits.", "modelUnavailable": "The selected model is unavailable.", "contextLimit": "The content exceeds this model's capacity. Shorten the input or choose a larger-context model.", + "requestTooLarge": "The model service rejected this request as too large. Reduce large attachments or start a new chat with a short summary before trying again.", + "outboundRequestTooLarge": "OpenSquilla stopped this request because the content to send exceeds its size limit. Reduce large attachments or start a new chat with a short summary before trying again.", "unsupported": "This model does not support this request.", "imageUnsupported": "This model does not support images.", "ensembleImageUnsupported": "This mode does not support images.", diff --git a/opensquilla-webui/src/locales/es.json b/opensquilla-webui/src/locales/es.json index 694ebba83e..62b64f92e0 100644 --- a/opensquilla-webui/src/locales/es.json +++ b/opensquilla-webui/src/locales/es.json @@ -4077,6 +4077,8 @@ "credits": "La cuenta del servicio del modelo no tiene saldo suficiente.", "modelUnavailable": "El modelo seleccionado no está disponible.", "contextLimit": "El contenido supera la capacidad del modelo. Acorta la entrada o elige un modelo con más contexto.", + "requestTooLarge": "El servicio del modelo rechazó esta solicitud por su tamaño. Reduce los archivos adjuntos grandes o inicia un chat nuevo con un resumen breve antes de volver a intentarlo.", + "outboundRequestTooLarge": "OpenSquilla detuvo esta solicitud porque el contenido que se iba a enviar supera su límite de tamaño. Reduce los archivos adjuntos grandes o inicia un chat nuevo con un resumen breve.", "unsupported": "Este modelo no admite esta solicitud.", "imageUnsupported": "Este modelo no admite imágenes.", "ensembleImageUnsupported": "Este modo no admite imágenes.", diff --git a/opensquilla-webui/src/locales/fr.json b/opensquilla-webui/src/locales/fr.json index 81171f7b72..f184184ceb 100644 --- a/opensquilla-webui/src/locales/fr.json +++ b/opensquilla-webui/src/locales/fr.json @@ -4077,6 +4077,8 @@ "credits": "Le solde du compte du service de modèle est insuffisant.", "modelUnavailable": "Le modèle sélectionné est indisponible.", "contextLimit": "Le contenu dépasse la capacité du modèle. Raccourcissez l’entrée ou choisissez un modèle avec un contexte plus grand.", + "requestTooLarge": "Le service du modèle a rejeté cette requête car elle est trop volumineuse. Réduisez les pièces jointes volumineuses ou démarrez une nouvelle conversation avec un bref résumé avant de réessayer.", + "outboundRequestTooLarge": "OpenSquilla a arrêté cette requête car le contenu à envoyer dépasse sa limite de taille. Réduisez les pièces jointes volumineuses ou démarrez une nouvelle conversation avec un bref résumé.", "unsupported": "Ce modèle ne prend pas en charge cette demande.", "imageUnsupported": "Ce modèle ne prend pas en charge les images.", "ensembleImageUnsupported": "Ce mode ne prend pas en charge les images.", diff --git a/opensquilla-webui/src/locales/ja.json b/opensquilla-webui/src/locales/ja.json index 1e4bff0db5..1a5fc59512 100644 --- a/opensquilla-webui/src/locales/ja.json +++ b/opensquilla-webui/src/locales/ja.json @@ -4077,6 +4077,8 @@ "credits": "モデルサービスのアカウント残高が不足しています。", "modelUnavailable": "選択したモデルは利用できません。", "contextLimit": "内容がモデルの容量を超えています。入力を短くするか、コンテキストの大きいモデルを選んでください。", + "requestTooLarge": "モデルサービスがリクエストのサイズを理由に拒否しました。大きな添付ファイルを減らすか、新しいチャットで短い要約を伝えてから再試行してください。", + "outboundRequestTooLarge": "送信する内容がサイズ制限を超えるため、OpenSquilla がリクエストを停止しました。大きな添付ファイルを減らすか、新しいチャットで短い要約を伝えてください。", "unsupported": "このモデルはこのリクエストに対応していません。", "imageUnsupported": "このモデルは画像に対応していません。", "ensembleImageUnsupported": "このモードは画像に対応していません。", diff --git a/opensquilla-webui/src/locales/zh-Hans.json b/opensquilla-webui/src/locales/zh-Hans.json index 615415ccdc..3fdd8a4c72 100644 --- a/opensquilla-webui/src/locales/zh-Hans.json +++ b/opensquilla-webui/src/locales/zh-Hans.json @@ -4179,6 +4179,8 @@ "credits": "模型服务账户余额不足", "modelUnavailable": "当前模型不可用", "contextLimit": "内容超出当前模型容量，请缩短输入或选择更大上下文的模型", + "requestTooLarge": "模型服务拒绝了过大的请求。请减少大附件，或在新对话中提供简短摘要后再试。", + "outboundRequestTooLarge": "OpenSquilla 已拦截超出发送大小限制的请求。请减少大附件，或在新对话中提供简短摘要后再试。", "unsupported": "当前模型不支持此请求", "imageUnsupported": "当前模型不支持图片", "ensembleImageUnsupported": "当前模式不支持图片", diff --git a/opensquilla-webui/src/utils/chat/chatErrorPresentation.test.ts b/opensquilla-webui/src/utils/chat/chatErrorPresentation.test.ts index c3d11e66f6..f4f9d11326 100644 --- a/opensquilla-webui/src/utils/chat/chatErrorPresentation.test.ts +++ b/opensquilla-webui/src/utils/chat/chatErrorPresentation.test.ts @@ -14,6 +14,8 @@ describe('Chat user error presentation', () => { ['provider_overloaded', 'busy'], ['auth_invalid', 'credentials', 'open-provider-settings'], ['context_overflow', 'contextLimit', 'choose-model'], + ['payload_too_large', 'requestTooLarge'], + ['wire_payload_too_large', 'outboundRequestTooLarge'], ['unsupported_feature', 'unsupported', 'choose-model'], ['insufficient_credits', 'credits', 'open-provider-settings'], ['model_not_found', 'modelUnavailable', 'choose-model'], @@ -35,6 +37,8 @@ describe('Chat user error presentation', () => { it.each([ ['no_provider', 'noProvider', 'open-model-settings'], ['provider_request_too_large', 'contextLimit', 'choose-model'], + ['provider_payload_too_large', 'requestTooLarge', undefined], + ['provider_wire_payload_too_large', 'outboundRequestTooLarge', undefined], ['provider_request_budget_exhausted', 'contextLimit', 'choose-model'], ['current_turn_context_exhausted', 'contextLimit', 'choose-model'], ['attachment_capacity_too_large', 'contextLimit', 'choose-model'], @@ -129,7 +133,7 @@ describe('Chat user error presentation', () => { expect(chatErrorPresentation({ terminalStatus: 'timeout', code: 'provider_retry_after_deadline' })).toEqual({ messageKey: 'chat.errorMessage.timeout' }) expect(chatErrorPresentation({ failureKind: 'insufficient_credits' })).toEqual({ messageKey: 'chat.errorMessage.credits', action: 'open-provider-settings' }) expect(chatErrorPresentation({ failureKind: 'usage_limit' })).toEqual({ messageKey: 'chat.errorMessage.usageLimit' }) - for (const code of ['400', '401', '403', '404', '429', '503']) expect(chatErrorPresentation({ code })).toEqual({ messageKey: 'chat.errorMessage.unknown' }) + for (const code of ['400', '401', '403', '404', '413', '429', '503']) expect(chatErrorPresentation({ code })).toEqual({ messageKey: 'chat.errorMessage.unknown' }) }) it('never uses a general retry hint to offer whole-turn replay', () => { diff --git a/opensquilla-webui/src/utils/chat/chatErrorPresentation.ts b/opensquilla-webui/src/utils/chat/chatErrorPresentation.ts index 0dd8cde742..0d15dbbc59 100644 --- a/opensquilla-webui/src/utils/chat/chatErrorPresentation.ts +++ b/opensquilla-webui/src/utils/chat/chatErrorPresentation.ts @@ -69,6 +69,8 @@ const PROVIDER_PRESENTATIONS = new Map([ ['provider_overloaded', presentation('busy')], ['auth_invalid', presentation('credentials', 'open-provider-settings')], ['context_overflow', presentation('contextLimit', 'choose-model')], + ['payload_too_large', presentation('requestTooLarge')], + ['wire_payload_too_large', presentation('outboundRequestTooLarge')], ['unsupported_feature', presentation('unsupported', 'choose-model')], ['insufficient_credits', presentation('credits', 'open-provider-settings')], ['usage_limit', presentation('usageLimit')], @@ -117,6 +119,8 @@ export function chatErrorPresentation(input: ChatErrorPresentationInput): ChatEr } if (causes.includes('ensemble_multimodal_unsupported')) return presentation('ensembleImageUnsupported', 'choose-model') if (causes.includes('image_input_unsupported')) return presentation('imageUnsupported', 'choose-model') + if (causes.includes('provider_wire_payload_too_large')) return presentation('outboundRequestTooLarge') + if (causes.includes('provider_payload_too_large')) return presentation('requestTooLarge') if (causes.some(value => CONTEXT_LIMITS.has(value))) return presentation('contextLimit', 'choose-model') if (causes.some(value => ['attachment_capacity_unknown', 'attachment_capacity_unavailable'].includes(value))) { return presentation('modelUnavailable', 'open-model-settings') diff --git a/opensquilla-webui/src/utils/chat/providerFailure.test.ts b/opensquilla-webui/src/utils/chat/providerFailure.test.ts index 88e06b1c0b..8c426accd4 100644 --- a/opensquilla-webui/src/utils/chat/providerFailure.test.ts +++ b/opensquilla-webui/src/utils/chat/providerFailure.test.ts @@ -14,6 +14,19 @@ describe('provider terminal metadata', () => { expect(normalizeTurnOutcome({ ...base, [key]: fields })).toEqual(normalizeTurnOutcome({ ...base, ...fields })) }) + it('keeps remote and local request-size causes distinct in persisted outcomes', () => { + const remote = normalizeTurnOutcome({ turn_id: 'remote', failure_kind: 'payload_too_large' }) + const local = normalizeTurnOutcome({ turn_id: 'local', failure_kind: 'wire_payload_too_large' }) + expect(remote?.failureKind).toBe('payload_too_large') + expect(local?.failureKind).toBe('wire_payload_too_large') + expect(chatErrorPresentation({ failureKind: remote?.failureKind })).toEqual({ + messageKey: 'chat.errorMessage.requestTooLarge', + }) + expect(chatErrorPresentation({ failureKind: local?.failureKind })).toEqual({ + messageKey: 'chat.errorMessage.outboundRequestTooLarge', + }) + }) + it.each(['ABCDEF01', 'abcdef0', 'abcdef012', 'abcdef01 ', 123, null, ''])('rejects invalid id %s', errorId => { const outcome = normalizeTurnOutcome({ turn_id: 't', error_id: errorId }) expect(outcome?.errorId).toBeNull() @@ -38,7 +51,7 @@ describe('provider terminal metadata', () => { it.each(['en', 'zh-Hans', 'de', 'es', 'fr', 'ja'] as const)('localizes allowlisted causes in %s', async locale => { await loadLocaleMessages(locale) i18n.global.locale.value = locale - for (const kind of ['rate_limited', 'provider_overloaded', 'auth_invalid', 'context_overflow', + for (const kind of ['rate_limited', 'provider_overloaded', 'auth_invalid', 'context_overflow', 'payload_too_large', 'wire_payload_too_large', 'unsupported_feature', 'insufficient_credits', 'model_not_found', 'transport_transient', 'policy_refusal', 'empty_response', 'malformed_response', 'bad_request'] as const) { const message = localizedChatErrorMessage('429', 'safe fallback', false, kind) @@ -97,6 +110,8 @@ describe('provider terminal metadata', () => { for (const [code, key] of [ ['timeout', 'timeout'], ['llm_timeout', 'timeout'], ['provider_output_truncated', 'responseIncomplete'], ['provider_request_too_large', 'contextLimit'], + ['provider_payload_too_large', 'requestTooLarge'], + ['provider_wire_payload_too_large', 'outboundRequestTooLarge'], ]) { expect(localizedChatErrorMessage(code, 'specific guidance', false, 'rate_limited')) .toBe(i18n.global.t(`chat.errorMessage.${key}`)) diff --git a/opensquilla-webui/src/utils/chat/providerFailure.ts b/opensquilla-webui/src/utils/chat/providerFailure.ts index 81b4a6e9cd..3e16905a47 100644 --- a/opensquilla-webui/src/utils/chat/providerFailure.ts +++ b/opensquilla-webui/src/utils/chat/providerFailure.ts @@ -1,6 +1,7 @@ /** Stable public taxonomy only; never infer a provider cause from prose or HTTP codes. */ const FAILURE_KINDS = new Set([ - 'rate_limited', 'provider_overloaded', 'auth_invalid', 'context_overflow', + 'rate_limited', 'provider_overloaded', 'auth_invalid', 'context_overflow', 'payload_too_large', + 'wire_payload_too_large', 'unsupported_feature', 'insufficient_credits', 'usage_limit', 'usage_limit_reached', 'model_not_found', 'transport_transient', 'policy_refusal', 'empty_response', 'malformed_response', 'bad_request', ]) @@ -17,6 +18,7 @@ export function diagnosticErrorId(value: unknown): string | undefined { const SPECIAL_TERMINALS = new Set([ 'timeout', 'llm_timeout', 'iteration_timeout', 'stream_idle_timeout', 'provider_output_truncated', 'provider_request_too_large', 'provider_request_budget_exhausted', 'current_turn_context_exhausted', + 'provider_payload_too_large', 'provider_wire_payload_too_large', 'model_repetition_loop_detected', 'tool_run_budget_exhausted', 'llm_budget_exhausted', 'turn_llm_call_budget_exceeded', 'turn_input_token_budget_exceeded', 'turn_output_token_budget_exceeded', 'turn_billed_cost_budget_exceeded', 'max_iterations', diff --git a/src/opensquilla/contracts/image_preparation.py b/src/opensquilla/contracts/image_preparation.py new file mode 100644 index 0000000000..5b0830fc4d --- /dev/null +++ b/src/opensquilla/contracts/image_preparation.py @@ -0,0 +1,154 @@ +"""Bound a tool image before it becomes provider-visible replay content.""" + +from __future__ import annotations + +import io +from dataclasses import dataclass + +from PIL import Image, ImageOps, features + +from opensquilla.contracts.attachments import normalize_attachment_mime +from opensquilla.contracts.image_validation import validate_image_bytes + +MAX_TOOL_IMAGE_DIMENSION = 2048 +MAX_TOOL_IMAGE_BASE64_CHARS = 1536 * 1024 +_LOSSY_QUALITIES = (88, 76, 64, 52, 40, 28) + + +@dataclass(frozen=True, slots=True) +class PreparedToolImage: + """Provider image bytes and the facts needed to describe any transformation.""" + + data: bytes + media_type: str + changed: bool + lossy: bool + original_size: int + encoded_chars: int + + +def _encoded_chars(payload: bytes) -> int: + return ((len(payload) + 2) // 3) * 4 + + +def _encode(image: Image.Image, *, format_name: str, **options: object) -> bytes: + output = io.BytesIO() + image.save(output, format=format_name, **options) + return output.getvalue() + + +def _has_alpha(image: Image.Image) -> bool: + return image.mode in {"RGBA", "LA", "PA"} or "transparency" in image.info + + +def _lossless_candidate(image: Image.Image, mime: str) -> tuple[bytes, str, bool]: + if mime == "image/jpeg": + return ( + _encode( + image.convert("RGB"), format_name="JPEG", quality=92, subsampling=0, optimize=True + ), + mime, + True, + ) + if mime == "image/webp": + mode = "RGBA" if _has_alpha(image) else "RGB" + return ( + _encode(image.convert(mode), format_name="WEBP", lossless=True, method=6), + mime, + False, + ) + # A transformed single-frame GIF is saved as PNG so palette quantization + # does not silently discard colors or transparency. + mode = "RGBA" if _has_alpha(image) else "RGB" + return _encode(image.convert(mode), format_name="PNG", optimize=True), "image/png", False + + +def _lossy_candidates(image: Image.Image) -> tuple[str, Image.Image]: + if _has_alpha(image): + return "image/webp", image.convert("RGBA") + return "image/jpeg", image.convert("RGB") + + +def prepare_tool_result_image(payload: bytes, mime: str) -> PreparedToolImage: + """Return a valid image bounded for a single model tool result. + + Source bytes are never mutated. A caller with a readable source path may + retain that path as the original while placing these prepared bytes in the + provider replay. Animated images remain unchanged or fail explicitly; a + first-frame conversion would misrepresent what the model received. + """ + + normalized_mime = normalize_attachment_mime(mime) + if normalized_mime is None: + raise ValueError("tool image media type is missing") + validate_image_bytes(payload, normalized_mime) + with Image.open(io.BytesIO(payload)) as source: + dimensions = source.size + needs_resize = max(dimensions) > MAX_TOOL_IMAGE_DIMENSION + if not needs_resize and _encoded_chars(payload) <= MAX_TOOL_IMAGE_BASE64_CHARS: + return PreparedToolImage( + data=payload, + media_type=normalized_mime, + changed=False, + lossy=False, + original_size=len(payload), + encoded_chars=_encoded_chars(payload), + ) + if getattr(source, "n_frames", 1) > 1: + raise ValueError("animated tool image exceeds the request image limit") + image = ImageOps.exif_transpose(source) + image.load() + if needs_resize: + image.thumbnail( + (MAX_TOOL_IMAGE_DIMENSION, MAX_TOOL_IMAGE_DIMENSION), + Image.Resampling.LANCZOS, + ) + + try: + lossless, lossless_mime, lossless_is_lossy = _lossless_candidate(image, normalized_mime) + except (KeyError, OSError, ValueError) as exc: + raise ValueError("tool image could not be prepared within the request limit") from exc + if _encoded_chars(lossless) <= MAX_TOOL_IMAGE_BASE64_CHARS: + validate_image_bytes(lossless, lossless_mime) + return PreparedToolImage( + data=lossless, + media_type=lossless_mime, + changed=True, + lossy=lossless_is_lossy, + original_size=len(payload), + encoded_chars=_encoded_chars(lossless), + ) + + output_mime, converted = _lossy_candidates(image) + if output_mime == "image/webp" and not features.check("webp"): + raise ValueError("transparent tool image exceeds the limit and WebP encoder is unavailable") + format_name = "WEBP" if output_mime == "image/webp" else "JPEG" + try: + for _ in range(12): + for quality in _LOSSY_QUALITIES: + options: dict[str, object] = {"quality": quality} + if format_name == "JPEG": + options["optimize"] = True + else: + # Fast WebP encoding avoids a long event-loop stall when a + # tool returns several contact sheets in one batch. + options["method"] = 0 + candidate = _encode(converted, format_name=format_name, **options) + if _encoded_chars(candidate) <= MAX_TOOL_IMAGE_BASE64_CHARS: + validate_image_bytes(candidate, output_mime) + return PreparedToolImage( + data=candidate, + media_type=output_mime, + changed=True, + lossy=True, + original_size=len(payload), + encoded_chars=_encoded_chars(candidate), + ) + width, height = converted.size + smaller = (max(1, width * 4 // 5), max(1, height * 4 // 5)) + if smaller == converted.size: + break + converted = converted.resize(smaller, Image.Resampling.LANCZOS) + except (KeyError, OSError, ValueError) as exc: + raise ValueError("tool image could not be prepared within the request limit") from exc + raise ValueError("tool image could not be prepared within the request limit") diff --git a/src/opensquilla/engine/agent.py b/src/opensquilla/engine/agent.py index 9269c06612..76e463d0ba 100644 --- a/src/opensquilla/engine/agent.py +++ b/src/opensquilla/engine/agent.py @@ -7,6 +7,7 @@ import asyncio import base64 +import binascii import contextlib import copy import functools @@ -41,7 +42,10 @@ ) from opensquilla.context_budget import ContextBudgetClass, ContextBudgetGovernor from opensquilla.contracts.attachments import IMAGE_ATTACHMENT_MIMES, normalize_attachment_mime -from opensquilla.contracts.image_validation import validate_image_bytes +from opensquilla.contracts.image_preparation import ( + MAX_TOOL_IMAGE_BASE64_CHARS, + prepare_tool_result_image, +) from opensquilla.contracts.turn_execution import TurnExecutionContext from opensquilla.engine.agent_injection import PendingInputProvider from opensquilla.engine.cache_break_monitor import ( @@ -128,6 +132,7 @@ bind_usage_accounting_scope, current_usage_accounting_scope, has_known_provider_usage_receipt, + mark_usage_call_not_sent, normalize_provider_usage, provider_accounts_physical_usage, start_usage_call, @@ -1587,6 +1592,7 @@ def _active_user_message_index_for_request( for index in range(len(messages) - 1, -1, -1): message = messages[index] if (message.role != "user" or _message_has_tool_result(message) + or message._request_only_pending_tool_image or isinstance(message, _LiveTurnContinuationMessage)): continue content = ( @@ -1600,6 +1606,7 @@ def _active_user_message_index_for_request( for index in range(len(messages) - 1, -1, -1): message = messages[index] if (message.role != "user" or _message_has_tool_result(message) + or message._request_only_pending_tool_image or isinstance(message, _LiveTurnContinuationMessage)): continue content = ( @@ -1657,6 +1664,13 @@ def _assistant_replay_tail(messages: list[Message], start: int) -> list[Message] return tail[first:] +def _tool_image_native_ack_indexes(messages: list[Message]) -> list[int]: + return [ + index for index, message in enumerate(messages) + if message.role == "assistant" and message._tool_image_native_ack is True + ] + + def _serialize_assistant_replay(messages: list[Message]) -> dict[str, Any]: """Persist image provenance without adding it to provider wire messages.""" serialized = [] @@ -1909,6 +1923,14 @@ def _safe_provider_terminal_message( ProviderFailureKind.CONTEXT_OVERFLOW: ( "The request exceeds the model provider's context window." ), + ProviderFailureKind.PAYLOAD_TOO_LARGE: ( + "The model provider rejected the request size. Reduce or remove attachments " + "and try again." + ), + ProviderFailureKind.WIRE_PAYLOAD_TOO_LARGE: ( + "OpenSquilla could not send this request within its size limit. " + "Reduce or remove attachments and try again." + ), ProviderFailureKind.UNSUPPORTED_FEATURE: ( "The model provider does not support this request." ), @@ -2253,6 +2275,171 @@ def _classify_provider_attempt( ) +def _clear_unverified_tool_image_retention(value: Any) -> None: + """Do not infer a readable original from a historical tool receipt alone.""" + if isinstance(value, ContentBlockImage): + value.durable_retained = None + elif isinstance(value, ContentBlockToolResult): + _clear_unverified_tool_image_retention(value.content) + elif isinstance(value, list): + for item in value: + _clear_unverified_tool_image_retention(item) + elif isinstance(value, dict): + block_type = str(value.get("type") or "").strip().lower() + if block_type == "image": + value.pop("durable_retained", None) + elif block_type == "tool_result": + _clear_unverified_tool_image_retention(value.get("content")) + + +def _has_historical_tool_image(message: Message) -> bool: + if ( + not isinstance(message, Message) + or not isinstance(message.content, list) + or not count_projected_image_blocks([message]) + ): + return False + if any(isinstance(block, ContentBlockToolResult) for block in message.content): + return True + # Interrupted/unsigned tool chains are quoted as context during replay. + # Their private provenance follows request-local copies, unlike text that + # a user can reproduce in an image caption. + return message.is_recorded_tool_context + + +def _historical_tool_images(value: Any) -> Iterator[ContentBlockImage]: + if isinstance(value, ContentBlockImage): + yield value + elif isinstance(value, ContentBlockToolResult): + yield from _historical_tool_images(value.content) + elif isinstance(value, list): + for item in value: + yield from _historical_tool_images(item) + elif isinstance(value, dict): + if value.get("type") == "image": + try: + yield ContentBlockImage.model_validate(value) + except ValueError: + pass + elif value.get("type") == "tool_result": + yield from _historical_tool_images(value.get("content")) + + +_BOUNDED_TOOL_IMAGE_NOTE = ( + "[Recorded tool image moved to the bounded image context below; " + "the context states how many earlier images were omitted.]" +) + + +def _move_historical_tool_images_to_context(value: Any) -> Any: + """Replace only request-local pixels while keeping tool-result structure.""" + if isinstance(value, ContentBlockImage): + return ContentBlockText(text=_BOUNDED_TOOL_IMAGE_NOTE) + if isinstance(value, ContentBlockToolResult): + result = value.model_copy(deep=True) + result.content = _move_historical_tool_images_to_context(result.content) + return result + if isinstance(value, list): + return [_move_historical_tool_images_to_context(item) for item in value] + if isinstance(value, dict): + if value.get("type") == "image": + return {"type": "text", "text": _BOUNDED_TOOL_IMAGE_NOTE} + if value.get("type") == "tool_result": + copied_result = copy.deepcopy(value) + copied_result["content"] = _move_historical_tool_images_to_context( + value.get("content") + ) + return copied_result + return copy.deepcopy(value) + + +def _pending_tool_image_window_context(messages: list[Message]) -> Message | None: + """Carry a bounded unseen tool image tail across the ordinary turn window.""" + + # Keep the latest useful frames. Repeated failed turns must not accumulate + # unbounded image bodies in a rescued request prefix. + max_images = 4 + max_data_chars = 4 * MAX_TOOL_IMAGE_BASE64_CHARS + selected: list[ContentBlockImage] = [] + total_chars = 0 + candidates = [ + image + for message in messages + for image in _historical_tool_images(message.content) + ] + for image in reversed(candidates): + if len(selected) >= max_images: + break + if image.source_type != "base64": + continue + try: + prepared = prepare_tool_result_image( + base64.b64decode(image.data, validate=True), image.media_type + ) + except (ValueError, binascii.Error): + continue + encoded = base64.b64encode(prepared.data).decode("ascii") + data_chars = len(encoded) + if total_chars + data_chars > max_data_chars: + continue + selected.append(image.model_copy(update={ + "data": encoded, + "media_type": prepared.media_type, + "local_path": None, + "durable_retained": None, + }, deep=True)) + total_chars += data_chars + if not candidates: + return None + omitted_count = len(candidates) - len(selected) + context = Message( + role="user", + content=[ + ContentBlockText( + text=( + "[Recorded tool image context] " + f"{len(selected)} of {len(candidates)} tool images selected for " + "image-capable requests; " + f"{omitted_count} images omitted by recency or size limits. " + "Selected images may be resized or compressed, or removed for a model " + "without image support. Analyze only image blocks actually present below." + ) + ), + *(image.model_copy(deep=True) for image in reversed(selected)), + ], + ) + context._recorded_tool_context = True + context._request_only_pending_tool_image = True + return context + + +def _bound_provider_tool_image_batch(messages: list[Message]) -> list[Message]: + """Bound both new and historical tool images in the physical request view.""" + tool_messages = [message for message in messages if _has_historical_tool_image(message)] + images = [ + image for message in tool_messages + for image in _historical_tool_images(message.content) + ] + if ( + len(images) <= 4 + and all(len(image.data) <= MAX_TOOL_IMAGE_BASE64_CHARS for image in images) + and sum(len(image.data) for image in images) <= 4 * MAX_TOOL_IMAGE_BASE64_CHARS + ): + return messages + bounded: list[Message] = [] + for message in messages: + if _has_historical_tool_image(message): + detached = message.model_copy(deep=True) + detached.content = _move_historical_tool_images_to_context(detached.content) + bounded.append(detached) + else: + bounded.append(message) + context = _pending_tool_image_window_context(tool_messages) + if context is not None: + bounded.append(context) + return bounded + + def _strip_historical_image_blocks( messages: list[Message], *, @@ -2263,17 +2450,59 @@ def _strip_historical_image_blocks( Current-turn uploads arrive through ``extra_messages``. Model capability is checked separately on each physical request, including tool-loaded images. """ + # Tool images are delivered to the model immediately after the tool call. + # On later turns the recorded result remains useful, but replaying its + # pixels again can multiply the request body even when the user sends only + # text. Keep the canonical replay intact and omit only its provider view. + # A tool result can be persisted even when its following provider request + # fails. Keep the final unacknowledged batch available to the next vision + # request. New rows acknowledge only a committed reply to a native tool + # image request; older rows retain the prior substantive-reply heuristic. + pending_image_indexes: set[int] = set() if preserve_images: - # Keep the historical object graph intact for a vision-capable route. - # The outbound projection below still deep-copies it before a physical - # provider call, so callers cannot mutate the canonical transcript. - return messages + for index, message in enumerate(messages): + if message.role == "assistant" and ( + message._tool_image_native_ack is True + or ( + message._tool_image_native_ack is None + and not message.is_trusted_runtime_notice + and (_message_has_visible_text(message) or _message_has_tool_use(message)) + ) + ): + pending_image_indexes.clear() + elif ( + _has_historical_tool_image(message) + ): + pending_image_indexes.add(index) + + projected_history: list[Message] = [] + for index, message in enumerate(messages): + if ( + index not in pending_image_indexes + and _has_historical_tool_image(message) + ): + # A stored path alone does not prove that a local tool image still + # exists or remains readable. Do not claim durable retention in + # the marker without a verified attachment reference. + detached = message.model_copy(deep=True) + _clear_unverified_tool_image_retention(detached.content) + projected_history.extend(project_messages( + [detached], + mode=ImageProjectionMode.MARKER, + marker_state=ImageMarkerState.NOT_REREAD, + ).messages) + else: + projected_history.append(message) + if preserve_images: + # User images and the pending tool batch remain available to vision. + # Historical projection is request-local; canonical replay is unchanged. + return projected_history # Historical images are not silently deleted. Project them into a # truthful marker, recursively (including images nested in tool results), # while keeping the original transcript available for a later vision turn. return project_messages( - messages, + projected_history, mode=ImageProjectionMode.MARKER, marker_state=ImageMarkerState.NOT_REREAD, ).messages @@ -5798,6 +6027,16 @@ def current_assistant_replay(self) -> dict[str, Any] | None: return None return _serialize_assistant_replay(tail) + def current_tool_image_native_ack_indexes(self) -> list[int]: + """Snapshot committed image acknowledgements for cancellation persistence.""" + view: Callable[[], tuple[list[Message], int]] | None = getattr( + self, "_active_replay_view", None + ) + if view is None: + return [] + messages, start = view() + return _tool_image_native_ack_indexes(_assistant_replay_tail(messages, start)) + def _freeze_current_replay_view(self) -> None: """Release the live turn's history cells while keeping its accepted tail.""" view: Callable[[], tuple[list[Message], int]] | None = getattr( @@ -6189,7 +6428,42 @@ async def _turn_generator( historical_projection=historical_projection_result.__dict__, ) history = limit_turns(sanitized_history, self.config.max_history_turns) + dropped_history_count = len(sanitized_history) - len(history) + pending_tool_messages = [ + item for item in sanitized_history if _has_historical_tool_image(item) + ] if self.config.preserve_historical_images else [] + pending_tool_images = [ + image + for item in pending_tool_messages + for image in _historical_tool_images(item.content) + ] + bound_pending_batch = ( + len(pending_tool_images) > 4 + or any(len(image.data) > MAX_TOOL_IMAGE_BASE64_CHARS + for image in pending_tool_images) + or sum(len(image.data) for image in pending_tool_images) + > 4 * MAX_TOOL_IMAGE_BASE64_CHARS + ) + if bound_pending_batch: + bounded_history: list[Message] = [] + for item in history: + if _has_historical_tool_image(item): + detached = item.model_copy(deep=True) + detached.content = _move_historical_tool_images_to_context(detached.content) + bounded_history.append(detached) + else: + bounded_history.append(item) + history = bounded_history + pending_window_context = _pending_tool_image_window_context( + pending_tool_messages if bound_pending_batch else [ + item + for item in sanitized_history[:dropped_history_count] + if _has_historical_tool_image(item) + ] + ) if self.config.preserve_historical_images else None history = repair_tool_pairing(history) + if pending_window_context is not None: + history.append(pending_window_context) initial_provider_history = tuple(history) canonical_history = repair_tool_pairing( limit_turns( @@ -6197,6 +6471,8 @@ async def _turn_generator( self.config.max_history_turns, ) ) + if pending_window_context is not None: + canonical_history.append(pending_window_context) self._write_context_stage( "session:limited", history, @@ -7184,8 +7460,13 @@ def _finish_reasoning_block( ) rebound_config = None if callable(prepare_continuation): + bounded_admission_messages = await asyncio.to_thread( + _bound_provider_tool_image_batch, + request_turn_messages, + ) rebound_config = prepare_continuation( - request_turn_messages, chat_cfg + bounded_admission_messages, + chat_cfg, ) if inspect.isawaitable(rebound_config): rebound_config = await rebound_config @@ -7926,6 +8207,7 @@ def _notify_call_outcome(*, ok: bool, failure_kind: str = "") -> None: usage_call: UsageCallStart | None = None usage_call_terminal = False + usage_call_not_sent = False usage_unknown_reason = "provider_stream_ended_without_usage" if usage_scope is not None and not provider_accounts_physical_usage( self.provider @@ -8242,10 +8524,10 @@ def _active_stream_deadline() -> float | None: if isinstance(row, dict) ] turn_model_usage_breakdown.extend(terminal_rows) - for item in terminal_usage.items: + for usage_item in terminal_usage.items: if ( - item.cost_source == "provider_billed" - or item.billed_cost_nanos > 0 + usage_item.cost_source == "provider_billed" + or usage_item.billed_cost_nanos > 0 ): total_provider_billed_entries += 1 else: @@ -8253,16 +8535,16 @@ def _active_stream_deadline() -> float | None: if self._usage_tracker and self._session_key: self._usage_tracker.add( self._session_key, - input_tokens=item.input_tokens, - output_tokens=item.output_tokens, - model_id=item.model, - cache_read_tokens=item.cache_read_tokens, - cache_write_tokens=item.cache_write_tokens, + input_tokens=usage_item.input_tokens, + output_tokens=usage_item.output_tokens, + model_id=usage_item.model, + cache_read_tokens=usage_item.cache_read_tokens, + cache_write_tokens=usage_item.cache_write_tokens, billed_cost=( - item.billed_cost_nanos / 1_000_000_000 + usage_item.billed_cost_nanos / 1_000_000_000 ), - provider=item.provider, - cost_source=item.cost_source, + provider=usage_item.provider, + cost_source=usage_item.cost_source, ) _accumulate_turn_cost( raw_ev, @@ -9047,6 +9329,11 @@ def _active_stream_deadline() -> float | None: pending_tool_events.clear() usage_unknown_reason = provider_error_usage_reason(raw_ev.code) known_usage_receipt = has_known_provider_usage_receipt(raw_ev) + usage_call_not_sent = ( + getattr(raw_ev, "_opensquilla_request_not_dispatched", False) + is True + and not known_usage_receipt + ) error_usage: UsageCallResult | None = None if known_usage_receipt and not cost_receipt_counted: usage_default_provider = ( @@ -9158,6 +9445,7 @@ def _active_stream_deadline() -> float | None: elif ( raw_ev.tool_argument_rejection is not None and not known_usage_receipt + and not usage_call_not_sent ): # This completed, rejected generation is still # a physical request with unknown cost. Keep @@ -9347,10 +9635,13 @@ def _active_stream_deadline() -> float | None: finally: stream_close_evidence.reset(close_token) if usage_call is not None and not usage_call_terminal: - await self._usage_call_unknown( - usage_call, - usage_unknown_reason, - ) + if usage_call_not_sent and usage_scope is not None: + await mark_usage_call_not_sent(usage_scope, usage_call) + else: + await self._usage_call_unknown( + usage_call, + usage_unknown_reason, + ) if reasoning_only_act_now_for_call is not None and not bool( getattr( self.config, @@ -9361,7 +9652,7 @@ def _active_stream_deadline() -> float | None: reasoning_only_act_now_message = None if _got_error: - if not cost_receipt_counted and ( + if not cost_receipt_counted and not usage_call_not_sent and ( provider_error is None or provider_error.tool_argument_rejection is None ): @@ -12267,14 +12558,28 @@ def capacity_fallback_event() -> ProviderActivityEvent | None: if native_content is not None: assistant_content = native_content if assistant_content: - turn_messages.append( - Message( - role="assistant", - content=assistant_content, - reasoning_content=iter_reasoning_content, - provider_replay=iter_provider_replay, - ) + accepted_assistant = Message( + role="assistant", + content=assistant_content, + reasoning_content=iter_reasoning_content, + provider_replay=iter_provider_replay, ) + if any( + _has_historical_tool_image(item) + for item in canonical_request_messages + ): + native_tool_image_sent = ( + selector_image_provider.last_image_request_had_native_tool_images + if selector_projects_images + else any( + _has_historical_tool_image(item) + for item in request_messages + ) + ) + accepted_assistant._tool_image_native_ack = bool( + native_tool_image_sent + ) + turn_messages.append(accepted_assistant) # Detect incomplete tool calls (stream interrupted mid-generation) if pending_tools and not tool_calls: @@ -12362,7 +12667,8 @@ def _tool_result_images( 1 <= len(data) <= 16 * 1024 * 1024 ): raise ValueError("image is empty or exceeds the size limit") - validate_image_bytes(base64.b64decode(data, validate=True), mime) + payload = base64.b64decode(data, validate=True) + prepared = prepare_tool_result_image(payload, mime) attachment_id = item.get("attachment_id") retained_metadata = { key: value for key in ("name", "local_path", "source_url") @@ -12370,8 +12676,8 @@ def _tool_result_images( } images.append( ContentBlockImage( - media_type=mime, - data=data, + media_type=prepared.media_type, + data=base64.b64encode(prepared.data).decode("ascii"), attachment_id=( attachment_id if isinstance(attachment_id, str) and attachment_id @@ -13424,6 +13730,10 @@ async def _run_after_key_lock() -> ToolResult: *turn_messages[:current_turn_start_index], *project_incomplete_tool_history(turn_messages[current_turn_start_index:]), ] + self._history = [ + item for item in self._history + if not item._request_only_pending_tool_image + ] self._write_context_stage("session:after", self._history) # ------ → DONE ------ @@ -13811,6 +14121,9 @@ async def _run_after_key_lock() -> ToolResult: _serialize_assistant_replay(replay_messages) if replay_messages else None ), + tool_image_native_ack_indexes=( + _tool_image_native_ack_indexes(replay_messages) + ), input_tokens=done_input_tokens, output_tokens=done_output_tokens, reasoning_tokens=done_reasoning_tokens, @@ -14515,6 +14828,7 @@ def _provider_request_messages_with_sanitize( record=not preview, ) source_messages = repair_tool_pairing(source_messages) + source_messages = _bound_provider_tool_image_batch(source_messages) request_messages, sanitize_result = sanitize_session_messages(source_messages) if not preview: self._remember_provider_visible_tool_results(request_messages) diff --git a/src/opensquilla/engine/history.py b/src/opensquilla/engine/history.py index 8c56866d62..3a2c47a4de 100644 --- a/src/opensquilla/engine/history.py +++ b/src/opensquilla/engine/history.py @@ -239,7 +239,10 @@ def project_history_replay( message_provenance.pop() unique_notices = tuple(terminal_notices) - messages.extend(Message(role="assistant", content=notice) for notice in unique_notices) + for notice in unique_notices: + trusted = Message(role="assistant", content=notice) + trusted._trusted_runtime_notice = True + messages.append(trusted) message_provenance.extend(terminal_notices[notice] for notice in unique_notices) return HistoryReplayProjection( messages=tuple(messages), @@ -408,6 +411,8 @@ def _project_value(value: Any) -> Any: def _is_synthetic_context_message(message: Message) -> bool: + if message._request_only_pending_tool_image: + return True content = message.content if isinstance(content, list): # Historical tool media is extracted beside the leading record text. @@ -424,7 +429,9 @@ def _is_real_user_turn(message: Message) -> bool: return False content = message.content if isinstance(content, list): - return not all(isinstance(block, ContentBlockToolResult) for block in content) + # Tool outcomes may carry adjacent image/document blocks. They belong + # to the preceding assistant call, not to a new user turn. + return not any(isinstance(block, ContentBlockToolResult) for block in content) return True @@ -633,6 +640,33 @@ class AssistantReplayError(ValueError): """A durable assistant replay envelope cannot be restored without data loss.""" +TOOL_IMAGE_NATIVE_ACK_CONTEXT_KEY = "tool_image_native_ack_indexes_v1" + + +def restore_tool_image_ack_context( + messages: list[Message], turn_context: Mapping[str, Any] | None, +) -> None: + """Restore accepted native-image evidence without changing provider messages. + + Older transcript rows have no such field and keep their historical replay + behavior. A malformed new field cannot claim that an image was seen. + """ + if ( + not isinstance(turn_context, Mapping) + or TOOL_IMAGE_NATIVE_ACK_CONTEXT_KEY not in turn_context + ): + return + raw_indexes = turn_context[TOOL_IMAGE_NATIVE_ACK_CONTEXT_KEY] + valid_indexes = ( + isinstance(raw_indexes, list) + and all(type(index) is int and 0 <= index < len(messages) for index in raw_indexes) + ) + acknowledged = set(raw_indexes) if valid_indexes else set() + for index, message in enumerate(messages): + if message.role == "assistant": + message._tool_image_native_ack = index in acknowledged + + def decode_assistant_replay(value: Mapping[str, Any]) -> list[Message]: """Restore versioned accepted messages without guessing missing native state.""" @@ -822,9 +856,11 @@ def reconstruct_messages_from_entry( return [] if role == "assistant" and assistant_replay is not None: + replayed_messages = decode_assistant_replay(assistant_replay) + restore_tool_image_ack_context(replayed_messages, turn_context) replayed_messages = project_incomplete_tool_history( _project_unsigned_silent_replay( - decode_assistant_replay(assistant_replay), turn_context, content, + replayed_messages, turn_context, content, ) ) artifact_context = artifact_history_context(content) diff --git a/src/opensquilla/engine/replay_compat.py b/src/opensquilla/engine/replay_compat.py index 14193d36f8..55b1831ee5 100644 --- a/src/opensquilla/engine/replay_compat.py +++ b/src/opensquilla/engine/replay_compat.py @@ -90,7 +90,9 @@ def record_block(block: Any) -> Any: ] records.append({"role": message.role, "content": content}) text = introduction + "\n" + json.dumps(records, ensure_ascii=False, separators=(",", ":")) - return Message( + recorded = Message( role="user", content=[ContentBlockText(text=text), *media] if media else text, ) + recorded._recorded_tool_context = True + return recorded diff --git a/src/opensquilla/engine/runtime.py b/src/opensquilla/engine/runtime.py index 71c3d5cf20..af675cb606 100644 --- a/src/opensquilla/engine/runtime.py +++ b/src/opensquilla/engine/runtime.py @@ -19,6 +19,7 @@ import math import os import platform +import stat import tempfile import time import uuid @@ -62,6 +63,9 @@ from opensquilla.contracts.attachments import ( ALLOWED_MEDIA_TYPES as _ALLOWED_ENGINE_MEDIA_TYPES, ) +from opensquilla.contracts.attachments import ( + IMAGE_ATTACHMENT_BYTES as _IMAGE_ATTACHMENT_BYTES, +) from opensquilla.contracts.attachments import ( IMAGE_ATTACHMENT_MIMES as _IMAGE_ATTACHMENT_MIMES, ) @@ -85,7 +89,12 @@ ) from opensquilla.contracts.image_validation import validate_image_bytes from opensquilla.contracts.turn_execution import TurnExecutionContext -from opensquilla.engine.agent import Agent, ToolHandler, _ToolArgumentProgress +from opensquilla.engine.agent import ( + Agent, + ToolHandler, + _has_historical_tool_image, + _ToolArgumentProgress, +) from opensquilla.engine.agent_injection import PendingInputProvider from opensquilla.engine.cache_break_monitor import notify_compaction from opensquilla.engine.fallback import FallbackPolicy, sleep_before_retry @@ -268,6 +277,9 @@ classify_image_failure, project_messages, ) +from opensquilla.provider.image_projection import ( + count_image_blocks as count_replay_image_blocks, +) from opensquilla.provider.model_catalog import ( resolve_effective_context_window, shared_catalog, @@ -1646,6 +1658,19 @@ def _should_use_selector_fallback(provider_name: str, event: ProviderErrorEvent) } +def _is_local_wire_body_rejection(provider_name: str, event: ProviderErrorEvent) -> bool: + """Keep the wire-body guard distinct from other unsent budget failures.""" + return bool( + event._opensquilla_request_not_dispatched + and classify_provider_error( + provider_name=provider_name, + status_code=event.http_status, + raw_code=event.provider_code or event.code, + message=event.message, + ) is ProviderFailureKind.WIRE_PAYLOAD_TOO_LARGE + ) + + def _report_credential_pool_failure( provider_name: str, turn_metadata: dict[str, Any] | None, @@ -2328,6 +2353,7 @@ def __init__( self._image_probe_forbidden = False self._stream_recovery_request: StreamRecoveryRequest | None = None self.last_image_request_had_native_images = False + self.last_image_request_had_native_tool_images = False self._fallback_limits: dict[tuple[str, str], tuple[int, int]] = {} self._fallback_deployment_limits: dict[_FallbackDeploymentIdentity, tuple[int, int]] = {} self._fallback_deployment_capabilities: dict[ @@ -2699,6 +2725,9 @@ def _project_image_messages_for_active_leg( ), ) self.last_image_request_had_native_images = projection.output_image_count > 0 + self.last_image_request_had_native_tool_images = any( + _has_historical_tool_image(message) for message in projection.messages + ) if projection.input_image_count and isinstance(turn_metadata, dict): # These fields describe the physical leg that is about to run. # A previous native probe must not leave stale metadata after a @@ -4031,6 +4060,12 @@ def primary_stream_factory() -> AsyncGenerator[Any, None]: code="incomplete_tool_stream", ) + local_wire_body_rejected = ( + isinstance(event, ProviderErrorEvent) + and _is_local_wire_body_rejection( + active_provider_id or self.provider_name, event, + ) + ) local_admission_fallback_index = ( self._local_admission_fallback_index( active_config, @@ -4038,6 +4073,7 @@ def primary_stream_factory() -> AsyncGenerator[Any, None]: ) if isinstance(event, ProviderErrorEvent) and event.code == "provider_request_budget_exhausted" + and not local_wire_body_rejected else 0 ) local_admission_escalation = local_admission_fallback_index > 0 @@ -4050,6 +4086,9 @@ def primary_stream_factory() -> AsyncGenerator[Any, None]: ) while ( isinstance(event, ProviderErrorEvent) + and not _is_local_wire_body_rejection( + active_provider_id or self.provider_name, event, + ) and not precise_image_capability_rejection and ( _should_use_selector_fallback(self.provider_name, event) @@ -8222,10 +8261,14 @@ def resolve_retained_attachment( final_text_parts[:] = [partial_text] if partial_text else [] reasoning_content = "".join(reasoning_parts).strip() assistant_replay = None + tool_image_native_ack_indexes: list[int] = [] replay_snapshot = getattr(agent, "current_assistant_replay", None) if callable(replay_snapshot): try: assistant_replay = replay_snapshot() + ack_snapshot = getattr(agent, "current_tool_image_native_ack_indexes", None) + if callable(ack_snapshot): + tool_image_native_ack_indexes = ack_snapshot() except Exception: log.warning( "turn_runner.cancelled_replay_snapshot_failed", @@ -8291,12 +8334,27 @@ def resolve_retained_attachment( if cancelled_turn_usage is not None else None ) - await _finish_required_cancel_cleanup( - self._append_session_message( - session_key, - **append_kwargs, - ) + from opensquilla.engine.history import TOOL_IMAGE_NATIVE_ACK_CONTEXT_KEY + from opensquilla.session.turn_context import ( + current_turn_context, + turn_context_scope, ) + + durable_context = current_turn_context() or {} + if assistant_replay is not None: + messages = assistant_replay.get("messages") + replay_count = len(messages) if isinstance(messages, list) else 0 + durable_context[TOOL_IMAGE_NATIVE_ACK_CONTEXT_KEY] = [ + index for index in tool_image_native_ack_indexes + if type(index) is int and 0 <= index < replay_count + ] + with turn_context_scope(durable_context): + await _finish_required_cancel_cleanup( + self._append_session_message( + session_key, + **append_kwargs, + ) + ) execution_context.publish_visible( text=body, generation_epoch=execution_context.generation_epoch, @@ -11904,7 +11962,9 @@ async def _router_previous_assistant_context( user_texts: list[str] = [] user_contents: list[str] = [] user_image_flags: list[bool] = [] + tool_image_flags: list[bool] = [] unanchored_image_replay = False + pending_tool_image = False for index, entry in enumerate(entries): if getattr(entry, "role", None) == "assistant": # A queued user row can precede the previous turn's completion. @@ -11915,23 +11975,73 @@ async def _router_previous_assistant_context( from opensquilla.engine.history import ( AssistantReplayError, decode_assistant_replay, + restore_tool_image_ack_context, + ) + from opensquilla.provider.types import ( + ContentBlockText, + ContentBlockToolResult, + ContentBlockToolUse, ) - from opensquilla.provider.types import ContentBlockImage try: replay_messages = decode_assistant_replay(replay) + turn_context = getattr(entry, "turn_context", None) + restore_tool_image_ack_context( + replay_messages, + turn_context if isinstance(turn_context, dict) else None, + ) except AssistantReplayError: replay_messages = [] - if any( - isinstance(block, ContentBlockImage) - for message in replay_messages - if isinstance(message.content, list) - for block in message.content - ): - if user_image_flags: - user_image_flags[-1] = True + for replay_message in replay_messages: + blocks = ( + replay_message.content + if isinstance(replay_message.content, list) + else [] + ) + if replay_message.role == "assistant" and ( + replay_message._tool_image_native_ack is True + or ( + replay_message._tool_image_native_ack is None + and ( + ( + isinstance(replay_message.content, str) + and bool(replay_message.content.strip()) + ) + or any( + isinstance(block, ContentBlockToolUse) + or ( + isinstance(block, ContentBlockText) + and bool(block.text.strip()) + ) + for block in blocks + ) + ) + ) + ): + pending_tool_image = False + elif ( + replay_message.role == "user" + and any(isinstance(block, ContentBlockToolResult) for block in blocks) + and count_replay_image_blocks([replay_message]) + ): + pending_tool_image = True + if count_replay_image_blocks(replay_messages): + if tool_image_flags: + tool_image_flags[-1] = True else: unanchored_image_replay = True + elif ( + isinstance(getattr(entry, "content", None), str) + and bool(entry.content.strip()) + ) or bool(getattr(entry, "tool_calls", None)): + from opensquilla.engine.history import TOOL_IMAGE_NATIVE_ACK_CONTEXT_KEY + + turn_context = getattr(entry, "turn_context", None) + if not ( + isinstance(turn_context, dict) + and TOOL_IMAGE_NATIVE_ACK_CONTEXT_KEY in turn_context + ): + pending_tool_image = False if getattr(entry, "role", None) != "user": continue if index in excluded_user_indexes: @@ -11943,6 +12053,7 @@ async def _router_previous_assistant_context( continue user_contents.append(content) user_image_flags.append(self._attachment_envelope_has_image(content)) + tool_image_flags.append(False) unpacked = self._maybe_unpack_attachments(content) text = unpacked.strip() if isinstance(unpacked, str) else content.strip() if len(text) > _ROUTER_HISTORY_USER_MAX_CHARS: @@ -11953,11 +12064,24 @@ async def _router_previous_assistant_context( if user_texts: context["history_user_texts"] = user_texts[-_ROUTER_HISTORY_USER_MAX_TURNS:] router_cfg = getattr(self._turn_config(), "squilla_router", None) + # A failed post-tool request leaves its final image batch pending even + # across later failed user turns. Route the next attempt to vision so + # the request-local replay can actually deliver those images. image_positions = [index for index, has_image in enumerate(user_image_flags) if has_image] - if image_positions or unanchored_image_replay: + recent_tool_image = bool(tool_image_flags and tool_image_flags[-1]) or pending_tool_image + if ( + recent_tool_image + and tool_image_flags + and (len(tool_image_flags) - 1) not in image_positions + ): + image_positions.append(len(tool_image_flags) - 1) + unanchored_recent_tool_image = ( + unanchored_image_replay or pending_tool_image + ) and not user_texts + if image_positions or unanchored_recent_tool_image: context["history_has_recent_image"] = True context["history_image_turn_count"] = ( - len(image_positions) + int(unanchored_image_replay) + len(image_positions) + int(unanchored_recent_tool_image) ) if image_positions: turns_since_last_image = len(user_contents) - image_positions[-1] - 1 @@ -14543,6 +14667,7 @@ async def _load_history( history, history_materializer if retain_history_material else None, attachment_replay_session_id, + workspace_dir, ) summary_markers.extend(replay.legacy_summary_markers) @@ -15021,21 +15146,129 @@ def _restore_retained_history_image_paths( messages: list[Message], materializer: AttachmentWorkspaceMaterializer | None, session_id: str, + workspace_dir: str | Path | None = None, ) -> None: """Rebind retained tool images to this session without rewriting unchanged replay.""" - from opensquilla.attachment_workspace import _safe_path_segment + from opensquilla.attachment_workspace import _safe_filename, _safe_path_segment from opensquilla.provider.types import ContentBlockImage, ContentBlockToolResult current_path_prefix = ( ".opensquilla", "attachments", _safe_path_segment(session_id, fallback="session"), ) + + def original_is_readable(path: str, name: str) -> bool: + """Accept only an intact immutable image in this session's workspace.""" + if workspace_dir is None: + return False + relative = Path(path) + parts = relative.parts + if ( + relative.is_absolute() + or len(parts) != 4 + or parts[:3] != current_path_prefix + or len(parts[3]) < 14 + or parts[3][12] != "-" + or any(char not in "0123456789abcdef" for char in parts[3][:12]) + or parts[3][13:] != _safe_filename(name) + ): + return False + original_mime = { + ".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", + ".gif": "image/gif", ".webp": "image/webp", + }.get(relative.suffix.lower()) + if original_mime is None: + return False + root = Path(workspace_dir).expanduser().resolve() + source = root.joinpath(*parts) + try: + if source.parent.resolve(strict=True) != source.parent: + return False + before = source.lstat() + if not stat.S_ISREG(before.st_mode) or before.st_size > _IMAGE_ATTACHMENT_BYTES: + return False + descriptor = os.open( + source, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) + | getattr(os, "O_NONBLOCK", 0) | getattr(os, "O_BINARY", 0), + ) + with os.fdopen(descriptor, "rb") as handle: + opened = os.fstat(handle.fileno()) + if not stat.S_ISREG(opened.st_mode) or ( + opened.st_dev, opened.st_ino + ) != (before.st_dev, before.st_ino): + return False + payload = handle.read(_IMAGE_ATTACHMENT_BYTES + 1) + after = os.fstat(handle.fileno()) + if ( + len(payload) != before.st_size + or (after.st_size, after.st_mtime_ns) != + (opened.st_size, opened.st_mtime_ns) + or not hashlib.sha256(payload).hexdigest().startswith(parts[3][:12]) + ): + return False + validate_image_bytes(payload, original_mime) + return True + except (OSError, ValueError): + return False + + def prepared_name(name: str, mime: str) -> str: + safe_name = _safe_filename(name) + suffix = { + "image/png": ".png", "image/jpeg": ".jpg", + "image/gif": ".gif", "image/webp": ".webp", + }.get(mime) + if suffix is None: + return safe_name + if safe_name.lower().endswith(suffix) or ( + mime == "image/jpeg" and safe_name.lower().endswith(".jpeg") + ): + return safe_name + return f"{Path(safe_name).stem or 'image'}{suffix}" + + def clear_nested_retention(value: Any) -> tuple[Any, bool]: + """Nested tool media has no verified workspace path at this boundary.""" + if isinstance(value, ContentBlockImage): + if value.local_path or value.durable_retained: + return value.model_copy(update={ + "local_path": None, "durable_retained": None, + }), True + return value, False + if isinstance(value, list): + changed = False + output: list[Any] = [] + for item in value: + rebound, item_changed = clear_nested_retention(item) + output.append(rebound) + changed = changed or item_changed + return (output if changed else value), changed + if isinstance(value, dict): + if value.get("type") == "image": + if "local_path" not in value and "durable_retained" not in value: + return value, False + mapped_value = dict(value) + mapped_value.pop("local_path", None) + mapped_value.pop("durable_retained", None) + return mapped_value, True + if value.get("type") == "tool_result" and "content" in value: + content, changed = clear_nested_retention(value["content"]) + if changed: + return {**value, "content": content}, True + return value, False + for message in messages: if not isinstance(message.content, list): continue - paths: dict[tuple[str, str | None], str | None] = {} + paths: dict[tuple[str, str | None], tuple[str | None, str | None, bool]] = {} blocks = list(message.content) for index, block in enumerate(blocks): + if isinstance(block, ContentBlockToolResult) and isinstance( + block.content, list + ): + nested_content, nested_changed = clear_nested_retention(block.content) + if nested_changed: + blocks[index] = block.model_copy(update={"content": nested_content}) + continue if not ( isinstance(block, ContentBlockImage) and block.source_type == "base64" @@ -15044,19 +15277,41 @@ def _restore_retained_history_image_paths( and block.name ): continue + if original_is_readable(block.local_path, block.name): + continue + derivative_name = prepared_name(block.name, block.media_type) + try: + source_hash_prefix = Path(block.local_path).name[:12] + original_bytes_copied = ( + derivative_name == _safe_filename(block.name) + and len(source_hash_prefix) == 12 + and all(char in "0123456789abcdef" for char in source_hash_prefix) + and hashlib.sha256(base64.b64decode(block.data, validate=True)) + .hexdigest().startswith(source_hash_prefix) + ) + except (ValueError, binascii.Error): + original_bytes_copied = False path = ( materializer.materialize_image_path( - {"mime": block.media_type, "data": block.data, "name": block.name}, + {"mime": block.media_type, "data": block.data, + "name": derivative_name}, session_id, ) if materializer is not None else None ) if path and path != block.local_path: - paths[(block.local_path, block.source_url)] = path - blocks[index] = block.model_copy(update={"local_path": path}) - elif not path and Path(block.local_path).parts[:3] != current_path_prefix: - paths[(block.local_path, block.source_url)] = None - blocks[index] = block.model_copy(update={"local_path": None}) + paths[(block.local_path, block.source_url)] = ( + path, derivative_name, original_bytes_copied, + ) + blocks[index] = block.model_copy(update={ + "local_path": path, "name": derivative_name, + }) + elif not path: + paths[(block.local_path, block.source_url)] = (None, None, False) + blocks[index] = block.model_copy(update={ + "local_path": None, "durable_retained": None, + }) if not paths: + message.content = blocks continue for index, block in enumerate(blocks): if not isinstance(block, ContentBlockToolResult) or not isinstance( @@ -15077,9 +15332,16 @@ def _restore_retained_history_image_paths( identity = (local_path, source_url) if identity not in paths: continue - path = paths[identity] - if path: - receipt["local_path"] = path + rebound_path, rebound_name, original_bytes_copied = paths[identity] + if rebound_path and rebound_name: + receipt["local_path"] = rebound_path + receipt["name"] = rebound_name + receipt["retention_note"] = ( + "An identical copy of the original image is readable in this session." + if original_bytes_copied else + "The original image is unavailable in this session; this readable " + "file is a prepared derivative used for model input." + ) else: receipt.pop("local_path", None) receipt["retention_note"] = ( diff --git a/src/opensquilla/engine/thinking.py b/src/opensquilla/engine/thinking.py index 7c651a745b..50ebd3b735 100644 --- a/src/opensquilla/engine/thinking.py +++ b/src/opensquilla/engine/thinking.py @@ -52,24 +52,23 @@ def drop_reasoning( if not filtered: filtered = [ContentBlockText(text="")] out.append( - Message( - role="assistant", - provider_replay=msg.provider_replay, - content=filtered, - reasoning_content=( + msg.model_copy( + update={ + "content": filtered, + "reasoning_content": ( msg.reasoning_content if keep_reasoning_content else None - ), + ), + }, ) ) else: out.append( - Message( - role="assistant", - provider_replay=msg.provider_replay, - content=msg.content, - reasoning_content=( + msg.model_copy( + update={ + "reasoning_content": ( msg.reasoning_content if keep_reasoning_content else None - ), + ), + }, ) ) continue diff --git a/src/opensquilla/engine/turn_runner/turn_finalizer_stage.py b/src/opensquilla/engine/turn_runner/turn_finalizer_stage.py index 5c85f5b429..214a9b8720 100644 --- a/src/opensquilla/engine/turn_runner/turn_finalizer_stage.py +++ b/src/opensquilla/engine/turn_runner/turn_finalizer_stage.py @@ -34,6 +34,7 @@ from __future__ import annotations +from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass, replace from typing import TYPE_CHECKING, Any, Literal, Protocol, runtime_checkable @@ -799,10 +800,26 @@ async def run( append_kwargs["expected_session_id"] = inp.expected_session_id if inp.expected_session_epoch is not None: append_kwargs["expected_session_epoch"] = inp.expected_session_epoch - append_result = await self._transcript_append.append_message( - inp.session_key, - **append_kwargs, - ) + append_scope: AbstractContextManager[None] = nullcontext() + ack_indexes = getattr(done_event, "tool_image_native_ack_indexes", None) + if assistant_replay is not None and isinstance(ack_indexes, list): + from opensquilla.engine.history import TOOL_IMAGE_NATIVE_ACK_CONTEXT_KEY + from opensquilla.session.turn_context import ( + current_turn_context, + turn_context_scope, + ) + + durable_context = current_turn_context() or {} + durable_context[TOOL_IMAGE_NATIVE_ACK_CONTEXT_KEY] = [ + index for index in ack_indexes + if type(index) is int and 0 <= index < len(assistant_replay.get("messages", [])) + ] + append_scope = turn_context_scope(durable_context) + with append_scope: + append_result = await self._transcript_append.append_message( + inp.session_key, + **append_kwargs, + ) if isinstance(append_result, TranscriptAppendResult): transcript_appended = append_result.appended if assistant_message_id is None: diff --git a/src/opensquilla/engine/types.py b/src/opensquilla/engine/types.py index 51e4f73a31..800f2847d5 100644 --- a/src/opensquilla/engine/types.py +++ b/src/opensquilla/engine/types.py @@ -452,6 +452,11 @@ class DoneEvent: assistant_replay: dict[str, Any] | None = field( default=None, repr=False, metadata={"private": True} ) + # Positions of committed assistant replies whose physical request carried + # native tool images. Persisted only in transcript turn_context. + tool_image_native_ack_indexes: list[int] | None = field( + default=None, repr=False, metadata={"private": True} + ) @property def upstream_cost_usd(self) -> float: @@ -462,8 +467,11 @@ def upstream_cost_usd(self) -> float: def public_agent_event_payload(event: Any) -> dict[str, Any]: """Serialize an event without copying private provider continuation state.""" if isinstance(event, DoneEvent): - payload = asdict(replace(event, assistant_replay=None)) + payload = asdict(replace( + event, assistant_replay=None, tool_image_native_ack_indexes=None, + )) payload.pop("assistant_replay", None) + payload.pop("tool_image_native_ack_indexes", None) return payload if isinstance(event, ToolResultEvent) and event.is_error: # A failed mutation must never advertise a successful file snapshot, diff --git a/src/opensquilla/engine/usage_accounting.py b/src/opensquilla/engine/usage_accounting.py index 919a328fde..8cc90cc618 100644 --- a/src/opensquilla/engine/usage_accounting.py +++ b/src/opensquilla/engine/usage_accounting.py @@ -374,6 +374,33 @@ async def mark_usage_call_unknown( ) +async def mark_usage_call_not_sent( + scope: UsageAccountingScope, + call: UsageCallStart, +) -> None: + """Close a reserved call rejected by a trusted adapter before HTTP dispatch.""" + + mark_not_sent = getattr(scope.sink, "mark_not_sent", None) + if not callable(mark_not_sent): + # External sinks implementing the original three-method contract keep + # their conservative unknown outcome until they support this terminal. + await mark_usage_call_unknown(scope, call, "cancelled_before_provider_request") + return + terminal_task = asyncio.create_task(mark_not_sent(call)) + try: + await asyncio.shield(terminal_task) + except asyncio.CancelledError: + with contextlib.suppress(Exception): + await terminal_task + raise + except Exception as exc: # noqa: BLE001 - preserve the local rejection + log.warning( + "usage_accounting.mark_not_sent_failed", + event_id=call.event_id, + error=str(exc), + ) + + def _consume_usage_task_result(task: asyncio.Future[Any]) -> None: """Consume a detached terminal-write result after bounded cancellation.""" @@ -470,7 +497,14 @@ async def account_provider_stream( unknown_reason = provider_error_usage_reason( getattr(event, "code", None) ) - if has_known_provider_usage_receipt(event) and not terminal: + if ( + getattr(event, "_opensquilla_request_not_dispatched", False) is True + and not terminal + and not has_known_provider_usage_receipt(event) + ): + terminal = True + await mark_usage_call_not_sent(scope, call) + elif has_known_provider_usage_receipt(event) and not terminal: terminal = True await finalize_usage_call(scope, call, event) yield event @@ -1039,6 +1073,7 @@ def normalize_provider_usage( "finalize_usage_call", "has_known_provider_usage_receipt", "mark_usage_call_unknown", + "mark_usage_call_not_sent", "normalize_provider_usage", "provider_accounts_physical_usage", "start_usage_call", diff --git a/src/opensquilla/gateway/usage_ledger_runtime.py b/src/opensquilla/gateway/usage_ledger_runtime.py index 1d3657b70a..e4209afc2b 100644 --- a/src/opensquilla/gateway/usage_ledger_runtime.py +++ b/src/opensquilla/gateway/usage_ledger_runtime.py @@ -408,6 +408,24 @@ async def mark_unknown(self, call: UsageCallStart, reason: str) -> None: ) raise + async def mark_not_sent(self, call: UsageCallStart) -> None: + """Close a reserved call that the local adapter never dispatched.""" + + completed_at_ms = max(call.started_at_ms, time.time_ns() // 1_000_000) + try: + record = await self._storage.mark_usage_event_not_sent( + call.event_id, + completed_at_ms=completed_at_ms, + ) + await self._notify_goal_usage(record) + except Exception: + self._schedule_retry( + self._retry_not_sent(call.event_id, completed_at_ms), + event_id=call.event_id, + operation="mark_not_sent", + ) + raise + def _schedule_retry( self, coroutine: Any, @@ -477,6 +495,20 @@ async def _retry_unknown( continue log.warning("usage.ledger_unknown_retry_exhausted", event_id=event_id) + async def _retry_not_sent(self, event_id: str, completed_at_ms: int) -> None: + for delay in self._retry_delays: + await asyncio.sleep(delay) + try: + record = await self._storage.mark_usage_event_not_sent( + event_id, + completed_at_ms=completed_at_ms, + ) + await self._notify_goal_usage(record) + return + except Exception: # noqa: BLE001 - bounded retry; recovery handles residue. + continue + log.warning("usage.ledger_not_sent_retry_exhausted", event_id=event_id) + async def close(self, *, drain_timeout: float | None = None) -> None: tasks = tuple(self._tasks) if not tasks: diff --git a/src/opensquilla/gateway/usage_query.py b/src/opensquilla/gateway/usage_query.py index 844f5e48fb..c57fbcf28e 100644 --- a/src/opensquilla/gateway/usage_query.py +++ b/src/opensquilla/gateway/usage_query.py @@ -24,6 +24,7 @@ ) from opensquilla.provider.fx import canonical_native_per_usd_rates from opensquilla.session.cost_rollup import rollup_cost_source +from opensquilla.session.usage_ledger import LOCAL_PRE_DISPATCH_ORIGIN _PRESET_DAYS = { "today": 1, @@ -1273,7 +1274,11 @@ async def _query_usage_ledger_once( ) # Storage primary keys already guarantee this; keep adapters/replays from # multiplying the same physical event across any reporting dimension. - events = list({_text(_field(event, "event_id")): event for event in events}.values()) + events = list({ + _text(_field(event, "event_id")): event + for event in events + if _event_origin(event) != LOCAL_PRE_DISPATCH_ORIGIN + }.values()) # Partition by status in one pass. Equality-based list membership here is # quadratic for the common all-finalized case and becomes visible at only # a few thousand ledger rows. diff --git a/src/opensquilla/provider/anthropic.py b/src/opensquilla/provider/anthropic.py index 5c5703db88..c683cfe144 100644 --- a/src/opensquilla/provider/anthropic.py +++ b/src/opensquilla/provider/anthropic.py @@ -29,6 +29,7 @@ from .registry import AuthHeaderStyle from .request_proof import ( ProviderRequestBudgetExceededError, + local_budget_rejection_event, project_final_request_payload, protected_tool_result_indexes, prove_provider_payload_from_env, @@ -559,10 +560,7 @@ async def _stream( if budget_decision.action == "budget_limited": proof = budget_decision.proof or {} log.warning("provider.request_budget_exhausted", **proof) - yield ErrorEvent( - message=json.dumps(proof, ensure_ascii=False, sort_keys=True), - code="provider_request_budget_exhausted", - ) + yield local_budget_rejection_event(proof) return if budget_decision.action == "invalid_request": log.warning("provider.request_serialization_failed") @@ -585,10 +583,7 @@ async def _stream( ) except ProviderRequestBudgetExceededError as exc: log.warning("provider.request_budget_exhausted", **exc.proof) - yield ErrorEvent( - message=json.dumps(exc.proof, ensure_ascii=False, sort_keys=True), - code="provider_request_budget_exhausted", - ) + yield local_budget_rejection_event(exc.proof) return headers = { diff --git a/src/opensquilla/provider/failures.py b/src/opensquilla/provider/failures.py index d0a41e8f36..ac4c8d6938 100644 --- a/src/opensquilla/provider/failures.py +++ b/src/opensquilla/provider/failures.py @@ -74,6 +74,8 @@ class ProviderFailureKind(StrEnum): PROVIDER_OVERLOADED = "provider_overloaded" AUTH_INVALID = "auth_invalid" CONTEXT_OVERFLOW = "context_overflow" + PAYLOAD_TOO_LARGE = "payload_too_large" + WIRE_PAYLOAD_TOO_LARGE = "wire_payload_too_large" UNSUPPORTED_FEATURE = "unsupported_feature" INSUFFICIENT_CREDITS = "insufficient_credits" USAGE_LIMIT = "usage_limit" @@ -223,6 +225,30 @@ def _is_input_length_range_error(signal: FailureSignal) -> bool: )) +def _is_local_wire_budget_exhausted(status_code: int | None, raw_code: str, message: str) -> bool: + """Use only a bounded local proof, never arbitrary upstream size prose.""" + if status_code is not None or raw_code != "provider_request_budget_exhausted": + return False + if not isinstance(message, str) or len(message) > 65_536: + return False + try: + proof = json.loads(message) + except ValueError: + return False + if not isinstance(proof, dict): + return False + used = proof.get("wire_json_bytes") + limit = proof.get("wire_body_max_bytes") + return ( + proof.get("wire_byte_budget_exhausted") is True + and proof.get("fits_wire_byte_budget") is False + and type(used) is int + and type(limit) is int + and limit > 0 + and used > limit + ) + + def _is_anthropic_missing_model(signal: FailureSignal) -> bool: """Anthropic's missing-resource type needs an explicit model target. @@ -508,6 +534,29 @@ def classify_provider_error( ) -> ProviderFailureKind: """Classify a provider error into a stable runtime failure kind.""" + normalized_code = (raw_code or "").strip().lower() + # A provider's explicit token-window code is stronger evidence than the + # transport status. A bare 413, however, only proves a request-size + # rejection; compacting token history or retrying another model blindly + # could resend the same oversized body. + if normalized_code in { + "context_length_exceeded", + "context_window_exceeded", + "context_overflow", + }: + return ProviderFailureKind.CONTEXT_OVERFLOW + if _is_local_wire_budget_exhausted(status_code, normalized_code, message): + return ProviderFailureKind.WIRE_PAYLOAD_TOO_LARGE + if status_code == 413 or normalized_code in { + "413", + "payload_too_large", + "request_entity_too_large", + "request_too_large", + "body_too_large", + "content_too_large", + }: + return ProviderFailureKind.PAYLOAD_TOO_LARGE + # Explicit business codes win over contradictory HTTP envelopes/prose. structured_kind = { "usage_limit": ProviderFailureKind.USAGE_LIMIT, @@ -520,14 +569,14 @@ def classify_provider_error( "rate_limit_error": ProviderFailureKind.RATE_LIMITED, "invalid_api_key": ProviderFailureKind.AUTH_INVALID, "authentication_error": ProviderFailureKind.AUTH_INVALID, - }.get((raw_code or "").strip().lower()) + }.get(normalized_code) if structured_kind is not None: return structured_kind provider = (provider_name or "").lower() family = _failure_family(provider) signal = FailureSignal( status_code=status_code, - raw_code=(raw_code or "").strip().lower(), + raw_code=normalized_code, message=message or "", text=_joined(status_code, raw_code, message), ) diff --git a/src/opensquilla/provider/ollama.py b/src/opensquilla/provider/ollama.py index 9407c09ec0..2161b0baf2 100644 --- a/src/opensquilla/provider/ollama.py +++ b/src/opensquilla/provider/ollama.py @@ -33,6 +33,7 @@ from .recovery_origin import transport_recovery_origin from .request_proof import ( ProviderRequestBudgetExceededError, + local_budget_rejection_event, project_final_request_payload, protected_tool_result_indexes, prove_provider_payload_from_env, @@ -326,10 +327,7 @@ async def _stream( if budget_decision.action == "budget_limited": proof = budget_decision.proof or {} log.warning("provider.request_budget_exhausted", **proof) - yield ErrorEvent( - message=json.dumps(proof, ensure_ascii=False, sort_keys=True), - code="provider_request_budget_exhausted", - ) + yield local_budget_rejection_event(proof) return if budget_decision.action == "invalid_request": log.warning("provider.request_serialization_failed") @@ -351,10 +349,7 @@ async def _stream( ) except ProviderRequestBudgetExceededError as exc: log.warning("provider.request_budget_exhausted", **exc.proof) - yield ErrorEvent( - message=json.dumps(exc.proof, ensure_ascii=False, sort_keys=True), - code="provider_request_budget_exhausted", - ) + yield local_budget_rejection_event(exc.proof) return endpoint = f"{self._base_url}/api/chat" diff --git a/src/opensquilla/provider/openai.py b/src/opensquilla/provider/openai.py index 0e469e5d44..c768815f43 100644 --- a/src/opensquilla/provider/openai.py +++ b/src/opensquilla/provider/openai.py @@ -81,6 +81,7 @@ from .request_proof import ( ProviderRequestBudgetExceededError, fit_payload_generation_budget, + local_budget_rejection_event, project_final_request_payload, protected_tool_result_indexes, prove_provider_payload_from_env, @@ -3885,10 +3886,7 @@ async def _stream( if budget_decision.action == "budget_limited": proof = budget_decision.proof or {} log.warning("provider.request_budget_exhausted", **proof) - yield ErrorEvent( - message=json.dumps(proof, ensure_ascii=False, sort_keys=True), - code="provider_request_budget_exhausted", - ) + yield local_budget_rejection_event(proof) return if budget_decision.action == "invalid_request": log.warning("provider.request_serialization_failed") @@ -3913,10 +3911,7 @@ async def _stream( except ProviderRequestBudgetExceededError as exc: exc.proof.update(generation_budget_proof) log.warning("provider.request_budget_exhausted", **exc.proof) - yield ErrorEvent( - message=json.dumps(exc.proof, ensure_ascii=False, sort_keys=True), - code="provider_request_budget_exhausted", - ) + yield local_budget_rejection_event(exc.proof) return retained_replay_units = _retained_reasoning_replay_units( payload, diff --git a/src/opensquilla/provider/openai_codex.py b/src/opensquilla/provider/openai_codex.py index 26b8a27cc3..920df6acaa 100644 --- a/src/opensquilla/provider/openai_codex.py +++ b/src/opensquilla/provider/openai_codex.py @@ -50,6 +50,7 @@ from .request_proof import ( RESPONSES_REQUEST_ENVELOPE, ProviderRequestBudgetExceededError, + local_budget_rejection_event, project_final_request_payload, prove_provider_payload_from_env, provider_request_character_budget, @@ -325,10 +326,7 @@ async def _stream( if budget_decision.action == "budget_limited": proof = budget_decision.proof or {} log.warning("provider.request_budget_exhausted", **proof) - yield ErrorEvent( - message=json.dumps(proof, ensure_ascii=False, sort_keys=True), - code="provider_request_budget_exhausted", - ) + yield local_budget_rejection_event(proof) return if budget_decision.action == "invalid_request": log.warning("provider.request_serialization_failed") @@ -351,10 +349,7 @@ async def _stream( ) except ProviderRequestBudgetExceededError as exc: log.warning("provider.request_budget_exhausted", **exc.proof) - yield ErrorEvent( - message=json.dumps(exc.proof, ensure_ascii=False, sort_keys=True), - code="provider_request_budget_exhausted", - ) + yield local_budget_rejection_event(exc.proof) return # A coordinator-issued physical attempt owns its retry decision. Do diff --git a/src/opensquilla/provider/openai_responses.py b/src/opensquilla/provider/openai_responses.py index 7251947709..071b314b87 100644 --- a/src/opensquilla/provider/openai_responses.py +++ b/src/opensquilla/provider/openai_responses.py @@ -45,6 +45,7 @@ from .request_proof import ( RESPONSES_REQUEST_ENVELOPE, ProviderRequestBudgetExceededError, + local_budget_rejection_event, project_final_request_payload, prove_provider_payload_from_env, provider_request_character_budget, @@ -394,10 +395,7 @@ async def _complete_items( if budget_decision.action == "budget_limited": proof = budget_decision.proof or {} log.warning("provider.request_budget_exhausted", **proof) - yield ErrorEvent( - message=json.dumps(proof, ensure_ascii=False, sort_keys=True), - code="provider_request_budget_exhausted", - ) + yield local_budget_rejection_event(proof) return if budget_decision.action == "invalid_request": log.warning("provider.request_serialization_failed") @@ -420,10 +418,7 @@ async def _complete_items( ) except ProviderRequestBudgetExceededError as exc: log.warning("provider.request_budget_exhausted", **exc.proof) - yield ErrorEvent( - message=json.dumps(exc.proof, ensure_ascii=False, sort_keys=True), - code="provider_request_budget_exhausted", - ) + yield local_budget_rejection_event(exc.proof) return endpoint = self._api_url("/v1/responses") diff --git a/src/opensquilla/provider/request_proof.py b/src/opensquilla/provider/request_proof.py index 99eacd6a5d..b2904a06d6 100644 --- a/src/opensquilla/provider/request_proof.py +++ b/src/opensquilla/provider/request_proof.py @@ -19,7 +19,7 @@ from opensquilla.context_budget import hard_input_tokens -from .types import ChatConfig, ContentBlockText, ProviderFinalRequestProjection +from .types import ChatConfig, ContentBlockText, ErrorEvent, ProviderFinalRequestProjection _COMPACTED_STRING_MAX_CHARS = 1200 _COMPACTED_TAIL_STRING_MAX_CHARS = 640 @@ -34,6 +34,10 @@ # encoding differ, so this is not an upper bound or a model-specific price. _IMAGE_MEDIA_TOKEN_FLOOR = 1_024 _IMAGE_ESTIMATE_GRID_EDGE = 32 +# A transport limit is independent of the model's image/token budget. Base64 +# media can have a small token reserve while making the JSON body enormous. +_DEFAULT_WIRE_BODY_MAX_BYTES = 24 * 1024 * 1024 +_WIRE_BODY_MAX_BYTES_ENV = "OPENSQUILLA_PROVIDER_REQUEST_MAX_BODY_BYTES" # PDF expansion includes page text and images; keep its existing byte reserve. _PDF_MEDIA_TOKEN_FLOOR = 4_096 _PDF_MEDIA_BYTES_PER_TOKEN = 128 @@ -148,6 +152,21 @@ def __init__(self, proof: dict[str, Any]) -> None: ProviderRequestBudgetExceeded = ProviderRequestBudgetExceededError +def local_budget_rejection_event(proof: dict[str, Any]) -> ErrorEvent: + """Build a budget error for an adapter rejection before HTTP dispatch.""" + + event = ErrorEvent( + message=json.dumps(proof, ensure_ascii=False, sort_keys=True), + code="provider_request_budget_exhausted", + ) + # This factory is called only from an adapter's local budget branches, + # before opening the transport. Character and token rejections are just + # as unsent as wire-byte rejections. Upstream errors are constructed by + # the HTTP handlers and cannot acquire this private dispatch evidence. + event._opensquilla_request_not_dispatched = True + return event + + @dataclass(frozen=True) class ProviderRequestEnvelopeShape: """Describe where one provider's final request stores model-visible input. @@ -377,6 +396,19 @@ def provider_request_character_budget(payload: dict[str, Any], config: ChatConfi return budget +def _wire_body_max_bytes() -> int: + """Return the JSON transport guard; zero explicitly disables it.""" + + raw = os.environ.get(_WIRE_BODY_MAX_BYTES_ENV) + if raw is None or not raw.strip(): + return _DEFAULT_WIRE_BODY_MAX_BYTES + try: + parsed = int(raw) + except ValueError: + return _DEFAULT_WIRE_BODY_MAX_BYTES + return parsed if parsed >= 0 else _DEFAULT_WIRE_BODY_MAX_BYTES + + def _serialized_token_estimate(serialized_payload: str) -> tuple[int, str]: from opensquilla.token_estimation import estimate_tokens_with_source @@ -1779,6 +1811,10 @@ def project_provider_payload( ) wire_json_chars = len(wire_json) wire_json_bytes = len(wire_json.encode("utf-8")) + wire_body_max_bytes = _wire_body_max_bytes() + fits_wire_byte_budget = ( + wire_body_max_bytes == 0 or wire_json_bytes <= wire_body_max_bytes + ) budget_payload, media = _budget_projection( payload, envelope_shape, @@ -1818,7 +1854,7 @@ def project_provider_payload( (token_budget is None and proof_budget <= 0) or estimated_tokens <= effective_token_budget ) - fits = fits_char_budget and fits_token_budget + fits = fits_char_budget and fits_token_budget and fits_wire_byte_budget proof: dict[str, Any] = { "projection_adapter": projection_adapter, "execution_status_version": 1, @@ -1926,6 +1962,12 @@ def project_provider_payload( if media.excluded_blocks: proof["media_chars_excluded"] = media.excluded_chars proof["media_blocks_excluded"] = media.excluded_blocks + if not fits_wire_byte_budget: + proof["wire_json_bytes"] = wire_json_bytes + proof["wire_body_max_bytes"] = wire_body_max_bytes + proof["fits_wire_byte_budget"] = False + proof["wire_byte_budget_exhausted"] = True + proof["wire_body_excess_bytes"] = wire_json_bytes - wire_body_max_bytes if not fits: proof["fallback_reason"] = "provider_request_budget_exhausted" return proof @@ -2044,13 +2086,18 @@ def prove_or_compact_provider_payload( # transport's JSON contract. HTTPX rejects NaN/Infinity and other # non-JSON values, so validate with the same strictness here and let # the coordinator map failures to ``invalid_request``. - json.dumps( + wire_json = json.dumps( payload, ensure_ascii=False, separators=(",", ":"), allow_nan=False, ) - return payload, None + wire_body_max_bytes = _wire_body_max_bytes() + if ( + wire_body_max_bytes == 0 + or len(wire_json.encode("utf-8")) <= wire_body_max_bytes + ): + return payload, None if not envelope_shape.allow_request_compaction: proof = prove_provider_payload( payload, @@ -2079,6 +2126,15 @@ def prove_or_compact_provider_payload( ) except ProviderRequestBudgetExceededError as first_error: first_chars = int(first_error.proof["estimated_chars"]) + # The request-only reducer does not rewrite image/PDF blocks. When + # media alone exceeds the transport cap, avoid repeated deep copies + # and fail with the exact size evidence for history-level recovery. + if ( + first_error.proof.get("wire_byte_budget_exhausted") + and int(first_error.proof.get("media_chars_excluded") or 0) + >= int(first_error.proof["wire_body_max_bytes"]) + ): + raise else: if scrubbed_projection: proof["compact_needed"] = True @@ -2263,7 +2319,18 @@ def prove_provider_payload_from_env( except ValueError: proof_budget = 0 if proof_budget <= 0 and token_budget is None: - return None + wire_json = json.dumps( + payload, + ensure_ascii=False, + separators=(",", ":"), + allow_nan=False, + ) + wire_body_max_bytes = _wire_body_max_bytes() + if ( + wire_body_max_bytes == 0 + or len(wire_json.encode("utf-8")) <= wire_body_max_bytes + ): + return None return prove_provider_payload( payload, projection_adapter=projection_adapter, diff --git a/src/opensquilla/provider/types.py b/src/opensquilla/provider/types.py index 269a919fd0..01b7cf7ab9 100644 --- a/src/opensquilla/provider/types.py +++ b/src/opensquilla/provider/types.py @@ -262,6 +262,9 @@ class ErrorEvent: # Trusted transport evidence, distinct from broad retry classifications. # Unknown extensions and protocol/consumer failures cannot authorize replay. recovery_origin: RecoveryOrigin = "unknown" + # Set only by a local adapter after its final request-budget proof rejects + # the request before HTTP dispatch. Never inferred from provider response text. + _opensquilla_request_not_dispatched: bool = field(default=False, repr=False) @dataclass(frozen=True, slots=True) @@ -863,6 +866,12 @@ class Message(BaseModel): role: Literal["user", "assistant"] content: MessageContent _execution_identity_span: tuple[int, int] | None = PrivateAttr(default=None) + _recorded_tool_context: bool = PrivateAttr(default=False) + _trusted_runtime_notice: bool = PrivateAttr(default=False) + # Request-local acknowledgement for a tool image in a native provider + # request. None denotes replay written before this evidence was recorded. + _tool_image_native_ack: bool | None = PrivateAttr(default=None) + _request_only_pending_tool_image: bool = PrivateAttr(default=False) reasoning_content: str | None = None provider_replay: ProviderReplayState | None = None @@ -870,6 +879,14 @@ class Message(BaseModel): def execution_identity_span(self) -> tuple[int, int] | None: return self._execution_identity_span + @property + def is_recorded_tool_context(self) -> bool: + return self._recorded_tool_context + + @property + def is_trusted_runtime_notice(self) -> bool: + return self._trusted_runtime_notice + # --------------------------------------------------------------------------- # Offline failure injection (test-only seam) @@ -897,6 +914,12 @@ def synthetic_failure_event(kind: ProviderFailureKind) -> ErrorEvent: _Kind.PROVIDER_OVERLOADED: ("529", "injected upstream overloaded"), _Kind.AUTH_INVALID: ("401", "injected invalid api key"), _Kind.CONTEXT_OVERFLOW: ("", "injected context window overflow"), + _Kind.PAYLOAD_TOO_LARGE: ("413", "injected request entity too large"), + _Kind.WIRE_PAYLOAD_TOO_LARGE: ( + "provider_request_budget_exhausted", + '{"wire_byte_budget_exhausted":true,"fits_wire_byte_budget":false,' + '"wire_json_bytes":2,"wire_body_max_bytes":1}', + ), _Kind.UNSUPPORTED_FEATURE: ("", "injected unsupported feature"), _Kind.INSUFFICIENT_CREDITS: ("402", "injected insufficient credits"), _Kind.USAGE_LIMIT: ("usage_limit_reached", "injected usage limit reached"), diff --git a/src/opensquilla/session/storage.py b/src/opensquilla/session/storage.py index b1c4694330..e21e41cdf1 100644 --- a/src/opensquilla/session/storage.py +++ b/src/opensquilla/session/storage.py @@ -140,6 +140,8 @@ read_usage_attributions_on_conn, ) from opensquilla.session.usage_ledger import ( + LOCAL_PRE_DISPATCH_ORIGIN, + LOCAL_WIRE_BODY_REJECTION_REASON, UsageBackfillBatch, UsageBackfillCursor, UsageBackfillEntry, @@ -4143,6 +4145,10 @@ async def _finalize_usage_event_on_conn( persisted = await self._get_usage_event_on_conn(conn, event_id=event_id) if persisted is None: raise KeyError(f"usage event not found: {event_id}") + if persisted.origin == LOCAL_PRE_DISPATCH_ORIGIN: + raise UsageLedgerConflictError( + "a locally rejected request cannot receive provider usage" + ) if completion.completed_at_ms < persisted.started_at_ms: raise ValueError("completed_at_ms must not precede started_at_ms") @@ -4271,7 +4277,8 @@ async def _refresh_goal_usage_coverage_on_conn( """UPDATE session_goals SET usage_coverage = CASE WHEN EXISTS (SELECT 1 FROM usage_events WHERE usage_events.goal_id = session_goals.goal_id - AND (status = 'unknown' OR coverage_status = 'usage_missing' + AND ((status = 'unknown' AND origin != 'local_pre_dispatch') + OR coverage_status = 'usage_missing' OR (status = 'started' AND EXISTS ( SELECT 1 FROM agent_tasks AS executor WHERE executor.task_id = COALESCE( @@ -4376,6 +4383,58 @@ async def mark_usage_event_unknown( await self._refresh_goal_usage_coverage_on_conn(conn, record.goal_id) return replace(record, transition_applied=True) + async def mark_usage_event_not_sent( + self, + event_id: str, + *, + completed_at_ms: int, + ) -> UsageEventRecord: + """Terminalize a reserved request rejected by local wire-body admission. + + The row remains auditable, but its origin excludes it from provider + usage and incomplete-receipt projections. Repeating this exact terminal + transition is idempotent; conflicting terminal outcomes are rejected. + """ + + if not event_id: + raise ValueError("event_id must not be empty") + if completed_at_ms < 0: + raise ValueError("completed_at_ms must be non-negative") + async with self._write_transaction("mark_usage_event_not_sent") as conn: + persisted = await self._get_usage_event_on_conn(conn, event_id=event_id) + if persisted is None: + raise KeyError(f"usage event not found: {event_id}") + if ( + persisted.status == "unknown" + and persisted.origin == LOCAL_PRE_DISPATCH_ORIGIN + and persisted.unknown_reason == LOCAL_WIRE_BODY_REJECTION_REASON + ): + return persisted + if persisted.status != "started" or persisted.origin != "live_provider": + raise UsageLedgerConflictError("usage event has a conflicting terminal outcome") + if completed_at_ms < persisted.started_at_ms: + raise ValueError("completed_at_ms must not precede started_at_ms") + await conn.execute( + """ + UPDATE usage_events + SET completed_at_ms = ?, status = 'unknown', + origin = ?, coverage_status = 'not_sent', + missing_cost_entries = 0, unknown_reason = ? + WHERE event_id = ? AND status = 'started' AND origin = 'live_provider' + """, + ( + completed_at_ms, + LOCAL_PRE_DISPATCH_ORIGIN, + LOCAL_WIRE_BODY_REJECTION_REASON, + event_id, + ), + ) + record = await self._get_usage_event_on_conn(conn, event_id=event_id) + assert record is not None + if record.goal_id is not None: + await self._refresh_goal_usage_coverage_on_conn(conn, record.goal_id) + return replace(record, transition_applied=True) + async def recover_started_usage_events( self, *, diff --git a/src/opensquilla/session/terminal_reply.py b/src/opensquilla/session/terminal_reply.py index a5b87b0433..540d2f27e0 100644 --- a/src/opensquilla/session/terminal_reply.py +++ b/src/opensquilla/session/terminal_reply.py @@ -12,6 +12,19 @@ ) CONTEXT_PAYLOAD_TOO_LARGE_CODE = "provider_request_too_large" +PROVIDER_PAYLOAD_TOO_LARGE_CODE = "provider_payload_too_large" +PROVIDER_WIRE_PAYLOAD_TOO_LARGE_CODE = "provider_wire_payload_too_large" +_REQUEST_SIZE_MESSAGES = { + PROVIDER_PAYLOAD_TOO_LARGE_CODE: ( + "The model provider rejected this request as too large. Reduce large " + "attachments or start a new chat with a concise summary before trying again." + ), + PROVIDER_WIRE_PAYLOAD_TOO_LARGE_CODE: ( + "OpenSquilla stopped this request because the content to send exceeds " + "its size limit. Reduce large attachments or start a new chat with a " + "concise summary before trying again." + ), +} CONTEXT_PAYLOAD_TOO_LARGE_MESSAGES = { "provider_system_prompt_too_large": ( "The fixed system instructions exceed the provider request budget. " @@ -78,6 +91,8 @@ ), "auth_invalid": "The model provider rejected the configured credentials.", "context_overflow": "The request exceeds the model provider's context window.", + "payload_too_large": _REQUEST_SIZE_MESSAGES[PROVIDER_PAYLOAD_TOO_LARGE_CODE], + "wire_payload_too_large": _REQUEST_SIZE_MESSAGES[PROVIDER_WIRE_PAYLOAD_TOO_LARGE_CODE], "unsupported_feature": "The model provider does not support this request.", "insufficient_credits": "The model provider account has insufficient credits.", "usage_limit": "The model provider account has reached its usage limit.", @@ -115,6 +130,8 @@ "provider_retry_after_deadline", "provider_request_budget_exhausted", "provider_request_too_large", + PROVIDER_PAYLOAD_TOO_LARGE_CODE, + PROVIDER_WIRE_PAYLOAD_TOO_LARGE_CODE, "request_error", "response_incomplete", "synthetic_upstream_failure", @@ -151,16 +168,35 @@ def safe_provider_failure_code(raw_code: str | None, failure_kind: str | None) - """Return a bounded terminal code without relaying provider-controlled text.""" normalized_code = str(raw_code or "").strip().lower().replace("-", "_") + normalized_kind = str(failure_kind or "").strip().lower().replace("-", "_") + if normalized_kind == "wire_payload_too_large": + return PROVIDER_WIRE_PAYLOAD_TOO_LARGE_CODE + if normalized_kind == "payload_too_large": + return PROVIDER_PAYLOAD_TOO_LARGE_CODE + if normalized_code == PROVIDER_WIRE_PAYLOAD_TOO_LARGE_CODE: + return "provider_error" if normalized_code.isascii() and normalized_code.isdigit() and len(normalized_code) <= 3: return normalized_code if normalized_code in _SAFE_PROVIDER_TERMINAL_CODES: return normalized_code - normalized_kind = str(failure_kind or "").strip().lower().replace("-", "_") if normalized_kind in _SAFE_PROVIDER_FAILURE_MESSAGES: return f"provider_{normalized_kind}" return "provider_error" +def _request_size_code(record_or_payload: Any) -> str | None: + for field in ("error_class", "terminal_reason"): + code = _normalize(_read_value(record_or_payload, field)) + if code in _REQUEST_SIZE_MESSAGES: + return code + kind = _normalize(_read_value(record_or_payload, "failure_kind")) + if kind == "wire_payload_too_large": + return PROVIDER_WIRE_PAYLOAD_TOO_LARGE_CODE + if kind == "payload_too_large": + return PROVIDER_PAYLOAD_TOO_LARGE_CODE + return None + + def safe_error_id(value: object) -> str | None: """Accept only the existing durable diagnostic reference format.""" if isinstance(value, str) and len(value) == 8 and all(c in "0123456789abcdef" for c in value): @@ -185,10 +221,12 @@ def build_terminal_reply( del surface, locale # Reserved for future surface/locale-specific phrasing. existing = _read_value(record_or_payload, "terminal_message") + request_size_code = _request_size_code(record_or_payload) if ( isinstance(existing, str) and existing.strip() and not _contains_context_payload_marker(existing) + and request_size_code is None ): return existing.strip() @@ -206,6 +244,8 @@ def build_terminal_reply( or "stream idle" in error_message ): return "The task timed out before it could finish." + if request_size_code is not None: + return _REQUEST_SIZE_MESSAGES[request_size_code] if is_context_payload_too_large(record_or_payload) or ( isinstance(existing, str) and _contains_context_payload_marker(existing) ): @@ -342,6 +382,9 @@ def sanitize_agent_error( fallback_error_class: str | None = None, fallback_error_message: str = "Agent error", ) -> tuple[str | None, str]: + request_size_code = _request_size_code(record_or_payload) + if request_size_code is not None: + return request_size_code, _REQUEST_SIZE_MESSAGES[request_size_code] if is_context_payload_too_large(record_or_payload): return CONTEXT_PAYLOAD_TOO_LARGE_CODE, build_terminal_reply(record_or_payload) if _is_provider_output_truncated(record_or_payload): diff --git a/src/opensquilla/session/usage_ledger.py b/src/opensquilla/session/usage_ledger.py index ed5baba0e3..1cb468f792 100644 --- a/src/opensquilla/session/usage_ledger.py +++ b/src/opensquilla/session/usage_ledger.py @@ -14,6 +14,8 @@ USD_NANO_SCALE = 1_000_000_000 _MAX_SQLITE_INTEGER = (1 << 63) - 1 +LOCAL_PRE_DISPATCH_ORIGIN = "local_pre_dispatch" +LOCAL_WIRE_BODY_REJECTION_REASON = "wire_body_budget_not_sent" UsageEventStatus = Literal["started", "finalized", "unknown"] UsageBackfillStatus = Literal["pending", "running", "complete", "partial", "failed"] diff --git a/tests/test_contracts/test_image_preparation.py b/tests/test_contracts/test_image_preparation.py new file mode 100644 index 0000000000..1387bf31de --- /dev/null +++ b/tests/test_contracts/test_image_preparation.py @@ -0,0 +1,149 @@ +"""Synthetic tool-image preparation checks.""" + +from __future__ import annotations + +import io +import random + +import pytest +from PIL import Image, features + +from opensquilla.contracts import image_preparation +from opensquilla.contracts.image_preparation import ( + MAX_TOOL_IMAGE_BASE64_CHARS, + MAX_TOOL_IMAGE_DIMENSION, + prepare_tool_result_image, +) +from opensquilla.contracts.image_validation import validate_image_bytes + + +def _image_bytes(image: Image.Image, format_name: str) -> bytes: + output = io.BytesIO() + image.save(output, format=format_name) + return output.getvalue() + + +def test_small_image_keeps_exact_original_bytes() -> None: + original = _image_bytes(Image.new("RGB", (24, 16), "#365a9c"), "PNG") + + prepared = prepare_tool_result_image(original, "image/png") + + assert prepared.data == original + assert prepared.media_type == "image/png" + assert prepared.changed is False + assert prepared.lossy is False + assert prepared.encoded_chars == ((len(original) + 2) // 3) * 4 + + +def test_large_dimension_is_resized_and_remains_valid() -> None: + original = _image_bytes(Image.new("RGBA", (3072, 1024), "#365a9c88"), "PNG") + + prepared = prepare_tool_result_image(original, "image/png") + + assert prepared.changed is True + assert prepared.lossy is False + assert prepared.media_type == "image/png" + assert prepared.encoded_chars <= MAX_TOOL_IMAGE_BASE64_CHARS + validate_image_bytes(prepared.data, prepared.media_type) + with Image.open(io.BytesIO(prepared.data)) as image: + assert max(image.size) <= MAX_TOOL_IMAGE_DIMENSION + assert image.getpixel((0, 0))[3] == 136 + + +def test_over_budget_static_png_uses_bounded_lossy_image() -> None: + pixels = random.Random(47).randbytes(1024 * 1024 * 3) + original = _image_bytes(Image.frombytes("RGB", (1024, 1024), pixels), "PNG") + assert ((len(original) + 2) // 3) * 4 > MAX_TOOL_IMAGE_BASE64_CHARS + + prepared = prepare_tool_result_image(original, "image/png") + + assert prepared.changed is True + assert prepared.lossy is True + assert prepared.media_type == "image/jpeg" + assert prepared.encoded_chars <= MAX_TOOL_IMAGE_BASE64_CHARS + validate_image_bytes(prepared.data, prepared.media_type) + with Image.open(io.BytesIO(prepared.data)) as image: + assert image.size == (1024, 1024) + + +def test_over_budget_transparent_png_preserves_alpha() -> None: + if not features.check("webp"): + pytest.skip("Pillow WebP encoder unavailable") + pixels = random.Random(71).randbytes(768 * 768 * 4) + original = _image_bytes(Image.frombytes("RGBA", (768, 768), pixels), "PNG") + assert ((len(original) + 2) // 3) * 4 > MAX_TOOL_IMAGE_BASE64_CHARS + + prepared = prepare_tool_result_image(original, "image/png") + + assert prepared.media_type == "image/webp" + assert prepared.lossy is True + assert prepared.encoded_chars <= MAX_TOOL_IMAGE_BASE64_CHARS + validate_image_bytes(prepared.data, prepared.media_type) + with Image.open(io.BytesIO(prepared.data)) as image: + assert image.mode == "RGBA" + assert image.getchannel("A").getextrema() == (0, 255) + + +def test_transparent_image_fails_explicitly_without_webp_encoder( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pixels = random.Random(72).randbytes(32 * 32 * 4) + original = _image_bytes(Image.frombytes("RGBA", (32, 32), pixels), "PNG") + monkeypatch.setattr(image_preparation, "MAX_TOOL_IMAGE_BASE64_CHARS", 256) + monkeypatch.setattr(image_preparation.features, "check", lambda _name: False) + + with pytest.raises(ValueError, match="WebP encoder is unavailable"): + prepare_tool_result_image(original, "image/png") + + +@pytest.mark.parametrize("encoder_error", [KeyError("WEBP"), OSError("encoder unavailable")]) +def test_transparent_image_wraps_webp_encoder_failure( + monkeypatch: pytest.MonkeyPatch, + encoder_error: Exception, +) -> None: + pixels = random.Random(73).randbytes(32 * 32 * 4) + original = _image_bytes(Image.frombytes("RGBA", (32, 32), pixels), "PNG") + monkeypatch.setattr(image_preparation, "MAX_TOOL_IMAGE_BASE64_CHARS", 256) + monkeypatch.setattr(image_preparation.features, "check", lambda _name: True) + encode = image_preparation._encode + + def failing_webp(image: Image.Image, *, format_name: str, **options: object) -> bytes: + if format_name == "WEBP": + raise encoder_error + return encode(image, format_name=format_name, **options) + + monkeypatch.setattr(image_preparation, "_encode", failing_webp) + + with pytest.raises(ValueError, match="could not be prepared"): + prepare_tool_result_image(original, "image/png") + + +def test_oversized_animation_is_not_silently_flattened() -> None: + first = Image.new("RGB", (2049, 2), "red") + second = Image.new("RGB", (2049, 2), "blue") + output = io.BytesIO() + first.save(output, format="GIF", save_all=True, append_images=[second], duration=100) + + with pytest.raises(ValueError, match="animated"): + prepare_tool_result_image(output.getvalue(), "image/gif") + + +def test_small_animation_keeps_all_frames() -> None: + first = Image.new("RGB", (16, 16), "red") + second = Image.new("RGB", (16, 16), "blue") + output = io.BytesIO() + first.save(output, format="GIF", save_all=True, append_images=[second], duration=100) + + prepared = prepare_tool_result_image(output.getvalue(), "image/gif") + + assert prepared.changed is False + assert prepared.data == output.getvalue() + with Image.open(io.BytesIO(prepared.data)) as image: + assert image.n_frames == 2 + + +def test_mismatched_mime_is_rejected() -> None: + original = _image_bytes(Image.new("RGB", (8, 8), "red"), "PNG") + + with pytest.raises(ValueError, match="mismatched"): + prepare_tool_result_image(original, "image/jpeg") diff --git a/tests/test_engine/test_agent_file_context.py b/tests/test_engine/test_agent_file_context.py index 94a46f39d8..ee15875462 100644 --- a/tests/test_engine/test_agent_file_context.py +++ b/tests/test_engine/test_agent_file_context.py @@ -416,5 +416,5 @@ async def capture(tool_call: Any) -> ToolResult: for message in requests[2] if isinstance(message.content, list) for block in message.content - ) + ) is not image_expected assert not any(event.kind == "error" for event in followup) diff --git a/tests/test_engine/test_agent_image_capability_guard.py b/tests/test_engine/test_agent_image_capability_guard.py index fca2fb70aa..f96928d2ca 100644 --- a/tests/test_engine/test_agent_image_capability_guard.py +++ b/tests/test_engine/test_agent_image_capability_guard.py @@ -2,6 +2,7 @@ import base64 import io +import json from collections.abc import AsyncIterator from pathlib import Path from types import SimpleNamespace @@ -1020,6 +1021,74 @@ async def load_images(call: ToolCall) -> ToolResult: assert request_config.provider_request_max_chars == 120_000 +@pytest.mark.asyncio +async def test_large_tool_image_is_prepared_before_provider_call_and_replay( + tmp_path: Path, +) -> None: + from opensquilla.attachment_workspace import AttachmentWorkspaceMaterializer + + provider = _ToolImageContinuationProvider() + context = ToolContext() + output = io.BytesIO() + Image.new("RGB", (1200, 1200), "#2468ac").save( + output, format="PNG", compress_level=0, + ) + original = output.getvalue() + assert len(base64.b64encode(original)) > 1_500_000 + materializer = AttachmentWorkspaceMaterializer( + media_root=tmp_path / "media", workspace_dir=tmp_path / "workspace", + ) + retained = materializer.materialize_bytes( + original, name="frame.png", mime="image/png", session_id="test-session", + ) + assert retained.rel_path is not None + + async def load_image(call: ToolCall) -> ToolResult: + context.tool_result_media[call.tool_use_id] = [{ + "mime": "image/png", "data": base64.b64encode(original).decode("ascii"), + "name": "frame.png", "local_path": retained.rel_path, + }] + return ToolResult( + tool_use_id=call.tool_use_id, tool_name=call.tool_name, + content=json.dumps({ + "status": "loaded", "path": "frame.png", "name": "frame.png", + "local_path": retained.rel_path, + }), + ) + + agent = Agent( + provider=provider, + config=AgentConfig(max_iterations=2, model_vision_support="supported"), + tool_context=context, tool_handler=load_image, + tool_definitions=[ToolDefinition( + name="observe", description="Load the frame.", + input_schema=ToolInputSchema(properties={}), + )], + ) + events = [event async for event in agent.run_turn("Inspect the frame.")] + + assert any(isinstance(event, DoneEvent) for event in events) + assert len(provider.calls) == 2 + provider_image = _tool_images(provider.calls[1]["messages"])[0] + replay_image = _tool_images(decode_assistant_replay( + next(event.assistant_replay for event in events if isinstance(event, DoneEvent)) + ))[0] + assert provider_image.data == replay_image.data + assert provider_image.media_type == replay_image.media_type + assert len(provider_image.data) <= 1536 * 1024 + assert base64.b64decode(provider_image.data, validate=True) != original + assert replay_image.local_path == retained.rel_path + assert replay_image.name == "frame.png" + assert (tmp_path / "workspace" / retained.rel_path).read_bytes() == original + assert any( + isinstance(block, ContentBlockToolResult) + and json.loads(block.content).get("local_path") == retained.rel_path + for message in provider.calls[1]["messages"] + if isinstance(message.content, list) + for block in message.content + ) + + @pytest.mark.asyncio async def test_loaded_tool_images_do_not_silently_drop_invalid_items_or_capacity() -> None: provider = _ToolImageContinuationProvider(reject_multiple=True) diff --git a/tests/test_engine/test_historical_image_followup.py b/tests/test_engine/test_historical_image_followup.py index c5e2604317..d7ae4311c2 100644 --- a/tests/test_engine/test_historical_image_followup.py +++ b/tests/test_engine/test_historical_image_followup.py @@ -1,7 +1,9 @@ from __future__ import annotations import base64 +import io import json +import random from collections.abc import AsyncIterator from dataclasses import dataclass from pathlib import Path @@ -10,9 +12,12 @@ from unittest.mock import MagicMock import pytest +from PIL import Image from opensquilla.attachment_refs import write_transcript_material +from opensquilla.contracts.image_preparation import prepare_tool_result_image from opensquilla.engine import Agent, AgentConfig +from opensquilla.engine.agent import _serialize_assistant_replay from opensquilla.engine.pipeline import TurnContext from opensquilla.engine.runtime import TurnRunner, _SelectorFallbackProvider from opensquilla.engine.steps.squilla_router import apply_squilla_router @@ -25,7 +30,12 @@ ModelCapabilities, TextDeltaEvent, ) -from opensquilla.provider.types import ContentBlockImage, ContentBlockText +from opensquilla.provider.types import ( + ContentBlockImage, + ContentBlockText, + ContentBlockToolResult, + ContentBlockToolUse, +) from opensquilla.session.attachment_manifest import ( ATTACHMENT_MANIFEST_PROVIDER, ATTACHMENT_MANIFEST_STATE_KIND, @@ -91,7 +101,7 @@ async def get_context_states(self, session_key: str) -> list[Any]: # noqa: ARG0 @pytest.mark.parametrize("lookback", [0, 1, 2]) @pytest.mark.parametrize("queued_before_completion", [False, True]) @pytest.mark.parametrize("prior_user_retained", [False, True]) -async def test_tool_loaded_active_images_ignore_the_legacy_turn_window( +async def test_tool_loaded_images_route_only_the_immediate_followup( lookback: int, queued_before_completion: bool, prior_user_retained: bool, ) -> None: manager = _FakeSessionManager() @@ -122,10 +132,14 @@ async def test_tool_loaded_active_images_ignore_the_legacy_turn_window( key, exclude_last_user=True, bound_user_message_id="current", ) - assert context["history_has_recent_image"] is True - if prior_user_retained: - assert context["turns_since_last_image"] == (0 if queued_before_completion else 1) - assert context["history_image_turn_count"] == 1 + assert context.get("history_has_recent_image") is (True if queued_before_completion else None) + if prior_user_retained and queued_before_completion: + assert context["turns_since_last_image"] == 0 + else: + assert "turns_since_last_image" not in context + assert context.get("history_image_turn_count") == ( + 1 if queued_before_completion else None + ) class _CanonicalSessionManager(_FakeSessionManager): @@ -211,6 +225,66 @@ def _message_has_marker(message: Message, marker: str) -> bool: ) or isinstance(message.content, str) and marker in message.content +@pytest.mark.asyncio +async def test_prior_tool_images_are_markers_while_user_image_remains_available() -> None: + manager = _FakeSessionManager() + key = "agent:main:tool-image-replay-projection" + config = GatewayConfig(llm={"provider": "openai"}) + runner = TurnRunner(provider_selector=MagicMock(), session_manager=manager, config=config) + await manager.create(key) + await manager.append_message(key, "user", _inline_image_envelope("Reference picture.")) + image_data = _b64(image_bytes(color="#123456")) + replay = _serialize_assistant_replay([ + Message(role="assistant", content=[ + ContentBlockToolUse(id="view-1", name="image", input={"path": "/tmp/frame.png"}), + ]), + Message(role="user", content=[ + ContentBlockToolResult(tool_use_id="view-1", content='{"status":"loaded"}'), + *[ + ContentBlockImage( + media_type="image/png", data=image_data, + local_path=f"/tmp/frame-{index}.png", durable_retained=True, + ) + for index in range(12) + ], + ]), + Message(role="assistant", content="I checked the frame."), + ]) + entry = await manager.append_message(key, "assistant", "I checked the frame.") + entry.assistant_replay = replay + await manager.append_message(key, "user", "Continue the edit.", "current") + + provider = _CapturingProvider() + agent = Agent(provider=provider, config=AgentConfig( + max_iterations=1, model_vision_support="supported", + preserve_historical_images=True, + metadata={"image_context_has_images": True}, + )) + await runner._load_history(agent, key, bound_user_message_id="current") + assert sum( + isinstance(block, ContentBlockImage) + for message in agent._history if isinstance(message.content, list) + for block in message.content + ) == 13 + events = [event async for event in agent.run_turn("Continue the edit.")] + + assert any(event.kind == "done" for event in events) + sent = provider.calls[0]["messages"] + sent_images = [ + block for message in sent if isinstance(message.content, list) + for block in message.content if isinstance(block, ContentBlockImage) + ] + assert len(sent_images) == 1 + assert sent_images[0].data == _b64(image_bytes()) + assert sum( + isinstance(block, ContentBlockText) and "历史图片本回合未重新读取" in block.text + for message in sent if isinstance(message.content, list) + for block in message.content + ) == 12 + assert not any(_message_has_marker(message, "原图已保留") for message in sent) + assert entry.assistant_replay == replay + + @pytest.mark.parametrize("legacy_lookback", [0, 1, 8]) async def test_image_route_keeps_old_text_and_image_paths_until_compaction( tmp_path: Path, legacy_lookback: int, @@ -357,6 +431,10 @@ async def no_sandbox(*args, **kwargs): assert _build_openai_messages(replayed) == _build_openai_messages( decode_assistant_replay(canonical)[-1] ) + else: + assert "identical copy of the original image" in json.loads( + tool_result.content + )["retention_note"] token = current_tool_context.set(context) try: await filesystem.read_file(image.local_path, _tool_use_id="reread") @@ -368,6 +446,97 @@ async def no_sandbox(*args, **kwargs): assert context.tool_result_media["reread"][0]["data"] == _b64(payload) +@pytest.mark.parametrize("source_case", ["same_session", "changed_source", "fork", "write_deny"]) +async def test_prepared_tool_image_replay_keeps_original_or_labels_derivative( + tmp_path: Path, source_case: str, +) -> None: + from opensquilla.attachment_workspace import AttachmentWorkspaceMaterializer + from opensquilla.tools.types import ToolContext + + manager = _FakeSessionManager() + key = "agent:main:prepared-tool-image-history" + node = await manager.create(key) + config = GatewayConfig(workspace_dir=str(tmp_path / "workspace")) + materializer = AttachmentWorkspaceMaterializer( + media_root=tmp_path / "media", workspace_dir=config.workspace_dir, + ) + output = io.BytesIO() + pixels = random.Random(42).randbytes(900 * 900 * 3) + Image.frombytes("RGB", (900, 900), pixels).save(output, format="PNG", compress_level=0) + original = output.getvalue() + prepared = prepare_tool_result_image(original, "image/png") + assert prepared.media_type == "image/jpeg" + parent_id = node.session_id if source_case in {"same_session", "changed_source"} else "parent" + source = materializer.materialize_bytes( + original, name="frame.png", mime="image/png", session_id=parent_id, + ) + assert source.rel_path is not None + original_path = Path(config.workspace_dir) / source.rel_path + if source_case == "changed_source": + original_path.write_bytes(image_bytes()) + source_url = "https://images.example.test/frame.png" + receipt = json.dumps({ + "status": "loaded", "path": source_url, "source_url": source_url, + "local_path": source.rel_path, "name": "frame.png", + }) + canonical = _serialize_assistant_replay([ + Message(role="assistant", content=[ContentBlockToolUse( + id="image-read", name="image", input={"path": source_url}, + )]), + Message(role="user", content=[ + ContentBlockToolResult(tool_use_id="image-read", content=receipt), + ContentBlockImage( + media_type=prepared.media_type, data=_b64(prepared.data), + name="frame.png", source_url=source_url, + local_path=source.rel_path, durable_retained=True, + ), + ]), + ]) + await manager.append_message(key, "user", "Load the frame.") + saved = await manager.append_message(key, "assistant", "Frame loaded.") + saved.assistant_replay = canonical + await manager.append_message(key, "user", "Read the source image again.", "current") + context = ToolContext( + workspace_dir=config.workspace_dir, workspace_strict=True, + artifact_session_id=node.session_id, artifact_media_root=str(tmp_path / "media"), + sandbox_gateway_config=config, + ) + if source_case == "write_deny": + context.workspace_write_deny_globs = [".opensquilla/attachments/**"] + agent = Agent(provider=_CapturingProvider(), tool_context=context, config=AgentConfig( + model_vision_support="supported", workspace_dir=config.workspace_dir, + )) + runner = TurnRunner(provider_selector=MagicMock(), session_manager=manager, config=config) + await runner._load_history(agent, key, bound_user_message_id="current") + + replayed = agent._history[-1] + image = next(block for block in replayed.content if isinstance(block, ContentBlockImage)) + tool_result = next( + block for block in replayed.content if isinstance(block, ContentBlockToolResult) + ) + rewritten_receipt = json.loads(tool_result.content) + assert image.data == _b64(prepared.data) + assert image.media_type == "image/jpeg" + assert json.dumps(saved.assistant_replay) == json.dumps(canonical) + if source_case == "same_session": + assert image.local_path == source.rel_path + assert image.name == "frame.png" + assert tool_result.content == receipt + assert original_path.read_bytes() == original + elif source_case == "write_deny": + assert image.local_path is None + assert "local_path" not in rewritten_receipt + assert "No readable local copy" in rewritten_receipt["retention_note"] + else: + assert image.local_path != source.rel_path + assert image.local_path is not None + assert image.local_path.endswith("-frame.jpg") + assert image.name == rewritten_receipt["name"] == "frame.jpg" + assert rewritten_receipt["local_path"] == image.local_path + assert "prepared derivative" in rewritten_receipt["retention_note"] + assert (Path(config.workspace_dir) / image.local_path).read_bytes() == prepared.data + + @pytest.mark.parametrize("image_source", ["active", "archive", "bound", "agent_history"]) async def test_text_primary_fallback_recovers_selected_historical_original( image_source: str, diff --git a/tests/test_engine/test_local_wire_rejection_usage.py b/tests/test_engine/test_local_wire_rejection_usage.py new file mode 100644 index 0000000000..2579f80e49 --- /dev/null +++ b/tests/test_engine/test_local_wire_rejection_usage.py @@ -0,0 +1,299 @@ +from __future__ import annotations + +from collections.abc import AsyncIterator +from dataclasses import replace +from typing import Any + +import httpx +import pytest + +from opensquilla.engine import Agent, AgentConfig +from opensquilla.engine.usage_accounting import ( + UsageAccountingScope, + UsageCallResult, + UsageCallStart, + UsageExecutionContext, + account_provider_stream, + bind_usage_accounting_scope, +) +from opensquilla.provider.openai import OpenAIProvider +from opensquilla.provider.request_proof import local_budget_rejection_event +from opensquilla.provider.types import ChatConfig, ErrorEvent, Message + +_WIRE_PROOF = { + "wire_byte_budget_exhausted": True, + "fits_wire_byte_budget": False, + "wire_json_bytes": 101, + "wire_body_max_bytes": 100, +} + + +class _Sink: + def __init__(self) -> None: + self.started: list[UsageCallStart] = [] + self.not_sent: list[UsageCallStart] = [] + self.unknown: list[tuple[UsageCallStart, str]] = [] + self.finalized: list[tuple[UsageCallStart, UsageCallResult]] = [] + + async def start(self, call: UsageCallStart) -> None: + self.started.append(call) + + async def mark_not_sent(self, call: UsageCallStart) -> None: + self.not_sent.append(call) + + async def mark_unknown(self, call: UsageCallStart, reason: str) -> None: + self.unknown.append((call, reason)) + + async def finalize(self, call: UsageCallStart, result: UsageCallResult) -> None: + self.finalized.append((call, result)) + + +class _LegacySink: + """An external sink may still implement the original accounting contract.""" + + def __init__(self) -> None: + self.started: list[UsageCallStart] = [] + self.unknown: list[tuple[UsageCallStart, str]] = [] + + async def start(self, call: UsageCallStart) -> None: + self.started.append(call) + + async def finalize(self, call: UsageCallStart, result: UsageCallResult) -> None: + del call, result + + async def mark_unknown(self, call: UsageCallStart, reason: str) -> None: + self.unknown.append((call, reason)) + + +def _scope(sink: _Sink) -> UsageAccountingScope: + return UsageAccountingScope( + sink=sink, + context=UsageExecutionContext( + execution_id="synthetic-wire-rejection", + agent_run_id="synthetic-wire-rejection", + ), + ) + + +class _RejectedProvider: + provider_name = "synthetic-provider" + retry_failed_call_safe = False + + def __init__(self, error: ErrorEvent) -> None: + self.error = error + + def chat( + self, + messages: list[Message], + tools: list[Any] | None = None, + config: ChatConfig | None = None, + ) -> AsyncIterator[ErrorEvent]: + del messages, tools, config + return self._stream() + + async def _stream(self) -> AsyncIterator[ErrorEvent]: + yield self.error + + +@pytest.mark.parametrize("upstream", [False, True]) +async def test_direct_agent_accounts_only_dispatched_failures_as_unknown( + upstream: bool, +) -> None: + error = ( + ErrorEvent(code="413", http_status=413, message="synthetic upstream rejection") + if upstream + else local_budget_rejection_event(_WIRE_PROOF) + ) + sink = _Sink() + agent = Agent( + provider=_RejectedProvider(error), + config=AgentConfig( + max_iterations=1, + max_provider_retries=0, + provider_id="synthetic-provider", + model_id="synthetic-model", + ), + ) + with bind_usage_accounting_scope(_scope(sink)): + _ = [event async for event in agent.run_turn("synthetic input")] + + assert len(sink.started) == 1 + assert not sink.finalized + if upstream: + assert sink.not_sent == [] + assert [call.event_id for call, _ in sink.unknown] == [sink.started[0].event_id] + else: + assert [call.event_id for call in sink.not_sent] == [sink.started[0].event_id] + assert sink.unknown == [] + + +async def test_physical_stream_adaptor_rejects_wire_body_without_http_or_unknown_usage( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("OPENSQUILLA_PROVIDER_REQUEST_MAX_BODY_BYTES", "100") + provider = OpenAIProvider( + api_key="synthetic-key", + model="synthetic-model", + base_url="https://example.test/v1", + ) + + def unexpected_stream(*args: Any, **kwargs: Any) -> None: + del args, kwargs + raise AssertionError("local rejection must precede HTTP dispatch") + + monkeypatch.setattr(httpx.AsyncClient, "stream", unexpected_stream) + sink = _Sink() + with bind_usage_accounting_scope(_scope(sink)): + events = [ + event + async for event in account_provider_stream( + lambda: provider.chat( + [Message(role="user", content="synthetic input")], + config=ChatConfig(provider_request_max_chars=20_000), + ), + provider="synthetic-provider", + model="synthetic-model", + ) + ] + + assert len(events) == 1 + assert isinstance(events[0], ErrorEvent) + assert events[0]._opensquilla_request_not_dispatched is True + assert [call.event_id for call in sink.not_sent] == [sink.started[0].event_id] + assert sink.unknown == [] + + +async def test_local_character_budget_rejection_is_not_a_provider_usage_gap( + monkeypatch: pytest.MonkeyPatch, +) -> None: + provider = OpenAIProvider( + api_key="synthetic-key", + model="synthetic-model", + base_url="https://example.test/v1", + ) + + def unexpected_stream(*args: Any, **kwargs: Any) -> None: + del args, kwargs + raise AssertionError("local budget rejection must precede HTTP dispatch") + + monkeypatch.setattr(httpx.AsyncClient, "stream", unexpected_stream) + sink = _Sink() + with bind_usage_accounting_scope(_scope(sink)): + events = [ + event + async for event in account_provider_stream( + lambda: provider.chat( + [Message(role="user", content="synthetic input")], + config=ChatConfig(provider_request_max_chars=1), + ), + provider="synthetic-provider", + model="synthetic-model", + ) + ] + + assert len(events) == 1 + assert isinstance(events[0], ErrorEvent) + assert events[0]._opensquilla_request_not_dispatched is True + assert [call.event_id for call in sink.not_sent] == [sink.started[0].event_id] + assert sink.unknown == [] + + +async def test_physical_stream_does_not_trust_upstream_budget_prose() -> None: + sink = _Sink() + + async def upstream() -> AsyncIterator[ErrorEvent]: + # A provider may echo a local-looking code and proof. It cannot mark + # the adapter's private pre-dispatch evidence. + yield ErrorEvent( + code="provider_request_budget_exhausted", + message=local_budget_rejection_event(_WIRE_PROOF).message, + http_status=413, + ) + + with bind_usage_accounting_scope(_scope(sink)): + _ = [ + event + async for event in account_provider_stream( + upstream, + provider="synthetic-provider", + model="synthetic-model", + ) + ] + + assert sink.not_sent == [] + assert [call.event_id for call, _ in sink.unknown] == [sink.started[0].event_id] + + +async def test_external_legacy_sink_keeps_conservative_terminal_outcome() -> None: + sink = _LegacySink() + + async def rejected() -> AsyncIterator[ErrorEvent]: + yield local_budget_rejection_event(_WIRE_PROOF) + + with bind_usage_accounting_scope(_scope(sink)): + _ = [ + event + async for event in account_provider_stream( + rejected, provider="synthetic-provider", model="synthetic-model", + ) + ] + + assert len(sink.started) == 1 + assert sink.unknown == [(sink.started[0], "cancelled_before_provider_request")] + + +async def test_real_http_413_remains_unknown_after_provider_dispatch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("OPENSQUILLA_PROVIDER_REQUEST_MAX_BODY_BYTES", "0") + requests: list[httpx.Request] = [] + + def respond(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(413, json={ + "error": { + "code": "provider_request_budget_exhausted", + "message": local_budget_rejection_event(_WIRE_PROOF).message, + }, + }) + + provider = OpenAIProvider( + api_key="synthetic-key", + model="synthetic-model", + base_url="https://example.test/v1", + ) + + async def build_client(*, timeout: Any) -> httpx.AsyncClient: + return httpx.AsyncClient(transport=httpx.MockTransport(respond), timeout=timeout) + + monkeypatch.setattr(provider, "_build_http_client", build_client) + sink = _Sink() + with bind_usage_accounting_scope(_scope(sink)): + events = [ + event + async for event in account_provider_stream( + lambda: provider.chat( + [Message(role="user", content="synthetic input")], + config=ChatConfig(provider_request_max_chars=20_000), + ), + provider="synthetic-provider", + model="synthetic-model", + ) + ] + + assert len(requests) == 1 + assert len(events) == 1 + assert isinstance(events[0], ErrorEvent) + assert events[0].http_status == 413 + assert events[0]._opensquilla_request_not_dispatched is False + assert sink.not_sent == [] + assert [call.event_id for call, _ in sink.unknown] == [sink.started[0].event_id] + + +def test_local_budget_factory_marks_all_pre_dispatch_rejections() -> None: + character_proof = {"fits": False, "fits_char_budget": False, "proof_budget": 1} + assert local_budget_rejection_event(character_proof)._opensquilla_request_not_dispatched + assert local_budget_rejection_event(_WIRE_PROOF)._opensquilla_request_not_dispatched is True + assert replace( + local_budget_rejection_event(_WIRE_PROOF), generation_epoch=2 + )._opensquilla_request_not_dispatched is True diff --git a/tests/test_engine/test_pending_tool_image_history_window.py b/tests/test_engine/test_pending_tool_image_history_window.py new file mode 100644 index 0000000000..71a005a083 --- /dev/null +++ b/tests/test_engine/test_pending_tool_image_history_window.py @@ -0,0 +1,300 @@ +"""A failed image handoff remains usable across the ordinary turn window.""" + +from __future__ import annotations + +import base64 +import io +import random +from collections.abc import AsyncIterator +from typing import Any + +import pytest +from PIL import Image + +from opensquilla.contracts.image_preparation import MAX_TOOL_IMAGE_BASE64_CHARS +from opensquilla.engine import Agent, AgentConfig +from opensquilla.engine.agent import ( + _active_user_message_index_for_request, + _pending_tool_image_window_context, + _serialize_assistant_replay, +) +from opensquilla.engine.history import limit_turns, reconstruct_messages_from_entry +from opensquilla.engine.runtime import TurnRunner +from opensquilla.engine.types import DoneEvent +from opensquilla.provider import ChatConfig, Message +from opensquilla.provider import DoneEvent as ProviderDoneEvent +from opensquilla.provider import TextDeltaEvent as ProviderTextDeltaEvent +from opensquilla.provider.types import ( + ContentBlockImage, + ContentBlockToolResult, + ContentBlockToolUse, +) +from tests.helpers.image_bytes import image_bytes + + +class _CapturingProvider: + provider_name = "synthetic" + + def __init__(self) -> None: + self.messages: list[Message] = [] + + def chat( + self, + messages: list[Message], + tools: list[Any] | None = None, # noqa: ARG002 + config: ChatConfig | None = None, # noqa: ARG002 + ) -> AsyncIterator[Any]: + self.messages = messages + return self._answer() + + async def _answer(self) -> AsyncIterator[Any]: + yield ProviderTextDeltaEvent(text="Frame analyzed.") + yield ProviderDoneEvent(stop_reason="end_turn", input_tokens=3, output_tokens=2) + + async def list_models(self) -> list[Any]: + return [] + + +def test_request_only_image_context_is_not_the_active_user_turn() -> None: + image = ContentBlockImage( + media_type="image/png", data=base64.b64encode(image_bytes()).decode("ascii") + ) + tool_result = Message( + role="user", + content=[ContentBlockToolResult(tool_use_id="view-1", content="loaded"), image], + ) + context = _pending_tool_image_window_context([tool_result]) + assert context is not None + request = [Message(role="user", content="Inspect this frame."), context] + + assert _active_user_message_index_for_request(request) == 0 + assert limit_turns(request, max_turns=1) == request + + +@pytest.mark.asyncio +async def test_pending_tool_image_survives_one_failed_user_turn_with_tiny_window() -> None: + encoded = base64.b64encode(image_bytes()).decode("ascii") + replay = _serialize_assistant_replay( + [ + Message( + role="assistant", + content=[ContentBlockToolUse(id="view-1", name="observe", input={})], + ), + Message( + role="user", + content=[ + ContentBlockToolResult(tool_use_id="view-1", content="frame loaded"), + ContentBlockImage(media_type="image/png", data=encoded), + ], + ), + ] + ) + history = [Message(role="user", content="Inspect this frame.")] + history.extend( + reconstruct_messages_from_entry("assistant", "", None, assistant_replay=replay) + ) + history.append(Message(role="user", content="Please try again after the failed request.")) + provider = _CapturingProvider() + agent = Agent( + provider=provider, + config=AgentConfig( + max_iterations=1, + max_history_turns=1, + model_vision_support="supported", + preserve_historical_images=True, + ), + ) + agent.set_history(history) + + events = [event async for event in agent.run_turn("Continue analyzing that frame.")] + + assert any(isinstance(event, DoneEvent) for event in events) + images = [ + block + for message in provider.messages + if isinstance(message.content, list) + for block in message.content + if isinstance(block, ContentBlockImage) + ] + assert [image.data for image in images] == [encoded] + assert all( + not message._request_only_pending_tool_image + for message in agent.history_snapshot() + ) + + +def test_nested_tool_images_drop_unverified_retained_path() -> None: + image = ContentBlockImage( + media_type="image/png", + data=base64.b64encode(image_bytes()).decode("ascii"), + name="frame.png", + local_path=".opensquilla/attachments/other/123456789abc-frame.png", + durable_retained=True, + ) + message = Message( + role="user", + content=[ + ContentBlockToolResult( + tool_use_id="view-1", + content=[image, {"type": "image", "local_path": "stale.png", + "durable_retained": True, "data": image.data, + "media_type": "image/png"}], + ) + ], + ) + + TurnRunner._restore_retained_history_image_paths( + [message], materializer=None, session_id="current" + ) + + assert isinstance(message.content, list) + result = message.content[0] + assert isinstance(result, ContentBlockToolResult) + assert isinstance(result.content, list) + native_image = result.content[0] + assert isinstance(native_image, ContentBlockImage) + assert native_image.local_path is None + assert native_image.durable_retained is None + mapped_image = result.content[1] + assert isinstance(mapped_image, dict) + assert "local_path" not in mapped_image + assert "durable_retained" not in mapped_image + + +@pytest.mark.asyncio +@pytest.mark.parametrize("max_history_turns", [1, 20]) +@pytest.mark.parametrize("image_count", [6, 12]) +async def test_pending_batch_caps_frames_and_describes_omissions( + max_history_turns: int, + image_count: int, +) -> None: + images = [ + ContentBlockImage( + media_type="image/png", + data=base64.b64encode( + image_bytes(color=f"#{index:02x}{index:02x}{index:02x}") + ).decode("ascii"), + attachment_id=f"frame-{index}", + ) + for index in range(image_count) + ] + replay = _serialize_assistant_replay( + [ + Message( + role="assistant", + content=[ContentBlockToolUse(id="view-1", name="observe", input={})], + ), + Message( + role="user", + content=[ + ContentBlockToolResult(tool_use_id="view-1", content="frames loaded"), + *images, + ], + ), + ] + ) + history = [Message(role="user", content="Inspect the video frames.")] + history.extend( + reconstruct_messages_from_entry("assistant", "", None, assistant_replay=replay) + ) + history.append(Message(role="user", content="Try again after the failure.")) + provider = _CapturingProvider() + agent = Agent( + provider=provider, + config=AgentConfig( + max_iterations=1, + max_history_turns=max_history_turns, + model_vision_support="supported", + preserve_historical_images=True, + ), + ) + agent.set_history(history) + + events = [event async for event in agent.run_turn("Continue.")] + assert any(isinstance(event, DoneEvent) for event in events) + + request_images = [ + block + for message in provider.messages + if isinstance(message.content, list) + for block in message.content + if isinstance(block, ContentBlockImage) + ] + assert [image.data for image in request_images] == [ + image.data for image in images[-4:] + ] + assert any( + f"{image_count - 4} images omitted by recency or size limits" in block.text + for message in provider.messages + if isinstance(message.content, list) + for block in message.content + if hasattr(block, "text") + ) + assert all( + not message._request_only_pending_tool_image + for message in agent.history_snapshot() + ) + if max_history_turns == 20: + assert sum( + isinstance(block, ContentBlockImage) + for message in agent.history_snapshot() + if isinstance(message.content, list) + for block in message.content + ) == image_count + + +@pytest.mark.asyncio +async def test_legacy_oversized_pending_image_is_reprepared_only_for_request() -> None: + pixels = random.Random(7).randbytes(1300 * 1300 * 3) + output = io.BytesIO() + Image.frombytes("RGB", (1300, 1300), pixels).save(output, format="PNG") + encoded = base64.b64encode(output.getvalue()).decode("ascii") + assert len(encoded) > MAX_TOOL_IMAGE_BASE64_CHARS + replay = _serialize_assistant_replay([ + Message( + role="assistant", + content=[ContentBlockToolUse(id="view-1", name="observe", input={})], + ), + Message( + role="user", + content=[ + ContentBlockToolResult(tool_use_id="view-1", content="large frame loaded"), + ContentBlockImage(media_type="image/png", data=encoded), + ], + ), + ]) + history = [Message(role="user", content="Inspect the frame.")] + history.extend(reconstruct_messages_from_entry( + "assistant", "", None, assistant_replay=replay, + )) + provider = _CapturingProvider() + agent = Agent( + provider=provider, + config=AgentConfig( + max_iterations=1, + max_history_turns=20, + model_vision_support="supported", + preserve_historical_images=True, + ), + ) + agent.set_history(history) + + events = [event async for event in agent.run_turn("Continue.")] + + assert any(isinstance(event, DoneEvent) for event in events) + request_images = [ + block + for message in provider.messages + if isinstance(message.content, list) + for block in message.content + if isinstance(block, ContentBlockImage) + ] + assert len(request_images) == 1 + assert len(request_images[0].data) <= MAX_TOOL_IMAGE_BASE64_CHARS + assert request_images[0].data != encoded + assert any( + isinstance(block, ContentBlockImage) and block.data == encoded + for message in agent.history_snapshot() + if isinstance(message.content, list) + for block in message.content + ) diff --git a/tests/test_engine/test_runtime_artifacts.py b/tests/test_engine/test_runtime_artifacts.py index b7af55d475..1fa1b978d6 100644 --- a/tests/test_engine/test_runtime_artifacts.py +++ b/tests/test_engine/test_runtime_artifacts.py @@ -1258,7 +1258,11 @@ def set_request_image_context(self, messages) -> None: from opensquilla.engine.history import decode_assistant_replay captured = decode_assistant_replay(assistant.assistant_replay) - assert history_capture.history[:-1] == captured + # Reload also restores private tool-image acknowledgement state from + # turn context; the replay payload carries the public provider messages. + assert [message.model_dump() for message in history_capture.history[:-1]] == [ + message.model_dump() for message in captured + ] assert history_capture.history[-1].role == "user" if native_replay: native_assistants = [message for message in captured if message.role == "assistant"] diff --git a/tests/test_engine/test_runtime_router_fallback.py b/tests/test_engine/test_runtime_router_fallback.py index 527d5f5765..cc8718520d 100644 --- a/tests/test_engine/test_runtime_router_fallback.py +++ b/tests/test_engine/test_runtime_router_fallback.py @@ -170,7 +170,7 @@ async def slow_router(ctx: TurnContext) -> TurnContext: "system prompt", [], ), - timeout=0.25, + timeout=1.0, ) assert resolved_provider is provider @@ -211,7 +211,7 @@ async def test_squilla_router_timeout_does_not_late_append_history( "system prompt", [], ), - timeout=0.25, + timeout=1.0, ) await asyncio.sleep(0.1) @@ -262,7 +262,7 @@ async def test_squilla_router_timeout_does_not_late_mutate_history_entries( "system prompt", [], ), - timeout=0.25, + timeout=1.0, ) await asyncio.sleep(0.1) diff --git a/tests/test_engine/test_selector_fallback_routed_model.py b/tests/test_engine/test_selector_fallback_routed_model.py index 5287a13857..d9b90f54ef 100644 --- a/tests/test_engine/test_selector_fallback_routed_model.py +++ b/tests/test_engine/test_selector_fallback_routed_model.py @@ -49,9 +49,11 @@ from opensquilla.provider.model_catalog import ModelCatalog from opensquilla.provider.openai import OpenAIProvider from opensquilla.provider.protocol import IMAGE_INPUT_UNSUPPORTED_CODE +from opensquilla.provider.request_proof import local_budget_rejection_event from opensquilla.provider.selector import ModelSelector, ProviderConfig, SelectorConfig from opensquilla.provider.types import ( ContentBlockImage, + ContentBlockText, ContentBlockToolResult, ContentBlockToolUse, ) @@ -2636,8 +2638,10 @@ def next_fallback_after_failure_matching(self, _exc, *, predicate): ] -async def test_local_admission_failure_escalates_to_larger_authorized_leg( +@pytest.mark.parametrize("wire_rejection", [False, True]) +async def test_local_admission_only_escalates_context_budget_to_larger_leg( monkeypatch: Any, + wire_rejection: bool, ) -> None: class _AdmissionProvider: provider_name = "openai" @@ -2650,10 +2654,17 @@ def __init__(self, model: str, *, fits: bool) -> None: async def chat(self, messages, tools=None, config=None): del messages, tools, config if not self.fits: - yield ErrorEvent( - message='{"reason":"provider_request_budget_exhausted"}', - code="provider_request_budget_exhausted", + proof = ( + { + "wire_byte_budget_exhausted": True, + "fits_wire_byte_budget": False, + "wire_json_bytes": 101, + "wire_body_max_bytes": 100, + } + if wire_rejection + else {"fits": False, "fits_char_budget": False, "proof_budget": 1} ) + yield local_budget_rejection_event(proof) return self.network_calls += 1 yield TextDeltaEvent(text="fallback answer") @@ -2717,10 +2728,14 @@ def _resolve_context_window( ] assert selector.primary.network_calls == 0 - assert selector.fallback.network_calls == 1 - assert any(isinstance(event, TextDeltaEvent) for event in events) - assert metadata["executed_model"] == "large-model" - assert metadata["router_fallback_reason"] == "local_admission_escalation" + if wire_rejection: + assert selector.fallback.network_calls == 0 + assert any(isinstance(event, ErrorEvent) for event in events) + else: + assert selector.fallback.network_calls == 1 + assert any(isinstance(event, TextDeltaEvent) for event in events) + assert metadata["executed_model"] == "large-model" + assert metadata["router_fallback_reason"] == "local_admission_escalation" async def test_local_admission_skips_larger_fallback_with_explicit_tool_denial( @@ -3321,6 +3336,7 @@ def _image_continuation_runtime( automatic: bool = True, vision_support: str = "supported", target_provider: str = "openrouter", + image_context_length: int = 96_000, ) -> SimpleNamespace: """Exercise real selection and routing with credential-free physical legs.""" @@ -3349,7 +3365,7 @@ def _image_continuation_runtime( }, { "id": image_model, - "context_length": 96_000, + "context_length": image_context_length, "top_provider": {"max_completion_tokens": 16_000}, "supported_parameters": ["tools"], **({"architecture": {"input_modalities": ["text", "image"]}} @@ -3722,3 +3738,84 @@ async def forbidden_analysis(*args): assert runtime.metadata["image_input_output_count"] == 0 assert runtime.metadata["image_input_marker_count"] == 1 assert context.tool_result_media == {} + + +@pytest.mark.asyncio +async def test_twelve_tool_frames_use_bounded_capacity_for_vision_route( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from opensquilla.engine.history import HistoryReplayProjection, project_history_replay_capacity + + runtime = _image_continuation_runtime(monkeypatch, image_context_length=24_000) + runtime.primary._streams = [[ + ToolUseStartEvent(tool_use_id="frames", tool_name="observe"), + ToolUseEndEvent(tool_use_id="frames", tool_name="observe", arguments={}), + DoneEvent(stop_reason="tool_use"), + ]] + context = ToolContext(is_owner=True) + frames = [ + b64encode(image_bytes(color=f"#{index:02x}{index:02x}{index:02x}")).decode("ascii") + for index in range(12) + ] + + async def load_frames(call: Any) -> ToolResult: + context.tool_result_media[call.tool_use_id] = [ + {"mime": "image/png", "data": data, "attachment_id": f"frame-{index}"} + for index, data in enumerate(frames) + ] + return ToolResult( + tool_use_id=call.tool_use_id, + tool_name=call.tool_name, + content="Twelve synthetic frames loaded", + ) + + agent = Agent( + provider=runtime.wrapper, + config=AgentConfig( + model_id="synthetic/text-small", + max_iterations=2, + max_provider_retries=0, + model_vision_support="unsupported", + model_capabilities=ModelCapabilities(supports_vision=False, supports_tools=True), + ), + tool_context=context, + tool_handler=load_frames, + tool_definitions=[ToolDefinition( + name="observe", + description="Load synthetic frames.", + input_schema=ToolInputSchema(properties={}), + )], + ) + + events = [event async for event in agent.run_turn("Inspect the frames.")] + + assert not [event for event in events if event.kind == "error"], events + assert any(event.kind == "done" for event in events) + assert len(runtime.primary.calls) == 1 + assert runtime.wrapper.active_deployment_config().model == runtime.image_model + call = runtime.built[-1].calls[0] + sent_frames = [ + block.data + for message in call["messages"] + if isinstance(message.content, list) + for block in message.content + if isinstance(block, ContentBlockImage) + ] + assert sent_frames == frames[-4:] + assert any( + "8 images omitted by recency or size limits" in block.text + for message in call["messages"] + if isinstance(message.content, list) + for block in message.content + if isinstance(block, ContentBlockText) + ) + canonical_history = agent.history_snapshot() + assert project_history_replay_capacity( + HistoryReplayProjection(messages=tuple(canonical_history)), + ).estimated_tokens > 8_000 + assert sum( + isinstance(block, ContentBlockImage) + for message in canonical_history + if isinstance(message.content, list) + for block in message.content + ) == 12 diff --git a/tests/test_engine/test_unseen_tool_image_replay.py b/tests/test_engine/test_unseen_tool_image_replay.py new file mode 100644 index 0000000000..b91501dfd3 --- /dev/null +++ b/tests/test_engine/test_unseen_tool_image_replay.py @@ -0,0 +1,715 @@ +from __future__ import annotations + +import base64 +from collections.abc import AsyncIterator +from copy import deepcopy +from dataclasses import dataclass +from typing import Any +from unittest.mock import MagicMock + +import pytest + +from opensquilla.engine import Agent, AgentConfig, ToolResult +from opensquilla.engine.agent import _serialize_assistant_replay +from opensquilla.engine.history import ( + TOOL_IMAGE_NATIVE_ACK_CONTEXT_KEY, + limit_turns, + reconstruct_messages_from_entry, + repair_tool_pairing, +) +from opensquilla.engine.runtime import TurnRunner +from opensquilla.engine.types import DoneEvent, ErrorEvent, ToolCall +from opensquilla.gateway.config import GatewayConfig +from opensquilla.provider import ( + ChatConfig, + Message, + ToolDefinition, + ToolInputSchema, + ToolUseEndEvent, + ToolUseStartEvent, +) +from opensquilla.provider import DoneEvent as ProviderDoneEvent +from opensquilla.provider import ErrorEvent as ProviderErrorEvent +from opensquilla.provider import TextDeltaEvent as ProviderTextDeltaEvent +from opensquilla.provider.types import ( + ContentBlockImage, + ContentBlockText, + ContentBlockToolResult, + ContentBlockToolUse, +) +from opensquilla.tools.types import ToolContext +from tests.helpers.image_bytes import image_bytes + + +def _image(image_id: str) -> ContentBlockImage: + color = "#112233" if image_id == "first-frame" else "#445566" + return ContentBlockImage( + media_type="image/png", + data=base64.b64encode(image_bytes(color=color)).decode("ascii"), + attachment_id=image_id, + ) + + +def _images(messages: list[Message]) -> list[ContentBlockImage]: + return [ + block + for message in messages + if isinstance(message.content, list) + for block in message.content + if isinstance(block, ContentBlockImage) + ] + + +def _tool_replay(*, final_answer: bool, second_image: bool = False) -> dict[str, Any]: + messages = [ + Message( + role="assistant", + content=[ + ContentBlockToolUse(id="read-1", name="observe", input={}), + ], + ), + Message( + role="user", + content=[ + ContentBlockToolResult(tool_use_id="read-1", content="loaded first frame"), + _image("first-frame"), + ], + ), + ] + if second_image: + messages.extend( + [ + Message( + role="assistant", + content=[ + ContentBlockToolUse(id="read-2", name="observe", input={}), + ], + ), + Message( + role="user", + content=[ + ContentBlockToolResult(tool_use_id="read-2", content="loaded second frame"), + _image("second-frame"), + ], + ), + ] + ) + elif final_answer: + messages.append(Message(role="assistant", content="The first frame is blue.")) + return _serialize_assistant_replay(messages) + + +def _interrupted_batch_replay(*, final_answer: bool) -> dict[str, Any]: + messages = [ + Message( + role="assistant", + content=[ + ContentBlockToolUse(id="read-1", name="observe", input={}), + ContentBlockToolUse(id="read-2", name="observe", input={}), + ], + ), + Message( + role="user", + content=[ + ContentBlockToolResult( + tool_use_id="read-1", + content=[_image("first-frame")], + ), + ], + ), + ] + if final_answer: + messages.append(Message(role="assistant", content="The frame is blue.")) + return _serialize_assistant_replay(messages) + + +class _ToolThen413Provider: + provider_name = "synthetic" + + def __init__(self) -> None: + self.calls: list[list[Message]] = [] + + def chat( + self, + messages: list[Message], + tools: list[Any] | None = None, # noqa: ARG002 + config: ChatConfig | None = None, # noqa: ARG002 + ) -> AsyncIterator[Any]: + self.calls.append(messages) + return self._tool_stream() if len(self.calls) == 1 else self._error_stream() + + async def _tool_stream(self) -> AsyncIterator[Any]: + yield ToolUseStartEvent(tool_use_id="read-1", tool_name="observe") + yield ToolUseEndEvent(tool_use_id="read-1", tool_name="observe", arguments={}) + yield ProviderDoneEvent(stop_reason="tool_use", input_tokens=1, output_tokens=1) + + async def _error_stream(self) -> AsyncIterator[Any]: + yield ProviderErrorEvent(code="413", message="HTTP 413") + + async def list_models(self) -> list[Any]: + return [] + + +class _ToolThenAnswerProvider(_ToolThen413Provider): + async def _error_stream(self) -> AsyncIterator[Any]: + yield ProviderTextDeltaEvent(text="Frame understood.") + yield ProviderDoneEvent(stop_reason="end_turn", input_tokens=3, output_tokens=2) + + +class _AnswerProvider: + provider_name = "synthetic" + + def __init__(self) -> None: + self.calls: list[list[Message]] = [] + + def chat( + self, + messages: list[Message], + tools: list[Any] | None = None, # noqa: ARG002 + config: ChatConfig | None = None, # noqa: ARG002 + ) -> AsyncIterator[Any]: + self.calls.append(messages) + return self._stream() + + async def _stream(self) -> AsyncIterator[Any]: + yield ProviderTextDeltaEvent(text="Frame understood.") + yield ProviderDoneEvent(stop_reason="end_turn", input_tokens=3, output_tokens=2) + + async def list_models(self) -> list[Any]: + return [] + + +@dataclass +class _TranscriptEntry: + role: str + content: str + message_id: str | None = None + assistant_replay: dict[str, Any] | None = None + turn_context: dict[str, Any] | None = None + + +@dataclass +class _SessionNode: + session_key: str + session_id: str + + +class _SessionManager: + def __init__(self, entries: list[_TranscriptEntry]) -> None: + self.entries = entries + + async def get_transcript(self, session_key: str) -> list[_TranscriptEntry]: # noqa: ARG002 + return list(self.entries) + + async def get_session(self, session_key: str) -> _SessionNode: + return _SessionNode(session_key=session_key, session_id="synthetic-session") + + async def get_context_states(self, session_key: str) -> list[Any]: # noqa: ARG002 + return [] + + +async def _continue_replay( + replay: dict[str, Any], + *, + extra_history: list[Message] | None = None, + turn_context: dict[str, Any] | None = None, +) -> tuple[list[Message], dict[str, Any]]: + stored_replay = deepcopy(replay) + history = [Message(role="user", content="Inspect the frames.")] + history.extend( + reconstruct_messages_from_entry( + "assistant", + "", + None, + assistant_replay=stored_replay, + turn_context=turn_context, + ) + ) + history.extend(extra_history or []) + provider = _AnswerProvider() + agent = Agent( + provider=provider, + config=AgentConfig( + max_iterations=1, + model_vision_support="supported", + preserve_historical_images=True, + ), + ) + agent.set_history(history) + events = [event async for event in agent.run_turn("Continue the analysis.")] + assert any(isinstance(event, DoneEvent) for event in events) + assert not any(isinstance(event, ErrorEvent) for event in events) + assert stored_replay == replay + return provider.calls[0], stored_replay + + +@pytest.mark.asyncio +async def test_failed_post_tool_request_replays_unseen_image_on_next_vision_turn() -> None: + provider = _ToolThen413Provider() + context = ToolContext() + + async def load_frame(call: ToolCall) -> ToolResult: + context.tool_result_media[call.tool_use_id] = [ + { + "mime": "image/png", + "data": base64.b64encode(image_bytes()).decode("ascii"), + "attachment_id": "first-frame", + } + ] + return ToolResult( + tool_use_id=call.tool_use_id, + tool_name=call.tool_name, + content="loaded first frame", + ) + + agent = Agent( + provider=provider, + config=AgentConfig( + max_iterations=2, + max_provider_retries=0, + max_overflow_retries=0, + compaction_enabled=False, + model_vision_support="supported", + ), + tool_handler=load_frame, + tool_context=context, + tool_definitions=[ + ToolDefinition( + name="observe", + description="Load a synthetic frame.", + input_schema=ToolInputSchema(properties={}), + ) + ], + ) + events = [event async for event in agent.run_turn("Inspect the frame.")] + assert any(isinstance(event, ErrorEvent) for event in events) + assert len(provider.calls) == 2 + assert [image.attachment_id for image in _images(provider.calls[1])] == ["first-frame"] + done = next(event for event in events if isinstance(event, DoneEvent)) + assert done.assistant_replay is not None + + next_request, _ = await _continue_replay(done.assistant_replay) + assert [image.data for image in _images(next_request)] == [ + base64.b64encode(image_bytes()).decode("ascii") + ] + assert any( + isinstance(block, ContentBlockToolResult) and block.tool_use_id == "read-1" + for message in next_request + if isinstance(message.content, list) + for block in message.content + ) + + +@pytest.mark.asyncio +async def test_first_post_tool_request_bounds_twelve_images_without_losing_replay() -> None: + provider = _ToolThenAnswerProvider() + context = ToolContext() + encoded = [ + base64.b64encode(image_bytes(color=f"#{index:02x}{index:02x}{index:02x}")) + .decode("ascii") + for index in range(12) + ] + + async def load_frames(call: ToolCall) -> ToolResult: + context.tool_result_media[call.tool_use_id] = [ + {"mime": "image/png", "data": data, "attachment_id": f"frame-{index}"} + for index, data in enumerate(encoded) + ] + return ToolResult( + tool_use_id=call.tool_use_id, + tool_name=call.tool_name, + content="loaded synthetic frames", + ) + + agent = Agent( + provider=provider, + config=AgentConfig(max_iterations=2, model_vision_support="supported"), + tool_handler=load_frames, + tool_context=context, + tool_definitions=[ + ToolDefinition( + name="observe", description="Load synthetic frames.", + input_schema=ToolInputSchema(properties={}), + ) + ], + ) + events = [event async for event in agent.run_turn("Inspect the frames.")] + + assert any(isinstance(event, DoneEvent) for event in events) + assert len(provider.calls) == 2 + assert [image.data for image in _images(provider.calls[1])] == encoded[-4:] + assert any( + "8 images omitted by recency or size limits" in block.text + for message in provider.calls[1] + if isinstance(message.content, list) + for block in message.content + if isinstance(block, ContentBlockText) + ) + assert sum( + isinstance(block, ContentBlockImage) + for message in agent.history_snapshot() + if isinstance(message.content, list) + for block in message.content + ) == 12 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("vision_support", "native_ack"), + [("supported", True), ("unsupported", False)], +) +async def test_only_committed_native_tool_image_response_records_ack( + vision_support: str, native_ack: bool, +) -> None: + provider = _ToolThenAnswerProvider() + context = ToolContext() + + async def load_frame(call: ToolCall) -> ToolResult: + context.tool_result_media[call.tool_use_id] = [ + {"mime": "image/png", "data": _image("first-frame").data} + ] + return ToolResult( + tool_use_id=call.tool_use_id, + tool_name=call.tool_name, + content="loaded first frame", + ) + + agent = Agent( + provider=provider, + config=AgentConfig( + max_iterations=2, + model_vision_support=vision_support, + ), + tool_handler=load_frame, + tool_context=context, + tool_definitions=[ + ToolDefinition( + name="observe", description="Load a synthetic frame.", + input_schema=ToolInputSchema(properties={}), + ) + ], + ) + events = [event async for event in agent.run_turn("Inspect the frame.")] + done = next(event for event in events if isinstance(event, DoneEvent)) + assert done.assistant_replay is not None + assert bool(_images(provider.calls[1])) is native_ack + assert done.tool_image_native_ack_indexes == ([2] if native_ack else []) + + next_request, _ = await _continue_replay( + done.assistant_replay, + turn_context={ + TOOL_IMAGE_NATIVE_ACK_CONTEXT_KEY: done.tool_image_native_ack_indexes, + }, + ) + assert bool(_images(next_request)) is not native_ack + + +@pytest.mark.asyncio +async def test_completed_post_tool_answer_keeps_prior_image_as_marker() -> None: + replay = _tool_replay(final_answer=True) + next_request, stored_replay = await _continue_replay(replay) + + assert not _images(next_request) + assert any( + isinstance(block, ContentBlockText) and "历史图片本回合未重新读取" in block.text + for message in next_request + if isinstance(message.content, list) + for block in message.content + ) + assert stored_replay == replay + + +@pytest.mark.asyncio +async def test_accepted_text_only_marker_answer_does_not_acknowledge_tool_image() -> None: + replay = _tool_replay(final_answer=True) + next_request, _ = await _continue_replay( + replay, + turn_context={TOOL_IMAGE_NATIVE_ACK_CONTEXT_KEY: []}, + ) + + assert [image.data for image in _images(next_request)] == [_image("first-frame").data] + + +@pytest.mark.asyncio +async def test_native_acknowledges_only_its_own_tool_image_batch() -> None: + replay = _tool_replay(final_answer=False, second_image=True) + next_request, _ = await _continue_replay( + replay, + turn_context={TOOL_IMAGE_NATIVE_ACK_CONTEXT_KEY: [2]}, + ) + + assert [image.data for image in _images(next_request)] == [_image("second-frame").data] + + +@pytest.mark.asyncio +async def test_marker_reply_after_native_tool_use_keeps_later_batch_pending() -> None: + replay = _tool_replay(final_answer=False, second_image=True) + replay["messages"].append( + Message(role="assistant", content="I can only read the text marker.").model_dump( + mode="json" + ) + ) + next_request, _ = await _continue_replay( + replay, + turn_context={TOOL_IMAGE_NATIVE_ACK_CONTEXT_KEY: [2]}, + ) + + assert [image.data for image in _images(next_request)] == [_image("second-frame").data] + + +@pytest.mark.asyncio +async def test_only_unseen_tail_image_replays_after_two_tool_batches() -> None: + replay = _tool_replay(final_answer=False, second_image=True) + next_request, stored_replay = await _continue_replay(replay) + + assert [image.data for image in _images(next_request)] == [_image("second-frame").data] + assert any( + isinstance(block, ContentBlockText) and "历史图片本回合未重新读取" in block.text + for message in next_request + if isinstance(message.content, list) + for block in message.content + ) + assert stored_replay == replay + + +@pytest.mark.asyncio +async def test_empty_recovery_scaffold_does_not_acknowledge_tool_image() -> None: + replay = _tool_replay(final_answer=False) + replay["messages"].extend( + [ + Message(role="assistant", content=[ContentBlockText(text="")]).model_dump(mode="json"), + Message(role="user", content="Continue after the empty answer.").model_dump( + mode="json" + ), + ] + ) + + next_request, stored_replay = await _continue_replay(replay) + + assert [image.data for image in _images(next_request)] == [_image("first-frame").data] + assert stored_replay == replay + + +@pytest.mark.asyncio +async def test_image_pending_after_failed_followup_still_routes_and_replays() -> None: + replay = _tool_replay(final_answer=False) + session_key = "agent:main:synthetic-image-pending" + manager = _SessionManager( + [ + _TranscriptEntry(role="user", content="Inspect the first frame."), + _TranscriptEntry(role="assistant", content="", assistant_replay=deepcopy(replay)), + _TranscriptEntry(role="system", content="Error: HTTP 413"), + _TranscriptEntry(role="user", content="Continue once."), + _TranscriptEntry(role="system", content="Error: HTTP 413"), + _TranscriptEntry(role="user", content="Continue again.", message_id="current"), + ] + ) + runner = TurnRunner( + provider_selector=MagicMock(), + session_manager=manager, + config=GatewayConfig(), + ) + + route_context = await runner._router_previous_assistant_context( + session_key, + exclude_last_user=True, + bound_user_message_id="current", + ) + assert route_context.get("history_has_recent_image") is True + + history = [Message(role="user", content="Inspect the first frame.")] + history.extend( + reconstruct_messages_from_entry( + "assistant", + "", + None, + assistant_replay=deepcopy(replay), + ) + ) + history.append(Message(role="user", content="Continue once.")) + provider = _AnswerProvider() + agent = Agent( + provider=provider, + config=AgentConfig( + max_iterations=1, + model_vision_support="supported", + preserve_historical_images=True, + ), + ) + agent.set_history(history) + events = [event async for event in agent.run_turn("Continue again.")] + + assert any(isinstance(event, DoneEvent) for event in events) + assert [image.data for image in _images(provider.calls[0])] == [_image("first-frame").data] + assert manager.entries[1].assistant_replay == replay + + +@pytest.mark.asyncio +@pytest.mark.parametrize("native_followup", [False, True]) +async def test_router_uses_durable_native_ack_for_later_text_reply( + native_followup: bool, +) -> None: + followup_replay = _serialize_assistant_replay( + [Message(role="assistant", content="Continuing from the prior observation.")] + ) + manager = _SessionManager( + [ + _TranscriptEntry(role="user", content="Inspect the frame."), + _TranscriptEntry( + role="assistant", + content="I only received a text marker.", + assistant_replay=_tool_replay(final_answer=True), + turn_context={TOOL_IMAGE_NATIVE_ACK_CONTEXT_KEY: []}, + ), + _TranscriptEntry(role="user", content="Continue once."), + _TranscriptEntry( + role="assistant", + content="Continuing from the prior observation.", + assistant_replay=followup_replay, + turn_context={ + TOOL_IMAGE_NATIVE_ACK_CONTEXT_KEY: [0] if native_followup else [], + }, + ), + _TranscriptEntry(role="user", content="Continue again.", message_id="current"), + ] + ) + runner = TurnRunner( + provider_selector=MagicMock(), session_manager=manager, config=GatewayConfig() + ) + + route_context = await runner._router_previous_assistant_context( + "agent:main:synthetic-native-image-ack", + exclude_last_user=True, + bound_user_message_id="current", + ) + + assert (route_context.get("history_has_recent_image") is True) is (not native_followup) + + +@pytest.mark.asyncio +async def test_unanchored_pending_image_routes_without_a_prior_user_row() -> None: + manager = _SessionManager( + [ + _TranscriptEntry( + role="assistant", + content="", + assistant_replay=_tool_replay(final_answer=False), + ), + _TranscriptEntry(role="user", content="Continue.", message_id="current"), + ] + ) + runner = TurnRunner( + provider_selector=MagicMock(), + session_manager=manager, + config=GatewayConfig(), + ) + + context = await runner._router_previous_assistant_context( + "agent:main:synthetic-unanchored-image", + exclude_last_user=True, + bound_user_message_id="current", + ) + + assert context["history_has_recent_image"] is True + assert context["history_image_turn_count"] == 1 + assert "last_image_turn_text" not in context + + +@pytest.mark.asyncio +async def test_trusted_runtime_notice_does_not_acknowledge_unseen_image() -> None: + replay = _tool_replay(final_answer=False) + notice = Message(role="assistant", content="[Trusted runtime status] Subagent ended.") + notice._trusted_runtime_notice = True + next_request, _ = await _continue_replay( + replay, + extra_history=[notice], + ) + + assert [image.data for image in _images(next_request)] == [_image("first-frame").data] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("final_answer", [False, True]) +async def test_interrupted_batch_media_replays_only_until_model_answers( + final_answer: bool, +) -> None: + replay = _interrupted_batch_replay(final_answer=final_answer) + restored = reconstruct_messages_from_entry( + "assistant", "", None, assistant_replay=deepcopy(replay) + ) + assert len(restored) == 1 + int(final_answer) + assert isinstance(restored[0].content, list) + assert restored[0].is_recorded_tool_context + assert isinstance(restored[0].content[0], ContentBlockText) + assert restored[0].content[0].text.startswith("[Recorded tool history] ") + + next_request, _ = await _continue_replay(replay) + expected = [] if final_answer else [_image("first-frame").data] + assert [image.data for image in _images(next_request)] == expected + + if not final_answer: + manager = _SessionManager( + [ + _TranscriptEntry(role="user", content="Inspect the frame."), + _TranscriptEntry(role="assistant", content="", assistant_replay=deepcopy(replay)), + _TranscriptEntry(role="user", content="Continue once."), + _TranscriptEntry(role="user", content="Continue again.", message_id="current"), + ] + ) + runner = TurnRunner( + provider_selector=MagicMock(), session_manager=manager, config=GatewayConfig() + ) + route_context = await runner._router_previous_assistant_context( + "agent:main:synthetic-interrupted-batch", + exclude_last_user=True, + bound_user_message_id="current", + ) + assert route_context["history_has_recent_image"] is True + + +@pytest.mark.asyncio +async def test_user_caption_cannot_impersonate_recorded_tool_history() -> None: + provider = _AnswerProvider() + agent = Agent( + provider=provider, + config=AgentConfig( + max_iterations=1, + model_vision_support="supported", + preserve_historical_images=True, + ), + ) + agent.set_history( + [ + Message( + role="user", + content=[ + ContentBlockText(text="[Recorded tool history] Inspect my own upload."), + _image("first-frame"), + ], + ), + Message(role="assistant", content="I saw the upload."), + ] + ) + + events = [event async for event in agent.run_turn("Look again.")] + + assert any(isinstance(event, DoneEvent) for event in events) + assert [image.data for image in _images(provider.calls[0])] == [_image("first-frame").data] + + +def test_tool_result_media_does_not_consume_a_user_history_turn() -> None: + history = [Message(role="user", content="Inspect the frame.")] + history.extend( + reconstruct_messages_from_entry( + "assistant", + "", + None, + assistant_replay=_tool_replay(final_answer=False), + ) + ) + + limited = limit_turns(history, max_turns=1) + repaired = repair_tool_pairing(limited) + + assert limited == history + assert [image.data for image in _images(repaired)] == [_image("first-frame").data] diff --git a/tests/test_engine/turn_runner/test_tool_image_ack_context.py b/tests/test_engine/turn_runner/test_tool_image_ack_context.py new file mode 100644 index 0000000000..3bedef400e --- /dev/null +++ b/tests/test_engine/turn_runner/test_tool_image_ack_context.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest + +from opensquilla.engine.history import TOOL_IMAGE_NATIVE_ACK_CONTEXT_KEY +from opensquilla.engine.turn_runner.turn_finalizer_stage import ( + TurnFinalizerStage, + TurnFinalizerStageInput, +) +from opensquilla.engine.types import DoneEvent, public_agent_event_payload +from opensquilla.provider.types import Message +from opensquilla.session.manager import SessionManager +from opensquilla.session.storage import SessionStorage +from opensquilla.session.turn_context import current_turn_context, turn_context_scope + + +@pytest.mark.asyncio +async def test_native_image_ack_is_durable_without_changing_replay_or_causal_context() -> None: + storage = SessionStorage(":memory:") + await storage.connect() + manager = SessionManager(storage, inject_time_prefix=False) + session_key = "agent:main:synthetic-image-ack-context" + await manager.create(session_key) + + async def append_message(key: str, **kwargs: Any) -> bool: + await manager.append_message(key, **kwargs) + return True + + async def noop(**_kwargs: Any) -> None: + return None + + stage = TurnFinalizerStage( + transcript_append=SimpleNamespace(append_message=append_message), + turn_memory_capture=SimpleNamespace(capture_turn=noop), + session_totals=SimpleNamespace(rollup=noop), + turn_error_persist=SimpleNamespace(persist_error=noop), + ) + replay = { + "version": 1, + "messages": [ + Message(role="assistant", content="The frame is blue.").model_dump(mode="json") + ], + } + done = DoneEvent( + text="The frame is blue.", + assistant_replay=replay, + tool_image_native_ack_indexes=[0], + ) + inp = TurnFinalizerStageInput( + final_text_parts=[done.text], + turn_segments=[], + turn_artifacts=[], + error_message=None, + pending_error_event=None, + done_event=done, + runtime_message="Inspect the frame.", + input_mode="user", + input_provenance=None, + resolved_model="synthetic-vision-model", + agent_id="main", + session_key=session_key, + tool_context=None, + run_kind="default", + heartbeat_ack_max_chars=300, + no_memory_capture=True, + ) + try: + with turn_context_scope({"turn_id": "synthetic-turn"}): + await stage.run(inp) + assert current_turn_context() == {"turn_id": "synthetic-turn"} + entries = await manager.get_transcript(session_key) + assert len(entries) == 1 + assert entries[0].assistant_replay == replay + assert entries[0].turn_context == { + "turn_id": "synthetic-turn", + TOOL_IMAGE_NATIVE_ACK_CONTEXT_KEY: [0], + } + assert "tool_image_native_ack_indexes" not in public_agent_event_payload(done) + finally: + await storage.close() diff --git a/tests/test_gateway/test_terminal_reply.py b/tests/test_gateway/test_terminal_reply.py index 70d5a26f6f..938b054a87 100644 --- a/tests/test_gateway/test_terminal_reply.py +++ b/tests/test_gateway/test_terminal_reply.py @@ -8,7 +8,10 @@ from opensquilla.session.models import AgentTaskStatus from opensquilla.session.terminal_reply import ( CONTEXT_PAYLOAD_TOO_LARGE_MESSAGES, + PROVIDER_PAYLOAD_TOO_LARGE_CODE, + PROVIDER_WIRE_PAYLOAD_TOO_LARGE_CODE, build_terminal_reply, + is_context_payload_too_large, safe_provider_failure_code, safe_provider_failure_message, sanitize_agent_error, @@ -478,3 +481,34 @@ def test_safe_provider_failure_message_is_allowlisted() -> None: "PRIVATE_PROVIDER_CODE_BODY", "rate_limited", ) == "provider_rate_limited" + + +def test_request_size_rejections_have_distinct_safe_codes_and_guidance() -> None: + assert safe_provider_failure_code("413", "payload_too_large") == ( + PROVIDER_PAYLOAD_TOO_LARGE_CODE + ) + assert safe_provider_failure_code( + PROVIDER_WIRE_PAYLOAD_TOO_LARGE_CODE, "wire_payload_too_large" + ) == PROVIDER_WIRE_PAYLOAD_TOO_LARGE_CODE + assert safe_provider_failure_code( + "provider_request_budget_exhausted", "wire_payload_too_large" + ) == PROVIDER_WIRE_PAYLOAD_TOO_LARGE_CODE + assert safe_provider_failure_code( + PROVIDER_WIRE_PAYLOAD_TOO_LARGE_CODE, "unknown" + ) == "provider_error" + assert safe_provider_failure_message("wire_payload_too_large").startswith("OpenSquilla") + assert safe_provider_failure_code("413", "context_overflow") == "413" + + for code in (PROVIDER_PAYLOAD_TOO_LARGE_CODE, PROVIDER_WIRE_PAYLOAD_TOO_LARGE_CODE): + payload = { + "status": "failed", + "error_class": code, + "error_message": "private upstream body", + "terminal_message": "private upstream body", + } + expected = build_terminal_reply(payload) + assert "private upstream body" not in expected + assert "large" in expected or "size" in expected + assert "context window" not in expected + assert is_context_payload_too_large(payload) is False + assert sanitize_agent_error(payload) == (code, expected) diff --git a/tests/test_gateway/test_wire_rejection_usage_ledger.py b/tests/test_gateway/test_wire_rejection_usage_ledger.py new file mode 100644 index 0000000000..7d609122fd --- /dev/null +++ b/tests/test_gateway/test_wire_rejection_usage_ledger.py @@ -0,0 +1,119 @@ +from __future__ import annotations + +import time + +import pytest + +from opensquilla.engine.usage_accounting import UsageCallStart +from opensquilla.gateway.usage_ledger_runtime import SessionUsageEventSink +from opensquilla.gateway.usage_query import query_usage_ledger +from opensquilla.session.models import SessionNode +from opensquilla.session.storage import SessionStorage +from opensquilla.session.usage_ledger import ( + LOCAL_PRE_DISPATCH_ORIGIN, + LOCAL_WIRE_BODY_REJECTION_REASON, + UsageEventCompletion, + UsageLedgerConflictError, +) + + +def _call(index: int, started_at_ms: int) -> UsageCallStart: + return UsageCallStart( + event_id=f"synthetic-event-{index}", + execution_id="synthetic-execution", + call_index=index, + agent_run_id="synthetic-run", + turn_id=f"synthetic-turn-{index}", + parent_turn_id=None, + session_id="synthetic-session", + agent_id="main", + run_kind="agent", + provider="synthetic-provider", + model="synthetic-model", + started_at_ms=started_at_ms, + ) + + +async def test_local_wire_rejection_is_durable_idempotent_and_excluded_from_usage() -> None: + storage = SessionStorage(":memory:") + await storage.connect() + sink = SessionUsageEventSink(storage) + try: + await storage.upsert_session(SessionNode( + session_key="agent:main:webchat:synthetic", + session_id="synthetic-session", + )) + started_at = int(time.time() * 1_000) - 100 + await storage.initialize_usage_ledger(now_ms=started_at - 10) + local = _call(1, started_at) + await sink.start(local) + await sink.mark_not_sent(local) + await sink.mark_not_sent(local) + + rows = await storage.query_usage_events( + None, None, statuses=("unknown",), session_id="synthetic-session", + ) + assert len(rows) == 1 + assert rows[0].origin == LOCAL_PRE_DISPATCH_ORIGIN + assert rows[0].unknown_reason == LOCAL_WIRE_BODY_REJECTION_REASON + assert rows[0].missing_cost_entries == 0 + assert rows[0].transition_applied is False + assert await storage.get_turn_usage_projection( + session_id="synthetic-session", + session_epoch=0, + turn_id=local.turn_id or "", + ) is None + assert await storage.get_turn_usage_projections( + session_id="synthetic-session", + session_epoch=0, + turn_ids=[local.turn_id or ""], + ) == {} + session_totals = await storage.reconcile_session_usage_totals_from_ledger( + session_key="agent:main:webchat:synthetic", + expected_epoch=0, + expected_session_id="synthetic-session", + ) + assert session_totals is not None + assert session_totals.missing_cost_entries == 0 + + local_only = await query_usage_ledger( + storage, + {"range": {"fromMs": started_at - 20, "toMs": started_at + 2_000}}, + now_ms=started_at + 3_000, + ) + assert local_only["eventCount"] == 0 + assert local_only["missingCostEntries"] == 0 + + upstream = _call(2, started_at + 1) + await sink.start(upstream) + await sink.mark_unknown(upstream, "provider_error:413") + await sink.mark_unknown(upstream, "provider_error:413") + upstream_projection = await storage.get_turn_usage_projection( + session_id="synthetic-session", + session_epoch=0, + turn_id=upstream.turn_id or "", + ) + assert upstream_projection is not None + assert upstream_projection["unknown_usage_events"] == 1 + assert upstream_projection["missing_cost_entries"] == 1 + + report = await query_usage_ledger( + storage, + {"range": {"fromMs": started_at - 20, "toMs": started_at + 2_000}}, + now_ms=started_at + 3_000, + ) + assert report["eventCount"] == 1 + assert report["missingCostEntries"] == 1 + assert report["totals"]["costNanos"] == 0 + assert report["totals"]["eventCount"] == 1 + + with pytest.raises(UsageLedgerConflictError): + await storage.finalize_usage_event( + local.event_id, + # A late receipt must not silently turn an unsubmitted call + # into a billed provider request. + UsageEventCompletion(completed_at_ms=started_at + 1), + ) + finally: + await sink.close() + await storage.close() diff --git a/tests/test_provider/test_failure_classification_parity.py b/tests/test_provider/test_failure_classification_parity.py index 0e4c5750e0..99d6766c34 100644 --- a/tests/test_provider/test_failure_classification_parity.py +++ b/tests/test_provider/test_failure_classification_parity.py @@ -42,6 +42,25 @@ ("openai", 401, "", "context length exceeded", K.CONTEXT_OVERFLOW), # Context overflow outranks policy refusal within the pre-pass. ("openai", None, "", "context window hit; flagged by moderation", K.CONTEXT_OVERFLOW), + ("openai", 413, "context_length_exceeded", "", K.CONTEXT_OVERFLOW), +] + +# A verified transport status or structured payload code proves request-size +# rejection. Response prose alone cannot establish an HTTP status. +_PAYLOAD_TOO_LARGE_CASES: list[Case] = [ + ("openai", 413, "", "", K.PAYLOAD_TOO_LARGE), + ("openai", None, "413", "", K.PAYLOAD_TOO_LARGE), + ("anthropic", 413, "", "input exceeds the gateway body limit", K.PAYLOAD_TOO_LARGE), + ("ollama", 413, "", "", K.PAYLOAD_TOO_LARGE), + (_UNREGISTERED, 413, "rate_limit_exceeded", "", K.PAYLOAD_TOO_LARGE), + ("openrouter", None, "payload_too_large", "", K.PAYLOAD_TOO_LARGE), + ("openai", None, "request_entity_too_large", "", K.PAYLOAD_TOO_LARGE), + ( + "openai", None, "provider_request_budget_exhausted", + '{"wire_byte_budget_exhausted":true,"fits_wire_byte_budget":false,' + '"wire_json_bytes":900,"wire_body_max_bytes":800}', + K.WIRE_PAYLOAD_TOO_LARGE, + ), ] # --- shared pre-pass: policy refusal --- @@ -295,12 +314,15 @@ (_UNREGISTERED, None, "", "the provider sent 520 tokens", K.UNKNOWN), # A bare "429 Too Many Requests" body without a status code has no trigger. (_UNREGISTERED, None, "", "429 Too Many Requests", K.UNKNOWN), + (_UNREGISTERED, None, "", "HTTP 413 Request Entity Too Large", K.UNKNOWN), + ("openai", None, "provider_wire_payload_too_large", "", K.UNKNOWN), # Near-miss on the exact-match empty-response shapes. (_UNREGISTERED, None, "", "the empty response body was discarded", K.UNKNOWN), ] CORPUS: list[Case] = ( _CONTEXT_OVERFLOW_CASES + + _PAYLOAD_TOO_LARGE_CASES + _POLICY_REFUSAL_CASES + _EMPTY_RESPONSE_CASES + _OPENAI_COMPAT_CASES diff --git a/tests/test_provider_failures.py b/tests/test_provider_failures.py index 3bb583d901..5fbd4b34e2 100644 --- a/tests/test_provider_failures.py +++ b/tests/test_provider_failures.py @@ -1,5 +1,7 @@ from __future__ import annotations +import json + import pytest import structlog @@ -64,6 +66,56 @@ def test_provider_request_budget_exhausted_is_context_overflow() -> None: ) +@pytest.mark.parametrize("provider_name", ["openai", "anthropic", "ollama", "unknown-provider"]) +def test_http_413_is_request_size_rejection_without_retry_or_model_hop( + provider_name: str, +) -> None: + kind = classify_provider_error( + provider_name, 413, raw_code="413", message="input exceeds maximum body size" + ) + assert kind is ProviderFailureKind.PAYLOAD_TOO_LARGE + assert decide_recovery_action(kind) is ProviderRecoveryAction.SURFACE + assert FallbackPolicy().should_retry(kind, attempt=0) is False + + +@pytest.mark.parametrize("raw_code", ["context_length_exceeded", "context_window_exceeded"]) +def test_explicit_token_window_code_on_413_keeps_context_recovery(raw_code: str) -> None: + kind = classify_provider_error("openai", 413, raw_code=raw_code) + assert kind is ProviderFailureKind.CONTEXT_OVERFLOW + assert decide_recovery_action(kind) is ProviderRecoveryAction.COMPACT_AND_RETRY + + +def test_local_wire_budget_proof_surfaces_without_context_compaction() -> None: + proof = { + "fallback_reason": "provider_request_budget_exhausted", + "wire_byte_budget_exhausted": True, + "fits_wire_byte_budget": False, + "wire_json_bytes": 900, + "wire_body_max_bytes": 800, + } + kind = classify_provider_error( + "openai", None, "provider_request_budget_exhausted", json.dumps(proof) + ) + assert kind is ProviderFailureKind.WIRE_PAYLOAD_TOO_LARGE + assert decide_recovery_action(kind) is ProviderRecoveryAction.SURFACE + assert FallbackPolicy().should_retry(kind, attempt=0) is False + + +@pytest.mark.parametrize("proof", [ + {"wire_byte_budget_exhausted": "true"}, + {"wire_byte_budget_exhausted": True}, + {"wire_byte_budget_exhausted": True, "fits_wire_byte_budget": False, + "wire_json_bytes": 500, "wire_body_max_bytes": 800}, +]) +def test_incomplete_wire_proof_preserves_legacy_context_classification( + proof: dict[str, object], +) -> None: + proof["fallback_reason"] = "provider_request_budget_exhausted" + assert classify_provider_error( + "openai", None, "provider_request_budget_exhausted", json.dumps(proof) + ) is ProviderFailureKind.CONTEXT_OVERFLOW + + @pytest.mark.parametrize("provider_name", ["ensemble", "openrouter"]) def test_ensemble_multimodal_rejection_is_surfaced_without_fallback( provider_name: str, diff --git a/tests/test_provider_request_proof.py b/tests/test_provider_request_proof.py index 78fb51741a..428a569e2e 100644 --- a/tests/test_provider_request_proof.py +++ b/tests/test_provider_request_proof.py @@ -1,11 +1,13 @@ from __future__ import annotations +import asyncio import base64 import io import json from copy import deepcopy from typing import Any +import httpx import pytest from PIL import Image @@ -27,7 +29,12 @@ provider_request_character_budget, provider_request_token_budget, ) -from opensquilla.provider.types import ChatConfig, ContentBlockToolResult, Message +from opensquilla.provider.types import ( + ChatConfig, + ContentBlockImage, + ContentBlockToolResult, + Message, +) @pytest.mark.parametrize( @@ -162,6 +169,145 @@ def test_provider_request_proof_allows_payload_within_budget() -> None: assert proof["tool_schema_too_large"] is False +@pytest.mark.parametrize( + ("environment", "expected"), + [ + (None, 24 * 1024 * 1024), + ("1024", 1024), + ("0", 0), + ("invalid", 24 * 1024 * 1024), + ("-1", 24 * 1024 * 1024), + ], +) +def test_wire_body_limit_defaults_and_environment_override( + monkeypatch: pytest.MonkeyPatch, + environment: str | None, + expected: int, +) -> None: + if environment is None: + monkeypatch.delenv("OPENSQUILLA_PROVIDER_REQUEST_MAX_BODY_BYTES", raising=False) + else: + monkeypatch.setenv("OPENSQUILLA_PROVIDER_REQUEST_MAX_BODY_BYTES", environment) + assert request_proof._wire_body_max_bytes() == expected + + +def test_wire_byte_limit_is_independent_of_image_token_budget( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("OPENSQUILLA_PROVIDER_REQUEST_MAX_BODY_BYTES", "600") + payload = {"messages": [{"role": "user", "content": [{ + "type": "image_url", + "image_url": {"url": "data:image/png;base64," + "a" * 700}, + }]}]} + proof = project_provider_payload( + payload, + projection_adapter="openai", + proof_budget=20_000, + token_budget=10_000, + ) + request = httpx.Request("POST", "https://example.test/chat", json=payload) + assert proof["wire_json_bytes"] == len(request.content) + assert proof["fits_char_budget"] is True + assert proof["fits_token_budget"] is True + assert proof["fits_wire_byte_budget"] is False + assert proof["wire_byte_budget_exhausted"] is True + assert proof["wire_body_max_bytes"] == 600 + assert proof["wire_body_excess_bytes"] == len(request.content) - 600 + assert proof["fits"] is False + assert proof["fallback_reason"] == "provider_request_budget_exhausted" + + +def test_wire_byte_limit_is_inclusive_and_can_be_disabled( + monkeypatch: pytest.MonkeyPatch, +) -> None: + payload = {"messages": [{"role": "user", "content": "中文 input"}]} + body_bytes = len(httpx.Request("POST", "https://example.test/chat", json=payload).content) + monkeypatch.setenv("OPENSQUILLA_PROVIDER_REQUEST_MAX_BODY_BYTES", str(body_bytes)) + assert prove_provider_payload( + payload, projection_adapter="openai", proof_budget=10_000, + )["fits"] is True + monkeypatch.setenv("OPENSQUILLA_PROVIDER_REQUEST_MAX_BODY_BYTES", str(body_bytes - 1)) + with pytest.raises(ProviderRequestBudgetExceeded) as exc_info: + prove_provider_payload(payload, projection_adapter="openai", proof_budget=10_000) + assert exc_info.value.proof["wire_json_bytes"] == body_bytes + monkeypatch.setenv("OPENSQUILLA_PROVIDER_REQUEST_MAX_BODY_BYTES", "0") + assert prove_provider_payload( + payload, projection_adapter="openai", proof_budget=10_000, + )["fits"] is True + + +def test_disabled_character_and_token_proof_still_blocks_oversized_media( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("OPENSQUILLA_PROVIDER_REQUEST_MAX_BODY_BYTES", "900") + payload = {"messages": [{"role": "user", "content": [ + {"type": "image_url", "image_url": {"url": "data:image/png;base64," + "a" * 500}}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64," + "a" * 500}}, + ]}]} + before = deepcopy(payload) + with pytest.raises(ProviderRequestBudgetExceeded) as exc_info: + prove_or_compact_provider_payload( + payload, projection_adapter="openai", proof_budget=0, + ) + assert exc_info.value.proof["fits_char_budget"] is True + assert exc_info.value.proof["fits_token_budget"] is True + assert exc_info.value.proof["retry_count"] == 0 + assert exc_info.value.proof["wire_byte_budget_exhausted"] is True + assert payload == before + with pytest.raises(ProviderRequestBudgetExceeded): + request_proof.prove_provider_payload_from_env( + payload, projection_adapter="openai", + ) + + +def test_wire_byte_limit_allows_request_compaction_of_text( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("OPENSQUILLA_PROVIDER_REQUEST_MAX_BODY_BYTES", "2000") + payload = {"messages": [ + {"role": "user", "content": "synthetic request"}, + {"role": "tool", "content": "x" * 5000}, + ]} + compacted, proof = prove_or_compact_provider_payload( + payload, projection_adapter="openai", proof_budget=30_000, + ) + assert proof is not None + assert proof["fits"] is True + assert proof["retry_count"] == 1 + assert len(httpx.Request("POST", "https://example.test/chat", json=compacted).content) <= 2000 + assert payload["messages"][1]["content"] == "x" * 5000 + + +def test_provider_does_not_send_request_exceeding_wire_byte_limit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("OPENSQUILLA_PROVIDER_REQUEST_MAX_BODY_BYTES", "800") + + def unexpected_stream(*args: Any, **kwargs: Any) -> None: + del args, kwargs + raise AssertionError("Oversized request reached HTTP transport") + + monkeypatch.setattr(httpx.AsyncClient, "stream", unexpected_stream) + provider = OpenAIProvider( + api_key="synthetic", model="synthetic-model", base_url="https://example.test", + ) + messages = [Message(role="user", content=[ + ContentBlockImage(media_type="image/png", data="a" * 1_000), + ])] + + async def collect_events() -> list[Any]: + return [ + event async for event in provider.chat( + messages, config=ChatConfig(provider_request_max_chars=20_000), + ) + ] + + events = asyncio.run(collect_events()) + assert len(events) == 1 + assert events[0].code == "provider_request_budget_exhausted" + assert json.loads(events[0].message)["wire_byte_budget_exhausted"] is True + + def test_character_cap_does_not_restrict_independent_token_capacity( monkeypatch: pytest.MonkeyPatch, ) -> None: