From 5e2b9b807d1d4a9a4294434315aa2e5199586dd6 Mon Sep 17 00:00:00 2001 From: Sertac Ozercan Date: Fri, 11 Sep 2026 13:08:12 -0700 Subject: [PATCH 1/9] feat(runtime): support skills and governed tool workflows Signed-off-by: Sertac Ozercan --- docs/agentkitfile.md | 6 + docs/foundry-hosted-brokered.md | 77 +++-- docs/instruction-skills.md | 80 +++++ docs/orka.md | 6 +- docs/runtime-adapters.md | 4 +- runtimes/common/README.md | 6 +- runtimes/common/agentkit_serve_common/acp.py | 45 ++- .../common/agentkit_serve_common/config.py | 11 +- .../common/agentkit_serve_common/foundry.py | 233 +++++++++++-- .../foundry_model_loop.py | 77 ++++- .../common/agentkit_serve_common/skills.py | 235 +++++++++++++ runtimes/common/tests/test_acp_protocol.py | 98 +++++- .../tests/test_foundry_brokered_protocol.py | 7 +- .../tests/test_foundry_tool_workflows.py | 312 ++++++++++++++++++ .../common/tests/test_instruction_skills.py | 263 +++++++++++++++ .../agentkit_serve/agent_factory.py | 25 +- .../tests/test_instruction_skills.py | 84 +++++ 17 files changed, 1486 insertions(+), 83 deletions(-) create mode 100644 docs/instruction-skills.md create mode 100644 runtimes/common/agentkit_serve_common/skills.py create mode 100644 runtimes/common/tests/test_foundry_tool_workflows.py create mode 100644 runtimes/common/tests/test_instruction_skills.py create mode 100644 runtimes/microsoft-agent-framework/tests/test_instruction_skills.py diff --git a/docs/agentkitfile.md b/docs/agentkitfile.md index 25de313..dd76821 100644 --- a/docs/agentkitfile.md +++ b/docs/agentkitfile.md @@ -237,6 +237,12 @@ use filesystem skills, stage them under `/agent/skills` in the runtime/deploymen image or prefer MCP-backed skills. Memory providers require an explicit `AGENTKIT_MEMORY_SCOPE` runtime env var; choose a per-user/session-safe scope. +Orka harness v2 and the hosted brokered model loop support the narrower +[bundled instruction-only skills](instruction-skills.md) mode. It exposes +`load_skill` over a startup snapshot of `SKILL.md` documents. Resources, scripts, +remote skill sources, search providers, and memory providers are not available +through that mode. + ### `expose` ```yaml diff --git a/docs/foundry-hosted-brokered.md b/docs/foundry-hosted-brokered.md index 9bc93d5..a8216a1 100644 --- a/docs/foundry-hosted-brokered.md +++ b/docs/foundry-hosted-brokered.md @@ -164,7 +164,7 @@ payloads. Ordinary client requests must not be able to submit tool results. Unknown response IDs, unknown call IDs, orphan tool outputs, duplicate conflicts, malformed outputs, missing or wrong continuation auth, and multiple tool outputs all return deterministic error envelopes. An identical duplicate continuation -from the broker path returns the same final response idempotently. +returns its original next response, even if later tool rounds have completed. The body field is an AgentKit/Orka compatibility extension rather than a standard OpenAI Responses field. Prove that the target Foundry gateway accepts and forwards @@ -367,22 +367,61 @@ deploy/foundry/scripts/local_brokered_conformance_container.sh \ --transcript-dir ./foundry-brokered-agentkit-transcript ``` -## Lower-level model-loop fallback - -Phase A4/A5 has an opt-in fallback when a high-level framework cannot prove -pause/resume: set `AGENTKIT_FOUNDRY_BROKERED_MODEL_LOOP=1` in Foundry brokered -mode. AgentKit then calls the configured OpenAI-compatible chat-completions -model directly with the static safe `brokeredTools` as function schemas. If the -model requests exactly one configured tool, AgentKit rewrites the model's tool -call id to a stable hosted Responses `call__` id and -returns a `function_call` output item for Orka. On the Orka-authenticated -`function_call_output` continuation, AgentKit resumes the model with a `tool` -message and returns the final assistant message. - -In this mode AgentKit-owned MCP/direct tools remain disabled; only the static -safe brokered schemas are model-visible. The first implementation intentionally -limits each turn to one brokered tool call and rejects unknown, multiple, or -repeated model tool calls deterministically. +## Model-driven tool workflows + +Set `AGENTKIT_FOUNDRY_BROKERED_MODEL_LOOP=1` to let the model choose tools and +work through a task. AgentKit calls the configured OpenAI-compatible Chat +Completions model with the static `brokeredTools` schemas. For example, the +agent can inspect telemetry, use the result to look up an incident, and then +explain what it found. + +Each operational tool call returns a `function_call` for Orka to execute. After +Orka sends the matching `function_call_output`, AgentKit resumes the model. It +can request another tool or return an answer. Every round gets a fresh response +ID and call ID; the next result must match both. AgentKit validates each new +call's name, arguments, and schema before returning it. Retrying the same result +returns the cached next response without another model request. File-backed +state preserves these completed rounds across restarts. + +The model can make up to 16 sequential tool calls per user turn. At the limit, +AgentKit asks for a final answer without tools and rejects any further tool +call. Parallel tool batches are unsupported. AgentKit-owned MCP and direct +operational tools remain disabled. + +Agents can also use [bundled instruction skills](instruction-skills.md). With +filesystem skills configured under `/agent/skills`, AgentKit advertises the +local `load_skill` tool alongside the brokered schemas. It returns the packaged +instructions from an immutable startup snapshot. Skill loads count toward the +same 16-call budget; they do not invoke Orka or execute scripts. All operational +calls described by those instructions still go through Orka. + +For Orka harness v2, use the companion +[agent-runtime-foundry broker](https://github.com/orka-agents/agent-runtime-foundry/blob/main/docs/harness-v2.md). +Give its `ORKA_FOUNDRY_BROKER_AGENTKIT_CONTINUATION_PROOF` the same value as the +hosted agent's `AGENTKIT_FOUNDRY_BROKERED_CONTINUATION_PROOF`, using a secret of +at least 32 bytes without whitespace. The broker adds this value only after +checking ownership, the active lease, and the expected response and call. It +also translates MCP results into AgentKit's approved/error envelope. Keep the +proof out of the ACP child configuration. The gateway must forward the proof +field; local tests cannot establish that a public Foundry deployment does so. +See the shared-proof limitations above. Human tool approvals for external v2 +runtimes remain unsupported by Orka. + +## Hosted follow-up questions + +In model-loop mode, a follow-up user message can reference the final response ID +from the previous turn. It must use the same effective hosted session identity. +AgentKit retains recent user and assistant messages so the model can understand +questions such as "Which incident was that?" Raw tool messages and system +instructions are excluded from this retained dialogue. An assistant answer can +still contain information derived from a tool result. + +History keeps complete recent exchanges within half the model-message byte +limit or one eighth of the response-state byte limit, whichever is smaller. +It shares the response store's TTL, capacity, and optional file persistence. +Missing, expired, evicted, or mismatched session history is rejected. Requests +without a session identity remain stateless. Pending tool calls must finish +before a user can continue their response. ## Implementation status and evidence @@ -398,9 +437,9 @@ requires deployed Foundry/Orka/Fibey state. | A2 static schemas and drift control | Implemented in Go writer/validator and Python runtime; export CLI added | `brokeredTools` ABI, `agentkit-brokered-tools`, `tests/test_config_validation.py`, `tests/test_brokered_schema.py`, Go config/ABI tests | Orka Tool CRDs must be exported during deployment and current digests verified before live runs. | | A3 deterministic brokered runtime | Implemented for local/fake hosted protocol integration | deterministic `/responses` brokered path and tests | Live Orka deterministic read/write smoke still required. | | A4 framework pause/resume decision | Lower-level OpenAI-compatible fallback implemented; high-level framework native hooks remain gated | `agentkit_serve_common.foundry_model_loop`, `AGENTKIT_FOUNDRY_BROKERED_MODEL_LOOP=1`, model-loop tests | Live model smoke for brokered read/write prompts. | -| A5 first real model adapter brokered mode | Fallback model loop can emit/resume brokered calls from static safe schemas | model-loop tests in `tests/test_foundry_brokered_protocol.py` | Deployed real model read and write prompts, including declined/policy/error outcomes. | +| A5 first real model adapter brokered mode | Model loop supports sequential brokered calls, local instruction skills, and bounded hosted dialogue | `tests/test_foundry_brokered_protocol.py`, `tests/test_foundry_tool_workflows.py` | Deployed real model read and write prompts, including policy/error outcomes. | | A6 live Orka integration | Not proven in this repo state | Local AgentKit/Foundry side helpers exist | Deploy AgentKit and Orka hosted-Responses adapter; run brokered read/write approval smoke. | -| A7 Fibey | Not started; gates not satisfied | N/A | Requires A3/A5/A6 live gates first, then Fibey schemas/instructions/scenario. | +| A7 Fibey | Runtime supports packaged skills and sequential tool workflows | Catalog/ACP/MAF skill tests and hosted workflow tests | Package the Fibey skills and schemas, then validate the deployed scenario. | | A8 hardening/docs/review | Local docs/tests/autoreview complete for current patch | This doc, `docs/agent-abi.md`, `docs/runtime-capabilities.md`; full tests/lint; `$autoreview` clean | Record live transcript and Orka/Fibey validation evidence before final completion. | Local verification commands used for the current AgentKit patch: diff --git a/docs/instruction-skills.md b/docs/instruction-skills.md new file mode 100644 index 0000000..3a53e9a --- /dev/null +++ b/docs/instruction-skills.md @@ -0,0 +1,80 @@ +# Bundled skills in governed mode + +Skills let an agent load written guidance for a task before using its tools. +The Microsoft Agent Framework runtime can use the same bundled skills when +running under Orka harness v2 or the Foundry hosted brokered model loop. +Loading a skill reads instructions from the agent image; Orka still controls +operational tool calls and approvals. + +Declare the skill directory in your AgentKitfile: + +```yaml +runtime: microsoft-agent-framework +context: + providers: + - type: skills + source: filesystem + path: /agent/skills +``` + +Arrange each skill in a directory with a matching name: + +```text +skills/ + inspection/ + SKILL.md + parts-lookup/ + SKILL.md +``` + +For example, `skills/inspection/SKILL.md`: + +```markdown +--- +name: inspection +description: Prepare an equipment inspection using the current work order. +--- +Retrieve the work order using the authorized lookup tool. Summarize the required +checks, identify missing information, and cite the returned record. +``` + +AgentKit resolves an `instructions.file` into the agent configuration, but does +not copy skill directories automatically. Add them to the built agent image +before composing its Orka runtime or deploying it to Foundry: + +```dockerfile +FROM ghcr.io/acme/inspection-agent@sha256: +COPY --chown=0:0 skills/ /agent/skills/ +``` + +Make the directories and documents readable by the runtime user. Pin the image +that contains both the agent configuration and skills. For Orka ACP composition, +use that image's digest as `AGENTKIT_ADAPTER_DIGEST`; the existing +`agentConfigurationDigest` continues to cover the exact `agent.yaml` bytes. +Keep `/agent/skills` image-owned instead of mounting user or task workspace files +there. Changing a skill requires building and registering the updated image. + +The runtime lists each skill's name and description for the model. The model +calls `load_skill` with `{"skill_name":"inspection"}` to receive the original +`SKILL.md` text, including its frontmatter. This works with upstream instruction +skills that already use `load_skill`. + +Only `SKILL.md` documents are loaded. Files are snapshotted during startup and +tool calls use that immutable snapshot. Sibling files and scripts are not made +available, and the runtime does not expose `read_skill_resource` or +`run_skill_script`. A skill's text cannot grant tool access or bypass an +approval. Remote/search/memory context providers remain rejected in governed +mode. + +Each skill needs a matching lowercase name of at most 64 characters, a +description of at most 1024 characters, and nonempty instructions. Discovery +covers the selected directory and two directory levels below it. Duplicate +names, symlinked paths, nonregular or hardlinked documents, invalid UTF-8, and +invalid frontmatter fail startup. Each document is limited to 128 KiB; a catalog +is limited to 64 skills and 1 MiB of document text. A directory may contain at +most 256 entries. + +The hosted brokered loop must reserve `load_skill` for the local catalog when +skills are configured. Give operational tools their own names and schemas. +For the rest of the hosted setup, see +[Foundry hosted brokered tools](foundry-hosted-brokered.md). diff --git a/docs/orka.md b/docs/orka.md index e0eeb13..6f89574 100644 --- a/docs/orka.md +++ b/docs/orka.md @@ -27,8 +27,10 @@ supervisor starts this child under a private UID/GID and session tree: The v2 path is strict: -- `/agent/agent.yaml` must not contain direct `tools`, `brokeredTools`, or - context providers; +- `/agent/agent.yaml` must not contain direct `tools` or `brokeredTools`; +- the Microsoft Agent Framework runtime also supports bundled, instruction-only + filesystem skills; other context providers remain prohibited. See + [Bundled skills in governed mode](instruction-skills.md); - the registered model must equal `model.name` in the baked config; - `agentConfigurationDigest` is `sha256:` plus the SHA-256 of the exact `/agent/agent.yaml` bytes; diff --git a/docs/runtime-adapters.md b/docs/runtime-adapters.md index f1d31fa..8a543c1 100644 --- a/docs/runtime-adapters.md +++ b/docs/runtime-adapters.md @@ -49,7 +49,9 @@ acceptance, SSE replay, and cancel endpoints. ACP mode opens no listener. It speaks newline-delimited ACP JSON-RPC on stdin and stdout. The child verifies the configured model and SHA-256 digest of the exact `/agent/agent.yaml` bytes before accepting a session. It rejects baked -direct tools, `brokeredTools`, and context providers. At session creation it +direct tools and `brokeredTools`. The Microsoft Agent Framework adapter can +load [packaged skill instructions](instruction-skills.md); other context +providers remain prohibited. At session creation it accepts at most one loopback HTTP MCP server with bearer authentication, which is the prompt-scoped broker created by the Orka supervisor. diff --git a/runtimes/common/README.md b/runtimes/common/README.md index 3e0097b..4c6e244 100644 --- a/runtimes/common/README.md +++ b/runtimes/common/README.md @@ -57,8 +57,10 @@ The child accepts one ACP session, text and resource-link prompt blocks, cancellation, and at most one loopback HTTP MCP server carrying a bearer Authorization header. Resource links are added to the model prompt as labeled text and are never fetched by the child. The runtime keeps successful user and -assistant turns for later prompts. It rejects baked `tools`, `brokeredTools`, -and context providers. Orka owns process and workspace isolation, prompt-scoped +assistant turns for later prompts. It rejects baked `tools` and `brokeredTools`. +The Microsoft Agent Framework adapter accepts +[packaged skill instructions](../../docs/instruction-skills.md); other context +providers remain prohibited. Orka owns process and workspace isolation, prompt-scoped MCP authority, provider proxying, and cleanup proof. ## Adding a runtime adapter diff --git a/runtimes/common/agentkit_serve_common/acp.py b/runtimes/common/agentkit_serve_common/acp.py index 1264ed3..a72d6da 100644 --- a/runtimes/common/agentkit_serve_common/acp.py +++ b/runtimes/common/agentkit_serve_common/acp.py @@ -27,6 +27,11 @@ from .config import AgentSpec, load_with_bytes from .conversation import ConversationTurn, RunRequest, ToolCallEvent from .runtime import AgentRunError, RuntimeFactory, RuntimeSession +from .skills import ( + SkillCatalog, + SkillConfigurationError, + validate_packaged_skill_providers, +) ACP_PROTOCOL_VERSION = 1 ACP_AGENT_CONFIGURATION_DIGEST_ENV = "AGENTKIT_ACP_AGENT_CONFIGURATION_DIGEST" @@ -223,6 +228,19 @@ def _factory_supports_http_mcp(factory: RuntimeFactory) -> bool: return bool(capability()) if callable(capability) else False +def _factory_supports_packaged_skills(factory: RuntimeFactory) -> bool: + capability = getattr(factory, "supports_acp_packaged_skills", None) + return bool(capability()) if callable(capability) else False + + +def _only_packaged_skill_providers(spec: AgentSpec) -> bool: + try: + validate_packaged_skill_providers(spec) + except SkillConfigurationError: + return False + return True + + def _request_key(value: Any) -> str: if isinstance(value, bool) or not isinstance(value, (int, str)): raise ACPProtocolError(_INVALID_REQUEST, "JSON-RPC id must be a string or integer") @@ -327,8 +345,10 @@ def validate_acp_runtime_binding(config_bytes: bytes, spec: AgentSpec) -> None: raise ACPConfigurationError("ACP strict mode rejects baked direct tools") if spec.brokered_tools: raise ACPConfigurationError("ACP strict mode rejects baked brokeredTools") - if spec.context.providers: - raise ACPConfigurationError("ACP strict mode rejects baked context providers") + try: + validate_packaged_skill_providers(spec) + except SkillConfigurationError as exc: + raise ACPConfigurationError(str(exc)) from exc expected_digest = _required_environment(ACP_AGENT_CONFIGURATION_DIGEST_ENV) prefix = "sha256:" @@ -364,6 +384,10 @@ def validate_acp_runtime_binding(config_bytes: bytes, spec: AgentSpec) -> None: except ACPProtocolError as exc: raise ACPConfigurationError(exc.message) from exc _required_environment(ACP_PROVIDER_TOKEN_ENV) + try: + spec._packaged_skill_catalog = SkillCatalog.from_spec(spec) + except SkillConfigurationError as exc: + raise ACPConfigurationError(str(exc)) from exc class ACPStdioServer: @@ -378,7 +402,10 @@ def __init__(self, spec: AgentSpec, factory: RuntimeFactory, send: _MessageSende _factory_supports_http_mcp(factory) and not spec.tools and not spec.brokered_tools - and not spec.context.providers + and ( + not spec.context.providers + or (_factory_supports_packaged_skills(factory) and _only_packaged_skill_providers(spec)) + ) ) self.sessions: dict[str, _SessionState] = {} self.requests: dict[str, asyncio.Task[None]] = {} @@ -596,8 +623,15 @@ async def _create_session(self, params: Any) -> dict[str, Any]: raise ACPProtocolError(_INVALID_PARAMS, "ACP strict mode rejects baked direct tools") if self.spec.brokered_tools: raise ACPProtocolError(_INVALID_PARAMS, "ACP strict mode rejects baked brokeredTools") - if self.spec.context.providers: - raise ACPProtocolError(_INVALID_PARAMS, "ACP strict mode rejects baked context providers") + try: + validate_packaged_skill_providers(self.spec) + if self.spec.context.providers: + if not _factory_supports_packaged_skills(self.factory): + raise SkillConfigurationError("this runtime does not support ACP packaged skills") + if self.spec._packaged_skill_catalog is None: + self.spec._packaged_skill_catalog = SkillCatalog.from_spec(self.spec) + except SkillConfigurationError as exc: + raise ACPProtocolError(_INVALID_PARAMS, str(exc)) from exc mcp_servers = request.get("mcpServers", []) if not isinstance(mcp_servers, list): @@ -720,6 +754,7 @@ def _project_spec( projected = AgentSpec.model_validate(data) except ValueError as exc: raise ACPProtocolError(_INVALID_PARAMS, "ACP MCP server configuration is invalid") from exc + projected._packaged_skill_catalog = self.spec._packaged_skill_catalog return projected, environment async def _prompt(self, request_key: str, params: Any) -> dict[str, str]: diff --git a/runtimes/common/agentkit_serve_common/config.py b/runtimes/common/agentkit_serve_common/config.py index 531dcec..81fefea 100644 --- a/runtimes/common/agentkit_serve_common/config.py +++ b/runtimes/common/agentkit_serve_common/config.py @@ -19,13 +19,16 @@ import re import sys from pathlib import Path -from typing import Any, Literal, Mapping +from typing import TYPE_CHECKING, Any, Literal, Mapping import yaml -from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator, model_validator +from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, ValidationError, field_validator, model_validator from .yaml_support import safe_load_lossless +if TYPE_CHECKING: + from .skills import SkillCatalog + # The ABI schema version this reader understands (agent-abi.md: ``abiVersion: v0``). ABI_VERSION = "v0" _PROVIDER_OPENAI_COMPATIBLE = "openai-compatible" @@ -1163,6 +1166,10 @@ def _port_range(cls, value: int) -> int: class AgentSpec(_Strict): """The whole baked ``agent.yaml``.""" + # Resolved by the ACP adapter after binding the exact config bytes. This is + # process-local data, never an ABI field or request-controlled configuration. + _packaged_skill_catalog: SkillCatalog | None = PrivateAttr(default=None) + abi_version: str = Field(alias="abiVersion") metadata: Metadata model: ModelSpec diff --git a/runtimes/common/agentkit_serve_common/foundry.py b/runtimes/common/agentkit_serve_common/foundry.py index bbf80ae..ba7e16c 100644 --- a/runtimes/common/agentkit_serve_common/foundry.py +++ b/runtimes/common/agentkit_serve_common/foundry.py @@ -31,7 +31,7 @@ from fractions import Fraction from pathlib import Path from contextlib import asynccontextmanager -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from typing import Any, Mapping from fastapi import Depends, FastAPI, Request @@ -418,6 +418,11 @@ class _HostedResponseState: model_messages: list[dict[str, Any]] | None = None initial_usage: dict[str, int] = field(default_factory=dict) final_persistence_pending: bool = False + # All rounds share one bounded store entry. A completed round retains its + # response/call pairing and exact reply so retries cannot repeat model work. + response_calls: dict[str, str] = field(default_factory=dict) + continuation_payloads: dict[str, dict[str, Any]] = field(default_factory=dict) + conversation_history: list[dict[str, str]] = field(default_factory=list) class _StateExpired(KeyError): @@ -525,6 +530,12 @@ def _state_to_payload(state: _HostedResponseState) -> dict[str, Any]: } if state.terminal_error is not None: payload["terminalError"] = state.terminal_error + if state.response_calls: + payload["responseCalls"] = dict(state.response_calls) + if state.continuation_payloads: + payload["continuationPayloads"] = dict(state.continuation_payloads) + if state.conversation_history: + payload["conversationHistory"] = list(state.conversation_history) return payload @@ -554,6 +565,19 @@ def _state_from_payload(data: Mapping[str, Any]) -> _HostedResponseState: terminal_error = data.get("terminalError") model_messages = data.get("modelMessages") initial_usage = data.get("initialUsage", {}) + response_calls = data.get("responseCalls", {}) + continuation_payloads = data.get("continuationPayloads", {}) + conversation_history = data.get("conversationHistory", []) + if not isinstance(response_calls, dict) or not all(isinstance(k, str) and isinstance(v, str) for k, v in response_calls.items()): + raise ValueError("stored responseCalls must map response IDs to call IDs") + if not isinstance(continuation_payloads, dict) or not all(isinstance(k, str) and isinstance(v, dict) for k, v in continuation_payloads.items()): + raise ValueError("stored continuationPayloads must map call IDs to response objects") + if not isinstance(conversation_history, list) or not all( + isinstance(turn, dict) and set(turn) == {"role", "content"} + and turn["role"] in {"user", "assistant"} and isinstance(turn["content"], str) + for turn in conversation_history + ): + raise ValueError("stored conversationHistory must contain user and assistant text") if final_payload is not None and not isinstance(final_payload, dict): raise ValueError("stored finalPayload must be an object") if terminal_error is not None and not isinstance(terminal_error, str): @@ -582,8 +606,8 @@ def _state_from_payload(data: Mapping[str, Any]) -> _HostedResponseState: # Clear it on load so Orka can retry instead of being stuck behind # duplicate_continuation_in_progress until TTL expiry. status = "pending" - accepted = {} - accepted_sizes = {} + accepted = {key: value for key, value in accepted.items() if key in continuation_payloads} + accepted_sizes = {key: value for key, value in accepted_sizes.items() if key in accepted} return _HostedResponseState( response_id=str(data.get("responseID") or ""), session_id=str(data["sessionID"]) if data.get("sessionID") is not None else None, @@ -596,11 +620,14 @@ def _state_from_payload(data: Mapping[str, Any]) -> _HostedResponseState: terminal_error=terminal_error, model_messages=model_messages, initial_usage={str(key): int(value or 0) for key, value in initial_usage.items()}, + response_calls=response_calls, + continuation_payloads=continuation_payloads, + conversation_history=conversation_history, ) def _state_has_replay(state: _HostedResponseState) -> bool: - return state.final_payload is not None or state.terminal_error is not None + return state.final_payload is not None or state.terminal_error is not None or bool(state.continuation_payloads) class _FoundryResponseStateStore: @@ -764,7 +791,7 @@ def _commit( pending_finals = [ response_id for response_id, state in states.items() - if state.final_persistence_pending and state.status == "completed" and _state_has_replay(state) + if state.final_persistence_pending and _state_has_replay(state) ] if pending_finals: installed = dict(states) @@ -1008,20 +1035,25 @@ def evict_completed_to_capacity(self, *, reserve_slots: int = 0) -> None: def get(self, response_id: str) -> _HostedResponseState: with self._lock: state = self._states.get(response_id) + if state is None: + state = next((entry for entry in self._states.values() if response_id in entry.response_calls or ( + entry.session_id and entry.final_payload and entry.final_payload.get("id") == response_id + )), None) if state is None: raise KeyError(response_id) + root_id = state.response_id states = dict(self._states) now = time.time() active_resume_ids = self._active_resume_ids() target_expired = state.expires_at <= now and not ( - state.status == "resuming" and response_id in active_resume_ids + state.status == "resuming" and root_id in active_resume_ids ) changed = self._purge_expired_from(states, now=now, active_resume_ids=active_resume_ids) if changed: self._commit(states) if target_expired: raise _StateExpired(response_id, deepcopy(state)) - current = self._states.get(response_id) + current = self._states.get(root_id) if current is None: raise KeyError(response_id) return deepcopy(current) @@ -2149,6 +2181,100 @@ def _function_call_response_payload( return payload +def _model_pending_call( + result: ModelLoopToolRequest, + *, + model_loop: BrokeredChatModelLoop, + response_id: str, + call_id: str, + max_argument_bytes: int, + max_messages_bytes: int, +) -> _PendingCall: + tool = model_loop.tools_by_name.get(result.name) + if tool is None: + raise AgentRunError("model requested unknown brokered tool", status=400, code="unknown_brokered_tool") + _validate_model_brokered_arguments(result.arguments) + _validate_model_arguments_for_tool(result.arguments, tool) + if len(_canonical_output_json(result.arguments).encode("utf-8")) > max_argument_bytes: + raise AgentRunError("brokered function_call arguments are too large for pending state", status=413, code="brokered_arguments_too_large") + try: + _bounded_json_bytes(result.messages, max_bytes=max_messages_bytes) + except _SerializedPayloadTooLarge as exc: + raise AgentRunError("model loop messages are too large for pending state", status=413, code="brokered_model_messages_too_large") from exc + except (_InvalidUnicodeValue, TypeError, ValueError) as exc: + raise AgentRunError("model loop messages are invalid", status=502, code="InvalidModelResponse") from exc + return _PendingCall(call_id=call_id, item_id=_new_function_call_id(response_id), tool=tool, arguments=result.arguments) + + +def _bounded_conversation_history(messages: list[dict[str, Any]], *, max_bytes: int) -> list[dict[str, str]]: + """Retain complete recent exchanges, without system prompts or tool data.""" + turns = [ + {"role": message["role"], "content": message["content"]} + for message in messages + if message.get("role") in {"user", "assistant"} + and isinstance(message.get("content"), str) and message["content"] + and not message.get("tool_calls") + ] + while turns: + try: + _bounded_json_bytes(turns, max_bytes=max_bytes) + return turns + except _SerializedPayloadTooLarge: + # Drop an entire exchange rather than leave an orphan answer. + next_user = next((i for i in range(1, len(turns)) if turns[i]["role"] == "user"), len(turns)) + turns = turns[next_user:] + return [] + + +def _advance_brokered_state( + *, + spec: AgentSpec, + store: _FoundryResponseStateStore, + state: _HostedResponseState, + previous_response_id: str, + call_id: str, + result: ModelLoopToolRequest, + model_loop: BrokeredChatModelLoop, + response_id: str, + next_call_id: str, +) -> JSONResponse: + try: + call = _model_pending_call( + result, + model_loop=model_loop, + response_id=response_id, + call_id=next_call_id, + max_argument_bytes=model_loop.max_argument_bytes, + max_messages_bytes=model_loop.max_messages_bytes, + ) + except AgentRunError as exc: + if not _reset_unfinalized_continuation(store, state, call_id=call_id): + return _state_storage_error() + return _error(str(exc), status=exc.status, code=exc.code) + payload = _function_call_response_payload(spec, response_id=response_id, call=call, previous_response_id=previous_response_id, usage=result.usage) + following = deepcopy(state) + if not following.response_calls: + following.response_calls[state.response_id] = call_id + following.response_calls[response_id] = next_call_id + following.pending_calls[next_call_id] = call + following.continuation_payloads[call_id] = payload + following.model_messages = result.messages + following.initial_usage = _combine_usage(state.initial_usage, result.usage) + following.status = "pending" + following.expires_at = time.time() + store.ttl_seconds + try: + store.save(following) + except (_StateStoreFull, _StateSizeLimitExceeded) as exc: + if not _reset_unfinalized_continuation(store, state, call_id=call_id): + return _state_storage_error() + return _state_full_error() if isinstance(exc, _StateStoreFull) else _model_response_too_large_error() + except _StatePersistenceError: + following.final_persistence_pending = True + store.cache_in_memory(following) + return _state_storage_error() + return JSONResponse(payload) + + def _final_text_from_tool_output(call: _PendingCall, output: dict[str, Any]) -> str: if not output.get("approved"): error = output.get("error") if isinstance(output.get("error"), dict) else {} @@ -2171,6 +2297,7 @@ def _reset_unfinalized_continuation( state.final_payload = None state.terminal_error = None state.final_persistence_pending = False + state.conversation_history = [] state.expires_at = time.time() + store.ttl_seconds try: store.save(state) @@ -2194,6 +2321,7 @@ def _complete_with_state_full( state.terminal_error = _TERMINAL_STATE_FULL state.model_messages = None state.initial_usage = {} + state.conversation_history = [] state.final_persistence_pending = False state.expires_at = time.time() + store.ttl_seconds try: @@ -2292,6 +2420,10 @@ async def _handle_brokered_continuation( call_id = item.get("call_id") if not isinstance(call_id, str) or not call_id: return _error("function_call_output.call_id is required", status=400, code="missing_call_id") + if state.response_calls and state.response_calls.get(str(previous_response_id)) != call_id: + return _error("function_call_output does not match previous_response_id", status=400, code="unknown_call_id") + if not state.response_calls and previous_response_id != state.response_id: + return _error("function_call_output does not match previous_response_id", status=400, code="unknown_call_id") call = state.pending_calls.get(call_id) if call is None: return _error("unknown function_call_output call_id", status=400, code="unknown_call_id") @@ -2343,7 +2475,8 @@ async def _handle_brokered_continuation( output_digest = _output_digest(output_json) if existing_output_digest is not None: - if existing_output_digest == output_digest and _state_has_replay(state): + replay_payload = state.continuation_payloads.get(call_id, state.final_payload) + if existing_output_digest == output_digest and (replay_payload is not None or state.terminal_error is not None): if state.final_persistence_pending: state.final_persistence_pending = False try: @@ -2359,10 +2492,11 @@ async def _handle_brokered_continuation( state.final_persistence_pending = True store.cache_in_memory(state) return _state_storage_error() - if state.terminal_error == _TERMINAL_STATE_FULL: + # A later workflow failure must not replace an earlier round's reply. + if replay_payload is None and state.terminal_error == _TERMINAL_STATE_FULL: return _state_full_error() - assert state.final_payload is not None - return JSONResponse(state.final_payload) + assert replay_payload is not None + return JSONResponse(replay_payload) if existing_output_digest == output_digest: return _error( "matching function_call_output is already being processed", @@ -2410,7 +2544,9 @@ async def _handle_brokered_continuation( store.mark_resume_active(state.response_id) try: try: - model_result = await model_loop.resume(state.model_messages, call_id=call_id, output=output_json) + next_response_id = _new_response_id(str(previous_response_id)) + next_call_id = f"call_{next_response_id}_1" + model_result = await model_loop.resume(state.model_messages, call_id=call_id, output=output_json, next_call_id=next_call_id) except AgentRunError as exc: if not _reset_unfinalized_continuation(store, state, call_id=call_id): return _state_storage_error() @@ -2429,13 +2565,17 @@ async def _handle_brokered_continuation( if not _reset_unfinalized_continuation(store, state, call_id=call_id): return _state_storage_error() return _error("model resume failed", status=502, code="ModelResumeError") - if not isinstance(model_result, ModelLoopFinal): - if not _reset_unfinalized_continuation(store, state, call_id=call_id): - return _state_storage_error() - return _error( - "model requested another brokered tool after resume", - status=400, - code="tool_loop_limit_exceeded", + if isinstance(model_result, ModelLoopToolRequest): + return _advance_brokered_state( + spec=spec, + store=store, + state=state, + previous_response_id=str(previous_response_id), + call_id=call_id, + result=model_result, + model_loop=model_loop, + response_id=next_response_id, + next_call_id=next_call_id, ) result = RunResult(text=model_result.text, usage=_combine_usage(state.initial_usage, model_result.usage)) finally: @@ -2446,7 +2586,7 @@ async def _handle_brokered_continuation( resume_initial_usage = dict(state.initial_usage) has_resume_transcript = resume_model_messages is not None used_model_resume = has_resume_transcript and model_loop is not None - final_payload = _responses_payload(spec, result, previous_response_id=state.response_id) + final_payload = _responses_payload(spec, result, previous_response_id=str(previous_response_id)) state.accepted_output_digests[call_id] = output_digest state.accepted_output_sizes[call_id] = accepted_output_size state.status = "completed" @@ -2455,6 +2595,10 @@ async def _handle_brokered_continuation( if has_resume_transcript: state.model_messages = None state.initial_usage = {} + if state.session_id and isinstance(model_result, ModelLoopFinal): + state.conversation_history = _bounded_conversation_history( + model_result.messages, max_bytes=min(model_loop.max_messages_bytes // 2, store.max_bytes // 8), + ) state.final_persistence_pending = False state.expires_at = time.time() + store.ttl_seconds try: @@ -2537,6 +2681,7 @@ def max_brokered_request_body_bytes() -> int: max_argument_bytes=max_argument_bytes, max_output_bytes=max_output_bytes, max_response_bytes=response_states.max_bytes, + max_messages_bytes=model_messages_limit, ) if brokered_tools and _brokered_model_loop_enabled(brokered_model_loop_enabled) else None @@ -2684,6 +2829,7 @@ async def responses(request: Request): max_output_bytes=max_output_bytes, model_loop=model_loop, ) + previous_state = None if brokered_tools and isinstance(previous_response_id, str) and previous_response_id: try: previous_state = response_states.get(previous_response_id) @@ -2691,7 +2837,7 @@ async def responses(request: Request): logger.warning("failed to access Foundry brokered response state: %s", exc) return _state_storage_error() except _StateExpired as exc: - if exc.state.status in {"pending", "resuming"}: + if exc.state.status in {"pending", "resuming"} or (model_loop is not None and session_id): return _error("previous_response_id state has expired", status=410, code="response_state_expired") previous_state = None except KeyError: @@ -2702,6 +2848,15 @@ async def responses(request: Request): status=409, code="response_pending_function_call_output", ) + if model_loop is not None and (session_id or (previous_state is not None and previous_state.session_id)): + if previous_state is None: + return _error("unknown previous_response_id", status=404, code="unknown_previous_response_id") + if previous_state.session_id != session_id: + return _error("previous_response_id requires the same effective Foundry session", status=409, code="response_session_mismatch") + if not previous_state.final_payload or previous_state.final_payload.get("id") != previous_response_id: + return _error("follow-up requires the completed response ID", status=409, code="response_not_completed") + if not previous_state.conversation_history: + return _error("previous conversation no longer fits retained history", status=409, code="response_history_unavailable") try: run_request = _responses_input_to_run_request( @@ -2710,6 +2865,11 @@ async def responses(request: Request): ) except ValueError as exc: return _error(str(exc), status=400, code="invalid_input") + if model_loop is not None and session_id and previous_state is not None: + run_request = replace(run_request, history=( + *(ConversationTurn(role=turn["role"], text=turn["content"]) for turn in previous_state.conversation_history), + *run_request.history, + )) if brokered_tools: if not continuation_proof: @@ -2745,13 +2905,32 @@ async def responses(request: Request): return _model_response_too_large_error() return _error(str(exc), status=exc.status, code=exc.code) if isinstance(model_result, ModelLoopFinal): - return JSONResponse( - _responses_payload( - spec, - RunResult(text=model_result.text, usage=model_result.usage), - previous_response_id=previous_response_id_for_output, - ) + payload = _responses_payload( + spec, + RunResult(text=model_result.text, usage=model_result.usage), + previous_response_id=previous_response_id_for_output, ) + if run_request.session_id: + completed = _HostedResponseState( + response_id=response_id, + session_id=run_request.session_id, + pending_calls={}, + expires_at=time.time() + response_states.ttl_seconds, + status="completed", + final_payload=payload, + conversation_history=_bounded_conversation_history( + model_result.messages, max_bytes=min(model_messages_limit // 2, response_states.max_bytes // 8), + ), + ) + try: + response_states.add_reserved(completed) + except _StateStoreFull: + return _state_full_error() + except _StateSizeLimitExceeded: + return _model_response_too_large_error() + except _StatePersistenceError: + return _state_storage_error() + return JSONResponse(payload) tool = {tool.name: tool for tool in brokered_tools}.get(model_result.name) if tool is None: return _error("model requested unknown brokered tool", status=400, code="unknown_brokered_tool") diff --git a/runtimes/common/agentkit_serve_common/foundry_model_loop.py b/runtimes/common/agentkit_serve_common/foundry_model_loop.py index 7cbba5b..be8a778 100644 --- a/runtimes/common/agentkit_serve_common/foundry_model_loop.py +++ b/runtimes/common/agentkit_serve_common/foundry_model_loop.py @@ -3,8 +3,8 @@ This is the Phase A4 fallback path: when high-level frameworks cannot suspend and resume externally brokered tool calls, AgentKit can drive a minimal model loop itself. The loop exposes only static safe brokered schemas to the model, converts -one model tool request into a hosted Responses function_call, and later resumes -with Orka's function_call_output to obtain the final assistant message. +one model tool request at a time into a hosted Responses function_call, and +resumes with Orka's function_call_output until the assistant finishes. """ from __future__ import annotations @@ -13,6 +13,7 @@ import json import math import os +import uuid from dataclasses import dataclass, field from decimal import Decimal, InvalidOperation from typing import Any, Mapping, Sequence @@ -23,6 +24,7 @@ from .config import AgentSpec from .conversation import FORWARDED_ROLES, RunRequest from .runtime import AgentRunError, BrokeredToolDefinition +from .skills import SkillCatalog _MAX_ARGUMENT_DEPTH = 128 @@ -31,6 +33,7 @@ class ModelLoopFinal: text: str usage: dict[str, int] = field(default_factory=dict) + messages: list[dict[str, Any]] = field(default_factory=list) @dataclass(frozen=True) @@ -42,7 +45,7 @@ class ModelLoopToolRequest: class BrokeredChatModelLoop: - """Explicit one-tool brokered model loop over OpenAI Chat Completions.""" + """Bounded sequential tool loop over OpenAI Chat Completions.""" def __init__( self, @@ -53,6 +56,8 @@ def __init__( max_argument_bytes: int = 8192, max_output_bytes: int = 64 * 1024, max_response_bytes: int = 4 * 1024 * 1024, + max_messages_bytes: int = 1024 * 1024, + max_tool_calls: int = 16, ) -> None: self.spec = spec self.tools = list(tools) @@ -60,19 +65,55 @@ def __init__( self.max_argument_bytes = max_argument_bytes self.max_output_bytes = max_output_bytes self.max_response_bytes = max_response_bytes + self.max_messages_bytes = max_messages_bytes + self.max_tool_calls = max_tool_calls self.tools_by_name = {tool.name: tool for tool in self.tools} + self.skills = SkillCatalog.from_spec(spec) + if self.skills and "load_skill" in self.tools_by_name: + raise AgentBuildError("load_skill is reserved for packaged skills") async def start(self, request: RunRequest, *, call_id: str) -> ModelLoopFinal | ModelLoopToolRequest: - messages = self._initial_messages(request) - data = await self._chat(messages, tools=self._tool_payloads()) + return await self._advance(self._initial_messages(request), call_id=call_id) + + async def _advance(self, messages: list[dict[str, Any]], *, call_id: str) -> ModelLoopFinal | ModelLoopToolRequest: + usage: dict[str, int] = {} + while True: + if len(json.dumps(messages, ensure_ascii=True).encode("utf-8")) > self.max_messages_bytes: + raise AgentRunError("model loop messages are too large", status=413, code="brokered_model_messages_too_large") + result = await self._step(messages, call_id=call_id) + for key, value in result.usage.items(): + usage[key] = usage.get(key, 0) + value + if isinstance(result, ModelLoopFinal): + return ModelLoopFinal(text=result.text, usage=usage, messages=result.messages) + if not self.skills or result.name != "load_skill": + return ModelLoopToolRequest(name=result.name, arguments=result.arguments, messages=result.messages, usage=usage) + if set(result.arguments) != {"skill_name"} or not isinstance(result.arguments["skill_name"], str): + raise AgentRunError("load_skill requires a skill_name string", status=400, code="InvalidToolArguments") + try: + content = self.skills.load_skill(result.arguments["skill_name"]) + except (KeyError, ValueError) as exc: + raise AgentRunError("unknown packaged skill", status=400, code="InvalidToolArguments") from exc + messages = result.messages + skill_call_id = f"skill_{uuid.uuid4().hex}" + messages[-1]["tool_calls"][0]["id"] = skill_call_id + messages.append({"role": "tool", "tool_call_id": skill_call_id, "content": content}) + + async def _step(self, messages: list[dict[str, Any]], *, call_id: str) -> ModelLoopFinal | ModelLoopToolRequest: + current_turn = next((i for i in range(len(messages) - 1, -1, -1) if messages[i].get("role") == "user"), 0) + tool_count = sum(len(message.get("tool_calls") or []) for message in messages[current_turn:]) + exhausted = tool_count >= self.max_tool_calls + data = await self._chat(messages, tools=[] if exhausted else self._tool_payloads()) message = _choice_message(data) usage = _usage(data) tool_calls = message.get("tool_calls") if not tool_calls: - return ModelLoopFinal(text=_message_text(message, max_bytes=self.max_output_bytes), usage=usage) + text = _message_text(message, max_bytes=self.max_output_bytes) + return ModelLoopFinal(text=text, usage=usage, messages=[*messages, {"role": "assistant", "content": text}]) + if exhausted: + raise AgentRunError("model exceeded the tool call limit", status=400, code="tool_loop_limit_exceeded") if not isinstance(tool_calls, list) or len(tool_calls) != 1: raise AgentRunError( - "model requested multiple brokered tools; deterministic brokered mode supports one call per turn", + "model requested multiple tools at once; brokered mode requires sequential calls", status=400, code="multiple_tool_calls_unsupported", ) @@ -83,7 +124,7 @@ async def start(self, request: RunRequest, *, call_id: str) -> ModelLoopFinal | if not isinstance(function, Mapping): raise AgentRunError("model tool call is missing function payload", status=400, code="invalid_tool_call") name = function.get("name") - if not isinstance(name, str) or name not in self.tools_by_name: + if not isinstance(name, str) or (name not in self.tools_by_name and not (self.skills and name == "load_skill")): raise AgentRunError(f"model requested unknown brokered tool {name!r}", status=400, code="unknown_brokered_tool") raw_arguments = function.get("arguments", "{}") if isinstance(raw_arguments, str): @@ -112,7 +153,14 @@ async def start(self, request: RunRequest, *, call_id: str) -> ModelLoopFinal | } return ModelLoopToolRequest(name=name, arguments=arguments, messages=[*messages, assistant_message], usage=usage) - async def resume(self, messages: Sequence[Mapping[str, Any]], *, call_id: str, output: str) -> ModelLoopFinal: + async def resume( + self, + messages: Sequence[Mapping[str, Any]], + *, + call_id: str, + output: str, + next_call_id: str | None = None, + ) -> ModelLoopFinal | ModelLoopToolRequest: if len(output) > self.max_output_bytes: raise AgentRunError("brokered tool output is too large for model resume", status=413, code="brokered_output_too_large") try: @@ -127,11 +175,7 @@ async def resume(self, messages: Sequence[Mapping[str, Any]], *, call_id: str, o raise AgentRunError("brokered tool output is too large for model resume", status=413, code="brokered_output_too_large") resumed = [dict(message) for message in messages] resumed.append({"role": "tool", "tool_call_id": call_id, "content": output}) - data = await self._chat(resumed, tools=[]) - message = _choice_message(data) - if message.get("tool_calls"): - raise AgentRunError("model requested another brokered tool after resume", status=400, code="tool_loop_limit_exceeded") - return ModelLoopFinal(text=_message_text(message, max_bytes=self.max_output_bytes), usage=_usage(data)) + return await self._advance(resumed, call_id=next_call_id or f"call_{uuid.uuid4().hex}") async def validate_credentials(self) -> None: await self._auth_headers() @@ -148,6 +192,8 @@ def _initial_messages(self, request: RunRequest) -> list[dict[str, Any]]: messages: list[dict[str, Any]] = [] if self.spec.instructions: messages.append({"role": "system", "content": self.spec.instructions}) + if self.skills: + messages.append({"role": "system", "content": self.skills.instructions}) for turn in request.history: if turn.role in FORWARDED_ROLES and turn.text: messages.append({"role": turn.role, "content": turn.text}) @@ -155,7 +201,7 @@ def _initial_messages(self, request: RunRequest) -> list[dict[str, Any]]: return messages def _tool_payloads(self) -> list[dict[str, Any]]: - payloads: list[dict[str, Any]] = [] + payloads: list[dict[str, Any]] = [self.skills.tool_schema()] if self.skills else [] for tool in self.tools: description = f"Brokered class: {tool.brokered_class}. {tool.description}".strip() payloads.append( @@ -175,6 +221,7 @@ async def _chat(self, messages: Sequence[Mapping[str, Any]], *, tools: Sequence[ if tools: payload["tools"] = list(tools) payload["tool_choice"] = "auto" + payload["parallel_tool_calls"] = False headers = await self._auth_headers() client = self.http_client close_client = False diff --git a/runtimes/common/agentkit_serve_common/skills.py b/runtimes/common/agentkit_serve_common/skills.py new file mode 100644 index 0000000..5c91fc8 --- /dev/null +++ b/runtimes/common/agentkit_serve_common/skills.py @@ -0,0 +1,235 @@ +"""Image-bundled skill instructions for governed runtimes. + +Only SKILL.md documents are loaded. The resulting catalog is immutable and its +load_skill function does no I/O: it cannot read references, execute scripts, or +contact the services mentioned in a skill. Operational tools remain brokered. +""" + +from __future__ import annotations + +import os +import re +import stat +from dataclasses import dataclass +from html import escape +from pathlib import Path +from typing import Any + +import yaml + +from .config import AgentSpec +from .yaml_support import safe_load_lossless + +_SKILLS_ROOT = Path("/agent/skills") +_MAX_SKILL_BYTES = 128 * 1024 +_MAX_CATALOG_BYTES = 1024 * 1024 +_MAX_SKILLS = 64 +_MAX_DIRECTORY_ENTRIES = 256 +_MAX_DISCOVERY_DEPTH = 2 +_NAME_PATTERN = re.compile(r"[a-z0-9]+(?:-[a-z0-9]+)*") + + +class SkillConfigurationError(ValueError): + """A skill configuration or document cannot be used in governed mode.""" + + +@dataclass(frozen=True) +class _Skill: + name: str + description: str + content: str + + +@dataclass(frozen=True) +class SkillCatalog: + """Immutable instructions with a local load_skill tool that never performs I/O.""" + + documents: tuple[_Skill, ...] = () + + @classmethod + def from_spec(cls, spec: AgentSpec) -> SkillCatalog: + """Snapshot bounded, UTF-8 SKILL.md files from the image's skill directories.""" + validate_packaged_skill_providers(spec) + documents: list[_Skill] = [] + total_bytes = 0 + names: set[str] = set() + try: + for provider in spec.context.providers: + path = Path(provider.path or "") + directory_fd = _open_directory(path) + try: + discovered = _read_directory(directory_fd, path.name, depth=0) + finally: + os.close(directory_fd) + if not discovered: + raise SkillConfigurationError("filesystem skill directory contains no SKILL.md documents") + for skill in discovered: + if skill.name in names: + raise SkillConfigurationError("duplicate bundled skill name") + names.add(skill.name) + total_bytes += len(skill.content.encode("utf-8")) + if len(names) > _MAX_SKILLS or total_bytes > _MAX_CATALOG_BYTES: + raise SkillConfigurationError("bundled skill catalog is too large") + documents.append(skill) + except OSError as exc: + # Keep filesystem exception details and any unexpected file paths out of + # protocol responses. Symlinks and non-directories fail here as well. + raise SkillConfigurationError( + "filesystem skills must be readable image directories without symlinks" + ) from exc + return cls(tuple(sorted(documents, key=lambda skill: skill.name))) + + def __bool__(self) -> bool: + return bool(self.documents) + + @property + def names(self) -> tuple[str, ...]: + return tuple(skill.name for skill in self.documents) + + @property + def instructions(self) -> str: + if not self.documents: + return "" + entries = "\n".join( + f"{skill.name}{escape(skill.description)}" + for skill in self.documents + ) + return ( + "The agent image includes these instruction-only skills:\n" + f"\n{entries}\n\n" + "When a skill applies, use load_skill with its skill_name to retrieve its " + "instructions before following them. Skills supply guidance, not execution " + "permissions. Use the authorized tools for operational data and actions." + ) + + def tool_schema(self) -> dict[str, Any]: + """Return the local load_skill definition in OpenAI Chat tool format.""" + return { + "type": "function", + "function": { + "name": "load_skill", + "description": "Loads the full instructions for a bundled skill.", + "parameters": { + "type": "object", + "properties": { + "skill_name": { + "type": "string", + "description": "The name of the skill to load.", + "enum": list(self.names), + } + }, + "required": ["skill_name"], + "additionalProperties": False, + }, + }, + } + + def load_skill(self, skill_name: str) -> str: + """Read an already-loaded document, never a model-supplied file path.""" + if isinstance(skill_name, str): + for skill in self.documents: + if skill.name == skill_name: + return skill.content + raise SkillConfigurationError("unknown bundled skill") + + +def validate_packaged_skill_providers(spec: AgentSpec) -> None: + """Accept only filesystem skill declarations without external provider fields.""" + for provider in spec.context.providers: + if provider.type != "skills" or provider.source != "filesystem": + raise SkillConfigurationError("governed context providers must be instruction-only filesystem skills") + if any( + value is not None + for value in ( + provider.tool_ref, + provider.endpoint_env, + provider.index, + provider.index_env, + provider.store_name_env, + provider.auth, + ) + ): + raise SkillConfigurationError("filesystem skills must not configure external providers") + path = Path(provider.path or "") + if not path.is_absolute() or ".." in path.parts or not path.is_relative_to(_SKILLS_ROOT): + raise SkillConfigurationError("filesystem skills must be under /agent/skills") + + +def _open_directory(path: Path) -> int: + """Open each path component without following symlinks, including ancestors.""" + flags = os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW + directory_fd = os.open(path.anchor, flags) + try: + for component in path.parts[1:]: + child_fd = os.open(component, flags, dir_fd=directory_fd) + os.close(directory_fd) + directory_fd = child_fd + return directory_fd + except BaseException: + os.close(directory_fd) + raise + + +def _read_directory(directory_fd: int, name: str, *, depth: int) -> list[_Skill]: + entries: list[os.DirEntry[str]] = [] + with os.scandir(directory_fd) as scan: + for entry in scan: + if len(entries) >= _MAX_DIRECTORY_ENTRIES: + raise SkillConfigurationError("bundled skill directory is too large") + entries.append(entry) + if any(entry.name == "SKILL.md" for entry in entries): + return [_read_document(directory_fd, name)] + documents: list[_Skill] = [] + for entry in sorted(entries, key=lambda value: value.name): + if entry.is_symlink(): + raise SkillConfigurationError("bundled skill directories must not contain symlinks") + if depth < _MAX_DISCOVERY_DEPTH and entry.is_dir(follow_symlinks=False): + child_fd = os.open(entry.name, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=directory_fd) + try: + documents.extend(_read_directory(child_fd, entry.name, depth=depth + 1)) + finally: + os.close(child_fd) + if len(documents) > _MAX_SKILLS: + raise SkillConfigurationError("bundled skill catalog is too large") + return documents + + +def _read_document(directory_fd: int, directory_name: str) -> _Skill: + document_fd = os.open("SKILL.md", os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK, dir_fd=directory_fd) + try: + info = os.fstat(document_fd) + if not stat.S_ISREG(info.st_mode) or info.st_nlink != 1: + raise SkillConfigurationError("SKILL.md must be a regular file without hard links") + with os.fdopen(document_fd, "rb", closefd=False) as source: + raw = source.read(_MAX_SKILL_BYTES + 1) + finally: + os.close(document_fd) + if len(raw) > _MAX_SKILL_BYTES: + raise SkillConfigurationError("SKILL.md exceeds the document size limit") + try: + content = raw.decode("utf-8") + except UnicodeDecodeError as exc: + raise SkillConfigurationError("SKILL.md must be UTF-8") from exc + match = re.match(r"\A---[ \t]*\r?\n(.*?)\r?\n---[ \t]*(?:\r?\n|\Z)", content, re.DOTALL) + if match is None: + raise SkillConfigurationError("SKILL.md must contain YAML frontmatter") + try: + metadata = safe_load_lossless(match.group(1)) + except (yaml.YAMLError, ValueError, RecursionError) as exc: + raise SkillConfigurationError("SKILL.md frontmatter is invalid") from exc + if not isinstance(metadata, dict): + raise SkillConfigurationError("SKILL.md frontmatter must be a mapping") + name, description = metadata.get("name"), metadata.get("description") + if not isinstance(name, str) or len(name) > 64 or not _NAME_PATTERN.fullmatch(name) or name != directory_name: + raise SkillConfigurationError( + "SKILL.md name must match its directory and use lowercase letters, digits, and single hyphens" + ) + if not isinstance(description, str) or not description.strip() or len(description) > 1024: + raise SkillConfigurationError("SKILL.md description must contain 1 to 1024 characters") + try: + description.encode("utf-8") + except UnicodeEncodeError as exc: + raise SkillConfigurationError("SKILL.md description must be valid Unicode") from exc + if not content[match.end() :].strip() or "\x00" in content: + raise SkillConfigurationError("SKILL.md must contain nonempty text instructions") + return _Skill(name=name, description=description, content=content) diff --git a/runtimes/common/tests/test_acp_protocol.py b/runtimes/common/tests/test_acp_protocol.py index a24421d..190dca7 100644 --- a/runtimes/common/tests/test_acp_protocol.py +++ b/runtimes/common/tests/test_acp_protocol.py @@ -18,6 +18,7 @@ from agentkit_serve_common.acp import ACPConfigurationError, ACPStdioServer from agentkit_serve_common.config import AgentSpec from agentkit_serve_common.conversation import ConversationTurn, RunRequest, ToolCallEvent +from agentkit_serve_common.skills import _Skill, SkillCatalog from agentkit_serve_common.runtime import ( AgentRunError, OfflineEchoRuntimeFactory, @@ -1726,10 +1727,10 @@ def test_runtime_binding_rejects_profile_or_provider_mismatch( "context": { "providers": [ { - "name": "skills", - "type": "skills", - "source": "filesystem", - "path": "/agent/skills", + "name": "search", + "type": "search", + "endpointEnv": "SEARCH_ENDPOINT", + "indexEnv": "SEARCH_INDEX", } ] } @@ -1744,3 +1745,92 @@ def test_runtime_binding_rejects_baked_tool_and_context_paths(tmp_path, override with pytest.raises(ACPConfigurationError, match=message): acp.validate_acp_runtime_binding(b"unused", _spec(**override)) + + +def test_acp_binds_instruction_skill_snapshot_and_preserves_http_mcp(monkeypatch): + spec = _spec(context={"providers": [{ + "type": "skills", "source": "filesystem", "path": "/agent/skills", + }]}) + catalog = SkillCatalog((_Skill("inspection", "Inspect equipment.", "original instructions"),)) + monkeypatch.setattr(SkillCatalog, "from_spec", classmethod(lambda cls, _: catalog)) + monkeypatch.setenv(acp.ACP_AGENT_CONFIGURATION_DIGEST_ENV, "sha256:" + hashlib.sha256(b"exact").hexdigest()) + monkeypatch.setenv(acp.ACP_MODEL_ENV, "test-model") + _set_provider_environment(monkeypatch) + acp.validate_acp_runtime_binding(b"exact", spec) + assert spec._packaged_skill_catalog is catalog + assert "_packaged_skill_catalog" not in spec.model_dump() + monkeypatch.setattr(SkillCatalog, "from_spec", classmethod(lambda cls, _: pytest.fail("snapshot must not be reloaded"))) + + class SkillsFactory(RecordingFactory): + def supports_acp_packaged_skills(self): + return True + + async def exercise(): + messages = [] + + async def send(message): + messages.append(dict(message)) + + runtime = RecordingRuntime([RunResult("done")]) + factory = SkillsFactory(lambda: runtime, supports_http_mcp=True) + server = ACPStdioServer(spec, factory, send) + try: + initialized = await _send_to(server, messages, _initialize()) + assert initialized["result"]["agentCapabilities"]["mcpCapabilities"] == {"http": True} + request = _new_session(mcp_servers=[_orka_mcp_server()]) + request["params"]["context"] = {"providers": [{"type": "skills", "source": "filesystem", "path": "/workspace"}]} + request["params"]["_packaged_skill_catalog"] = {"inspection": "request instructions"} + session = await _send_to(server, messages, request) + assert "sessionId" in session["result"] + projected = factory.specs[0] + assert projected._packaged_skill_catalog is catalog + assert projected._packaged_skill_catalog.load_skill("inspection") == "original instructions" + assert projected.context.providers[0].path == "/agent/skills" + assert len(projected.tools) == 1 + assert projected.tools[0].url_env.startswith("AGENTKIT_ACP_SESSION_") + finally: + await server.close() + + asyncio.run(exercise()) + + +def test_acp_checks_config_binding_before_reading_skill_files(monkeypatch): + spec = _spec(context={"providers": [{ + "type": "skills", "source": "filesystem", "path": "/agent/skills", + }]}) + monkeypatch.setattr(SkillCatalog, "from_spec", classmethod(lambda cls, _: pytest.fail("unverified skills must not be read"))) + monkeypatch.setenv(acp.ACP_AGENT_CONFIGURATION_DIGEST_ENV, "sha256:" + hashlib.sha256(b"expected").hexdigest()) + with pytest.raises(ACPConfigurationError, match="exact agent config bytes"): + acp.validate_acp_runtime_binding(b"changed", spec) + + +def test_agent_config_cannot_supply_a_packaged_skill_snapshot(): + with pytest.raises(ValueError, match="Extra inputs are not permitted"): + _spec(_packaged_skill_catalog={"inspection": "unverified instructions"}) + + +def test_acp_rejects_skills_on_an_unsupported_runtime(monkeypatch): + _set_provider_environment(monkeypatch) + spec = _spec(context={"providers": [{ + "type": "skills", "source": "filesystem", "path": "/agent/skills", + }]}) + + async def exercise(): + messages = [] + + async def send(message): + messages.append(dict(message)) + + factory = RecordingFactory(lambda: RecordingRuntime([]), supports_http_mcp=True) + server = ACPStdioServer(spec, factory, send) + try: + initialized = await _send_to(server, messages, _initialize()) + assert initialized["result"]["agentCapabilities"]["mcpCapabilities"] == {} + session = await _send_to(server, messages, _new_session()) + assert session["error"]["code"] == -32602 + assert "does not support ACP packaged skills" in session["error"]["message"] + assert factory.specs == [] + finally: + await server.close() + + asyncio.run(exercise()) diff --git a/runtimes/common/tests/test_foundry_brokered_protocol.py b/runtimes/common/tests/test_foundry_brokered_protocol.py index f2acf92..2b1036c 100644 --- a/runtimes/common/tests/test_foundry_brokered_protocol.py +++ b/runtimes/common/tests/test_foundry_brokered_protocol.py @@ -2258,7 +2258,7 @@ def __init__(self) -> None: async def handle_async_request(self, request: httpx.Request) -> httpx.Response: payload = json.loads(request.content.decode("utf-8")) self.requests.append(payload) - if "tools" in payload: + if not any(message.get("role") == "tool" for message in payload["messages"]): self.initial_calls += 1 return httpx.Response( 200, @@ -2390,7 +2390,8 @@ def test_foundry_brokered_model_loop_emits_model_requested_tool_and_resumes_to_f assert fake.requests[0]["tools"][0]["function"]["description"].startswith("Brokered class: read.") assert fake.requests[1]["messages"][-1]["role"] == "tool" assert fake.requests[1]["messages"][-1]["tool_call_id"] == call["call_id"] - assert "tools" not in fake.requests[1] + assert fake.requests[1]["tools"] == fake.requests[0]["tools"] + assert fake.requests[1]["parallel_tool_calls"] is False assert CONTINUATION_PROOF not in json.dumps(fake.requests, sort_keys=True) @@ -4853,7 +4854,7 @@ def test_foundry_brokered_bounds_persisted_model_messages_and_releases_reservati assert oversized.json()["error"]["code"] == "brokered_model_messages_too_large" assert accepted.status_code == 200, accepted.text assert _call(accepted.json()) - assert len(fake.requests) == 2 + assert len(fake.requests) == 1 # Oversized model input is rejected before inference. def test_foundry_brokered_model_loop_reserves_capacity_before_model_call(): fake = _FakeChatTransport( diff --git a/runtimes/common/tests/test_foundry_tool_workflows.py b/runtimes/common/tests/test_foundry_tool_workflows.py new file mode 100644 index 0000000..fc12876 --- /dev/null +++ b/runtimes/common/tests/test_foundry_tool_workflows.py @@ -0,0 +1,312 @@ +"""Public Responses regressions for sequential governed tool workflows.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +from fastapi.testclient import TestClient + +import agentkit_serve_common.foundry as foundry_module +import agentkit_serve_common.skills as skills_module +from agentkit_serve_common.config import AgentSpec +from test_foundry_brokered_protocol import ( + CONTINUATION_AUTH, + _FakeChatTransport, + _call, + _chat_response, + _continuation, + _message_text, + _model_loop_app, + _multi_tool_spec, + _spec, +) + + +def _tool(name, arguments=None): + return _chat_response({ + "role": "assistant", + "content": None, + "tool_calls": [{"id": "untrusted-model-id", "type": "function", "function": { + "name": name, "arguments": json.dumps(arguments or {}), + }}], + }) + + +def _result(response, output): + return _continuation(response["id"], _call(response)["call_id"], output) + + +@pytest.fixture +def packaged_skill(tmp_path, monkeypatch): + root = tmp_path.resolve() / "skills" + directory = root / "inspection" + directory.mkdir(parents=True) + document = directory / "SKILL.md" + text = ( + "---\nname: inspection\ndescription: Inspect the affected site.\n---\n" + "Read telemetry, then retrieve active incidents using the authorized tools.\n" + ) + document.write_text(text, encoding="utf-8") + open_directory = skills_module._open_directory + + def open_test_directory(path: Path): + return open_directory(root / path.relative_to("/agent/skills")) + + monkeypatch.setattr(skills_module, "_open_directory", open_test_directory) + data = _multi_tool_spec().model_dump(by_alias=True) + data["context"] = {"providers": [{"type": "skills", "source": "filesystem", "path": "/agent/skills"}]} + return AgentSpec.model_validate(data), document, text + + +def test_hosted_skills_load_locally_between_governed_tool_rounds(packaged_skill): + spec, document, text = packaged_skill + fake = _FakeChatTransport([ + _tool("load_skill", {"skill_name": "inspection"}), + _tool("check-network-telemetry", {"site": "site-a"}), + _tool("load_skill", {"skill_name": "inspection"}), + _tool("get-active-incidents"), + _chat_response({"role": "assistant", "content": "Incident INC-17 explains the signal loss."}), + ]) + app = _model_loop_app(spec, fake) + document.unlink() # Both skill calls must use the startup snapshot. + with TestClient(app) as client: + first = client.post("/responses", json={"input": "Inspect site-a"}) + assert first.status_code == 200, first.text + assert _call(first.json())["name"] == "check-network-telemetry" + payload = _result(first.json(), {"approved": True, "output": {"signal": "low"}}) + assert client.post("/responses", json=payload).status_code == 403 + assert len(fake.requests) == 2 + second = client.post("/responses", json=payload, headers=CONTINUATION_AUTH) + assert second.status_code == 200, second.text + assert _call(second.json())["name"] == "get-active-incidents" + final = client.post("/responses", json=_result(second.json(), {"approved": True, "output": {"incident": "INC-17"}}), headers=CONTINUATION_AUTH) + assert final.status_code == 200, final.text + assert _message_text(final.json()) == "Incident INC-17 explains the signal loss." + assert final.json()["usage"] == {"input_tokens": 5, "output_tokens": 5, "total_tokens": 10} + + advertised = {tool["function"]["name"] for tool in fake.requests[0]["tools"]} + assert advertised == {"load_skill", "check-network-telemetry", "get-active-incidents"} + assert "inspection" in fake.requests[0]["messages"][1]["content"] + messages = fake.requests[-1]["messages"] + calls = [call for message in messages for call in message.get("tool_calls", [])] + outputs = [message for message in messages if message["role"] == "tool"] + assert len(calls) == len({call["id"] for call in calls}) == 4 + assert [call["id"] for call in calls] == [output["tool_call_id"] for output in outputs] + assert outputs[0]["content"] == outputs[2]["content"] == text + assert calls[1]["id"] == _call(first.json())["call_id"] + assert calls[3]["id"] == _call(second.json())["call_id"] + assert all(request["parallel_tool_calls"] is False for request in fake.requests) + + +@pytest.mark.parametrize("arguments", [ + {"skill_name": "missing"}, + {"skill_name": "../inspection"}, + {"skill_name": "inspection", "path": "/etc/passwd"}, + {"skill_name": 42}, +]) +def test_hosted_skill_load_rejects_unadvertised_input(packaged_skill, arguments): + spec, _, _ = packaged_skill + fake = _FakeChatTransport([_tool("load_skill", arguments)]) + with TestClient(_model_loop_app(spec, fake)) as client: + response = client.post("/responses", json={"input": "Inspect site-a"}) + assert response.status_code == 400 + assert response.json()["error"]["code"] == "InvalidToolArguments" + assert len(fake.requests) == 1 + + +def test_hosted_local_skill_calls_share_the_tool_budget(packaged_skill): + spec, _, _ = packaged_skill + fake = _FakeChatTransport([_tool("load_skill", {"skill_name": "inspection"}) for _ in range(17)]) + with TestClient(_model_loop_app(spec, fake)) as client: + response = client.post("/responses", json={"input": "Keep loading the skill"}) + assert response.status_code == 400 + assert response.json()["error"]["code"] == "tool_loop_limit_exceeded" + assert len(fake.requests) == 17 + assert "tools" not in fake.requests[-1] + + +@pytest.mark.parametrize("denied", [False, True]) +def test_chained_tool_calls_keep_pairing_replay_and_results_across_restart(tmp_path, denied): + state_file = tmp_path / "responses.json" + spec = _multi_tool_spec() + fake = _FakeChatTransport([ + _tool("check-network-telemetry", {"site": "site-a"}), + _tool("get-active-incidents"), + ]) + first_output = {"approved": False, "error": {"code": "policy_denied", "message": "Read denied"}} if denied else {"approved": True, "output": {"site": "site-a", "signal": "low"}} + with TestClient(_model_loop_app(spec, fake, max_pending_responses=1, response_state_file=state_file)) as client: + initial = client.post("/responses", json={"input": "Check telemetry, then incidents", "agent_session_id": "session-a"}).json() + first_result = _result(initial, first_output) + first_result["agent_session_id"] = "session-a" + unauthorized = client.post("/responses", json=first_result) + assert unauthorized.status_code == 403 + response = client.post("/responses", json=first_result, headers=CONTINUATION_AUTH) + assert response.status_code == 200, response.text + second = response.json() + assert second["id"] != initial["id"] + assert _call(second)["call_id"] != _call(initial)["call_id"] + assert second["previous_response_id"] == initial["id"] + assert _call(second)["name"] == "get-active-incidents" + assert client.post("/responses", json=first_result, headers=CONTINUATION_AUTH).json() == second + assert len(fake.requests) == 2 + assert json.loads(fake.requests[1]["messages"][-1]["content"]) == first_output + # Capacity is per workflow, so advancing remains possible with one entry. + assert len(json.loads(state_file.read_text())["states"]) == 1 + resumed_model = _FakeChatTransport([_chat_response({"role": "assistant", "content": "Incident INC-17 explains the outage."})]) + with TestClient(_model_loop_app(spec, resumed_model, max_pending_responses=1, response_state_file=state_file)) as client: + assert client.post("/responses", json=first_result, headers=CONTINUATION_AUTH).json() == second + second_result = _result(second, {"approved": True, "output": {"incident": "INC-17"}}) + second_result["agent_session_id"] = "session-a" + wrong_pair = dict(first_result, previous_response_id=second["id"]) + assert client.post("/responses", json=wrong_pair, headers=CONTINUATION_AUTH).status_code == 400 + wrong_session = dict(second_result, agent_session_id="session-b") + assert client.post("/responses", json=wrong_session, headers=CONTINUATION_AUTH).status_code == 409 + completed = client.post("/responses", json=second_result, headers=CONTINUATION_AUTH) + assert completed.status_code == 200, completed.text + assert _message_text(completed.json()) == "Incident INC-17 explains the outage." + assert completed.json()["previous_response_id"] == second["id"] + assert client.post("/responses", json=first_result, headers=CONTINUATION_AUTH).json() == second + assert client.post("/responses", json=second_result, headers=CONTINUATION_AUTH).json() == completed.json() + assert len(resumed_model.requests) == 1 + returned = [json.loads(m["content"]) for m in resumed_model.requests[0]["messages"] if m["role"] == "tool"] + assert returned == [first_output, {"approved": True, "output": {"incident": "INC-17"}}] + + +def test_every_new_round_validates_tool_and_arguments_before_export(): + fake = _FakeChatTransport([ + _tool("check-network-telemetry", {"site": "site-a"}), + _tool("check-network-telemetry", {"site": 42}), + _tool("get-active-incidents"), + ]) + with TestClient(_model_loop_app(_multi_tool_spec(), fake)) as client: + initial = client.post("/responses", json={"input": "Inspect the outage"}).json() + payload = _result(initial, {"approved": True, "output": {}}) + rejected = client.post("/responses", json=payload, headers=CONTINUATION_AUTH) + assert rejected.status_code == 400 + assert rejected.json()["error"]["code"] == "InvalidToolArguments" + retried = client.post("/responses", json=payload, headers=CONTINUATION_AUTH) + assert retried.status_code == 200, retried.text + assert _call(retried.json())["name"] == "get-active-incidents" + + +def test_chained_round_write_failure_retries_cached_transition_without_model_work(tmp_path, monkeypatch): + fake = _FakeChatTransport([_tool("conformance_read"), _tool("conformance_read")]) + original = foundry_module._FoundryResponseStateStore._persist + fail = False + + def persist(store, data): + if fail and b'"continuationPayloads"' in data: + raise foundry_module._StatePersistenceError("injected transition failure") + return original(store, data) + + monkeypatch.setattr(foundry_module._FoundryResponseStateStore, "_persist", persist) + with TestClient(_model_loop_app(_spec(), fake, response_state_file=tmp_path / "responses.json")) as client: + initial = client.post("/responses", json={"input": "Read twice"}).json() + payload = _result(initial, {"approved": True, "output": {}}) + fail = True + assert client.post("/responses", json=payload, headers=CONTINUATION_AUTH).status_code == 503 + fail = False + retried = client.post("/responses", json=payload, headers=CONTINUATION_AUTH) + assert retried.status_code == 200, retried.text + assert _call(retried.json())["call_id"] != _call(initial)["call_id"] + assert len(fake.requests) == 2 + + +def test_earlier_round_replays_after_later_capacity_failure_and_restart(tmp_path): + state_file = tmp_path / "responses.json" + fake = _FakeChatTransport([ + _tool("conformance_read"), + _tool("conformance_read"), + _tool("conformance_read"), + _chat_response({"role": "assistant", "content": "x" * 6_000}), + ]) + options = {"response_state_file": state_file, "max_response_state_bytes": 10_000} + with TestClient(_model_loop_app(_spec(), fake, **options)) as client: + initial = client.post("/responses", json={"input": "Read twice"}).json() + first_result = _result(initial, {"approved": True, "output": {}}) + second = client.post("/responses", json=first_result, headers=CONTINUATION_AUTH) + assert second.status_code == 200, second.text + unrelated = client.post("/responses", json={"input": "Another pending read"}) + assert unrelated.status_code == 200, unrelated.text + second_result = _result(second.json(), {"approved": True, "output": {}}) + failed = client.post("/responses", json=second_result, headers=CONTINUATION_AUTH) + assert failed.status_code == 429, failed.text + assert failed.json()["error"]["code"] == "brokered_response_state_full" + assert client.post("/responses", json=second_result, headers=CONTINUATION_AUTH).json() == failed.json() + replayed = client.post("/responses", json=first_result, headers=CONTINUATION_AUTH) + assert replayed.status_code == 200, replayed.text + assert replayed.json() == second.json() + assert len(fake.requests) == 4 + + restarted = _FakeChatTransport([]) + with TestClient(_model_loop_app(_spec(), restarted, **options)) as client: + replayed = client.post("/responses", json=first_result, headers=CONTINUATION_AUTH) + assert replayed.status_code == 200, replayed.text + assert replayed.json() == second.json() + assert client.post("/responses", json=second_result, headers=CONTINUATION_AUTH).json() == failed.json() + assert not restarted.requests + + +def test_brokered_workflow_has_a_finite_tool_budget(): + fake = _FakeChatTransport([_tool("conformance_read") for _ in range(17)]) + with TestClient(_model_loop_app(_spec(), fake, max_pending_responses=1)) as client: + response = client.post("/responses", json={"input": "Keep reading"}) + for _ in range(16): + assert response.status_code == 200, response.text + response = client.post("/responses", json=_result(response.json(), {"approved": True, "output": {}}), headers=CONTINUATION_AUTH) + assert response.status_code == 400 + assert response.json()["error"]["code"] == "tool_loop_limit_exceeded" + assert "tools" not in fake.requests[-1] + + +def test_hosted_followups_retain_dialogue_without_tool_data_across_restart(tmp_path): + state_file = tmp_path / "responses.json" + fake = _FakeChatTransport([ + _chat_response({"role": "assistant", "content": "I will inspect site-a."}), + _tool("conformance_read"), + _chat_response({"role": "assistant", "content": "Site-a has incident INC-17."}), + ]) + with TestClient(_model_loop_app(_spec(), fake, response_state_file=state_file)) as client: + first = client.post("/responses", json={"input": "My site is site-a", "agent_session_id": "session-a"}) + assert first.status_code == 200, first.text + second = client.post("/responses", json={"input": "Check it", "agent_session_id": "session-a", "previous_response_id": first.json()["id"]}) + assert second.status_code == 200, second.text + request = _result(second.json(), {"approved": True, "output": {"incident": "INC-17", "internal_note": "tool-data-only"}}) + request["agent_session_id"] = "session-a" + completed = client.post("/responses", json=request, headers=CONTINUATION_AUTH) + assert completed.status_code == 200, completed.text + context = fake.requests[1]["messages"] + assert context[-3:] == [ + {"role": "user", "content": "My site is site-a"}, + {"role": "assistant", "content": "I will inspect site-a."}, + {"role": "user", "content": "Check it"}, + ] + restarted = _FakeChatTransport([_chat_response({"role": "assistant", "content": "INC-17 is the incident at site-a."})]) + with TestClient(_model_loop_app(_spec(), restarted, response_state_file=state_file)) as client: + payload = {"input": "Which incident was that?", "agent_session_id": "session-a", "previous_response_id": completed.json()["id"]} + for session in [None, "session-b"]: + wrong = dict(payload) + if session is None: + del wrong["agent_session_id"] + else: + wrong["agent_session_id"] = session + assert client.post("/responses", json=wrong).status_code == 409 + followup = client.post("/responses", json=payload) + assert followup.status_code == 200, followup.text + assert _message_text(followup.json()) == "INC-17 is the incident at site-a." + messages = restarted.requests[0]["messages"] + assert any(message.get("content") == "Site-a has incident INC-17." for message in messages) + assert all(message["role"] != "tool" for message in messages) + assert "tool-data-only" not in json.dumps(messages) + assert len(restarted.requests) == 1 + + +def test_hosted_followup_with_unknown_history_fails_closed(): + fake = _FakeChatTransport([]) + with TestClient(_model_loop_app(_spec(), fake)) as client: + response = client.post("/responses", json={"input": "Continue", "agent_session_id": "session-a", "previous_response_id": "unknown-response"}) + assert response.status_code == 404 + assert response.json()["error"]["code"] == "unknown_previous_response_id" + assert not fake.requests diff --git a/runtimes/common/tests/test_instruction_skills.py b/runtimes/common/tests/test_instruction_skills.py new file mode 100644 index 0000000..0d98b82 --- /dev/null +++ b/runtimes/common/tests/test_instruction_skills.py @@ -0,0 +1,263 @@ +from __future__ import annotations + +import os +from pathlib import Path + +import pytest + +from agentkit_serve_common import skills +from agentkit_serve_common.config import AgentSpec +from agentkit_serve_common.skills import SkillCatalog, SkillConfigurationError + + +def _spec(*, providers=None) -> AgentSpec: + return AgentSpec.model_validate( + { + "abiVersion": "v0", + "metadata": {"name": "instruction-agent"}, + "model": { + "provider": "openai-compatible", + "baseURL": "https://model.example.invalid/v1", + "name": "test-model", + }, + "instructions": "Use the relevant skill.", + "context": { + "providers": providers + if providers is not None + else [ + { + "type": "skills", + "source": "filesystem", + "path": "/agent/skills", + } + ] + }, + "expose": {"openai": True, "port": 8080}, + } + ) + + +@pytest.fixture +def skill_root(tmp_path, monkeypatch): + root = tmp_path.resolve() / "skills" + root.mkdir() + open_directory = skills._open_directory + + # The packaged path is fixed; relocate it for tests without needing a host + # /agent directory. All actual traversal and file-read checks still run. + def open_test_directory(path: Path): + return open_directory(root / path.relative_to("/agent/skills")) + + monkeypatch.setattr(skills, "_open_directory", open_test_directory) + return root + + +def _write_skill(root, name="inspection", *, content=None): + directory = root / name + directory.mkdir(parents=True, exist_ok=True) + content = content or ( + f"---\nname: {name}\ndescription: Inspect the relevant equipment.\n---\n" + "Use the approved lookup tool to obtain the latest inspection record.\n" + ) + (directory / "SKILL.md").write_text(content, encoding="utf-8") + return content + + +def test_catalog_loads_exact_documents_and_never_reopens_files(skill_root): + original = _write_skill(skill_root) + _write_skill( + skill_root, + "parts", + content=( + "---\nname: parts\ndescription: Find parts & specifications.\n---\nRetrieve parts using authorized tools.\n" + ), + ) + (skill_root / "inspection" / "script.py").write_text("raise RuntimeError('must not execute')") + (skill_root / "inspection" / "secret.txt").write_text("not a skill document") + catalog = SkillCatalog.from_spec(_spec()) + + assert catalog.names == ("inspection", "parts") + assert "Find parts & specifications." in catalog.instructions + assert "secret.txt" not in catalog.instructions + assert "script.py" not in catalog.instructions + assert catalog.load_skill("inspection") == original + (skill_root / "inspection" / "SKILL.md").write_text("changed after startup") + (skill_root / "parts" / "SKILL.md").unlink() + assert catalog.load_skill("inspection") == original + assert "Retrieve parts" in catalog.load_skill("parts") + definition = catalog.tool_schema()["function"] + assert definition["name"] == "load_skill" + assert definition["parameters"]["properties"]["skill_name"]["enum"] == ["inspection", "parts"] + assert definition["parameters"]["additionalProperties"] is False + + +def test_nested_and_individual_skill_directories(skill_root): + _write_skill(skill_root / "operations", "inspection") + assert SkillCatalog.from_spec(_spec()).names == ("inspection",) + spec = _spec( + providers=[ + { + "type": "skills", + "source": "filesystem", + "path": "/agent/skills/operations/inspection", + } + ] + ) + assert SkillCatalog.from_spec(spec).names == ("inspection",) + + +@pytest.mark.parametrize("name", ["", "../inspection", "inspection/SKILL.md", None, {}, "unknown", "INSPECTION"]) +def test_model_cannot_request_file_paths_or_unknown_skills(skill_root, name): + _write_skill(skill_root) + catalog = SkillCatalog.from_spec(_spec()) + with pytest.raises(SkillConfigurationError, match="unknown bundled skill"): + catalog.load_skill(name) + + +def test_no_providers_needs_no_skill_directory(monkeypatch): + monkeypatch.setattr(skills, "_open_directory", lambda _: pytest.fail("unexpected filesystem access")) + catalog = SkillCatalog.from_spec(_spec(providers=[])) + assert not catalog + assert catalog.instructions == "" + + +@pytest.mark.parametrize( + "provider", + [ + {"type": "search", "endpointEnv": "SEARCH_ENDPOINT", "indexEnv": "SEARCH_INDEX"}, + {"type": "memory", "endpointEnv": "MEMORY_ENDPOINT", "storeNameEnv": "MEMORY_STORE"}, + {"type": "skills", "source": "filesystem", "path": "/agent/skills", "toolRef": "remote"}, + {"type": "skills", "source": "filesystem", "path": "/agent/skills", "endpointEnv": "REMOTE"}, + ], +) +def test_external_context_providers_remain_prohibited(provider): + with pytest.raises(SkillConfigurationError, match="providers"): + SkillCatalog.from_spec(_spec(providers=[provider])) + + +def test_path_traversal_remains_prohibited(): + spec = _spec( + providers=[ + { + "type": "skills", + "source": "filesystem", + "path": "/agent/skills/../skills", + } + ] + ) + with pytest.raises(SkillConfigurationError, match="under /agent/skills"): + SkillCatalog.from_spec(spec) + + +@pytest.mark.parametrize("target", ["root", "directory", "document"]) +def test_skill_path_symlinks_fail_closed(skill_root, target): + _write_skill(skill_root) + if target == "root": + link = skill_root + elif target == "directory": + link = skill_root / "inspection" + else: + link = skill_root / "inspection" / "SKILL.md" + moved = link.with_name(link.name + "-original") + link.rename(moved) + link.symlink_to(moved, target_is_directory=moved.is_dir()) + with pytest.raises(SkillConfigurationError, match="symlinks"): + SkillCatalog.from_spec(_spec()) + + +@pytest.mark.parametrize("kind", ["hard-link", "fifo", "directory"]) +def test_skill_document_must_be_an_unlinked_regular_file(skill_root, monkeypatch, kind): + _write_skill(skill_root) + path = skill_root / "inspection" / "SKILL.md" + if kind == "hard-link": + os.link(path, skill_root / "another-file") + else: + path.unlink() + if kind == "fifo": + os.mkfifo(path) + else: + path.mkdir() + open_file = os.open + opened = [] + + def open_and_capture(name, flags, **kwargs): + descriptor = open_file(name, flags, **kwargs) + if name == "SKILL.md": + opened.append(descriptor) + return descriptor + + monkeypatch.setattr(skills.os, "open", open_and_capture) + with pytest.raises(SkillConfigurationError): + SkillCatalog.from_spec(_spec()) + assert len(opened) == 1 + with pytest.raises(OSError): + os.fstat(opened[0]) + + +@pytest.mark.parametrize( + "content", + [ + "No frontmatter.", + "---\nname: another\ndescription: Description.\n---\nInstructions.", + "---\nname: inspection\ndescription: 42\n---\nInstructions.", + "---\nname: inspection\ndescription: Description.\n---\n ", + "---\nname: inspection\ndescription: Description.\n---\nInstructions.\x00", + "---\nname: inspection\ndescription: !!python/object:unsafe {}\n---\nInstructions.", + "---\n- inspection\n---\nInstructions.", + "---\nname: inspection\nname: another\ndescription: Description.\n---\nInstructions.", + '---\nname: inspection\ndescription: "\\ud800"\n---\nInstructions.', + ], +) +def test_invalid_skill_documents_fail_startup(skill_root, content): + _write_skill(skill_root, content=content) + with pytest.raises(SkillConfigurationError, match="SKILL.md"): + SkillCatalog.from_spec(_spec()) + + +def test_invalid_utf8_document_fails_startup(skill_root): + _write_skill(skill_root) + (skill_root / "inspection" / "SKILL.md").write_bytes(b"\xff") + with pytest.raises(SkillConfigurationError, match="UTF-8"): + SkillCatalog.from_spec(_spec()) + + +@pytest.mark.parametrize("limit", ["_MAX_SKILL_BYTES", "_MAX_CATALOG_BYTES", "_MAX_SKILLS", "_MAX_DIRECTORY_ENTRIES"]) +def test_skill_catalog_limits_are_enforced(skill_root, monkeypatch, limit): + _write_skill(skill_root) + _write_skill(skill_root, "parts") + monkeypatch.setattr(skills, limit, 1) + with pytest.raises(SkillConfigurationError, match="size limit|too large"): + SkillCatalog.from_spec(_spec()) + + +def test_duplicate_skill_names_fail_startup(skill_root): + _write_skill(skill_root / "first") + _write_skill(skill_root / "second") + with pytest.raises(SkillConfigurationError, match="duplicate"): + SkillCatalog.from_spec(_spec()) + + +def test_empty_or_missing_skill_directory_fails_startup(skill_root): + with pytest.raises(SkillConfigurationError, match="no SKILL.md"): + SkillCatalog.from_spec(_spec()) + skill_root.rmdir() + with pytest.raises(SkillConfigurationError, match="readable image directories"): + SkillCatalog.from_spec(_spec()) + + +def test_directory_swap_cannot_redirect_the_open_skill_read(skill_root, monkeypatch): + original = _write_skill(skill_root) + outside = skill_root.parent / "replacement" + _write_skill( + outside, content=("---\nname: inspection\ndescription: Replacement.\n---\nReplacement instructions.\n") + ) + read_document = skills._read_document + + def swap_before_read(directory_fd, directory_name): + directory = skill_root / "inspection" + directory.rename(skill_root / "original") + directory.symlink_to(outside / "inspection", target_is_directory=True) + return read_document(directory_fd, directory_name) + + monkeypatch.setattr(skills, "_read_document", swap_before_read) + assert SkillCatalog.from_spec(_spec()).load_skill("inspection") == original diff --git a/runtimes/microsoft-agent-framework/agentkit_serve/agent_factory.py b/runtimes/microsoft-agent-framework/agentkit_serve/agent_factory.py index 5518786..050e7b3 100644 --- a/runtimes/microsoft-agent-framework/agentkit_serve/agent_factory.py +++ b/runtimes/microsoft-agent-framework/agentkit_serve/agent_factory.py @@ -29,6 +29,7 @@ FileSkillsSource, FunctionInvocationContext, FunctionMiddleware, + FunctionTool, MCPSkillsSource, MCPStdioTool, MCPStreamableHTTPTool, @@ -412,11 +413,24 @@ def build_agent( client=None, ) -> Agent: """Assemble the MAF agent: client + system prompt + tools + context.""" + instructions = spec.instructions + tools = [build_tool(t, stack=stack) for t in spec.tools] + if spec._packaged_skill_catalog: + skills = spec._packaged_skill_catalog + instructions += "\n\n" + skills.instructions + definition = skills.tool_schema()["function"] + tools.append(FunctionTool( + name=definition["name"], + description=definition["description"], + input_model=definition["parameters"], + func=skills.load_skill, + approval_mode="never_require", + )) return Agent( client=client if client is not None else build_client(spec), - instructions=spec.instructions, + instructions=instructions, name=spec.metadata.name, - tools=[build_tool(t, stack=stack) for t in spec.tools], + tools=tools, context_providers=context_providers, middleware=[_ModelMessageMiddleware(), _MCPFailureMiddleware()], ) @@ -611,7 +625,8 @@ async def _build_context_providers(self): providers.append(await self._build_search_provider(provider)) elif provider.type == _CONTEXT_TYPE_SKILLS: if provider.source == _CONTEXT_SOURCE_FILESYSTEM: - providers.append(SkillsProvider(FileSkillsSource(provider.path))) + if self.spec._packaged_skill_catalog is None: + providers.append(SkillsProvider(FileSkillsSource(provider.path))) elif provider.source == _CONTEXT_SOURCE_MCP: providers.append(await self._build_mcp_skills_provider(provider)) elif provider.type == _CONTEXT_TYPE_MEMORY: @@ -721,6 +736,10 @@ def supports_acp_http_mcp() -> bool: return True +def supports_acp_packaged_skills() -> bool: + return True + + def build_runtime(spec: AgentSpec) -> RuntimeSession: """Build the runtime session consumed by the shared server.""" if offline_orka_echo_enabled(): diff --git a/runtimes/microsoft-agent-framework/tests/test_instruction_skills.py b/runtimes/microsoft-agent-framework/tests/test_instruction_skills.py new file mode 100644 index 0000000..01017e4 --- /dev/null +++ b/runtimes/microsoft-agent-framework/tests/test_instruction_skills.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +import asyncio + +from agent_framework import BaseChatClient, ChatResponse, Content, FunctionInvocationLayer, Message + +from agentkit_serve import agent_factory +from agentkit_serve_common.config import AgentSpec +from agentkit_serve_common.conversation import RunRequest +from agentkit_serve_common.skills import _Skill, SkillCatalog + + +def test_governed_skills_use_snapshot_and_expose_only_load_skill(monkeypatch): + content = "---\nname: inspection\ndescription: Inspect equipment.\n---\nUse the authorized lookup tool.\n" + spec = AgentSpec.model_validate( + { + "abiVersion": "v0", + "metadata": {"name": "inspection-agent"}, + "model": { + "provider": "openai-compatible", + "baseURL": "http://127.0.0.1:43123/v1", + "name": "test-model", + }, + "instructions": "Consult the relevant skill before answering.", + "context": { + "providers": [ + { + "type": "skills", + "source": "filesystem", + "path": "/agent/skills", + } + ] + }, + "expose": {"openai": True, "port": 8080}, + } + ) + spec._packaged_skill_catalog = SkillCatalog((_Skill("inspection", "Inspect equipment.", content),)) + + class SkillClient(FunctionInvocationLayer, BaseChatClient): + def __init__(self): + super().__init__() + self.requests = 0 + + async def _inner_get_response(self, *, messages, stream, options, **kwargs): + self.requests += 1 + assert [tool.name for tool in options.get("tools", [])] == ["load_skill"] + if self.requests == 1: + assert "inspection" in options["instructions"] + response = Content.from_function_call( + call_id="skill-call", name="load_skill", arguments={"skill_name": "inspection"} + ) + else: + assert self.requests == 2 + assert any( + item.type == "function_result" and content in str(item.result) + for message in messages + for item in message.contents + ) + response = Content.from_text("The inspection instructions are loaded.") + return ChatResponse(messages=[Message(role="assistant", contents=[response])]) + + def no_native_provider(*args, **kwargs): + raise AssertionError("governed skills must not use the filesystem or native SkillsProvider") + + monkeypatch.setattr(agent_factory, "FileSkillsSource", no_native_provider) + monkeypatch.setattr(agent_factory, "SkillsProvider", no_native_provider) + client = SkillClient() + monkeypatch.setattr(agent_factory, "build_client", lambda _: client) + + async def exercise(): + events = [] + + async def observe(event): + events.append(event) + + async with agent_factory.MAFRuntime(spec) as runtime: + result = await runtime.run(RunRequest("Inspect the equipment.", on_tool_event=observe)) + assert result.text == "The inspection instructions are loaded." + assert [(event.tool_name, event.status) for event in events] == [ + ("load_skill", "in_progress"), + ("load_skill", "completed"), + ] + + asyncio.run(exercise()) From c2e9cd53d1a592ab4631e53ae8c707e1fc1039fc Mon Sep 17 00:00:00 2001 From: Sertac Ozercan Date: Fri, 11 Sep 2026 13:18:04 -0700 Subject: [PATCH 2/9] fix(foundry): reject expired session follow-ups without identity Signed-off-by: Sertac Ozercan --- .../common/agentkit_serve_common/foundry.py | 4 +++- .../tests/test_foundry_tool_workflows.py | 18 ++++++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/runtimes/common/agentkit_serve_common/foundry.py b/runtimes/common/agentkit_serve_common/foundry.py index ba7e16c..428c894 100644 --- a/runtimes/common/agentkit_serve_common/foundry.py +++ b/runtimes/common/agentkit_serve_common/foundry.py @@ -2837,7 +2837,9 @@ async def responses(request: Request): logger.warning("failed to access Foundry brokered response state: %s", exc) return _state_storage_error() except _StateExpired as exc: - if exc.state.status in {"pending", "resuming"} or (model_loop is not None and session_id): + if exc.state.status in {"pending", "resuming"} or ( + model_loop is not None and (session_id or exc.state.session_id) + ): return _error("previous_response_id state has expired", status=410, code="response_state_expired") previous_state = None except KeyError: diff --git a/runtimes/common/tests/test_foundry_tool_workflows.py b/runtimes/common/tests/test_foundry_tool_workflows.py index fc12876..8002a49 100644 --- a/runtimes/common/tests/test_foundry_tool_workflows.py +++ b/runtimes/common/tests/test_foundry_tool_workflows.py @@ -310,3 +310,21 @@ def test_hosted_followup_with_unknown_history_fails_closed(): assert response.status_code == 404 assert response.json()["error"]["code"] == "unknown_previous_response_id" assert not fake.requests + + +@pytest.mark.parametrize("session_id", [None, "session-a", "session-b"]) +def test_expired_hosted_history_rejects_followup_even_without_session_id(session_id): + fake = _FakeChatTransport([ + _chat_response({"role": "assistant", "content": "Your site is site-a."}), + _chat_response({"role": "assistant", "content": "Must not run without the expired context."}), + ]) + with TestClient(_model_loop_app(_spec(), fake, state_ttl_seconds=0)) as client: + first = client.post("/responses", json={"input": "My site is site-a", "agent_session_id": "session-a"}) + assert first.status_code == 200, first.text + payload = {"input": "Which site was that?", "previous_response_id": first.json()["id"]} + if session_id is not None: + payload["agent_session_id"] = session_id + followup = client.post("/responses", json=payload) + assert followup.status_code == 410, followup.text + assert followup.json()["error"]["code"] == "response_state_expired" + assert len(fake.requests) == 1 From eaf549ff9b8040a52c8d6c91c2121de44fffea3e Mon Sep 17 00:00:00 2001 From: Sertac Ozercan Date: Fri, 11 Sep 2026 14:50:50 -0700 Subject: [PATCH 3/9] fix(foundry): acknowledge brokered streams before model work Signed-off-by: Sertac Ozercan --- docs/foundry-hosted-brokered.md | 15 +- .../common/agentkit_serve_common/foundry.py | 44 +- .../foundry_streaming.py | 162 ++++++++ .../common/tests/test_foundry_streaming.py | 388 ++++++++++++++++++ 4 files changed, 600 insertions(+), 9 deletions(-) create mode 100644 runtimes/common/agentkit_serve_common/foundry_streaming.py create mode 100644 runtimes/common/tests/test_foundry_streaming.py diff --git a/docs/foundry-hosted-brokered.md b/docs/foundry-hosted-brokered.md index a8216a1..5d4feb1 100644 --- a/docs/foundry-hosted-brokered.md +++ b/docs/foundry-hosted-brokered.md @@ -234,9 +234,18 @@ operations. ## Streaming -The current route is non-streaming. If clients send `stream: true`, AgentKit -returns the same normal JSON response rather than SSE. This keeps azd/direct curl -smokes deterministic while making streaming support an explicit future step. +In brokered mode, `/responses` honors `stream: true` with server-sent events. +AgentKit sends `response.created` before starting model work, then sends +`response.completed` or `response.failed` with the same response ID. These events +include the effective `agent_session_id` when one is present. This lets Orka +record which hosted response it accepted before waiting for the model. + +The stream contains acknowledgement and completion events, without token deltas. +Validation errors before acknowledgement keep their normal HTTP error status. +Disconnecting the stream cancels the active model request and releases its +in-progress state. Orka still uses authenticated session stop and idle checks to +confirm hosted cleanup. Requests without `stream: true` keep the JSON response +format. ## Troubleshooting diff --git a/runtimes/common/agentkit_serve_common/foundry.py b/runtimes/common/agentkit_serve_common/foundry.py index 428c894..8869d68 100644 --- a/runtimes/common/agentkit_serve_common/foundry.py +++ b/runtimes/common/agentkit_serve_common/foundry.py @@ -40,6 +40,7 @@ from .brokered import brokered_tool_definitions from .config import AgentSpec, _unsafe_brokered_key, _unsafe_brokered_text from .foundry_model_loop import BrokeredChatModelLoop, ModelLoopFinal, ModelLoopToolRequest +from .foundry_streaming import BrokeredResponseStream, brokered_stream_response from .conversation import FORWARDED_ROLES, ConversationTurn, RunRequest from .runtime import AgentRunError, BrokeredToolDefinition, RunResult, RuntimeFactory from .server import make_auth_dependency @@ -364,8 +365,14 @@ def _responses_input_to_run_request(value: Any, *, session_id: str | None) -> Ru return RunRequest(prompt=json.dumps(value, separators=(",", ":"), sort_keys=True), session_id=session_id) -def _responses_payload(spec: AgentSpec, result: RunResult, *, previous_response_id: str | None = None) -> dict[str, Any]: - response_id = _new_response_id(previous_response_id) +def _responses_payload( + spec: AgentSpec, + result: RunResult, + *, + previous_response_id: str | None = None, + response_id: str | None = None, +) -> dict[str, Any]: + response_id = response_id or _new_response_id(previous_response_id) message_id = _new_message_id(response_id) payload: dict[str, Any] = { "id": response_id, @@ -2355,6 +2362,7 @@ async def _handle_brokered_continuation( session_id: str | None, max_output_bytes: int, model_loop: BrokeredChatModelLoop | None = None, + stream: BrokeredResponseStream | None = None, ) -> JSONResponse: if not continuation_proof: return _error( @@ -2546,6 +2554,8 @@ async def _handle_brokered_continuation( try: next_response_id = _new_response_id(str(previous_response_id)) next_call_id = f"call_{next_response_id}_1" + if stream is not None: + await stream.accept(next_response_id) model_result = await model_loop.resume(state.model_messages, call_id=call_id, output=output_json, next_call_id=next_call_id) except AgentRunError as exc: if not _reset_unfinalized_continuation(store, state, call_id=call_id): @@ -2586,7 +2596,12 @@ async def _handle_brokered_continuation( resume_initial_usage = dict(state.initial_usage) has_resume_transcript = resume_model_messages is not None used_model_resume = has_resume_transcript and model_loop is not None - final_payload = _responses_payload(spec, result, previous_response_id=str(previous_response_id)) + final_payload = _responses_payload( + spec, + result, + previous_response_id=str(previous_response_id), + response_id=stream.created.result()["id"] if stream is not None and stream.created.done() else None, + ) state.accepted_output_digests[call_id] = output_digest state.accepted_output_sizes[call_id] = accepted_output_size state.status = "completed" @@ -2775,11 +2790,22 @@ async def responses(request: Request): except (UnicodeDecodeError, RecursionError, ValueError): return _error("Request body must be JSON", status=400, code="invalid_json") + if brokered_tools and isinstance(data, dict) and data.get("stream") is True: + return await brokered_stream_response( + spec.model.name, + lambda stream: execute_responses(request, data, request_body_size, stream), + ) + return await execute_responses(request, data, request_body_size) + + async def execute_responses( + request: Request, + data: Any, + request_body_size: int, + stream: BrokeredResponseStream | None = None, + ) -> JSONResponse: if not isinstance(data, dict): return _error("Request body must be a JSON object", status=400, code="invalid_request") - # Foundry/azd clients may include stream=true by default. The adapter is - # intentionally non-streaming, so tolerate the flag and return a normal - # completed response instead of failing readiness/e2e checks. + # Non-brokered adapters retain their existing buffered response behavior. if data.get("tools"): return _error( "request-supplied Responses tools are not allowed; hosted brokered mode uses static safe schemas", @@ -2808,6 +2834,8 @@ async def responses(request: Request): status=409, code="response_session_mismatch", ) + if stream is not None: + stream.session_id = session_id previous_response_id = data.get("previous_response_id") function_outputs = _function_call_outputs_from_input(data["input"]) if brokered_tools and not function_outputs and request_body_size > request_body_limit: @@ -2828,6 +2856,7 @@ async def responses(request: Request): session_id=session_id, max_output_bytes=max_output_bytes, model_loop=model_loop, + stream=stream, ) previous_state = None if brokered_tools and isinstance(previous_response_id, str) and previous_response_id: @@ -2901,6 +2930,8 @@ async def responses(request: Request): return _state_storage_error() try: try: + if stream is not None: + await stream.accept(response_id) model_result = await model_loop.start(run_request, call_id=call_id) except AgentRunError as exc: if exc.code == "ModelResponseTooLarge": @@ -2911,6 +2942,7 @@ async def responses(request: Request): spec, RunResult(text=model_result.text, usage=model_result.usage), previous_response_id=previous_response_id_for_output, + response_id=response_id if stream is not None else None, ) if run_request.session_id: completed = _HostedResponseState( diff --git a/runtimes/common/agentkit_serve_common/foundry_streaming.py b/runtimes/common/agentkit_serve_common/foundry_streaming.py new file mode 100644 index 0000000..c6e34d5 --- /dev/null +++ b/runtimes/common/agentkit_serve_common/foundry_streaming.py @@ -0,0 +1,162 @@ +"""Early Responses acknowledgements for the hosted brokered model loop.""" + +from __future__ import annotations + +import asyncio +import json +import time +from collections.abc import Awaitable, Callable +from typing import Any + +from fastapi.responses import JSONResponse, Response + + +class BrokeredResponseStream: + def __init__(self, model: str) -> None: + self.model = model + self.session_id: str | None = None + self.created: asyncio.Future[dict[str, Any]] = ( + asyncio.get_running_loop().create_future() + ) + self.delivered = asyncio.Event() + + def prepare(self, response_id: str) -> None: + if self.created.done(): + raise RuntimeError("hosted response was already acknowledged") + payload: dict[str, Any] = { + "id": response_id, + "object": "response", + "created_at": int(time.time()), + "status": "in_progress", + "model": self.model, + "output": [], + } + if self.session_id: + payload["agent_session_id"] = self.session_id + self.created.set_result(payload) + + async def accept(self, response_id: str) -> None: + self.prepare(response_id) + # Model work starts only after the response.created frame has been sent. + await self.delivered.wait() + + def terminal(self, result: JSONResponse | None) -> dict[str, Any]: + created = self.created.result() + payload = json.loads(result.body) if result is not None else {} + if ( + result is not None + and result.status_code < 400 + and payload.get("id") == created["id"] + ): + if self.session_id: + payload["agent_session_id"] = self.session_id + return {"type": "response.completed", "response": payload} + error = ( + payload.get("error") + if result is not None and result.status_code >= 400 + else None + ) + if not isinstance(error, dict): + error = {"code": "HostedResponseError", "message": "hosted response failed"} + return { + "type": "response.failed", + "response": {**created, "status": "failed", "error": error}, + } + + +def _frame(event: dict[str, Any]) -> bytes: + data = json.dumps(event, separators=(",", ":"), ensure_ascii=True) + return f"event: {event['type']}\ndata: {data}\n\n".encode("utf-8") + + +class _BrokeredStreamingResponse(Response): + media_type = "text/event-stream" + + def __init__( + self, stream: BrokeredResponseStream, operation: asyncio.Task[JSONResponse] + ) -> None: + super().__init__( + headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"} + ) + self.raw_headers = [ + (key, value) for key, value in self.raw_headers if key != b"content-length" + ] + self.stream = stream + self.operation = operation + + async def __call__(self, scope: Any, receive: Any, send: Any) -> None: + async def write_response() -> None: + await send( + { + "type": "http.response.start", + "status": self.status_code, + "headers": self.raw_headers, + } + ) + await send( + { + "type": "http.response.body", + "body": _frame( + { + "type": "response.created", + "response": self.stream.created.result(), + } + ), + "more_body": True, + } + ) + self.stream.delivered.set() + try: + result = await self.operation + except Exception: # Model failures must not leak exception text into SSE. + result = None + await send( + { + "type": "http.response.body", + "body": _frame(self.stream.terminal(result)), + "more_body": False, + } + ) + + async def wait_for_disconnect() -> None: + while True: + if (await receive())["type"] == "http.disconnect": + return + + writer = asyncio.create_task(write_response()) + disconnected = asyncio.create_task(wait_for_disconnect()) + try: + done, _ = await asyncio.wait( + {writer, disconnected}, return_when=asyncio.FIRST_COMPLETED + ) + if writer in done: + await writer + finally: + for task in (writer, disconnected, self.operation): + if not task.done(): + task.cancel() + await asyncio.gather( + writer, disconnected, self.operation, return_exceptions=True + ) + + +async def brokered_stream_response( + model: str, + operation: Callable[[BrokeredResponseStream], Awaitable[JSONResponse]], +) -> Response: + stream = BrokeredResponseStream(model) + task = asyncio.create_task(operation(stream)) + try: + await asyncio.wait({task, stream.created}, return_when=asyncio.FIRST_COMPLETED) + if not stream.created.done(): + # Rejected requests keep their HTTP error status. Deterministic tool + # responses and completed replays need no model work or early wait. + result = await task + if result.status_code >= 400: + return result + stream.prepare(json.loads(result.body)["id"]) + return _BrokeredStreamingResponse(stream, task) + except BaseException: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + raise diff --git a/runtimes/common/tests/test_foundry_streaming.py b/runtimes/common/tests/test_foundry_streaming.py new file mode 100644 index 0000000..08a26ac --- /dev/null +++ b/runtimes/common/tests/test_foundry_streaming.py @@ -0,0 +1,388 @@ +"""Hosted Responses acknowledgements, identity, and disconnect regressions.""" + +from __future__ import annotations + +import asyncio +import json +from contextlib import asynccontextmanager + +import httpx +import pytest + +from test_foundry_brokered_protocol import ( + CONTINUATION_AUTH, + _app, + _call, + _chat_response, + _continuation, + _spec, +) + + +def _tool_response(): + return _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "model-call", + "type": "function", + "function": { + "name": "conformance_read", + "arguments": '{"probe":true}', + }, + } + ], + } + ) + + +class HeldModel(httpx.AsyncBaseTransport): + def __init__(self, result="text", *, continuation=False): + self.result = result + self.held_call = 2 if continuation else 1 + self.calls = 0 + self.started = asyncio.Event() + self.release = asyncio.Event() + self.cancelled = asyncio.Event() + + async def handle_async_request(self, request): + self.calls += 1 + if self.calls < self.held_call: + return httpx.Response(200, request=request, json=_tool_response()) + if self.calls == self.held_call: + self.started.set() + try: + await self.release.wait() + except asyncio.CancelledError: + self.cancelled.set() + raise + if self.result == "error": + return httpx.Response( + 401, request=request, json={"error": "private-upstream-detail"} + ) + payload = ( + _tool_response() + if self.result == "tool" + else _chat_response({"role": "assistant", "content": "Verified response."}) + ) + return httpx.Response(200, request=request, json=payload) + + +@asynccontextmanager +async def _exchange(app, payload, *, hold_created=None): + incoming, outgoing = asyncio.Queue(), asyncio.Queue() + await incoming.put( + { + "type": "http.request", + "body": json.dumps(payload).encode(), + "more_body": False, + } + ) + + async def send(message): + await outgoing.put(message) + if ( + hold_created is not None + and message["type"] == "http.response.body" + and message.get("more_body") + ): + await hold_created.wait() + + scope = { + "type": "http", + "asgi": {"version": "3.0", "spec_version": "2.4"}, + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": "/responses", + "raw_path": b"/responses", + "query_string": b"", + "root_path": "", + "headers": [ + (b"content-type", b"application/json"), + *[ + (key.encode(), value.encode()) + for key, value in CONTINUATION_AUTH.items() + ], + ], + "client": ("127.0.0.1", 1234), + "server": ("testserver", 80), + } + task = asyncio.create_task(app(scope, incoming.get, send)) + try: + yield incoming, outgoing, task + finally: + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + +async def _created(outgoing): + headers = await asyncio.wait_for(outgoing.get(), 2) + assert headers["type"] == "http.response.start" and headers["status"] == 200 + assert dict(headers["headers"])[b"content-type"].startswith(b"text/event-stream") + assert b"content-length" not in dict(headers["headers"]) + return await _event(outgoing) + + +async def _event(outgoing): + message = await asyncio.wait_for(outgoing.get(), 2) + assert message["type"] == "http.response.body" + assert message["body"].endswith(b"\n\n") + data = next( + line.removeprefix(b"data: ") + for line in message["body"].splitlines() + if line.startswith(b"data: ") + ) + return json.loads(data) + + +@pytest.mark.parametrize("continuation", [False, True]) +@pytest.mark.parametrize("result", ["text", "tool"]) +def test_brokered_stream_ack_precedes_model_and_matches_completion_and_replay( + continuation, result +): + async def exercise(): + model = HeldModel(result, continuation=continuation) + async with httpx.AsyncClient(transport=model) as upstream: + app = _app( + _spec(), + brokered_model_loop_enabled=True, + brokered_model_http_client=upstream, + ) + payload = { + "input": "Read the synthetic data", + "agent_session_id": "session-a", + "stream": True, + } + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://testserver" + ) as client: + if continuation: + initial = ( + await client.post( + "/responses", + json={ + "input": "Read data", + "agent_session_id": "session-a", + }, + ) + ).json() + payload = { + **_continuation( + initial["id"], + _call(initial)["call_id"], + {"approved": True, "output": {"ok": True}}, + ), + "agent_session_id": "session-a", + "stream": True, + } + delivered = asyncio.Event() + async with _exchange(app, payload, hold_created=delivered) as ( + _, + outgoing, + request, + ): + created = await _created(outgoing) + assert created["type"] == "response.created" + assert created["response"]["status"] == "in_progress" + assert created["response"]["agent_session_id"] == "session-a" + assert created["response"]["output"] == [] + assert not model.started.is_set() + delivered.set() + await asyncio.wait_for(model.started.wait(), 2) + assert outgoing.empty() and not request.done() + model.release.set() + completed = await _event(outgoing) + await asyncio.wait_for(request, 2) + assert completed["type"] == "response.completed" + response = completed["response"] + assert response["id"] == created["response"]["id"] + assert response["agent_session_id"] == "session-a" + assert response["output"][0]["response_id"] == response["id"] + assert response["output"][0]["type"] == ( + "function_call" if result == "tool" else "message" + ) + if continuation: + async with _exchange(app, payload) as (_, outgoing, replay_request): + replay_created = await _created(outgoing) + replay_completed = await _event(outgoing) + await asyncio.wait_for(replay_request, 2) + assert replay_created["response"]["id"] == response["id"] + assert replay_completed == completed + buffered = await client.post( + "/responses", + json={**payload, "stream": False}, + headers=CONTINUATION_AUTH, + ) + assert buffered.status_code == 200 + assert buffered.json() == { + key: value + for key, value in response.items() + if key != "agent_session_id" + } + assert model.calls == 2 + + asyncio.run(exercise()) + + +@pytest.mark.parametrize("continuation", [False, True]) +def test_brokered_stream_error_retains_acknowledged_identity(continuation): + async def exercise(): + model = HeldModel("error", continuation=continuation) + async with httpx.AsyncClient(transport=model) as upstream: + app = _app( + _spec(), + brokered_model_loop_enabled=True, + brokered_model_http_client=upstream, + ) + payload = { + "input": "Read data", + "agent_session_id": "session-a", + "stream": True, + } + if continuation: + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://testserver" + ) as client: + initial = ( + await client.post( + "/responses", + json={ + "input": "Read data", + "agent_session_id": "session-a", + }, + ) + ).json() + payload = { + **_continuation( + initial["id"], + _call(initial)["call_id"], + {"approved": True, "output": {"ok": True}}, + ), + "agent_session_id": "session-a", + "stream": True, + } + async with _exchange(app, payload) as (_, outgoing, request): + created = await _created(outgoing) + await asyncio.wait_for(model.started.wait(), 2) + model.release.set() + failed = await _event(outgoing) + await asyncio.wait_for(request, 2) + assert failed["type"] == "response.failed" and "error" not in failed + assert failed["response"]["id"] == created["response"]["id"] + assert failed["response"]["agent_session_id"] == "session-a" + assert failed["response"]["status"] == "failed" + assert failed["response"]["error"]["code"] + assert "private-upstream-detail" not in json.dumps(failed) + + asyncio.run(exercise()) + + +@pytest.mark.parametrize("continuation", [False, True]) +def test_brokered_stream_disconnect_cancels_model_and_releases_initial_or_resume_state( + continuation, +): + async def exercise(): + model = HeldModel(continuation=continuation) + async with httpx.AsyncClient(transport=model) as upstream: + app = _app( + _spec(), + brokered_model_loop_enabled=True, + brokered_model_http_client=upstream, + max_pending_responses=1, + ) + payload = {"input": "Read data", "stream": True} + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://testserver" + ) as client: + if continuation: + initial = ( + await client.post("/responses", json={"input": "Read data"}) + ).json() + payload = { + **_continuation( + initial["id"], + _call(initial)["call_id"], + {"approved": True, "output": {"ok": True}}, + ), + "stream": True, + } + async with _exchange(app, payload) as (incoming, outgoing, request): + await _created(outgoing) + await asyncio.wait_for(model.started.wait(), 2) + await incoming.put({"type": "http.disconnect"}) + await asyncio.wait_for(request, 2) + assert model.cancelled.is_set() + assert outgoing.empty() + retry = await client.post( + "/responses", + json={**payload, "stream": False}, + headers=CONTINUATION_AUTH, + ) + assert retry.status_code == 200, retry.text + assert ( + retry.json()["output"][0]["content"][0]["text"] + == "Verified response." + ) + assert model.calls == (3 if continuation else 2) + + asyncio.run(exercise()) + + +def test_brokered_stream_disconnect_before_created_is_delivered_starts_no_model_work(): + async def exercise(): + model = HeldModel() + async with httpx.AsyncClient(transport=model) as upstream: + app = _app( + _spec(), + brokered_model_loop_enabled=True, + brokered_model_http_client=upstream, + max_pending_responses=1, + ) + async with _exchange( + app, + {"input": "Read data", "stream": True}, + hold_created=asyncio.Event(), + ) as (incoming, outgoing, request): + await _created(outgoing) + await incoming.put({"type": "http.disconnect"}) + await asyncio.wait_for(request, 2) + assert model.calls == 0 + model.release.set() + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://testserver" + ) as client: + retry = await client.post("/responses", json={"input": "Read data"}) + assert retry.status_code == 200, retry.text + + asyncio.run(exercise()) + + +def test_brokered_stream_validation_failure_keeps_http_error_without_ack(): + async def exercise(): + model = HeldModel() + async with httpx.AsyncClient(transport=model) as upstream: + app = _app( + _spec(), + brokered_model_loop_enabled=True, + brokered_model_http_client=upstream, + ) + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://testserver" + ) as client: + response = await client.post( + "/responses", + json={ + "input": "Read data", + "stream": True, + "tools": [{"type": "function"}], + }, + ) + assert response.status_code == 400 + assert response.json()["error"]["code"] == "tools_unsupported" + assert model.calls == 0 + + asyncio.run(exercise()) From 0e305cd092dbb4e3682d1a0bdd6d4922040f1217 Mon Sep 17 00:00:00 2001 From: Sertac Ozercan Date: Fri, 11 Sep 2026 15:03:08 -0700 Subject: [PATCH 4/9] fix(foundry): measure message limits with compact JSON Signed-off-by: Sertac Ozercan --- .../foundry_model_loop.py | 3 ++- .../tests/test_foundry_brokered_protocol.py | 19 +++++++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/runtimes/common/agentkit_serve_common/foundry_model_loop.py b/runtimes/common/agentkit_serve_common/foundry_model_loop.py index be8a778..3b6ad82 100644 --- a/runtimes/common/agentkit_serve_common/foundry_model_loop.py +++ b/runtimes/common/agentkit_serve_common/foundry_model_loop.py @@ -78,7 +78,8 @@ async def start(self, request: RunRequest, *, call_id: str) -> ModelLoopFinal | async def _advance(self, messages: list[dict[str, Any]], *, call_id: str) -> ModelLoopFinal | ModelLoopToolRequest: usage: dict[str, int] = {} while True: - if len(json.dumps(messages, ensure_ascii=True).encode("utf-8")) > self.max_messages_bytes: + encoded_messages = json.dumps(messages, ensure_ascii=True, separators=(",", ":")).encode("utf-8") + if len(encoded_messages) > self.max_messages_bytes: raise AgentRunError("model loop messages are too large", status=413, code="brokered_model_messages_too_large") result = await self._step(messages, call_id=call_id) for key, value in result.usage.items(): diff --git a/runtimes/common/tests/test_foundry_brokered_protocol.py b/runtimes/common/tests/test_foundry_brokered_protocol.py index 2b1036c..438fc07 100644 --- a/runtimes/common/tests/test_foundry_brokered_protocol.py +++ b/runtimes/common/tests/test_foundry_brokered_protocol.py @@ -4820,6 +4820,25 @@ def test_foundry_brokered_model_message_size_matches_persistence_encoding(): assert measured == persisted assert len(measured) > len(utf8_compact) + +def test_foundry_brokered_model_loop_accepts_exact_compact_message_limit(): + spec = _spec().model_copy(update={"instructions": ""}) + fake = _FakeChatTransport([ + _chat_response({"role": "assistant", "content": "Within the limit."}), + ]) + app = _model_loop_app(spec, fake, max_model_messages_bytes=512) + # A compact user message record adds 30 bytes around its ASCII content. + with TestClient(app) as client: + accepted = client.post("/responses", json={"input": "x" * 482}) + rejected = client.post("/responses", json={"input": "x" * 483}) + + assert accepted.status_code == 200, accepted.text + assert _message_text(accepted.json()) == "Within the limit." + assert rejected.status_code == 413 + assert rejected.json()["error"]["code"] == "brokered_model_messages_too_large" + assert len(fake.requests) == 1 + + def test_foundry_brokered_bounds_persisted_model_messages_and_releases_reservation(): tool_response = _chat_response( { From 604cd83118bb3d891a2243478f5e208381b145cb Mon Sep 17 00:00:00 2001 From: Sertac Ozercan Date: Fri, 11 Sep 2026 15:18:46 -0700 Subject: [PATCH 5/9] fix(foundry): acknowledge native responses before model work Signed-off-by: Sertac Ozercan --- docs/foundry-hosted-brokered.md | 12 +- .../common/agentkit_serve_common/foundry.py | 8 +- .../foundry_streaming.py | 2 +- .../common/tests/test_foundry_protocol.py | 5 +- .../common/tests/test_foundry_streaming.py | 255 +++++++++++++++++- 5 files changed, 268 insertions(+), 14 deletions(-) diff --git a/docs/foundry-hosted-brokered.md b/docs/foundry-hosted-brokered.md index 5d4feb1..d47102e 100644 --- a/docs/foundry-hosted-brokered.md +++ b/docs/foundry-hosted-brokered.md @@ -234,7 +234,8 @@ operations. ## Streaming -In brokered mode, `/responses` honors `stream: true` with server-sent events. +`/responses` honors `stream: true` with server-sent events, both with and without +`brokeredTools`. AgentKit sends `response.created` before starting model work, then sends `response.completed` or `response.failed` with the same response ID. These events include the effective `agent_session_id` when one is present. This lets Orka @@ -242,10 +243,11 @@ record which hosted response it accepted before waiting for the model. The stream contains acknowledgement and completion events, without token deltas. Validation errors before acknowledgement keep their normal HTTP error status. -Disconnecting the stream cancels the active model request and releases its -in-progress state. Orka still uses authenticated session stop and idle checks to -confirm hosted cleanup. Requests without `stream: true` keep the JSON response -format. +Disconnecting the stream cancels the active runtime or model request and releases +any in-progress brokered state. Orka still uses authenticated session stop and idle +checks to confirm hosted cleanup. Requests without `stream: true` keep the JSON +response format. Streaming does not add persistence or replay behavior to runtimes +without `brokeredTools`. ## Troubleshooting diff --git a/runtimes/common/agentkit_serve_common/foundry.py b/runtimes/common/agentkit_serve_common/foundry.py index 8869d68..0597569 100644 --- a/runtimes/common/agentkit_serve_common/foundry.py +++ b/runtimes/common/agentkit_serve_common/foundry.py @@ -2790,7 +2790,7 @@ async def responses(request: Request): except (UnicodeDecodeError, RecursionError, ValueError): return _error("Request body must be JSON", status=400, code="invalid_json") - if brokered_tools and isinstance(data, dict) and data.get("stream") is True: + if isinstance(data, dict) and data.get("stream") is True: return await brokered_stream_response( spec.model.name, lambda stream: execute_responses(request, data, request_body_size, stream), @@ -2805,7 +2805,6 @@ async def execute_responses( ) -> JSONResponse: if not isinstance(data, dict): return _error("Request body must be a JSON object", status=400, code="invalid_request") - # Non-brokered adapters retain their existing buffered response behavior. if data.get("tools"): return _error( "request-supplied Responses tools are not allowed; hosted brokered mode uses static safe schemas", @@ -3086,13 +3085,16 @@ async def execute_responses( ) ) + response_id = _new_response_id() if stream is not None else None try: + if stream is not None: + await stream.accept(response_id) result = await request.app.state.runtime.run(run_request) except AgentRunError as exc: return _non_brokered_agent_run_error(exc) except Exception: # noqa: BLE001 - deterministic protocol envelope. return _non_brokered_unexpected_runtime_error() - return JSONResponse(_responses_payload(spec, result)) + return JSONResponse(_responses_payload(spec, result, response_id=response_id)) return app diff --git a/runtimes/common/agentkit_serve_common/foundry_streaming.py b/runtimes/common/agentkit_serve_common/foundry_streaming.py index c6e34d5..27df09d 100644 --- a/runtimes/common/agentkit_serve_common/foundry_streaming.py +++ b/runtimes/common/agentkit_serve_common/foundry_streaming.py @@ -1,4 +1,4 @@ -"""Early Responses acknowledgements for the hosted brokered model loop.""" +"""Early Responses acknowledgements for hosted AgentKit runtimes.""" from __future__ import annotations diff --git a/runtimes/common/tests/test_foundry_protocol.py b/runtimes/common/tests/test_foundry_protocol.py index 146b4cc..6492faa 100644 --- a/runtimes/common/tests/test_foundry_protocol.py +++ b/runtimes/common/tests/test_foundry_protocol.py @@ -91,11 +91,12 @@ def test_foundry_non_brokered_ignores_brokered_state_file_env(monkeypatch, tmp_p assert resp.json()["output"][0]["content"][0]["text"] == "echo: hi" -def test_foundry_responses_tolerates_stream_flag_with_non_streaming_response(): +def test_foundry_responses_non_streaming_keeps_json_response(): app = create_foundry_app(_spec(), EchoFactory()) with TestClient(app) as client: - resp = client.post("/responses", json={"input": "hi", "stream": True}) + resp = client.post("/responses", json={"input": "hi", "stream": False}) assert resp.status_code == 200 + assert resp.headers["content-type"] == "application/json" assert resp.json()["status"] == "completed" assert resp.json()["output"][0]["content"][0]["text"] == "echo: hi" diff --git a/runtimes/common/tests/test_foundry_streaming.py b/runtimes/common/tests/test_foundry_streaming.py index 08a26ac..5140b9e 100644 --- a/runtimes/common/tests/test_foundry_streaming.py +++ b/runtimes/common/tests/test_foundry_streaming.py @@ -8,7 +8,8 @@ import httpx import pytest - +from agentkit_serve_common.foundry import create_foundry_app +from agentkit_serve_common.runtime import AgentRunError, RunResult from test_foundry_brokered_protocol import ( CONTINUATION_AUTH, _app, @@ -17,6 +18,8 @@ _continuation, _spec, ) +from test_foundry_protocol import EchoFactory, EchoRuntime +from test_foundry_protocol import _spec as native_spec def _tool_response(): @@ -62,6 +65,8 @@ async def handle_async_request(self, request): return httpx.Response( 401, request=request, json={"error": "private-upstream-detail"} ) + if self.result == "exception": + raise RuntimeError("private-upstream-detail") payload = ( _tool_response() if self.result == "tool" @@ -71,7 +76,7 @@ async def handle_async_request(self, request): @asynccontextmanager -async def _exchange(app, payload, *, hold_created=None): +async def _exchange(app, payload, *, hold_created=None, headers=None): incoming, outgoing = asyncio.Queue(), asyncio.Queue() await incoming.put( { @@ -104,7 +109,9 @@ async def send(message): (b"content-type", b"application/json"), *[ (key.encode(), value.encode()) - for key, value in CONTINUATION_AUTH.items() + for key, value in ( + CONTINUATION_AUTH if headers is None else headers + ).items() ], ], "client": ("127.0.0.1", 1234), @@ -386,3 +393,245 @@ async def exercise(): assert model.calls == 0 asyncio.run(exercise()) + + +class NativeModelRuntime(EchoRuntime): + def __init__(self, client, url): + super().__init__() + self.client = client + self.url = url + + async def run(self, request): + self.requests.append(request) + response = await self.client.post( + self.url, + json={"messages": [{"role": "user", "content": request.prompt}]}, + ) + if response.status_code >= 400: + raise AgentRunError(response.text, status=502, code="UpstreamFailure") + body = response.json() + return RunResult(text=body["choices"][0]["message"]["content"]) + + +def _native_app(client, *, url="http://model/v1/chat/completions", **options): + factory = EchoFactory() + factory.runtime = NativeModelRuntime(client, url) + return create_foundry_app(native_spec(), factory, **options), factory.runtime + + +@pytest.mark.parametrize("session_id", [None, "native-session"]) +def test_native_stream_ack_precedes_runtime_and_keeps_completion_identity(session_id): + async def exercise(): + model = HeldModel() + async with httpx.AsyncClient(transport=model) as upstream: + app, runtime = _native_app(upstream) + payload = {"input": "Summarize the evidence", "stream": True} + if session_id: + payload["agent_session_id"] = session_id + delivered = asyncio.Event() + async with ( + app.router.lifespan_context(app), + _exchange(app, payload, hold_created=delivered, headers={}) as ( + _, + outgoing, + request, + ), + ): + created = await _created(outgoing) + assert created["response"]["status"] == "in_progress" + assert created["response"].get("agent_session_id") == session_id + assert created["response"]["output"] == [] + assert runtime.requests == [] and model.calls == 0 + delivered.set() + await asyncio.wait_for(model.started.wait(), 2) + assert not request.done() and outgoing.empty() + model.release.set() + completed = await _event(outgoing) + await asyncio.wait_for(request, 2) + assert completed["type"] == "response.completed" + response = completed["response"] + assert response["id"] == created["response"]["id"] + assert response.get("agent_session_id") == session_id + assert response["output"][0]["response_id"] == response["id"] + assert response["output"][0]["content"][0]["text"] == "Verified response." + assert runtime.requests[0].session_id == session_id + assert model.calls == 1 + + asyncio.run(exercise()) + + +@pytest.mark.parametrize("result", ["error", "exception"]) +def test_native_stream_error_keeps_identity_and_sanitizes_runtime_details(result): + async def exercise(): + model = HeldModel(result) + async with httpx.AsyncClient(transport=model) as upstream: + app, _ = _native_app(upstream) + async with ( + app.router.lifespan_context(app), + _exchange( + app, + { + "input": "Summarize the evidence", + "agent_session_id": "native-session", + "stream": True, + }, + headers={}, + ) as (_, outgoing, request), + ): + created = await _created(outgoing) + await asyncio.wait_for(model.started.wait(), 2) + model.release.set() + failed = await _event(outgoing) + await asyncio.wait_for(request, 2) + assert failed["type"] == "response.failed" and "error" not in failed + assert failed["response"]["id"] == created["response"]["id"] + assert failed["response"]["agent_session_id"] == "native-session" + assert failed["response"]["status"] == "failed" + assert failed["response"]["error"]["code"] == "RuntimeFailure" + assert "private-upstream-detail" not in json.dumps(failed) + + asyncio.run(exercise()) + + +def test_native_stream_disconnect_before_ack_delivery_starts_no_runtime(): + async def exercise(): + model = HeldModel() + async with httpx.AsyncClient(transport=model) as upstream: + app, runtime = _native_app(upstream) + async with ( + app.router.lifespan_context(app), + _exchange( + app, + {"input": "Summarize the evidence", "stream": True}, + hold_created=asyncio.Event(), + headers={}, + ) as (incoming, outgoing, request), + ): + await _created(outgoing) + await incoming.put({"type": "http.disconnect"}) + await asyncio.wait_for(request, 2) + assert runtime.requests == [] and model.calls == 0 + + asyncio.run(exercise()) + + +def test_native_stream_disconnect_closes_model_http_socket_and_allows_next_request(): + async def exercise(): + model_started, model_disconnected = asyncio.Event(), asyncio.Event() + model_requests, handlers = [], set() + + async def serve_model(reader, writer): + handler = asyncio.current_task() + handlers.add(handler) + try: + headers = await reader.readuntil(b"\r\n\r\n") + length = next( + int(line.split(b":", 1)[1]) + for line in headers.splitlines() + if line.lower().startswith(b"content-length:") + ) + model_requests.append(await reader.readexactly(length)) + if len(model_requests) == 1: + model_started.set() + assert await reader.read(1) == b"" + model_disconnected.set() + else: + body = json.dumps( + _chat_response({"role": "assistant", "content": "Recovered."}) + ).encode() + writer.write( + b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n" + + f"Content-Length: {len(body)}\r\nConnection: close\r\n\r\n".encode() + + body + ) + await writer.drain() + finally: + writer.close() + await writer.wait_closed() + handlers.remove(handler) + + server = await asyncio.start_server(serve_model, "127.0.0.1", 0) + url = f"http://127.0.0.1:{server.sockets[0].getsockname()[1]}/chat/completions" + try: + async with ( + server, + httpx.AsyncClient(trust_env=False, timeout=5) as upstream, + ): + app, runtime = _native_app(upstream, url=url) + async with app.router.lifespan_context(app): + async with _exchange( + app, + {"input": "Wait for cancellation", "stream": True}, + headers={}, + ) as (incoming, outgoing, request): + await _created(outgoing) + await asyncio.wait_for(model_started.wait(), 2) + await incoming.put({"type": "http.disconnect"}) + await asyncio.wait_for(request, 2) + await asyncio.wait_for(model_disconnected.wait(), 2) + assert outgoing.empty() + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://hosted" + ) as client: + response = await client.post( + "/responses", + json={"input": "Next request", "stream": False}, + ) + assert response.status_code == 200 + assert ( + response.json()["output"][0]["content"][0]["text"] + == "Recovered." + ) + assert len(runtime.requests) == len(model_requests) == 2 + finally: + for handler in tuple(handlers): + handler.cancel() + await asyncio.gather(*handlers, return_exceptions=True) + + asyncio.run(exercise()) + + +@pytest.mark.parametrize( + ("payload", "authorized", "status", "code"), + [ + ({"input": "Read data", "stream": True}, False, 401, None), + ({"stream": True}, True, 400, "missing_input"), + ({"input": "x", "stream": True, "tools": [{}]}, True, 400, "tools_unsupported"), + ( + {"input": [{"role": "assistant", "content": "x"}], "stream": True}, + True, + 400, + "invalid_input", + ), + ({"input": "x" * 300, "stream": True}, True, 413, "request_too_large"), + ], +) +def test_native_stream_rejection_keeps_http_error_before_runtime( + payload, authorized, status, code +): + async def exercise(): + model = HeldModel() + async with httpx.AsyncClient(transport=model) as upstream: + app, runtime = _native_app( + upstream, auth_token="local-test", max_request_body_bytes=256 + ) + async with ( + app.router.lifespan_context(app), + httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://hosted" + ) as client, + ): + response = await client.post( + "/responses", + json=payload, + headers={"authorization": "Bearer local-test"} + if authorized + else {}, + ) + assert response.status_code == status + assert response.headers["content-type"] == "application/json" + if code: + assert response.json()["error"]["code"] == code + assert runtime.requests == [] and model.calls == 0 + + asyncio.run(exercise()) From e985d7083943d13ca553cb976a9bdb0ae0700ae8 Mon Sep 17 00:00:00 2001 From: Sertac Ozercan Date: Fri, 11 Sep 2026 15:33:42 -0700 Subject: [PATCH 6/9] fix(foundry): preserve response creation timestamps Signed-off-by: Sertac Ozercan --- .../common/agentkit_serve_common/foundry.py | 24 +++++++----- .../foundry_streaming.py | 14 ++++--- .../common/tests/test_foundry_streaming.py | 37 +++++++++++++++++-- 3 files changed, 57 insertions(+), 18 deletions(-) diff --git a/runtimes/common/agentkit_serve_common/foundry.py b/runtimes/common/agentkit_serve_common/foundry.py index 0597569..a1bfbf2 100644 --- a/runtimes/common/agentkit_serve_common/foundry.py +++ b/runtimes/common/agentkit_serve_common/foundry.py @@ -371,13 +371,14 @@ def _responses_payload( *, previous_response_id: str | None = None, response_id: str | None = None, + created_at: int | None = None, ) -> dict[str, Any]: response_id = response_id or _new_response_id(previous_response_id) message_id = _new_message_id(response_id) payload: dict[str, Any] = { "id": response_id, "object": "response", - "created_at": int(time.time()), + "created_at": int(time.time()) if created_at is None else created_at, "status": "completed", "model": spec.model.name, "output": [ @@ -2163,11 +2164,12 @@ def _function_call_response_payload( call: _PendingCall, previous_response_id: str | None = None, usage: Mapping[str, int] | None = None, + created_at: int | None = None, ) -> dict[str, Any]: payload: dict[str, Any] = { "id": response_id, "object": "response", - "created_at": int(time.time()), + "created_at": int(time.time()) if created_at is None else created_at, "status": "completed", "model": spec.model.name, "output": [ @@ -2244,6 +2246,7 @@ def _advance_brokered_state( model_loop: BrokeredChatModelLoop, response_id: str, next_call_id: str, + created_at: int | None = None, ) -> JSONResponse: try: call = _model_pending_call( @@ -2258,7 +2261,7 @@ def _advance_brokered_state( if not _reset_unfinalized_continuation(store, state, call_id=call_id): return _state_storage_error() return _error(str(exc), status=exc.status, code=exc.code) - payload = _function_call_response_payload(spec, response_id=response_id, call=call, previous_response_id=previous_response_id, usage=result.usage) + payload = _function_call_response_payload(spec, response_id=response_id, call=call, previous_response_id=previous_response_id, usage=result.usage, created_at=created_at) following = deepcopy(state) if not following.response_calls: following.response_calls[state.response_id] = call_id @@ -2531,6 +2534,7 @@ async def _handle_brokered_continuation( if state.status != "pending": return _error("previous response is not pending a tool result", status=409, code="response_not_pending") + created_at = None if state.model_messages is not None and model_loop is not None: try: model_loop.validate_static_credentials() @@ -2555,7 +2559,7 @@ async def _handle_brokered_continuation( next_response_id = _new_response_id(str(previous_response_id)) next_call_id = f"call_{next_response_id}_1" if stream is not None: - await stream.accept(next_response_id) + created_at = await stream.accept(next_response_id) model_result = await model_loop.resume(state.model_messages, call_id=call_id, output=output_json, next_call_id=next_call_id) except AgentRunError as exc: if not _reset_unfinalized_continuation(store, state, call_id=call_id): @@ -2586,6 +2590,7 @@ async def _handle_brokered_continuation( model_loop=model_loop, response_id=next_response_id, next_call_id=next_call_id, + created_at=created_at, ) result = RunResult(text=model_result.text, usage=_combine_usage(state.initial_usage, model_result.usage)) finally: @@ -2601,6 +2606,7 @@ async def _handle_brokered_continuation( result, previous_response_id=str(previous_response_id), response_id=stream.created.result()["id"] if stream is not None and stream.created.done() else None, + created_at=created_at, ) state.accepted_output_digests[call_id] = output_digest state.accepted_output_sizes[call_id] = accepted_output_size @@ -2929,8 +2935,7 @@ async def execute_responses( return _state_storage_error() try: try: - if stream is not None: - await stream.accept(response_id) + created_at = await stream.accept(response_id) if stream is not None else None model_result = await model_loop.start(run_request, call_id=call_id) except AgentRunError as exc: if exc.code == "ModelResponseTooLarge": @@ -2942,6 +2947,7 @@ async def execute_responses( RunResult(text=model_result.text, usage=model_result.usage), previous_response_id=previous_response_id_for_output, response_id=response_id if stream is not None else None, + created_at=created_at, ) if run_request.session_id: completed = _HostedResponseState( @@ -3026,6 +3032,7 @@ async def execute_responses( call=call, previous_response_id=previous_response_id_for_output, usage=model_result.usage, + created_at=created_at, ) ) finally: @@ -3087,14 +3094,13 @@ async def execute_responses( response_id = _new_response_id() if stream is not None else None try: - if stream is not None: - await stream.accept(response_id) + created_at = await stream.accept(response_id) if stream is not None else None result = await request.app.state.runtime.run(run_request) except AgentRunError as exc: return _non_brokered_agent_run_error(exc) except Exception: # noqa: BLE001 - deterministic protocol envelope. return _non_brokered_unexpected_runtime_error() - return JSONResponse(_responses_payload(spec, result, response_id=response_id)) + return JSONResponse(_responses_payload(spec, result, response_id=response_id, created_at=created_at)) return app diff --git a/runtimes/common/agentkit_serve_common/foundry_streaming.py b/runtimes/common/agentkit_serve_common/foundry_streaming.py index 27df09d..a8ce8bd 100644 --- a/runtimes/common/agentkit_serve_common/foundry_streaming.py +++ b/runtimes/common/agentkit_serve_common/foundry_streaming.py @@ -20,13 +20,13 @@ def __init__(self, model: str) -> None: ) self.delivered = asyncio.Event() - def prepare(self, response_id: str) -> None: + def prepare(self, response_id: str, *, created_at: int | None = None) -> None: if self.created.done(): raise RuntimeError("hosted response was already acknowledged") payload: dict[str, Any] = { "id": response_id, "object": "response", - "created_at": int(time.time()), + "created_at": int(time.time()) if created_at is None else created_at, "status": "in_progress", "model": self.model, "output": [], @@ -35,10 +35,11 @@ def prepare(self, response_id: str) -> None: payload["agent_session_id"] = self.session_id self.created.set_result(payload) - async def accept(self, response_id: str) -> None: + async def accept(self, response_id: str) -> int: self.prepare(response_id) # Model work starts only after the response.created frame has been sent. await self.delivered.wait() + return self.created.result()["created_at"] def terminal(self, result: JSONResponse | None) -> dict[str, Any]: created = self.created.result() @@ -66,7 +67,7 @@ def terminal(self, result: JSONResponse | None) -> dict[str, Any]: def _frame(event: dict[str, Any]) -> bytes: data = json.dumps(event, separators=(",", ":"), ensure_ascii=True) - return f"event: {event['type']}\ndata: {data}\n\n".encode("utf-8") + return f"event: {event['type']}\ndata: {data}\n\n".encode() class _BrokeredStreamingResponse(Response): @@ -108,7 +109,7 @@ async def write_response() -> None: self.stream.delivered.set() try: result = await self.operation - except Exception: # Model failures must not leak exception text into SSE. + except Exception: # noqa: BLE001 - never leak model exception text into SSE. result = None await send( { @@ -154,7 +155,8 @@ async def brokered_stream_response( result = await task if result.status_code >= 400: return result - stream.prepare(json.loads(result.body)["id"]) + payload = json.loads(result.body) + stream.prepare(payload["id"], created_at=payload["created_at"]) return _BrokeredStreamingResponse(stream, task) except BaseException: task.cancel() diff --git a/runtimes/common/tests/test_foundry_streaming.py b/runtimes/common/tests/test_foundry_streaming.py index 5140b9e..7591e77 100644 --- a/runtimes/common/tests/test_foundry_streaming.py +++ b/runtimes/common/tests/test_foundry_streaming.py @@ -149,8 +149,12 @@ async def _event(outgoing): @pytest.mark.parametrize("continuation", [False, True]) @pytest.mark.parametrize("result", ["text", "tool"]) def test_brokered_stream_ack_precedes_model_and_matches_completion_and_replay( - continuation, result + continuation, result, monkeypatch, tmp_path ): + clock = [1_700_000_000] + monkeypatch.setattr("agentkit_serve_common.foundry.time.time", lambda: clock[0]) + state_file = tmp_path / "responses.json" + async def exercise(): model = HeldModel(result, continuation=continuation) async with httpx.AsyncClient(transport=model) as upstream: @@ -158,6 +162,7 @@ async def exercise(): _spec(), brokered_model_loop_enabled=True, brokered_model_http_client=upstream, + response_state_file=state_file, ) payload = { "input": "Read the synthetic data", @@ -195,29 +200,47 @@ async def exercise(): created = await _created(outgoing) assert created["type"] == "response.created" assert created["response"]["status"] == "in_progress" + assert created["response"]["created_at"] == clock[0] assert created["response"]["agent_session_id"] == "session-a" assert created["response"]["output"] == [] assert not model.started.is_set() delivered.set() await asyncio.wait_for(model.started.wait(), 2) assert outgoing.empty() and not request.done() + clock[0] += 10 model.release.set() completed = await _event(outgoing) await asyncio.wait_for(request, 2) assert completed["type"] == "response.completed" response = completed["response"] assert response["id"] == created["response"]["id"] + assert response["created_at"] == created["response"]["created_at"] assert response["agent_session_id"] == "session-a" assert response["output"][0]["response_id"] == response["id"] assert response["output"][0]["type"] == ( "function_call" if result == "tool" else "message" ) if continuation: - async with _exchange(app, payload) as (_, outgoing, replay_request): + clock[0] += 10 + restarted_app = _app( + _spec(), + brokered_model_loop_enabled=True, + brokered_model_http_client=upstream, + response_state_file=state_file, + ) + async with _exchange(restarted_app, payload) as ( + _, + outgoing, + replay_request, + ): replay_created = await _created(outgoing) replay_completed = await _event(outgoing) await asyncio.wait_for(replay_request, 2) assert replay_created["response"]["id"] == response["id"] + assert ( + replay_created["response"]["created_at"] + == response["created_at"] + ) assert replay_completed == completed buffered = await client.post( "/responses", @@ -420,7 +443,12 @@ def _native_app(client, *, url="http://model/v1/chat/completions", **options): @pytest.mark.parametrize("session_id", [None, "native-session"]) -def test_native_stream_ack_precedes_runtime_and_keeps_completion_identity(session_id): +def test_native_stream_ack_precedes_runtime_and_keeps_completion_identity( + session_id, monkeypatch +): + clock = [1_700_000_000] + monkeypatch.setattr("agentkit_serve_common.foundry.time.time", lambda: clock[0]) + async def exercise(): model = HeldModel() async with httpx.AsyncClient(transport=model) as upstream: @@ -439,18 +467,21 @@ async def exercise(): ): created = await _created(outgoing) assert created["response"]["status"] == "in_progress" + assert created["response"]["created_at"] == clock[0] assert created["response"].get("agent_session_id") == session_id assert created["response"]["output"] == [] assert runtime.requests == [] and model.calls == 0 delivered.set() await asyncio.wait_for(model.started.wait(), 2) assert not request.done() and outgoing.empty() + clock[0] += 10 model.release.set() completed = await _event(outgoing) await asyncio.wait_for(request, 2) assert completed["type"] == "response.completed" response = completed["response"] assert response["id"] == created["response"]["id"] + assert response["created_at"] == created["response"]["created_at"] assert response.get("agent_session_id") == session_id assert response["output"][0]["response_id"] == response["id"] assert response["output"][0]["content"][0]["text"] == "Verified response." From 710e874020b289a190d129a26040b3469e410310 Mon Sep 17 00:00:00 2001 From: Sertac Ozercan Date: Fri, 11 Sep 2026 16:11:15 -0700 Subject: [PATCH 7/9] fix(foundry): retain safe brokered model diagnostics Signed-off-by: Sertac Ozercan --- .../common/agentkit_serve_common/foundry.py | 45 ++-- .../foundry_model_loop.py | 47 +++- .../tests/test_foundry_brokered_protocol.py | 13 +- .../common/tests/test_foundry_model_errors.py | 217 ++++++++++++++++++ 4 files changed, 300 insertions(+), 22 deletions(-) create mode 100644 runtimes/common/tests/test_foundry_model_errors.py diff --git a/runtimes/common/agentkit_serve_common/foundry.py b/runtimes/common/agentkit_serve_common/foundry.py index a1bfbf2..3e45f31 100644 --- a/runtimes/common/agentkit_serve_common/foundry.py +++ b/runtimes/common/agentkit_serve_common/foundry.py @@ -39,7 +39,12 @@ from .brokered import brokered_tool_definitions from .config import AgentSpec, _unsafe_brokered_key, _unsafe_brokered_text -from .foundry_model_loop import BrokeredChatModelLoop, ModelLoopFinal, ModelLoopToolRequest +from .foundry_model_loop import ( + BrokeredChatModelLoop, + ModelLoopFinal, + ModelLoopToolRequest, + normalized_model_error_details, +) from .foundry_streaming import BrokeredResponseStream, brokered_stream_response from .conversation import FORWARDED_ROLES, ConversationTurn, RunRequest from .runtime import AgentRunError, BrokeredToolDefinition, RunResult, RuntimeFactory @@ -193,6 +198,21 @@ def _model_response_too_large_error() -> JSONResponse: ) +def _brokered_model_run_error(exc: Exception) -> JSONResponse: + status = 502 + if isinstance(exc, AgentRunError): + normalized = normalized_model_error_details(exc) + if normalized is not None: + status, error = normalized + return JSONResponse({"error": error}, status_code=status) + if type(exc.status) is int: + if 400 <= exc.status < 500: + return _error(str(exc), status=exc.status, code=exc.code) + if 500 <= exc.status <= 599: + status = exc.status + return _error("model resume failed", status=status, code="ModelResumeError") + + def _non_brokered_agent_run_error(exc: AgentRunError) -> JSONResponse: if exc.status < 500: return _error(str(exc), status=exc.status, code=exc.code) @@ -2539,7 +2559,7 @@ async def _handle_brokered_continuation( try: model_loop.validate_static_credentials() except AgentRunError as exc: - return _error(str(exc), status=exc.status, code=exc.code) + return _brokered_model_run_error(exc) state.accepted_output_digests[call_id] = output_digest state.accepted_output_sizes[call_id] = accepted_output_size state.status = "resuming" @@ -2564,21 +2584,16 @@ async def _handle_brokered_continuation( except AgentRunError as exc: if not _reset_unfinalized_continuation(store, state, call_id=call_id): return _state_storage_error() - if exc.code == "ModelResponseTooLarge": - logger.warning("brokered model-loop response exceeded configured limits") - return _model_response_too_large_error() - if exc.status >= 500: - logger.warning("brokered model-loop resume failed: %s", exc) - return _error("model resume failed", status=exc.status, code="ModelResumeError") - return _error(str(exc), status=exc.status, code=exc.code) + logger.warning("brokered model-loop resume failed") + return _brokered_model_run_error(exc) except asyncio.CancelledError: _reset_unfinalized_continuation(store, state, call_id=call_id) raise except Exception as exc: # noqa: BLE001 - reset continuation state before surfacing unexpected model failures. - logger.exception("brokered model-loop resume failed unexpectedly") + logger.warning("brokered model-loop resume failed unexpectedly") if not _reset_unfinalized_continuation(store, state, call_id=call_id): return _state_storage_error() - return _error("model resume failed", status=502, code="ModelResumeError") + return _brokered_model_run_error(exc) if isinstance(model_result, ModelLoopToolRequest): return _advance_brokered_state( spec=spec, @@ -2919,7 +2934,7 @@ async def execute_responses( try: model_loop.validate_static_credentials() except AgentRunError as exc: - return _error(str(exc), status=exc.status, code=exc.code) + return _brokered_model_run_error(exc) response_id = _new_response_id(previous_response_id_for_output) call_id = f"call_{response_id}_1" try: @@ -2938,9 +2953,9 @@ async def execute_responses( created_at = await stream.accept(response_id) if stream is not None else None model_result = await model_loop.start(run_request, call_id=call_id) except AgentRunError as exc: - if exc.code == "ModelResponseTooLarge": - return _model_response_too_large_error() - return _error(str(exc), status=exc.status, code=exc.code) + return _brokered_model_run_error(exc) + except Exception as exc: # noqa: BLE001 - unknown model failures have no public diagnostic payload. + return _brokered_model_run_error(exc) if isinstance(model_result, ModelLoopFinal): payload = _responses_payload( spec, diff --git a/runtimes/common/agentkit_serve_common/foundry_model_loop.py b/runtimes/common/agentkit_serve_common/foundry_model_loop.py index 3b6ad82..35acae6 100644 --- a/runtimes/common/agentkit_serve_common/foundry_model_loop.py +++ b/runtimes/common/agentkit_serve_common/foundry_model_loop.py @@ -28,6 +28,38 @@ _MAX_ARGUMENT_DEPTH = 128 +_NORMALIZED_MODEL_ERRORS = { + "ModelAuthMissing": (503, "model authentication is not configured"), + "ModelAuthRejected": (503, "model service rejected configured credentials"), + "ModelUnavailable": (503, "model service is unavailable"), + "ModelUpstreamError": (502, "model service request failed"), + "InvalidModelResponse": (502, "model service returned an invalid response"), + "ModelResponseTooLarge": (502, "model response is too large to retain safely"), +} + + +class _ModelHTTPError(AgentRunError): + def __init__(self, message: str, *, status: int, code: str, upstream_status: int) -> None: + super().__init__(message, status=status, code=code) + self.upstream_status = upstream_status + + +def normalized_model_error_details(exc: AgentRunError) -> tuple[int, dict[str, Any]] | None: + """Project only runtime-owned model error definitions and bounded HTTP metadata.""" + definition = _NORMALIZED_MODEL_ERRORS.get(exc.code) if type(exc.code) is str else None + if definition is None: + return None + status, message = definition + error: dict[str, Any] = {"message": message, "code": exc.code} + # A similarly named attribute on a framework exception is not HTTP evidence. + if ( + isinstance(exc, _ModelHTTPError) + and type(exc.upstream_status) is int + and 400 <= exc.upstream_status <= 599 + ): + error["upstream_status"] = exc.upstream_status + return status, error + @dataclass(frozen=True) class ModelLoopFinal: @@ -319,22 +351,31 @@ def _chat_completions_url(base_url: str) -> str: def _normalized_model_http_error(status_code: int) -> AgentRunError: - if status_code in {401, 403}: + if type(status_code) is not int or not 400 <= status_code <= 599: return AgentRunError( + "model service request failed", + status=502, + code="ModelUpstreamError", + ) + if status_code in {401, 403}: + return _ModelHTTPError( "model service rejected configured credentials", status=503, code="ModelAuthRejected", + upstream_status=status_code, ) if status_code == 429 or status_code >= 500: - return AgentRunError( + return _ModelHTTPError( "model service is unavailable", status=503, code="ModelUnavailable", + upstream_status=status_code, ) - return AgentRunError( + return _ModelHTTPError( "model service request failed", status=502, code="ModelUpstreamError", + upstream_status=status_code, ) diff --git a/runtimes/common/tests/test_foundry_brokered_protocol.py b/runtimes/common/tests/test_foundry_brokered_protocol.py index 438fc07..1b0775f 100644 --- a/runtimes/common/tests/test_foundry_brokered_protocol.py +++ b/runtimes/common/tests/test_foundry_brokered_protocol.py @@ -2542,6 +2542,7 @@ def reject(request: httpx.Request, *, status: int = upstream_status) -> httpx.Re assert response.json()["error"] == { "message": "model service rejected configured credentials", "code": "ModelAuthRejected", + "upstream_status": upstream_status, } assert internal_url not in response.text assert "/chat/completions" not in response.text @@ -2597,7 +2598,7 @@ def respond(request: httpx.Request) -> httpx.Response: assert response.status_code == 502 assert response.json()["error"] == { - "message": "model response must be a JSON object", + "message": "model service returned an invalid response", "code": "InvalidModelResponse", } @@ -2707,7 +2708,7 @@ def respond(request: httpx.Request) -> httpx.Response: assert response.status_code == 502 assert response.json()["error"] == { - "message": "model service returned an invalid JSON response", + "message": "model service returned an invalid response", "code": "InvalidModelResponse", } @@ -2763,7 +2764,11 @@ def respond(request: httpx.Request) -> httpx.Response: retried = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) assert failed.status_code == 503 - assert failed.json()["error"] == {"message": "model resume failed", "code": "ModelResumeError"} + assert failed.json()["error"] == { + "message": "model service rejected configured credentials", + "code": "ModelAuthRejected", + "upstream_status": 401, + } assert internal_url not in failed.text assert retried.status_code == 200, retried.text assert _message_text(retried.json()) == "Retry after auth recovery." @@ -3166,7 +3171,7 @@ def handler(self, request: httpx.Request) -> httpx.Response: retried = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) assert failed.status_code == 502 - assert failed.json()["error"] == {"message": "model resume failed", "code": "ModelResumeError"} + assert failed.json()["error"] == {"message": "model service request failed", "code": "ModelUpstreamError"} assert retried.status_code == 200, retried.text assert _message_text(retried.json()) == "Recovered after retry." diff --git a/runtimes/common/tests/test_foundry_model_errors.py b/runtimes/common/tests/test_foundry_model_errors.py new file mode 100644 index 0000000..bc748d8 --- /dev/null +++ b/runtimes/common/tests/test_foundry_model_errors.py @@ -0,0 +1,217 @@ +"""Safe model diagnostics on the public hosted Responses boundary.""" + +from __future__ import annotations + +import json + +import httpx +import pytest +from fastapi.testclient import TestClient + +from agentkit_serve_common import foundry_model_loop +from agentkit_serve_common.runtime import AgentRunError +from test_foundry_brokered_protocol import ( + CONTINUATION_AUTH, + _app, + _call, + _continuation, +) +from test_foundry_streaming import _tool_response + + +def _streamed_error(app, *, continuation): + payload = { + "input": "Read synthetic data.", + "agent_session_id": "synthetic-session", + "stream": True, + } + with TestClient(app) as client: + if continuation: + initial = client.post("/responses", json={**payload, "stream": False}) + assert initial.status_code == 200 + payload = { + **_continuation( + initial.json()["id"], + _call(initial.json())["call_id"], + {"approved": True, "output": {"ok": True}}, + ), + "agent_session_id": "synthetic-session", + "stream": True, + } + response = client.post("/responses", json=payload, headers=CONTINUATION_AUTH) + + assert response.status_code == 200 + frames = [ + json.loads(line.removeprefix("data: ")) + for line in response.text.splitlines() + if line.startswith("data: ") + ] + assert [frame["type"] for frame in frames] == [ + "response.created", + "response.failed", + ] + created, failed = (frame["response"] for frame in frames) + assert failed["id"] == created["id"] + assert failed["created_at"] == created["created_at"] + assert failed["agent_session_id"] == created["agent_session_id"] == "synthetic-session" + assert failed["status"] == "failed" + assert "error" not in frames[-1] + assert "private-upstream-detail" not in response.text + return failed["error"] + + +@pytest.mark.parametrize("continuation", [False, True]) +@pytest.mark.parametrize( + ("status", "code", "message"), + [ + (401, "ModelAuthRejected", "model service rejected configured credentials"), + (403, "ModelAuthRejected", "model service rejected configured credentials"), + (429, "ModelUnavailable", "model service is unavailable"), + (500, "ModelUnavailable", "model service is unavailable"), + (599, "ModelUnavailable", "model service is unavailable"), + (400, "ModelUpstreamError", "model service request failed"), + ], +) +def test_model_http_error_retains_only_normalized_code_and_status( + continuation, status, code, message, caplog +): + calls = 0 + + def model(request): + nonlocal calls + calls += 1 + if continuation and calls == 1: + return httpx.Response(200, request=request, json=_tool_response()) + return httpx.Response( + status, + request=request, + headers={"x-request-id": "private-upstream-detail"}, + json={ + "error": { + "code": "private-upstream-detail", + "message": "private-upstream-detail", + } + }, + ) + + app = _app( + brokered_model_loop_enabled=True, + brokered_model_http_client=httpx.AsyncClient(transport=httpx.MockTransport(model)), + ) + error = _streamed_error(app, continuation=continuation) + assert error == {"code": code, "message": message, "upstream_status": status} + assert "private-upstream-detail" not in caplog.text + + +@pytest.mark.parametrize("continuation", [False, True]) +@pytest.mark.parametrize("kind", ["transport", "invalid-json"]) +def test_non_http_model_failures_have_no_upstream_status(continuation, kind): + calls = 0 + + def model(request): + nonlocal calls + calls += 1 + if continuation and calls == 1: + return httpx.Response(200, request=request, json=_tool_response()) + if kind == "transport": + raise httpx.ConnectError("private-upstream-detail", request=request) + return httpx.Response(200, request=request, content=b"private-upstream-detail") + + app = _app( + brokered_model_loop_enabled=True, + brokered_model_http_client=httpx.AsyncClient(transport=httpx.MockTransport(model)), + ) + error = _streamed_error(app, continuation=continuation) + assert set(error) == {"code", "message"} + assert error["code"] == ( + "ModelUpstreamError" if kind == "transport" else "InvalidModelResponse" + ) + + +def _raise_model_error(monkeypatch, *, continuation, error): + calls = 0 + + async def fail(self, messages, *, tools): + nonlocal calls + calls += 1 + if continuation and calls == 1: + return _tool_response() + raise error + + monkeypatch.setattr(foundry_model_loop.BrokeredChatModelLoop, "_chat", fail) + return _app(brokered_model_loop_enabled=True) + + +@pytest.mark.parametrize("continuation", [False, True]) +@pytest.mark.parametrize( + ("code", "message"), + [ + ("ModelAuthMissing", "model authentication is not configured"), + ("ModelAuthRejected", "model service rejected configured credentials"), + ("ModelUnavailable", "model service is unavailable"), + ("ModelUpstreamError", "model service request failed"), + ("InvalidModelResponse", "model service returned an invalid response"), + ("ModelResponseTooLarge", "model response is too large to retain safely"), + ], +) +def test_known_model_codes_use_fixed_messages_and_do_not_trust_status_attributes( + continuation, code, message, monkeypatch, caplog +): + error = AgentRunError("private-upstream-detail", status=999, code=code) + error.upstream_status = 401 + app = _raise_model_error(monkeypatch, continuation=continuation, error=error) + assert _streamed_error(app, continuation=continuation) == { + "code": code, + "message": message, + } + assert "private-upstream-detail" not in caplog.text + + +@pytest.mark.parametrize("continuation", [False, True]) +@pytest.mark.parametrize("upstream_status", [True, "401", 401.0, 399, 600, None]) +def test_normalized_http_error_rejects_invalid_upstream_status( + continuation, upstream_status, monkeypatch +): + error = foundry_model_loop._normalized_model_http_error(401) + error.upstream_status = upstream_status + app = _raise_model_error(monkeypatch, continuation=continuation, error=error) + assert _streamed_error(app, continuation=continuation) == { + "code": "ModelAuthRejected", + "message": "model service rejected configured credentials", + } + + +@pytest.mark.parametrize("continuation", [False, True]) +@pytest.mark.parametrize( + ("code", "status"), + [ + ("private-upstream-detail", 503), + (None, 999), + (["private-upstream-detail"], "503"), + ], +) +def test_unknown_model_codes_and_statuses_are_sanitized( + continuation, code, status, monkeypatch, caplog +): + error = AgentRunError("private-upstream-detail", status=status, code=code) + error.upstream_status = 401 + app = _raise_model_error(monkeypatch, continuation=continuation, error=error) + assert _streamed_error(app, continuation=continuation) == { + "code": "ModelResumeError", + "message": "model resume failed", + } + assert "private-upstream-detail" not in caplog.text + + +@pytest.mark.parametrize("continuation", [False, True]) +def test_unexpected_model_exception_remains_generic(continuation, monkeypatch, caplog): + app = _raise_model_error( + monkeypatch, + continuation=continuation, + error=RuntimeError("private-upstream-detail"), + ) + assert _streamed_error(app, continuation=continuation) == { + "code": "ModelResumeError", + "message": "model resume failed", + } + assert "private-upstream-detail" not in caplog.text From b451933ff848af6bcbf76eaecd2d2917515f9880 Mon Sep 17 00:00:00 2001 From: Sertac Ozercan Date: Fri, 11 Sep 2026 16:40:33 -0700 Subject: [PATCH 8/9] fix(foundry): retry confirmed model rate limits Signed-off-by: Sertac Ozercan --- docs/foundry-hosted-brokered.md | 15 +- .../foundry_model_loop.py | 71 +++- .../foundry_streaming.py | 13 +- .../tests/test_foundry_brokered_protocol.py | 5 +- .../tests/test_foundry_model_retries.py | 392 ++++++++++++++++++ .../common/tests/test_foundry_streaming.py | 5 +- 6 files changed, 480 insertions(+), 21 deletions(-) create mode 100644 runtimes/common/tests/test_foundry_model_retries.py diff --git a/docs/foundry-hosted-brokered.md b/docs/foundry-hosted-brokered.md index d47102e..e1ec034 100644 --- a/docs/foundry-hosted-brokered.md +++ b/docs/foundry-hosted-brokered.md @@ -396,8 +396,19 @@ state preserves these completed rounds across restarts. The model can make up to 16 sequential tool calls per user turn. At the limit, AgentKit asks for a final answer without tools and rejects any further tool -call. Parallel tool batches are unsupported. AgentKit-owned MCP and direct -operational tools remain disabled. +call. The model endpoint must support `parallel_tool_calls: false`, which asks +for one tool call at a time. AgentKit also rejects parallel tool batches if a +model ignores that setting. AgentKit-owned MCP and direct operational tools +remain disabled. + +An HTTP 429 from the model service is retried up to twice within the same +hosted response. AgentKit honors `retry-after-ms` or `Retry-After` delays of +up to 60 seconds each. A longer server delay ends the response with +`ModelUnavailable` and `upstream_status: 429`. Missing or malformed delay +headers use short exponential backoff with jitter. Disconnecting the hosted +stream cancels the wait. Other HTTP errors, transport failures, and invalid +model responses are not retried. These model retries do not repeat Orka tool +operations or submit a new hosted response. Agents can also use [bundled instruction skills](instruction-skills.md). With filesystem skills configured under `/agent/skills`, AgentKit advertises the diff --git a/runtimes/common/agentkit_serve_common/foundry_model_loop.py b/runtimes/common/agentkit_serve_common/foundry_model_loop.py index 35acae6..bcf1129 100644 --- a/runtimes/common/agentkit_serve_common/foundry_model_loop.py +++ b/runtimes/common/agentkit_serve_common/foundry_model_loop.py @@ -13,9 +13,12 @@ import json import math import os +import random +import time import uuid from dataclasses import dataclass, field from decimal import Decimal, InvalidOperation +from email.utils import mktime_tz, parsedate_tz from typing import Any, Mapping, Sequence import httpx @@ -27,6 +30,8 @@ from .skills import SkillCatalog _MAX_ARGUMENT_DEPTH = 128 +_MAX_RATE_LIMIT_RETRIES = 2 +_MAX_RETRY_AFTER_SECONDS = 60 _NORMALIZED_MODEL_ERRORS = { "ModelAuthMissing": (503, "model authentication is not configured"), @@ -259,20 +264,32 @@ async def _chat(self, messages: Sequence[Mapping[str, Any]], *, tools: Sequence[ client = self.http_client close_client = False if client is None: - client = httpx.AsyncClient(headers=headers, timeout=60) + client = httpx.AsyncClient(timeout=60) close_client = True try: - async with client.stream( - "POST", - _chat_completions_url(self.spec.model.base_url), - json=payload, - headers=headers or None, - ) as response: - response.raise_for_status() - response_body = await _read_response_body_bounded( - response, - max_bytes=self.max_response_bytes, - ) + for retry in range(_MAX_RATE_LIMIT_RETRIES + 1): + retry_delay = None + async with client.stream( + "POST", + _chat_completions_url(self.spec.model.base_url), + json=payload, + headers=headers or None, + ) as response: + # Only an explicit rate-limit rejection is safe to retry. + # Transport failures or accepted responses may have done work. + if response.status_code == 429 and retry < _MAX_RATE_LIMIT_RETRIES: + retry_delay = _rate_limit_retry_delay(response.headers, retry=retry) + if retry_delay is None: + response.raise_for_status() + response_body = await _read_response_body_bounded( + response, + max_bytes=self.max_response_bytes, + ) + break + # Release the HTTP response before waiting. Disconnect cancellation + # must stop this wait before another model request can be submitted. + await asyncio.sleep(retry_delay) + headers = await self._auth_headers() except httpx.HTTPStatusError as exc: raise _normalized_model_http_error(exc.response.status_code) from exc except AgentRunError: @@ -317,6 +334,36 @@ async def _auth_headers(self) -> dict[str, str]: return headers +def _rate_limit_retry_delay(headers: httpx.Headers, *, retry: int) -> float | None: + delay = _retry_after_seconds(headers) + if delay is None: + return (2.0 ** retry) * random.uniform(0.75, 1.0) + # Never shorten a longer server delay and retry before the requested window. + return delay if delay <= _MAX_RETRY_AFTER_SECONDS else None + + +def _retry_after_seconds(headers: httpx.Headers) -> float | None: + for name, divisor in (("retry-after-ms", 1000), ("retry-after", 1)): + value = headers.get(name) + if value is None: + continue + try: + delay = float(value) + except ValueError: + continue + if math.isfinite(delay) and delay >= 0: + return delay / divisor + value = headers.get("retry-after") + if value is not None: + try: + date = parsedate_tz(value) + if date is not None: + return max(0.0, mktime_tz(date) - time.time()) + except (TypeError, ValueError, OverflowError): + pass + return None + + async def _read_response_body_bounded(response: httpx.Response, *, max_bytes: int) -> bytearray: content_length = response.headers.get("content-length") if content_length is not None: diff --git a/runtimes/common/agentkit_serve_common/foundry_streaming.py b/runtimes/common/agentkit_serve_common/foundry_streaming.py index a8ce8bd..ebf1401 100644 --- a/runtimes/common/agentkit_serve_common/foundry_streaming.py +++ b/runtimes/common/agentkit_serve_common/foundry_streaming.py @@ -65,8 +65,12 @@ def terminal(self, result: JSONResponse | None) -> dict[str, Any]: } -def _frame(event: dict[str, Any]) -> bytes: - data = json.dumps(event, separators=(",", ":"), ensure_ascii=True) +def _frame(event: dict[str, Any], *, sequence_number: int) -> bytes: + data = json.dumps( + {**event, "sequence_number": sequence_number}, + separators=(",", ":"), + ensure_ascii=True, + ) return f"event: {event['type']}\ndata: {data}\n\n".encode() @@ -101,7 +105,8 @@ async def write_response() -> None: { "type": "response.created", "response": self.stream.created.result(), - } + }, + sequence_number=0, ), "more_body": True, } @@ -114,7 +119,7 @@ async def write_response() -> None: await send( { "type": "http.response.body", - "body": _frame(self.stream.terminal(result)), + "body": _frame(self.stream.terminal(result), sequence_number=1), "more_body": False, } ) diff --git a/runtimes/common/tests/test_foundry_brokered_protocol.py b/runtimes/common/tests/test_foundry_brokered_protocol.py index 1b0775f..b1db2f4 100644 --- a/runtimes/common/tests/test_foundry_brokered_protocol.py +++ b/runtimes/common/tests/test_foundry_brokered_protocol.py @@ -4603,13 +4603,14 @@ async def __aexit__(self, *args: Any) -> None: return None class FakeClient: - def __init__(self, *, headers: dict[str, str], timeout: int) -> None: + def __init__(self, *, headers: dict[str, str] | None = None, timeout: int) -> None: assert timeout == 60 - captured_headers.update(headers) + captured_headers.update(headers or {}) def stream(self, method: str, url: str, **kwargs: Any) -> FakeStream: assert method == "POST" assert url.endswith("/chat/completions") + captured_headers.update(kwargs.get("headers") or {}) return FakeStream() async def aclose(self) -> None: diff --git a/runtimes/common/tests/test_foundry_model_retries.py b/runtimes/common/tests/test_foundry_model_retries.py new file mode 100644 index 0000000..f3e22fe --- /dev/null +++ b/runtimes/common/tests/test_foundry_model_retries.py @@ -0,0 +1,392 @@ +"""Bounded rate-limit recovery without replaying hosted invocations.""" + +from __future__ import annotations + +import asyncio +import json +from email.utils import formatdate +from types import SimpleNamespace + +import httpx +import pytest +from fastapi.testclient import TestClient + +from agentkit_serve_common import foundry_model_loop +from test_foundry_brokered_protocol import ( + CONTINUATION_AUTH, + _app, + _call, + _chat_response, + _continuation, + _spec, +) +from test_foundry_streaming import _created, _event, _exchange, _tool_response + + +class RejectedBody(httpx.AsyncByteStream): + def __init__(self): + self.read = False + self.closed = False + + async def __aiter__(self): + self.read = True + yield b"private-upstream-detail" + + async def aclose(self): + self.closed = True + + +def _reject(request, bodies, *, headers=None): + body = RejectedBody() + bodies.append(body) + return httpx.Response( + 429, + request=request, + headers={"x-request-id": "private-upstream-detail", **(headers or {})}, + stream=body, + ) + + +def _capture_waits(monkeypatch, bodies): + waits = [] + + async def sleep(delay): + assert bodies and all(body.closed and not body.read for body in bodies) + waits.append(delay) + await asyncio.sleep(0) + + # Keep the real ASGI/event-loop scheduling untouched while advancing only + # this model client's retry waits without spending minutes in each test. + monkeypatch.setattr( + foundry_model_loop, + "asyncio", + SimpleNamespace(sleep=sleep, to_thread=asyncio.to_thread), + ) + return waits + + +async def _payload(client, *, continuation): + payload = { + "input": "Read synthetic data", + "agent_session_id": "session-a", + "stream": True, + } + if continuation: + initial = await client.post("/responses", json={**payload, "stream": False}) + assert initial.status_code == 200, initial.text + body = initial.json() + payload = { + **_continuation( + body["id"], + _call(body)["call_id"], + {"approved": True, "output": {"ok": True}}, + ), + "agent_session_id": "session-a", + "stream": True, + } + return payload + + +@pytest.mark.parametrize("continuation", [False, True]) +@pytest.mark.parametrize("completion", ["text", "tool", "exhausted"]) +def test_hosted_rate_limit_retries_keep_one_ack_and_unchanged_model_input( + continuation, completion, monkeypatch, caplog +): + bodies, requests, credentials = [], [], [] + waits = _capture_waits(monkeypatch, bodies) + first_attempt = 2 if continuation else 1 + + async def auth(self): + credentials.append(len(credentials) + 1) + return {"Authorization": f"Bearer synthetic-attempt-{len(credentials)}"} + + monkeypatch.setattr(foundry_model_loop.BrokeredChatModelLoop, "_auth_headers", auth) + + def model(request): + requests.append(request) + if len(requests) < first_attempt: + return httpx.Response(200, request=request, json=_tool_response()) + if len(requests) < first_attempt + 2 or completion == "exhausted": + return _reject(request, bodies, headers={"retry-after": "60"}) + result = ( + _tool_response() + if completion == "tool" + else _chat_response({"role": "assistant", "content": "Verified response."}) + ) + return httpx.Response(200, request=request, json=result) + + async def exercise(): + async with httpx.AsyncClient(transport=httpx.MockTransport(model)) as upstream: + app = _app( + brokered_model_loop_enabled=True, + brokered_model_http_client=upstream, + ) + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://hosted" + ) as client: + payload = await _payload(client, continuation=continuation) + async with _exchange(app, payload) as (_, outgoing, request): + created = await _created(outgoing) + terminal = await _event(outgoing) + await asyncio.wait_for(request, 2) + assert outgoing.empty() + assert created["type"] == "response.created" + assert terminal["response"]["id"] == created["response"]["id"] + assert ( + terminal["response"]["created_at"] + == created["response"]["created_at"] + ) + assert terminal["response"]["agent_session_id"] == "session-a" + assert "private-upstream-detail" not in json.dumps(terminal) + if completion == "exhausted": + assert terminal["type"] == "response.failed" + assert terminal["response"]["error"] == { + "code": "ModelUnavailable", + "message": "model service is unavailable", + "upstream_status": 429, + } + else: + assert terminal["type"] == "response.completed" + assert terminal["response"]["output"][0]["type"] == ( + "function_call" if completion == "tool" else "message" + ) + if continuation: + replay = await client.post( + "/responses", + json={**payload, "stream": False}, + headers=CONTINUATION_AUTH, + ) + assert replay.status_code == 200 + assert replay.json()["id"] == terminal["response"]["id"] + + asyncio.run(exercise()) + assert len(requests) == first_attempt + 2 + attempts = requests[first_attempt - 1 :] + assert attempts[0].content == attempts[1].content == attempts[2].content + assert [request.headers["authorization"] for request in requests] == [ + f"Bearer synthetic-attempt-{number}" for number in credentials + ] + assert waits == [60, 60] + assert all(body.closed and not body.read for body in bodies) + assert "private-upstream-detail" not in caplog.text + + +@pytest.mark.parametrize( + ("headers", "expected"), + [ + ({"retry-after-ms": "60000", "retry-after": "1"}, 60), + ({"retry-after-ms": "1250"}, 1.25), + ({"retry-after-ms": "invalid", "retry-after": "2"}, 2), + ({"retry-after-ms": "nan", "retry-after": "3"}, 3), + ({"retry-after-ms": "-1", "retry-after": "4"}, 4), + ({"retry-after": "60"}, 60), + ({"retry-after": "0"}, 0), + ({"retry-after": formatdate(1_700_000_060, usegmt=True)}, 60), + ({"retry-after": formatdate(1_699_999_999, usegmt=True)}, 0), + ({"retry-after": "61"}, None), + ({"retry-after-ms": "60001", "retry-after": "1"}, None), + ({"retry-after": formatdate(1_700_000_061, usegmt=True)}, None), + ], +) +def test_model_retry_respects_server_delay_and_never_shortens_long_windows( + headers, expected, monkeypatch +): + monkeypatch.setattr(foundry_model_loop.time, "time", lambda: 1_700_000_000) + bodies, requests = [], [] + waits = _capture_waits(monkeypatch, bodies) + + def model(request): + requests.append(request) + if len(requests) == 1: + return _reject(request, bodies, headers=headers) + return httpx.Response( + 200, + request=request, + json=_chat_response({"role": "assistant", "content": "Recovered."}), + ) + + app = _app( + brokered_model_loop_enabled=True, + brokered_model_http_client=httpx.AsyncClient( + transport=httpx.MockTransport(model) + ), + ) + with TestClient(app) as client: + result = client.post("/responses", json={"input": "Read data"}) + if expected is None: + assert len(requests) == 1 and waits == [] + assert result.status_code == 503 + assert result.json()["error"]["upstream_status"] == 429 + else: + assert len(requests) == 2 and waits == [expected] + assert result.status_code == 200 + assert result.json()["output"][0]["content"][0]["text"] == "Recovered." + assert all(body.closed and not body.read for body in bodies) + + +@pytest.mark.parametrize("header", [None, "invalid", "nan", "inf", "-1"]) +def test_model_missing_or_malformed_retry_header_uses_bounded_backoff( + header, monkeypatch +): + bodies = [] + waits = _capture_waits(monkeypatch, bodies) + headers = ( + {} if header is None else {"retry-after-ms": header, "retry-after": header} + ) + + def model(request): + return _reject(request, bodies, headers=headers) + + app = _app( + brokered_model_loop_enabled=True, + brokered_model_http_client=httpx.AsyncClient( + transport=httpx.MockTransport(model) + ), + ) + with TestClient(app) as client: + result = client.post("/responses", json={"input": "Read data"}) + assert result.status_code == 503 + assert result.json()["error"]["upstream_status"] == 429 + assert len(bodies) == 3 and len(waits) == 2 + assert 0.75 <= waits[0] <= 1 and 1.5 <= waits[1] <= 2 + + +@pytest.mark.parametrize( + "failure", [400, 401, 403, 408, 409, 500, 503, "connect", "read", "json"] +) +def test_model_does_not_retry_other_rejections_or_ambiguous_failures( + failure, monkeypatch +): + requests = [] + waits = _capture_waits(monkeypatch, []) + + class IncompleteBody(httpx.AsyncByteStream): + async def __aiter__(self): + yield b'{"choices":' + raise httpx.ReadError("private-upstream-detail") + + def model(request): + requests.append(request) + if failure == "connect": + raise httpx.ConnectError("private-upstream-detail", request=request) + if failure == "read": + return httpx.Response(200, request=request, stream=IncompleteBody()) + return httpx.Response( + 200 if failure == "json" else failure, + request=request, + headers={"retry-after-ms": "1"}, + content=b"private-upstream-detail", + ) + + app = _app( + brokered_model_loop_enabled=True, + brokered_model_http_client=httpx.AsyncClient( + transport=httpx.MockTransport(model) + ), + ) + with TestClient(app) as client: + result = client.post("/responses", json={"input": "Read data"}) + assert result.status_code >= 400 + assert len(requests) == 1 and waits == [] + assert "private-upstream-detail" not in result.text + + +@pytest.mark.parametrize("continuation", [False, True]) +def test_hosted_disconnect_during_real_rate_limit_wait_stops_retry_and_releases_state( + continuation, tmp_path +): + async def exercise(): + requests, handlers = [], set() + rejected_connection_closed = asyncio.Event() + reject_at = 2 if continuation else 1 + + async def serve_model(reader, writer): + handler = asyncio.current_task() + handlers.add(handler) + try: + headers = await reader.readuntil(b"\r\n\r\n") + length = next( + int(line.split(b":", 1)[1]) + for line in headers.splitlines() + if line.lower().startswith(b"content-length:") + ) + requests.append(await reader.readexactly(length)) + if len(requests) == reject_at: + # A definitive 429 is enough. Its body must not hold the + # connection open during the actual 60-second retry wait. + writer.write( + b"HTTP/1.1 429 Too Many Requests\r\nRetry-After: 60\r\n" + b"Content-Length: 65536\r\nConnection: close\r\n\r\n" + ) + await writer.drain() + assert await reader.read(1) == b"" + rejected_connection_closed.set() + else: + response = ( + _tool_response() + if len(requests) < reject_at + else _chat_response( + {"role": "assistant", "content": "Recovered."} + ) + ) + body = json.dumps(response).encode() + writer.write( + b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n" + + f"Content-Length: {len(body)}\r\nConnection: close\r\n\r\n".encode() + + body + ) + await writer.drain() + finally: + writer.close() + await writer.wait_closed() + handlers.remove(handler) + + server = await asyncio.start_server(serve_model, "127.0.0.1", 0) + spec = _spec() + spec.model.base_url = ( + f"http://127.0.0.1:{server.sockets[0].getsockname()[1]}/v1" + ) + state_file = tmp_path / "responses.json" + try: + async with ( + server, + httpx.AsyncClient(trust_env=False, timeout=5) as upstream, + ): + app = _app( + spec, + brokered_model_loop_enabled=True, + brokered_model_http_client=upstream, + response_state_file=state_file, + max_pending_responses=1, + ) + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://hosted" + ) as client: + payload = await _payload(client, continuation=continuation) + async with _exchange(app, payload) as (incoming, outgoing, request): + created = await _created(outgoing) + assert created["type"] == "response.created" + await asyncio.wait_for(rejected_connection_closed.wait(), 2) + assert not request.done() and outgoing.empty() + await incoming.put({"type": "http.disconnect"}) + await asyncio.wait_for(request, 2) + assert outgoing.empty() and len(requests) == reject_at + # A new client operation can use the released initial slot or + # resume the explicitly cancelled continuation. No automatic + # model retry remains alive after the disconnected request. + result = await client.post( + "/responses", + json={**payload, "stream": False}, + headers=CONTINUATION_AUTH, + ) + assert result.status_code == 200, result.text + assert ( + result.json()["output"][0]["content"][0]["text"] == "Recovered." + ) + assert len(requests) == reject_at + 1 + finally: + for handler in tuple(handlers): + handler.cancel() + await asyncio.gather(*handlers, return_exceptions=True) + + asyncio.run(exercise()) diff --git a/runtimes/common/tests/test_foundry_streaming.py b/runtimes/common/tests/test_foundry_streaming.py index 7591e77..8a26496 100644 --- a/runtimes/common/tests/test_foundry_streaming.py +++ b/runtimes/common/tests/test_foundry_streaming.py @@ -143,7 +143,10 @@ async def _event(outgoing): for line in message["body"].splitlines() if line.startswith(b"data: ") ) - return json.loads(data) + event = json.loads(data) + assert type(event["sequence_number"]) is int + assert event["sequence_number"] == (0 if event["type"] == "response.created" else 1) + return event @pytest.mark.parametrize("continuation", [False, True]) From cc6745a6cf27fb4c688ef67f78d41d6b8f48c479 Mon Sep 17 00:00:00 2001 From: Sertac Ozercan Date: Fri, 11 Sep 2026 16:57:07 -0700 Subject: [PATCH 9/9] fix(foundry): sanitize model validation errors Signed-off-by: Sertac Ozercan --- .../common/agentkit_serve_common/foundry.py | 30 +++++-- .../common/tests/test_foundry_model_errors.py | 79 ++++++++++++++++--- 2 files changed, 92 insertions(+), 17 deletions(-) diff --git a/runtimes/common/agentkit_serve_common/foundry.py b/runtimes/common/agentkit_serve_common/foundry.py index 3e45f31..e79c2e9 100644 --- a/runtimes/common/agentkit_serve_common/foundry.py +++ b/runtimes/common/agentkit_serve_common/foundry.py @@ -84,6 +84,20 @@ _STATE_FILE_ENV = "AGENTKIT_FOUNDRY_RESPONSE_STATE_FILE" _FOUNDRY_SESSION_ENV = "FOUNDRY_AGENT_SESSION_ID" _TERMINAL_STATE_FULL = "state_full" +_BROKERED_MODEL_VALIDATION_ERRORS = { + "InvalidToolArguments": (400, "model tool arguments are invalid"), + "InvalidToolOutput": (400, "brokered tool output is invalid"), + "UnsafeBrokeredArguments": (400, "model tool arguments are unsafe"), + "UnsupportedBrokeredSchema": (400, "brokered tool schema is unsupported"), + "unknown_brokered_tool": (400, "model requested unknown brokered tool"), + "invalid_tool_call": (400, "model tool call is invalid"), + "unsupported_tool_call": (400, "model returned an unsupported tool call"), + "multiple_tool_calls_unsupported": (400, "brokered mode requires sequential tool calls"), + "tool_loop_limit_exceeded": (400, "model exceeded the tool call limit"), + "brokered_arguments_too_large": (413, "model tool arguments are too large"), + "brokered_output_too_large": (413, "brokered tool output is too large for model resume"), + "brokered_model_messages_too_large": (413, "model loop messages are too large"), +} def _new_response_id(previous_response_id: str | None = None) -> str: @@ -205,11 +219,13 @@ def _brokered_model_run_error(exc: Exception) -> JSONResponse: if normalized is not None: status, error = normalized return JSONResponse({"error": error}, status_code=status) - if type(exc.status) is int: - if 400 <= exc.status < 500: - return _error(str(exc), status=exc.status, code=exc.code) - if 500 <= exc.status <= 599: - status = exc.status + # Validation exceptions may include model-supplied names or argument paths. + validation = _BROKERED_MODEL_VALIDATION_ERRORS.get(exc.code) if type(exc.code) is str else None + if validation is not None: + status, message = validation + return _error(message, status=status, code=exc.code) + if type(exc.status) is int and 400 <= exc.status <= 599: + status = exc.status return _error("model resume failed", status=status, code="ModelResumeError") @@ -2280,7 +2296,7 @@ def _advance_brokered_state( except AgentRunError as exc: if not _reset_unfinalized_continuation(store, state, call_id=call_id): return _state_storage_error() - return _error(str(exc), status=exc.status, code=exc.code) + return _brokered_model_run_error(exc) payload = _function_call_response_payload(spec, response_id=response_id, call=call, previous_response_id=previous_response_id, usage=result.usage, created_at=created_at) following = deepcopy(state) if not following.response_calls: @@ -2992,7 +3008,7 @@ async def execute_responses( _validate_model_brokered_arguments(model_result.arguments) _validate_model_arguments_for_tool(model_result.arguments, tool) except AgentRunError as exc: - return _error(str(exc), status=exc.status, code=exc.code) + return _brokered_model_run_error(exc) if len(_canonical_output_json(model_result.arguments).encode("utf-8")) > max_argument_bytes: return _error( "brokered function_call arguments are too large for pending state", diff --git a/runtimes/common/tests/test_foundry_model_errors.py b/runtimes/common/tests/test_foundry_model_errors.py index bc748d8..3ef7c50 100644 --- a/runtimes/common/tests/test_foundry_model_errors.py +++ b/runtimes/common/tests/test_foundry_model_errors.py @@ -15,15 +15,16 @@ _app, _call, _continuation, + _spec, ) from test_foundry_streaming import _tool_response -def _streamed_error(app, *, continuation): +def _response_error(app, *, continuation, stream=True, status=400): payload = { "input": "Read synthetic data.", "agent_session_id": "synthetic-session", - "stream": True, + "stream": stream, } with TestClient(app) as client: if continuation: @@ -36,10 +37,14 @@ def _streamed_error(app, *, continuation): {"approved": True, "output": {"ok": True}}, ), "agent_session_id": "synthetic-session", - "stream": True, + "stream": stream, } response = client.post("/responses", json=payload, headers=CONTINUATION_AUTH) + assert "private-upstream-detail" not in response.text + if not stream: + assert response.status_code == status + return response.json()["error"] assert response.status_code == 200 frames = [ json.loads(line.removeprefix("data: ")) @@ -56,10 +61,60 @@ def _streamed_error(app, *, continuation): assert failed["agent_session_id"] == created["agent_session_id"] == "synthetic-session" assert failed["status"] == "failed" assert "error" not in frames[-1] - assert "private-upstream-detail" not in response.text return failed["error"] +@pytest.mark.parametrize("continuation", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize( + ("kind", "code", "message"), + [ + ("unknown-name", "unknown_brokered_tool", "model requested unknown brokered tool"), + ("duplicate-key", "InvalidToolArguments", "model tool arguments are invalid"), + ("unicode-path", "InvalidToolArguments", "model tool arguments are invalid"), + ("schema-path", "InvalidToolArguments", "model tool arguments are invalid"), + ("unsafe-path", "UnsafeBrokeredArguments", "model tool arguments are unsafe"), + ], +) +def test_model_tool_validation_never_reflects_names_keys_or_paths( + continuation, stream, kind, code, message, caplog +): + calls = 0 + marker = "private-upstream-detail" + spec = _spec() + payload = _tool_response() + function = payload["choices"][0]["message"]["tool_calls"][0]["function"] + if kind == "unknown-name": + function["name"] = marker + elif kind == "duplicate-key": + function["arguments"] = f'{{"{marker}":true,"{marker}":false}}' + elif kind == "unicode-path": + function["arguments"] = json.dumps({marker: "\ud800"}) + elif kind == "schema-path": + function["arguments"] = json.dumps({marker: True}) + spec.brokered_tools[0].parameters["additionalProperties"] = False + else: + function["arguments"] = json.dumps({marker: {"api_key": "synthetic-value"}}) + + def model(request): + nonlocal calls + calls += 1 + result = _tool_response() if continuation and calls == 1 else payload + return httpx.Response(200, request=request, json=result) + + app = _app( + spec, + brokered_model_loop_enabled=True, + brokered_model_http_client=httpx.AsyncClient(transport=httpx.MockTransport(model)), + ) + assert _response_error(app, continuation=continuation, stream=stream) == { + "code": code, + "message": message, + } + assert calls == (2 if continuation else 1) + assert marker not in caplog.text + + @pytest.mark.parametrize("continuation", [False, True]) @pytest.mark.parametrize( ("status", "code", "message"), @@ -98,7 +153,7 @@ def model(request): brokered_model_loop_enabled=True, brokered_model_http_client=httpx.AsyncClient(transport=httpx.MockTransport(model)), ) - error = _streamed_error(app, continuation=continuation) + error = _response_error(app, continuation=continuation) assert error == {"code": code, "message": message, "upstream_status": status} assert "private-upstream-detail" not in caplog.text @@ -121,7 +176,7 @@ def model(request): brokered_model_loop_enabled=True, brokered_model_http_client=httpx.AsyncClient(transport=httpx.MockTransport(model)), ) - error = _streamed_error(app, continuation=continuation) + error = _response_error(app, continuation=continuation) assert set(error) == {"code", "message"} assert error["code"] == ( "ModelUpstreamError" if kind == "transport" else "InvalidModelResponse" @@ -152,6 +207,8 @@ async def fail(self, messages, *, tools): ("ModelUpstreamError", "model service request failed"), ("InvalidModelResponse", "model service returned an invalid response"), ("ModelResponseTooLarge", "model response is too large to retain safely"), + ("InvalidToolArguments", "model tool arguments are invalid"), + ("UnsafeBrokeredArguments", "model tool arguments are unsafe"), ], ) def test_known_model_codes_use_fixed_messages_and_do_not_trust_status_attributes( @@ -160,7 +217,7 @@ def test_known_model_codes_use_fixed_messages_and_do_not_trust_status_attributes error = AgentRunError("private-upstream-detail", status=999, code=code) error.upstream_status = 401 app = _raise_model_error(monkeypatch, continuation=continuation, error=error) - assert _streamed_error(app, continuation=continuation) == { + assert _response_error(app, continuation=continuation) == { "code": code, "message": message, } @@ -175,7 +232,7 @@ def test_normalized_http_error_rejects_invalid_upstream_status( error = foundry_model_loop._normalized_model_http_error(401) error.upstream_status = upstream_status app = _raise_model_error(monkeypatch, continuation=continuation, error=error) - assert _streamed_error(app, continuation=continuation) == { + assert _response_error(app, continuation=continuation) == { "code": "ModelAuthRejected", "message": "model service rejected configured credentials", } @@ -186,6 +243,8 @@ def test_normalized_http_error_rejects_invalid_upstream_status( ("code", "status"), [ ("private-upstream-detail", 503), + ("private-upstream-detail", 400), + (["private-upstream-detail"], 413), (None, 999), (["private-upstream-detail"], "503"), ], @@ -196,7 +255,7 @@ def test_unknown_model_codes_and_statuses_are_sanitized( error = AgentRunError("private-upstream-detail", status=status, code=code) error.upstream_status = 401 app = _raise_model_error(monkeypatch, continuation=continuation, error=error) - assert _streamed_error(app, continuation=continuation) == { + assert _response_error(app, continuation=continuation) == { "code": "ModelResumeError", "message": "model resume failed", } @@ -210,7 +269,7 @@ def test_unexpected_model_exception_remains_generic(continuation, monkeypatch, c continuation=continuation, error=RuntimeError("private-upstream-detail"), ) - assert _streamed_error(app, continuation=continuation) == { + assert _response_error(app, continuation=continuation) == { "code": "ModelResumeError", "message": "model resume failed", }