diff --git a/Makefile b/Makefile index 7f55e4d..114d950 100644 --- a/Makefile +++ b/Makefile @@ -30,10 +30,16 @@ ifneq ($(strip $(BUILDER)),) BUILDER_FLAG := --builder $(BUILDER) endif -# The runtime adapter image (agentkit-serve) is linux/amd64 (its uv base is -# amd64-only), so the test agent is built for the same platform. +# PLATFORM keeps the runtime adapter, generated agent image, and local run on +# one target architecture. linux/amd64 remains the compatibility default; +# override it (for example, PLATFORM=linux/arm64) for native ARM builds. PLATFORM ?= linux/amd64 +# The local run target binds the service across the container namespace, which +# requires bearer auth. Override this non-production token when desired, e.g. +# `make run-test-agent LOCAL_AUTH_TOKEN=my-local-token`. +LOCAL_AUTH_TOKEN ?= agentkit-local-dev-token + # RUNTIME selects which runtime adapter the test-agent targets: `pydantic-ai` # (default), Microsoft Agent Framework (`maf` alias or canonical name), or # LangGraph (`langgraph`). build-test-agent derives the adapter image, fixture, @@ -88,20 +94,23 @@ build-agentkit: # keeps the context small). .PHONY: build-serve build-serve: - docker buildx build . -f runtimes/pydantic-ai/Dockerfile -t agentkit-serve:$(TAG) --load + docker buildx build . -f runtimes/pydantic-ai/Dockerfile \ + --platform $(PLATFORM) -t agentkit-serve:$(TAG) --load # Build the Microsoft Agent Framework runtime adapter (agentkit-serve-maf) image. # This is the LLB base used when an agentkitfile selects # `runtime: microsoft-agent-framework` (alias `maf`). .PHONY: build-serve-maf build-serve-maf: - docker buildx build . -f runtimes/microsoft-agent-framework/Dockerfile -t agentkit-serve-maf:$(TAG) --load + docker buildx build . -f runtimes/microsoft-agent-framework/Dockerfile \ + --platform $(PLATFORM) -t agentkit-serve-maf:$(TAG) --load # Build the LangGraph runtime adapter (agentkit-serve-langgraph) image. # This is the LLB base used when an agentkitfile selects `runtime: langgraph`. .PHONY: build-serve-langgraph build-serve-langgraph: - docker buildx build . -f runtimes/langgraph/Dockerfile -t agentkit-serve-langgraph:$(TAG) --load + docker buildx build . -f runtimes/langgraph/Dockerfile \ + --platform $(PLATFORM) -t agentkit-serve-langgraph:$(TAG) --load # Build a test agent against the LOCAL frontend (BUILDKIT_SYNTAX) and the LOCAL # adapter (--build-arg adapter). The runtime, fixture, adapter image, and output @@ -116,8 +125,17 @@ build-test-agent: --platform $(PLATFORM) \ -t $(AGENT_IMAGE) --load --provenance=false -# Run the built test agent. Expects OPENAI_API_KEY in the environment; the agent -# binds 127.0.0.1 inside the container and serves the OpenAI /v1 façade on :8080. +# Run the built test agent. Expects OPENAI_API_KEY in the environment and forwards +# it by name so its value never appears in the command line. The service must bind +# 0.0.0.0 inside the container to cross the container network namespace, while the +# published host socket remains loopback-only. The expanded curl command records +# the configurable local bearer token required by that non-loopback container bind. .PHONY: run-test-agent run-test-agent: - docker run --rm --platform $(PLATFORM) -p 127.0.0.1:8080:8080 -e OPENAI_API_KEY=$$OPENAI_API_KEY $(AGENT_IMAGE) + @printf '%s\n' 'curl -fsS -H "Authorization: Bearer $(LOCAL_AUTH_TOKEN)" http://127.0.0.1:8080/v1/models' + docker run --rm --platform $(PLATFORM) \ + -p 127.0.0.1:8080:8080 \ + -e AGENTKIT_BIND=0.0.0.0 \ + -e AGENTKIT_AUTH_TOKEN="$(LOCAL_AUTH_TOKEN)" \ + -e OPENAI_API_KEY \ + $(AGENT_IMAGE) diff --git a/deploy/foundry/scripts/foundry_brokered_conformance.sh b/deploy/foundry/scripts/foundry_brokered_conformance.sh index 0c4f29e..ac4b42a 100755 --- a/deploy/foundry/scripts/foundry_brokered_conformance.sh +++ b/deploy/foundry/scripts/foundry_brokered_conformance.sh @@ -29,6 +29,8 @@ Optional: AGENTKIT_EXPECTED_CALL_ID_PREFIX Optional required call_id prefix, e.g. call_. AGENTKIT_EXPECTED_FINAL_TEXT Optional exact final assistant text; defaults are derived for SDK/AgentKit fixtures. AGENTKIT_CONTINUATION_PROOF Optional x-agentkit-brokered-continuation-proof header. + AGENTKIT_CONTINUATION_PROOF_BODY Optional proof sent only in the live continuation body. + It is omitted from the archived request transcript. EOF } @@ -100,7 +102,15 @@ continuation_response="$transcript_dir/04-continuation-response.json" summary_file="$transcript_dir/summary.json" expected_output_file="$transcript_dir/.expected-output.json" expected_final_text_file="$transcript_dir/.expected-final-text.txt" -trap 'rm -f -- "$expected_output_file" "$expected_final_text_file"' EXIT +rm -f "$continuation_response" "$summary_file" "$summary_file.tmp" "$expected_output_file" "$expected_final_text_file" +continuation_response_wire="$(mktemp "${TMPDIR:-/tmp}/agentkit-foundry-continuation-response.XXXXXX")" +continuation_request_wire="$(mktemp "${TMPDIR:-/tmp}/agentkit-foundry-continuation-request.XXXXXX")" + +cleanup() { + rm -f "$continuation_response_wire" "$continuation_request_wire" "$summary_file.tmp" \ + "$expected_output_file" "$expected_final_text_file" +} +trap cleanup EXIT printf '%s' "$conformance_output" >"$expected_output_file" printf '%s' "$expected_final_text" >"$expected_final_text_file" @@ -118,7 +128,12 @@ printf '%s\n' "$initial_curl_config" | curl -fsS \ -d "@$initial_request" >"$initial_response" unset initial_curl_config -read -r response_id call_id < <(EXPECTED_TOOL_NAME="$expected_tool_name" EXPECTED_ARGUMENTS="$expected_arguments" EXPECTED_CALL_ID="$expected_call_id" EXPECTED_CALL_ID_PREFIX="$expected_call_id_prefix" python3 - "$initial_response" <<'PY' +EXPECTED_TOOL_NAME="$expected_tool_name" \ +EXPECTED_ARGUMENTS="$expected_arguments" \ +EXPECTED_CALL_ID="$expected_call_id" \ +EXPECTED_CALL_ID_PREFIX="$expected_call_id_prefix" \ +CONFORMANCE_OUTPUT="$conformance_output" \ +python3 - "$initial_response" <<'PY' >"$continuation_request" import json import os import sys @@ -145,26 +160,40 @@ if expected_call_id != "auto": if expected_call_id_prefix: assert call_id.startswith(expected_call_id_prefix), call assert json.loads(call.get("arguments") or "{}") == expected_arguments, call -print(response_id, call_id) -PY -) -PREVIOUS_RESPONSE_ID="$response_id" CALL_ID="$call_id" CONFORMANCE_OUTPUT="$conformance_output" python3 - <<'PY' >"$continuation_request" -import json -import os -# Validate the configured output is JSON before placing it in the Responses item. -json.loads(os.environ["CONFORMANCE_OUTPUT"]) -print(json.dumps({ - "previous_response_id": os.environ["PREVIOUS_RESPONSE_ID"], +conformance_output = os.environ["CONFORMANCE_OUTPUT"] +json.loads(conformance_output) +continuation = { + "previous_response_id": response_id, "input": [{ "type": "function_call_output", - "call_id": os.environ["CALL_ID"], - "output": os.environ["CONFORMANCE_OUTPUT"], + "call_id": call_id, + "output": conformance_output, "status": "completed", }], -}, separators=(",", ":"))) +} +agent_session_id = body.get("agent_session_id") +if agent_session_id is not None: + assert isinstance(agent_session_id, str) and agent_session_id.strip(), body + continuation["agent_session_id"] = agent_session_id +print(json.dumps(continuation, separators=(",", ":"))) PY +cp "$continuation_request" "$continuation_request_wire" +if [[ -n "${AGENTKIT_CONTINUATION_PROOF_BODY:-}" ]]; then + AGENTKIT_CONTINUATION_PROOF_BODY="$AGENTKIT_CONTINUATION_PROOF_BODY" \ + python3 -c ' +import json +import os +import sys +from pathlib import Path + +body = json.loads(Path(sys.argv[1]).read_text(encoding="utf-8")) +body["brokered_continuation_proof"] = os.environ["AGENTKIT_CONTINUATION_PROOF_BODY"] +print(json.dumps(body, separators=(",", ":"))) +' "$continuation_request" >"$continuation_request_wire" +fi + continuation_curl_config="$({ write_curl_header "Authorization: Bearer ${token}" if [[ -n "${AGENTKIT_CONTINUATION_PROOF:-}" ]]; then @@ -175,9 +204,75 @@ printf '%s\n' "$continuation_curl_config" | curl -fsS \ --config - \ -H 'content-type: application/json' \ "$AGENT_RESPONSES_ENDPOINT" \ - -d "@$continuation_request" >"$continuation_response" + -d "@$continuation_request_wire" >"$continuation_response_wire" unset continuation_curl_config +AGENTKIT_CONTINUATION_PROOF="${AGENTKIT_CONTINUATION_PROOF:-}" \ +AGENTKIT_CONTINUATION_PROOF_BODY="${AGENTKIT_CONTINUATION_PROOF_BODY:-}" \ +python3 - "$continuation_response_wire" <<'PY' +import json +import os +import re +import sys +from pathlib import Path + +_JSON_ESCAPE_RE = re.compile(r'(?:\\u[0-9A-Fa-f]{4}){1,2}|\\["\\/bfnrt]') + +def decoded_fragment_contains_proof(value, proof): + seen = set() + while value not in seen: + seen.add(value) + if proof in value: + return True + + def decode_escape(match): + try: + return json.loads(f'"{match.group(0)}"') + except (ValueError, RecursionError): + return match.group(0) + + decoded = _JSON_ESCAPE_RE.sub(decode_escape, value) + if decoded == value: + return False + value = decoded + return False + + +def contains_proof(value, proof): + pending = [value] + decoded_strings = set() + while pending: + current = pending.pop() + if isinstance(current, str): + if decoded_fragment_contains_proof(current, proof): + return True + if current in decoded_strings: + continue + decoded_strings.add(current) + try: + pending.append(json.loads(current)) + except (ValueError, RecursionError): + pass + elif isinstance(current, list): + pending.extend(current) + elif isinstance(current, dict): + pending.extend(current.keys()) + pending.extend(current.values()) + return False + + +raw = Path(sys.argv[1]).read_text(encoding="utf-8") +try: + decoded = json.loads(raw) +except json.JSONDecodeError: + decoded = None +for name in ("AGENTKIT_CONTINUATION_PROOF", "AGENTKIT_CONTINUATION_PROOF_BODY"): + proof = os.environ.get(name, "") + if proof and (contains_proof(raw, proof) or contains_proof(decoded, proof)): + raise SystemExit("gateway response contains continuation proof; refusing to archive transcript") +PY +mv "$continuation_response_wire" "$continuation_response" + verifier_args=( "$transcript_dir" --expected-tool-name "$expected_tool_name" @@ -191,8 +286,7 @@ if [[ -n "$expected_call_id_prefix" ]]; then verifier_args+=(--expected-call-id-prefix "$expected_call_id_prefix") fi python3 deploy/foundry/scripts/verify_brokered_transcript.py "${verifier_args[@]}" >"$summary_file.tmp" -rm -f "$summary_file.tmp" -rm -f "$expected_output_file" "$expected_final_text_file" +cleanup trap - EXIT echo "Foundry brokered conformance passed. Sanitized transcript: ${transcript_dir}" diff --git a/deploy/foundry/scripts/verify_brokered_transcript.py b/deploy/foundry/scripts/verify_brokered_transcript.py index 9fbd5d4..ca7a356 100755 --- a/deploy/foundry/scripts/verify_brokered_transcript.py +++ b/deploy/foundry/scripts/verify_brokered_transcript.py @@ -122,7 +122,6 @@ def verify_transcript( expected_call_id: str = "call_conformance_1", expected_call_id_prefix: str | None = None, ) -> dict[str, Any]: - _require(expected_final_text is not None, "expected final text is required for transcript verification") root = Path(transcript_dir) expected_arguments = _parse_json_lossless(expected_arguments_json) expected_output = _parse_json_lossless(expected_output_json) @@ -132,6 +131,10 @@ def verify_transcript( continuation_response = _load_json(root / "04-continuation-response.json") _require(isinstance(initial_request, dict), "initial request must be a JSON object") + _require( + "brokered_continuation_proof" not in initial_request, + "sanitized initial request must not archive a continuation proof", + ) _require("tools" not in initial_request, "initial request must not contain request-level tools") _require("input" in initial_request, "initial request must contain input") @@ -140,6 +143,12 @@ def verify_transcript( initial_response_id = initial_response.get("id") _require(isinstance(initial_response_id, str) and initial_response_id.startswith("caresp_"), "initial response id must start with caresp_") _require(not initial_response_id.startswith("resp_"), "initial response id must not use old resp_ format") + agent_session_id = initial_response.get("agent_session_id") + if agent_session_id is not None: + _require( + isinstance(agent_session_id, str) and bool(agent_session_id.strip()), + "initial response agent_session_id must be a non-empty string", + ) output = initial_response.get("output") _require(isinstance(output, list) and len(output) == 1, "initial response output must contain exactly one item") call = output[0] @@ -161,7 +170,16 @@ def verify_transcript( _require(_strict_json_equal(parsed_arguments, expected_arguments), f"function_call arguments must be {expected_arguments}") _require(isinstance(continuation_request, dict), "continuation request must be a JSON object") + _require( + "brokered_continuation_proof" not in continuation_request, + "sanitized continuation request must not archive a continuation proof", + ) _require(continuation_request.get("previous_response_id") == initial_response_id, "continuation previous_response_id must match initial id") + if agent_session_id is not None: + _require( + continuation_request.get("agent_session_id") == agent_session_id, + "continuation agent_session_id must match the initial response", + ) continuation_input = continuation_request.get("input") _require(isinstance(continuation_input, list) and len(continuation_input) == 1, "continuation input must contain exactly one item") continuation_item = continuation_input[0] @@ -181,8 +199,8 @@ def verify_transcript( _require(isinstance(continuation_response_id, str) and continuation_response_id.startswith("caresp_"), "continuation response id must start with caresp_") _require(continuation_response_id != initial_response_id, "continuation response id must differ from initial response id") final_text = _message_text(continuation_response, response_id=continuation_response_id) - if expected_final_text is not None: - _require(final_text == expected_final_text, "final message text did not match expected conformance result") + _require(expected_final_text is not None, "expected final text is required for transcript verification") + _require(final_text == expected_final_text, "final message text did not match expected conformance result") return { "initial_response_id": initial_response_id, diff --git a/docs/agent-abi.md b/docs/agent-abi.md index 499bb9d..ff0df13 100644 --- a/docs/agent-abi.md +++ b/docs/agent-abi.md @@ -27,7 +27,7 @@ model: provider: openai-compatible baseURL: https://api.openai.com/v1 name: gpt-4o-mini - apiKeyEnv: OPENAI_API_KEY + apiKeyEnv: "OPENAI_API_KEY" # Optional future generic model auth. Capability-gated; apiKeyEnv remains the # normal v0 model-auth path. # auth: diff --git a/go.mod b/go.mod index 312353c..c52ed23 100644 --- a/go.mod +++ b/go.mod @@ -11,6 +11,7 @@ require ( github.com/opencontainers/image-spec v1.1.1 github.com/pkg/errors v0.9.1 github.com/sirupsen/logrus v1.9.4 + github.com/tonistiigi/fsutil v0.0.0-20251211185533-a2aa163d723f golang.org/x/sync v0.20.0 google.golang.org/grpc v1.81.1 ) @@ -43,7 +44,6 @@ require ( github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 // indirect github.com/secure-systems-lab/go-securesystemslib v0.10.0 // indirect github.com/shibumi/go-pathspec v1.3.0 // indirect - github.com/tonistiigi/fsutil v0.0.0-20251211185533-a2aa163d723f // indirect github.com/tonistiigi/go-csvvalue v0.0.0-20240814133006-030d3b2625d0 // indirect go.opentelemetry.io/auto/sdk v1.2.1 // indirect go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.68.0 // indirect diff --git a/pkg/agentkit/abi/render.go b/pkg/agentkit/abi/render.go index 05e0358..e30f361 100644 --- a/pkg/agentkit/abi/render.go +++ b/pkg/agentkit/abi/render.go @@ -3,6 +3,7 @@ package abi import ( + "bytes" "encoding/json" "math" "math/big" @@ -34,6 +35,26 @@ func (n yamlNumber) MarshalYAML() ([]byte, error) { return []byte(n), nil } +// yamlString centralizes scalar rendering for every string in the ABI. Always +// emitting a quoted scalar prevents YAML syntax and implicit type resolution +// from changing string data across the Go and Python readers. +type yamlString string + +func (s yamlString) MarshalYAML() ([]byte, error) { + return []byte(strconv.Quote(string(s))), nil +} + +func yamlStrings(values []string) []yamlString { + if len(values) == 0 { + return nil + } + out := make([]yamlString, len(values)) + for i, value := range values { + out[i] = yamlString(value) + } + return out +} + func yamlFloat(value float64, bitSize int) yamlNumber { rendered := strconv.FormatFloat(value, 'f', -1, bitSize) if value == 0 && math.Signbit(value) { @@ -43,65 +64,65 @@ func yamlFloat(value float64, bitSize int) yamlNumber { } type abiMetadata struct { - Name string `yaml:"name"` + Name yamlString `yaml:"name"` } type abiModel struct { - Provider string `yaml:"provider"` - BaseURL string `yaml:"baseURL"` - Name string `yaml:"name"` - APIKeyEnv string `yaml:"apiKeyEnv,omitempty"` - Auth *abiAuth `yaml:"auth,omitempty"` + Provider yamlString `yaml:"provider"` + BaseURL yamlString `yaml:"baseURL"` + Name yamlString `yaml:"name"` + APIKeyEnv yamlString `yaml:"apiKeyEnv,omitempty"` + Auth *abiAuth `yaml:"auth,omitempty"` } type abiToolHeader struct { - Name string `yaml:"name"` - Value string `yaml:"value,omitempty"` - ValueEnv string `yaml:"valueEnv,omitempty"` + Name yamlString `yaml:"name"` + Value yamlString `yaml:"value,omitempty"` + ValueEnv yamlString `yaml:"valueEnv,omitempty"` } type abiAuth struct { - Type string `yaml:"type"` - TokenEnv string `yaml:"tokenEnv,omitempty"` - Audience string `yaml:"audience,omitempty"` + Type yamlString `yaml:"type"` + TokenEnv yamlString `yaml:"tokenEnv,omitempty"` + Audience yamlString `yaml:"audience,omitempty"` } type abiTool struct { - Name string `yaml:"name"` - Type string `yaml:"type,omitempty"` - Transport string `yaml:"transport,omitempty"` - Command []string `yaml:"command,omitempty"` - URLEnv string `yaml:"urlEnv,omitempty"` + Name yamlString `yaml:"name"` + Type yamlString `yaml:"type,omitempty"` + Transport yamlString `yaml:"transport,omitempty"` + Command []yamlString `yaml:"command,omitempty"` + URLEnv yamlString `yaml:"urlEnv,omitempty"` Headers []abiToolHeader `yaml:"headers,omitempty"` Auth *abiAuth `yaml:"auth,omitempty"` - Approval string `yaml:"approval,omitempty"` - Env []string `yaml:"env,omitempty"` + Approval yamlString `yaml:"approval,omitempty"` + Env []yamlString `yaml:"env,omitempty"` } type abiBrokeredTool struct { - Name string `yaml:"name"` - Description string `yaml:"description"` - BrokeredClass string `yaml:"brokeredClass"` - Parameters map[string]any `yaml:"parameters"` - SchemaDigest string `yaml:"schemaDigest,omitempty"` + Name yamlString `yaml:"name"` + Description yamlString `yaml:"description"` + BrokeredClass yamlString `yaml:"brokeredClass"` + Parameters any `yaml:"parameters"` + SchemaDigest yamlString `yaml:"schemaDigest,omitempty"` } type abiEnvVar struct { - Name string `yaml:"name"` - Required bool `yaml:"required,omitempty"` + Name yamlString `yaml:"name"` + Required bool `yaml:"required,omitempty"` } type abiContextProvider struct { - Name string `yaml:"name,omitempty"` - Type string `yaml:"type"` - Source string `yaml:"source,omitempty"` - Path string `yaml:"path,omitempty"` - ToolRef string `yaml:"toolRef,omitempty"` - Index string `yaml:"index,omitempty"` - EndpointEnv string `yaml:"endpointEnv,omitempty"` - IndexEnv string `yaml:"indexEnv,omitempty"` - StoreNameEnv string `yaml:"storeNameEnv,omitempty"` - Auth *abiAuth `yaml:"auth,omitempty"` + Name yamlString `yaml:"name,omitempty"` + Type yamlString `yaml:"type"` + Source yamlString `yaml:"source,omitempty"` + Path yamlString `yaml:"path,omitempty"` + ToolRef yamlString `yaml:"toolRef,omitempty"` + Index yamlString `yaml:"index,omitempty"` + EndpointEnv yamlString `yaml:"endpointEnv,omitempty"` + IndexEnv yamlString `yaml:"indexEnv,omitempty"` + StoreNameEnv yamlString `yaml:"storeNameEnv,omitempty"` + Auth *abiAuth `yaml:"auth,omitempty"` } type abiContext struct { @@ -110,10 +131,10 @@ type abiContext struct { type abiObservability struct { OTel struct { - EndpointEnv string `yaml:"endpointEnv,omitempty"` + EndpointEnv yamlString `yaml:"endpointEnv,omitempty"` } `yaml:"otel,omitempty"` Logs struct { - LevelEnv string `yaml:"levelEnv,omitempty"` + LevelEnv yamlString `yaml:"levelEnv,omitempty"` } `yaml:"logs,omitempty"` } @@ -123,10 +144,10 @@ type abiExpose struct { } type abiAgent struct { - ABIVersion string `yaml:"abiVersion"` + ABIVersion yamlString `yaml:"abiVersion"` Metadata abiMetadata `yaml:"metadata"` Model abiModel `yaml:"model"` - Instructions string `yaml:"instructions"` + Instructions yamlString `yaml:"instructions"` Tools []abiTool `yaml:"tools"` BrokeredTools []abiBrokeredTool `yaml:"brokeredTools,omitempty"` Env []abiEnvVar `yaml:"env,omitempty"` @@ -193,60 +214,71 @@ func expandJSONNumber(value string) string { return sign + out } -func copyMap(in map[string]any) map[string]any { +func isNegativeJSONZero(value string) bool { + if !strings.HasPrefix(value, "-") { + return false + } + coefficient := value[1:] + if exponent := strings.IndexAny(coefficient, "eE"); exponent >= 0 { + coefficient = coefficient[:exponent] + } + sawZero := false + for _, char := range coefficient { + switch char { + case '0': + sawZero = true + case '.': + default: + return false + } + } + return sawZero +} + +func yamlJSONNumber(value json.Number) yamlNumber { + raw := value.String() + if isNegativeJSONZero(raw) { + return yamlNumber("-0.0") + } + return yamlNumber(expandJSONNumber(raw)) +} + +func copyMap(in map[string]any) map[any]any { if in == nil { return nil } - out := make(map[string]any, len(in)) - for k, v := range in { - out[k] = copyAny(v) + out := make(map[any]any, len(in)) + for key, value := range in { + out[yamlString(key)] = copyAny(value) } return out } +// normalizeJSONValue routes brokered schemas through encoding/json, matching the +// representation used by Go validation and digesting (including []byte base64). +func normalizeJSONValue(v any) (any, error) { + encoded, err := json.Marshal(v) + if err != nil { + return nil, err + } + decoder := json.NewDecoder(bytes.NewReader(encoded)) + decoder.UseNumber() + var normalized any + if err := decoder.Decode(&normalized); err != nil { + return nil, err + } + return copyAny(normalized), nil +} + func copyAny(v any) any { switch typed := v.(type) { + case string: + return yamlString(typed) case map[string]any: if typed == nil { return nil } return copyMap(typed) - case map[string]string: - if typed == nil { - return nil - } - out := make(map[string]string, len(typed)) - for key, value := range typed { - out[key] = value - } - return out - case map[string]int: - if typed == nil { - return nil - } - out := make(map[string]int, len(typed)) - for key, value := range typed { - out[key] = value - } - return out - case map[string]float64: - if typed == nil { - return nil - } - out := make(map[string]any, len(typed)) - for key, value := range typed { - out[key] = yamlFloat(value, 64) - } - return out - case map[string]bool: - if typed == nil { - return nil - } - out := make(map[string]bool, len(typed)) - for key, value := range typed { - out[key] = value - } - return out case []any: if typed == nil { return nil @@ -256,36 +288,8 @@ func copyAny(v any) any { out[i] = copyAny(item) } return out - case []string: - if typed == nil { - return nil - } - return append([]string(nil), typed...) - case []int: - if typed == nil { - return nil - } - return append([]int(nil), typed...) - case []float64: - if typed == nil { - return nil - } - out := make([]any, len(typed)) - for i, item := range typed { - out[i] = yamlFloat(item, 64) - } - return out - case []bool: - if typed == nil { - return nil - } - return append([]bool(nil), typed...) - case float32: - return yamlFloat(float64(typed), 32) - case float64: - return yamlFloat(typed, 64) case json.Number: - return yamlNumber(expandJSONNumber(typed.String())) + return yamlJSONNumber(typed) default: return copyReflectedJSON(typed) } @@ -313,10 +317,10 @@ func copyReflectedJSON(value any) any { if rv.Type().Key().Kind() != reflect.String { return value } - out := make(map[string]any, rv.Len()) + out := make(map[any]any, rv.Len()) iter := rv.MapRange() for iter.Next() { - out[iter.Key().String()] = copyAny(iter.Value().Interface()) + out[yamlString(iter.Key().String())] = copyAny(iter.Value().Interface()) } return out case reflect.Slice: @@ -350,58 +354,62 @@ func copyReflectedJSON(value any) any { // output is byte-compatible with agentkit-serve's strict (extra=forbid) reader. func Render(agent effective.Agent) ([]byte, error) { out := abiAgent{ - ABIVersion: Version, - Metadata: abiMetadata{Name: agent.Metadata.Name}, + ABIVersion: yamlString(Version), + Metadata: abiMetadata{Name: yamlString(agent.Metadata.Name)}, Model: abiModel{ - Provider: agent.Model.Provider, - BaseURL: agent.Model.BaseURL, - Name: agent.Model.Name, - APIKeyEnv: agent.Model.APIKeyEnv, + Provider: yamlString(agent.Model.Provider), + BaseURL: yamlString(agent.Model.BaseURL), + Name: yamlString(agent.Model.Name), + APIKeyEnv: yamlString(agent.Model.APIKeyEnv), }, - Instructions: agent.Instructions, + Instructions: yamlString(agent.Instructions), Tools: make([]abiTool, 0, len(agent.Tools)), Env: make([]abiEnvVar, 0, len(agent.Env)), Expose: abiExpose{OpenAI: agent.Expose.OpenAI, Port: agent.Expose.Port}, } if agent.Model.Auth != nil { - out.Model.Auth = &abiAuth{Type: agent.Model.Auth.Type, TokenEnv: agent.Model.Auth.TokenEnv, Audience: agent.Model.Auth.Audience} + out.Model.Auth = &abiAuth{Type: yamlString(agent.Model.Auth.Type), TokenEnv: yamlString(agent.Model.Auth.TokenEnv), Audience: yamlString(agent.Model.Auth.Audience)} } for _, t := range agent.Tools { tool := abiTool{ - Name: t.Name, - Type: t.Type, - Transport: t.Transport, - Command: t.Command, - URLEnv: t.URLEnv, + Name: yamlString(t.Name), + Type: yamlString(t.Type), + Transport: yamlString(t.Transport), + Command: yamlStrings(t.Command), + URLEnv: yamlString(t.URLEnv), Headers: make([]abiToolHeader, 0, len(t.Headers)), - Approval: t.Approval, - Env: t.Env, + Approval: yamlString(t.Approval), + Env: yamlStrings(t.Env), } for _, h := range t.Headers { - tool.Headers = append(tool.Headers, abiToolHeader{Name: h.Name, Value: h.Value, ValueEnv: h.ValueEnv}) + tool.Headers = append(tool.Headers, abiToolHeader{Name: yamlString(h.Name), Value: yamlString(h.Value), ValueEnv: yamlString(h.ValueEnv)}) } if len(tool.Headers) == 0 { tool.Headers = nil } if t.Auth != nil { - tool.Auth = &abiAuth{Type: t.Auth.Type, TokenEnv: t.Auth.TokenEnv, Audience: t.Auth.Audience} + tool.Auth = &abiAuth{Type: yamlString(t.Auth.Type), TokenEnv: yamlString(t.Auth.TokenEnv), Audience: yamlString(t.Auth.Audience)} } out.Tools = append(out.Tools, tool) } for _, t := range agent.BrokeredTools { + parameters, err := normalizeJSONValue(t.Parameters) + if err != nil { + return nil, err + } out.BrokeredTools = append(out.BrokeredTools, abiBrokeredTool{ - Name: t.Name, - Description: t.Description, - BrokeredClass: t.BrokeredClass, - Parameters: copyMap(t.Parameters), - SchemaDigest: t.SchemaDigest, + Name: yamlString(t.Name), + Description: yamlString(t.Description), + BrokeredClass: yamlString(t.BrokeredClass), + Parameters: parameters, + SchemaDigest: yamlString(t.SchemaDigest), }) } if len(out.BrokeredTools) == 0 { out.BrokeredTools = nil } for _, e := range agent.Env { - out.Env = append(out.Env, abiEnvVar{Name: e.Name, Required: e.Required}) + out.Env = append(out.Env, abiEnvVar{Name: yamlString(e.Name), Required: e.Required}) } if len(out.Env) == 0 { out.Env = nil @@ -410,18 +418,18 @@ func Render(agent effective.Agent) ([]byte, error) { ctx := &abiContext{Providers: make([]abiContextProvider, 0, len(agent.Context.Providers))} for _, provider := range agent.Context.Providers { p := abiContextProvider{ - Name: provider.Name, - Type: provider.Type, - Source: provider.Source, - Path: provider.Path, - ToolRef: provider.ToolRef, - Index: provider.Index, - EndpointEnv: provider.EndpointEnv, - IndexEnv: provider.IndexEnv, - StoreNameEnv: provider.StoreNameEnv, + Name: yamlString(provider.Name), + Type: yamlString(provider.Type), + Source: yamlString(provider.Source), + Path: yamlString(provider.Path), + ToolRef: yamlString(provider.ToolRef), + Index: yamlString(provider.Index), + EndpointEnv: yamlString(provider.EndpointEnv), + IndexEnv: yamlString(provider.IndexEnv), + StoreNameEnv: yamlString(provider.StoreNameEnv), } if provider.Auth != nil { - p.Auth = &abiAuth{Type: provider.Auth.Type, TokenEnv: provider.Auth.TokenEnv, Audience: provider.Auth.Audience} + p.Auth = &abiAuth{Type: yamlString(provider.Auth.Type), TokenEnv: yamlString(provider.Auth.TokenEnv), Audience: yamlString(provider.Auth.Audience)} } ctx.Providers = append(ctx.Providers, p) } @@ -429,8 +437,8 @@ func Render(agent effective.Agent) ([]byte, error) { } if agent.Observability.OTel.EndpointEnv != "" || agent.Observability.Logs.LevelEnv != "" { obs := &abiObservability{} - obs.OTel.EndpointEnv = agent.Observability.OTel.EndpointEnv - obs.Logs.LevelEnv = agent.Observability.Logs.LevelEnv + obs.OTel.EndpointEnv = yamlString(agent.Observability.OTel.EndpointEnv) + obs.Logs.LevelEnv = yamlString(agent.Observability.Logs.LevelEnv) out.Observability = obs } diff --git a/pkg/agentkit/abi/render_test.go b/pkg/agentkit/abi/render_test.go index 2d409db..b56b15e 100644 --- a/pkg/agentkit/abi/render_test.go +++ b/pkg/agentkit/abi/render_test.go @@ -16,17 +16,22 @@ import ( // so gosec's G101 string-literal credential heuristic does not false-positive on // the struct literal below. const ( - testAPIKeyEnvName = "OPENAI_API_KEY" //nolint:gosec // G101: env var NAME, not a credential - testInstructions = "Be helpful and cite sources." + testAPIKeyEnvName = "OPENAI_API_KEY" //nolint:gosec // G101: env var NAME, not a credential + testInstructions = "Be helpful and cite sources." + testHelloBase64 = "SGVsbG8=" + testYAMLPositiveInfinity = ".inf" ) const ( - jsonSchemaTypeKey = "type" - jsonSchemaTypeObject = "object" - jsonSchemaTypeNumber = "number" - jsonSchemaPropertiesKey = "properties" - jsonSchemaMinimumKey = "minimum" - jsonSchemaDefaultKey = "default" + jsonSchemaTypeKey = "type" + jsonSchemaDescriptionKey = "description" + jsonSchemaDefaultKey = "default" + jsonSchemaEnumKey = "enum" + jsonSchemaObject = "object" + jsonSchemaString = "string" + jsonSchemaNumber = "number" + jsonSchemaPropertiesKey = "properties" + jsonSchemaMinimumKey = "minimum" ) func sampleConfig() *config.AgentConfig { @@ -97,8 +102,8 @@ func TestRenderAgentYAMLShape(t *testing.T) { t.Errorf("rendered agent.yaml unexpectedly contains %q\n---\n%s", k, s) } } - if !strings.Contains(s, "abiVersion: v0") { - t.Errorf("expected abiVersion: v0\n---\n%s", s) + if !strings.Contains(s, "abiVersion: \"v0\"") { + t.Errorf("expected quoted abiVersion: v0\n---\n%s", s) } } @@ -166,7 +171,7 @@ func TestRenderAgentYAMLIncludesEnvRequirements(t *testing.T) { } s := string(out) - for _, want := range []string{"env:", "name: REQUIRED_FOO", "required: true", "name: OPTIONAL_BAR"} { + for _, want := range []string{"env:", "name: \"REQUIRED_FOO\"", "required: true", "name: \"OPTIONAL_BAR\""} { if !strings.Contains(s, want) { t.Fatalf("rendered agent.yaml missing %q\n---\n%s", want, s) } @@ -196,14 +201,14 @@ func TestRenderAgentYAMLIncludesRemoteMCPTool(t *testing.T) { s := string(out) for _, want := range []string{ - "type: mcp", - "transport: streamable-http", - "urlEnv: TOOLBOX_ENDPOINT", - "name: Foundry-Features", - "value: Toolboxes=V1Preview", + "type: \"mcp\"", + "transport: \"streamable-http\"", + "urlEnv: \"TOOLBOX_ENDPOINT\"", + "name: \"Foundry-Features\"", + "value: \"Toolboxes=V1Preview\"", "auth:", - "type: bearer", - "tokenEnv: TOOLBOX_TOKEN", + "type: \"bearer\"", + "tokenEnv: \"TOOLBOX_TOKEN\"", } { if !strings.Contains(s, want) { t.Fatalf("rendered agent.yaml missing %q\n---\n%s", want, s) @@ -229,12 +234,12 @@ func TestRenderAgentYAMLIncludesContextAndObservability(t *testing.T) { for _, want := range []string{ "context:", "providers:", - "name: knowledge", - "type: search", - "endpointEnv: SEARCH_ENDPOINT", - "indexEnv: SEARCH_INDEX", + "name: \"knowledge\"", + "type: \"search\"", + "endpointEnv: \"SEARCH_ENDPOINT\"", + "indexEnv: \"SEARCH_INDEX\"", "observability:", - "endpointEnv: OTEL_EXPORTER_OTLP_ENDPOINT", + "endpointEnv: \"OTEL_EXPORTER_OTLP_ENDPOINT\"", } { if !strings.Contains(s, want) { t.Fatalf("rendered agent.yaml missing %q\n---\n%s", want, s) @@ -250,7 +255,7 @@ func TestRenderAgentYAMLIncludesModelWorkloadIdentityAuth(t *testing.T) { t.Fatalf("render error: %v", err) } s := string(out) - for _, want := range []string{"auth:", "type: workload-identity-token", "audience: https://ai.azure.com/.default"} { + for _, want := range []string{"auth:", "type: \"workload-identity-token\"", "audience: \"https://ai.azure.com/.default\""} { if !strings.Contains(s, want) { t.Fatalf("rendered agent.yaml missing %q\n---\n%s", want, s) } @@ -265,9 +270,9 @@ func TestRenderAgentYAMLIncludesBrokeredTools(t *testing.T) { Description: "Read telemetry.", BrokeredClass: config.BrokeredClassRead, Parameters: map[string]any{ - jsonSchemaTypeKey: jsonSchemaTypeObject, + jsonSchemaTypeKey: jsonSchemaObject, jsonSchemaPropertiesKey: map[string]any{ - "site": map[string]any{jsonSchemaTypeKey: "string", jsonSchemaMinimumKey: 0.000001}, + "site": map[string]any{jsonSchemaTypeKey: jsonSchemaString, jsonSchemaMinimumKey: 0.000001}, "typedFloats": map[string]float64{jsonSchemaMinimumKey: 0.000001}, "empty": map[string]any{}, }, @@ -278,12 +283,12 @@ func TestRenderAgentYAMLIncludesBrokeredTools(t *testing.T) { t.Fatalf("render error: %v", err) } s := string(out) - for _, want := range []string{"brokeredTools:", "name: check-network-telemetry", "description: Read telemetry.", "brokeredClass: read", "properties:", "site:", "minimum: 0.000001", "typedFloats:", "empty: {}"} { + for _, want := range []string{"brokeredTools:", "name: \"check-network-telemetry\"", "description: \"Read telemetry.\"", "brokeredClass: \"read\"", "\"properties\":", "\"site\":", "\"minimum\": 0.000001", "\"typedFloats\":", "\"empty\": {}"} { if !strings.Contains(s, want) { t.Fatalf("rendered agent.yaml missing %q\n---\n%s", want, s) } } - for _, never := range []string{"url:", "secretRef:", "headers:", "auth" + ":", "1e-06"} { + for _, never := range []string{"\"url\":", "\"secretRef\":", "\"headers\":", "\"auth\":", "1e-06"} { if strings.Contains(s, never) { t.Fatalf("rendered brokered agent.yaml leaked %q\n---\n%s", never, s) } @@ -298,9 +303,9 @@ func TestRenderAgentYAMLFormatsJSONNumberBrokeredSchemaValuesAsNumbers(t *testin Description: "Read numeric data.", BrokeredClass: config.BrokeredClassRead, Parameters: map[string]any{ - jsonSchemaTypeKey: jsonSchemaTypeObject, + jsonSchemaTypeKey: jsonSchemaObject, jsonSchemaPropertiesKey: map[string]any{ - "small": map[string]any{jsonSchemaTypeKey: jsonSchemaTypeNumber, jsonSchemaMinimumKey: json.Number("1e-7")}, + "small": map[string]any{jsonSchemaTypeKey: jsonSchemaNumber, jsonSchemaMinimumKey: json.Number("1e-7")}, }, }, }} @@ -309,11 +314,469 @@ func TestRenderAgentYAMLFormatsJSONNumberBrokeredSchemaValuesAsNumbers(t *testin if err != nil { t.Fatalf("render error: %v", err) } - if !strings.Contains(string(out), "minimum: 0.0000001") || strings.Contains(string(out), "1e-7") { + if !strings.Contains(string(out), "\"minimum\": 0.0000001") || strings.Contains(string(out), "1e-7") { t.Fatalf("rendered agent.yaml did not preserve json.Number as fixed YAML number\n---\n%s", out) } } +func TestRenderAgentYAMLRejectsYAMLBinarySchemaValues(t *testing.T) { + const binaryAgentkitfile = `apiVersion: v1alpha1 +kind: Agent +metadata: + name: binary-schema +runtime: pydantic-ai +model: + provider: openai-compatible + baseURL: https://api.openai.com/v1 + name: gpt-4o-mini +instructions: Read binary fixtures. +brokeredTools: +- name: inspect-binary + description: Inspect binary metadata. + brokeredClass: read + parameters: + type: object + properties: + payload: + type: string + default: !!binary SGVsbG8= +expose: + openai: true + port: 8080 +` + + cfg, err := config.NewFromBytes([]byte(binaryAgentkitfile)) + if err != nil { + t.Fatalf("load binary agentkitfile: %v", err) + } + if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "unsupported JSON value []uint8") { + t.Fatalf("binary schema must be rejected rather than silently coerced, got: %v", err) + } +} + +func TestRenderAgentYAMLNormalizesNamedBrokeredSchemaTypes(t *testing.T) { + type schemaString string + type schemaBytes []byte + type schemaStrings []schemaString + type schemaMap map[schemaString]any + + const ( + propertyName = schemaString("line:\u2028break") + enumValue = schemaString("value\twith-tab") + ) + tool := config.BrokeredTool{ + Name: "typed-schema-tool", + Description: "Read typed schema data.", + BrokeredClass: config.BrokeredClassRead, + Parameters: map[string]any{ + jsonSchemaTypeKey: jsonSchemaObject, + jsonSchemaPropertiesKey: schemaMap{ + propertyName: map[string]any{ + jsonSchemaTypeKey: schemaString(jsonSchemaString), + jsonSchemaDefaultKey: schemaBytes("Hello"), + jsonSchemaEnumKey: schemaStrings{enumValue}, + }, + }, + }, + } + binaryDigest, err := config.BrokeredToolSchemaDigest(tool) + if err != nil { + t.Fatalf("digest named schema types: %v", err) + } + equivalent := tool + equivalent.Parameters = map[string]any{ + jsonSchemaTypeKey: jsonSchemaObject, + jsonSchemaPropertiesKey: map[string]any{ + string(propertyName): map[string]any{ + jsonSchemaTypeKey: jsonSchemaString, + jsonSchemaDefaultKey: testHelloBase64, + jsonSchemaEnumKey: []string{string(enumValue)}, + }, + }, + } + stringDigest, err := config.BrokeredToolSchemaDigest(equivalent) + if err != nil { + t.Fatalf("digest equivalent schema types: %v", err) + } + if binaryDigest != stringDigest { + t.Fatalf("named schema digest = %q, equivalent digest = %q", binaryDigest, stringDigest) + } + tool.SchemaDigest = binaryDigest + + agent := sampleAgent() + agent.Tools = nil + agent.BrokeredTools = []config.BrokeredTool{tool} + out, err := Render(agent) + if err != nil { + t.Fatalf("render named schema types: %v", err) + } + var got struct { + BrokeredTools []struct { + Parameters struct { + Properties map[string]struct { + Type string `yaml:"type"` + Default any `yaml:"default"` + Enum []string `yaml:"enum"` + } `yaml:"properties"` + } `yaml:"parameters"` + } `yaml:"brokeredTools"` + } + if err := yaml.Unmarshal(out, &got); err != nil { + t.Fatalf("parse rendered named schema types: %v\n---\n%s", err, out) + } + property, ok := got.BrokeredTools[0].Parameters.Properties[string(propertyName)] + if !ok { + t.Fatalf("rendered schema lost named property key %q: %#v", propertyName, got.BrokeredTools[0].Parameters.Properties) + } + if property.Type != jsonSchemaString { + t.Errorf("type = %q, want %q", property.Type, jsonSchemaString) + } + if property.Default != testHelloBase64 { + t.Errorf("default = %#v (%T), want base64 string", property.Default, property.Default) + } + if len(property.Enum) != 1 || property.Enum[0] != string(enumValue) { + t.Errorf("enum = %#v, want %q", property.Enum, enumValue) + } +} + +func TestRenderAgentYAMLRoundTripsYAMLSensitiveStrings(t *testing.T) { + tests := []struct { + name string + value string + }{ + {name: "document start", value: "---"}, + {name: "document end", value: "..."}, + {name: "document end prefix", value: "... value"}, + {name: "explicit key prefix", value: "? ask"}, + {name: "merge key", value: "<<"}, + {name: "value key", value: "="}, + {name: "positive infinity", value: testYAMLPositiveInfinity}, + {name: "negative infinity", value: "-.Inf"}, + {name: "not a number", value: ".NaN"}, + {name: "explicit key tab prefix", value: "?\task"}, + {name: "leading tab", value: "\tvalue"}, + {name: "embedded tab", value: "before\tafter"}, + {name: "trailing tab", value: "value\t"}, + {name: "control", value: "before\x01after"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + value := test.value + agent := sampleAgent() + agent.Instructions = value + agent.Tools = nil + agent.BrokeredTools = []config.BrokeredTool{{ + Name: "indicator-tool", + Description: "Read indicator values.", + BrokeredClass: config.BrokeredClassRead, + Parameters: map[string]any{ + jsonSchemaTypeKey: jsonSchemaObject, + jsonSchemaPropertiesKey: map[string]any{ + value: map[string]any{ + jsonSchemaTypeKey: jsonSchemaString, + jsonSchemaDefaultKey: value, + }, + }, + }, + }} + + out, err := Render(agent) + if err != nil { + t.Fatalf("render YAML-sensitive string %q: %v", value, err) + } + var got struct { + Instructions string `yaml:"instructions"` + BrokeredTools []struct { + Parameters struct { + Properties map[string]struct { + Default string `yaml:"default"` + } `yaml:"properties"` + } `yaml:"parameters"` + } `yaml:"brokeredTools"` + } + if err := yaml.Unmarshal(out, &got); err != nil { + t.Fatalf("parse rendered YAML-sensitive string %q: %v\n---\n%s", value, err, out) + } + if got.Instructions != value { + t.Errorf("instructions = %q, want %q", got.Instructions, value) + } + property, ok := got.BrokeredTools[0].Parameters.Properties[value] + if !ok { + t.Fatalf("rendered schema lost property key %q: %#v", value, got.BrokeredTools[0].Parameters.Properties) + } + if property.Default != value { + t.Errorf("schema default = %q, want %q", property.Default, value) + } + }) + } +} + +func TestRenderAgentYAMLRoundTripsOrdinaryMultilineStringsInNestedFields(t *testing.T) { + const multiline = "first line\nsecond line\rthird line" + agent := sampleAgent() + agent.Tools = nil + agent.BrokeredTools = []config.BrokeredTool{{ + Name: "multiline-tool", + Description: multiline, + BrokeredClass: config.BrokeredClassRead, + Parameters: map[string]any{ + jsonSchemaTypeKey: jsonSchemaObject, + jsonSchemaDescriptionKey: multiline, + jsonSchemaPropertiesKey: map[string]any{ + multiline: map[string]any{ + jsonSchemaTypeKey: jsonSchemaString, + jsonSchemaDescriptionKey: multiline, + }, + }, + }, + }} + + out, err := Render(agent) + if err != nil { + t.Fatalf("render multiline strings: %v", err) + } + var got struct { + BrokeredTools []struct { + Description string `yaml:"description"` + Parameters struct { + Description string `yaml:"description"` + Properties map[string]struct { + Description string `yaml:"description"` + } `yaml:"properties"` + } `yaml:"parameters"` + } `yaml:"brokeredTools"` + } + if err := yaml.Unmarshal(out, &got); err != nil { + t.Fatalf("rendered multiline agent.yaml did not parse: %v\n---\n%s", err, out) + } + if len(got.BrokeredTools) != 1 { + t.Fatalf("brokeredTools = %#v", got.BrokeredTools) + } + tool := got.BrokeredTools[0] + for field, value := range map[string]string{ + jsonSchemaDescriptionKey: tool.Description, + "parameters.description": tool.Parameters.Description, + "parameters.properties.description": tool.Parameters.Properties[multiline].Description, + } { + if value != multiline { + t.Errorf("%s = %q, want %q", field, value, multiline) + } + } +} + +func TestRenderAgentYAMLEscapesYAMLLineBreaksInEveryStringField(t *testing.T) { + const lineBreakText = "NEL:\u0085LS:\u2028PS:\u2029end" + expected := map[string]bool{} + marked := func(field string) string { + value := field + " " + lineBreakText + expected[value] = true + return value + } + + agent := effective.Agent{ + Metadata: config.Metadata{Name: marked("metadata.name")}, + Model: config.Model{ + Provider: marked("model.provider"), + BaseURL: marked("model.baseURL"), + Name: marked("model.name"), + APIKeyEnv: marked("model.apiKeyEnv"), + Auth: &config.Auth{ + Type: marked("model.auth.type"), + TokenEnv: marked("model.auth.tokenEnv"), + Audience: marked("model.auth.audience"), + }, + }, + Instructions: marked("instructions"), + Tools: []config.Tool{{ + Name: marked("tools.name"), + Type: marked("tools.type"), + Transport: marked("tools.transport"), + Command: []string{marked("tools.command[0]"), marked("tools.command[1]")}, + URLEnv: marked("tools.urlEnv"), + Headers: []config.ToolHeader{{ + Name: marked("tools.headers.name"), + Value: marked("tools.headers.value"), + ValueEnv: marked("tools.headers.valueEnv"), + }}, + Auth: &config.Auth{ + Type: marked("tools.auth.type"), + TokenEnv: marked("tools.auth.tokenEnv"), + Audience: marked("tools.auth.audience"), + }, + Approval: marked("tools.approval"), + Env: []string{marked("tools.env[0]"), marked("tools.env[1]")}, + }}, + BrokeredTools: []config.BrokeredTool{{ + Name: marked("brokeredTools.name"), + Description: marked("brokeredTools.description"), + BrokeredClass: marked("brokeredTools.brokeredClass"), + Parameters: map[string]any{ + marked("brokeredTools.parameters.key"): marked("brokeredTools.parameters.value"), + "slice": []string{marked("brokeredTools.parameters.slice")}, + }, + SchemaDigest: marked("brokeredTools.schemaDigest"), + }}, + Env: []config.EnvVar{{Name: marked("env.name")}}, + Context: config.Context{Providers: []config.ContextProvider{{ + Name: marked("context.providers.name"), + Type: marked("context.providers.type"), + Source: marked("context.providers.source"), + Path: marked("context.providers.path"), + ToolRef: marked("context.providers.toolRef"), + Index: marked("context.providers.index"), + EndpointEnv: marked("context.providers.endpointEnv"), + IndexEnv: marked("context.providers.indexEnv"), + StoreNameEnv: marked("context.providers.storeNameEnv"), + Auth: &config.Auth{ + Type: marked("context.providers.auth.type"), + TokenEnv: marked("context.providers.auth.tokenEnv"), + Audience: marked("context.providers.auth.audience"), + }, + }}}, + Observability: config.Observability{ + OTel: config.ObservabilityOTel{EndpointEnv: marked("observability.otel.endpointEnv")}, + Logs: config.ObservabilityLogs{LevelEnv: marked("observability.logs.levelEnv")}, + }, + Expose: config.Expose{OpenAI: true, Port: 8080}, + } + + out, err := Render(agent) + if err != nil { + t.Fatalf("render error: %v", err) + } + if strings.ContainsAny(string(out), "\u0085\u2028\u2029") { + t.Fatalf("rendered agent.yaml contains an unescaped YAML line-break code point\n---\n%s", out) + } + + var decoded any + if err := yaml.Unmarshal(out, &decoded); err != nil { + t.Fatalf("rendered agent.yaml did not parse: %v\n---\n%s", err, out) + } + seen := map[string]bool{} + var collectStrings func(any) + collectStrings = func(value any) { + switch typed := value.(type) { + case string: + seen[typed] = true + case map[string]any: + for key, item := range typed { + seen[key] = true + collectStrings(item) + } + case map[any]any: + for key, item := range typed { + collectStrings(key) + collectStrings(item) + } + case []any: + for _, item := range typed { + collectStrings(item) + } + } + } + collectStrings(decoded) + for value := range expected { + if !seen[value] { + t.Errorf("round-tripped agent.yaml lost %q\n---\n%s", value, out) + } + } +} + +func TestRenderAgentYAMLEdgeCasesMatchCrossLanguageGolden(t *testing.T) { + const ( + lineBreakText = "NEL:\u0085LS:\u2028PS:\u2029end" + propertyName = "line:\u2028break" + ) + + tool := config.BrokeredTool{ + Name: "unicode-numeric-tool", + Description: "description " + lineBreakText, + BrokeredClass: config.BrokeredClassRead, + Parameters: map[string]any{ + jsonSchemaTypeKey: jsonSchemaObject, + jsonSchemaDescriptionKey: "schema " + lineBreakText, + jsonSchemaPropertiesKey: map[string]any{ + "12:34:56": map[string]any{ + jsonSchemaTypeKey: jsonSchemaString, + jsonSchemaDefaultKey: "2001-12-14 21:59:43.10 -5", + }, + testYAMLPositiveInfinity: map[string]any{ + jsonSchemaTypeKey: jsonSchemaString, + jsonSchemaDefaultKey: testYAMLPositiveInfinity, + }, + "<<": map[string]any{ + jsonSchemaTypeKey: jsonSchemaString, + jsonSchemaDefaultKey: "<<", + }, + "=": map[string]any{ + jsonSchemaTypeKey: jsonSchemaString, + jsonSchemaDefaultKey: "=", + }, + "? ask": map[string]any{ + jsonSchemaTypeKey: jsonSchemaString, + jsonSchemaDefaultKey: "before\tafter", + }, + "binary": map[string]any{ + jsonSchemaTypeKey: jsonSchemaString, + jsonSchemaDefaultKey: testHelloBase64, + }, + propertyName: map[string]any{ + jsonSchemaTypeKey: "number", + jsonSchemaDescriptionKey: "property " + lineBreakText, + jsonSchemaMinimumKey: math.Copysign(0, -1), + }, + }, + "required": []any{propertyName}, + }, + } + digest, err := config.BrokeredToolSchemaDigest(tool) + if err != nil { + t.Fatalf("schema digest: %v", err) + } + tool.SchemaDigest = digest + + cfg := sampleConfig() + cfg.Tools = nil + cfg.Model.APIKeyEnv = "" + cfg.Instructions = config.Source{Inline: "instructions " + lineBreakText} + cfg.BrokeredTools = []config.BrokeredTool{tool} + if err := cfg.Validate(); err != nil { + t.Fatalf("edge-case config should pass Go validation: %v", err) + } + + out, err := Render(effective.FromConfig(cfg, cfg.Instructions.Inline)) + if err != nil { + t.Fatalf("render error: %v", err) + } + want, err := os.ReadFile("testdata/edge-cases.yaml") + if err != nil { + t.Fatalf("read edge-case golden: %v", err) + } + if string(out) != string(want) { + t.Fatalf("rendered edge-case agent.yaml drifted from cross-language golden\n--- got ---\n%s\n--- want ---\n%s", out, want) + } +} + +func TestRenderAgentYAMLRejectsUnderflowingJSONNumbers(t *testing.T) { + cfg := sampleConfig() + cfg.Tools = nil + cfg.Instructions = config.Source{Inline: testInstructions} + cfg.BrokeredTools = []config.BrokeredTool{{ + Name: "tiny-number-tool", + Description: "Read tiny numeric data.", + BrokeredClass: config.BrokeredClassRead, + Parameters: map[string]any{ + jsonSchemaTypeKey: jsonSchemaObject, + jsonSchemaPropertiesKey: map[string]any{ + "tiny": map[string]any{jsonSchemaTypeKey: jsonSchemaNumber, jsonSchemaMinimumKey: json.Number("-1e-350")}, + }, + }, + }} + if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "cannot be represented exactly") { + t.Fatalf("underflowing json.Number must fail exact representation validation, got: %v", err) + } +} + func TestRenderAgentYAMLPreservesNegativeZeroBrokeredSchemaFloats(t *testing.T) { cfg := sampleConfig() cfg.Tools = nil @@ -322,10 +785,10 @@ func TestRenderAgentYAMLPreservesNegativeZeroBrokeredSchemaFloats(t *testing.T) Description: "Preserve negative zero.", BrokeredClass: config.BrokeredClassRead, Parameters: map[string]any{ - jsonSchemaTypeKey: jsonSchemaTypeObject, + jsonSchemaTypeKey: jsonSchemaObject, jsonSchemaPropertiesKey: map[string]any{ - "offset": map[string]any{jsonSchemaTypeKey: jsonSchemaTypeNumber, jsonSchemaDefaultKey: math.Copysign(0, -1)}, - "jsonOffset": map[string]any{jsonSchemaTypeKey: jsonSchemaTypeNumber, jsonSchemaDefaultKey: json.Number("-0e0")}, + "offset": map[string]any{jsonSchemaTypeKey: jsonSchemaNumber, jsonSchemaDefaultKey: math.Copysign(0, -1)}, + "jsonOffset": map[string]any{jsonSchemaTypeKey: jsonSchemaNumber, jsonSchemaDefaultKey: json.Number("-0e0")}, "largeInteger": map[string]any{jsonSchemaTypeKey: "integer", jsonSchemaDefaultKey: json.Number("9007199254740995.0")}, "typedNested": map[string][]float64{"enum": {math.Copysign(0, -1)}}, "typedNil": map[string]any{jsonSchemaDefaultKey: map[string]string(nil)}, @@ -344,16 +807,16 @@ func TestRenderAgentYAMLPreservesNegativeZeroBrokeredSchemaFloats(t *testing.T) if err != nil { t.Fatalf("render error: %v", err) } - if strings.Count(string(out), "default: -0.0") != 2 { + if strings.Count(string(out), `"default": -0.0`) != 2 { t.Fatalf("rendered agent.yaml did not preserve negative zero as a float\n---\n%s", out) } - if !strings.Contains(string(out), "default: 9007199254740995") || strings.Contains(string(out), "9007199254740995.0") { + if !strings.Contains(string(out), `"default": 9007199254740995`) || strings.Contains(string(out), "9007199254740995.0") { t.Fatalf("rendered agent.yaml did not preserve a large integral decimal as an integer\n---\n%s", out) } if strings.Count(string(out), yamlNegativeZero) != 3 { t.Fatalf("rendered agent.yaml did not normalize negative zero in typed nested containers\n---\n%s", out) } - if strings.Count(string(out), "default: null") != 2 { + if strings.Count(string(out), `"default": null`) != 2 { t.Fatalf("rendered agent.yaml did not preserve typed nil containers as null\n---\n%s", out) } } diff --git a/pkg/agentkit/abi/testdata/agent.yaml b/pkg/agentkit/abi/testdata/agent.yaml index cd883d1..5a9fb8d 100644 --- a/pkg/agentkit/abi/testdata/agent.yaml +++ b/pkg/agentkit/abi/testdata/agent.yaml @@ -1,19 +1,19 @@ -abiVersion: v0 +abiVersion: "v0" metadata: - name: acme-support + name: "acme-support" model: - provider: openai-compatible - baseURL: https://api.openai.com/v1 - name: gpt-4o-mini - apiKeyEnv: OPENAI_API_KEY -instructions: Be helpful and cite sources. + provider: "openai-compatible" + baseURL: "https://api.openai.com/v1" + name: "gpt-4o-mini" + apiKeyEnv: "OPENAI_API_KEY" +instructions: "Be helpful and cite sources." tools: -- name: fetch +- name: "fetch" command: - - uvx - - mcp-server-fetch + - "uvx" + - "mcp-server-fetch" env: - - FETCH_TIMEOUT + - "FETCH_TIMEOUT" expose: openai: true port: 8080 diff --git a/pkg/agentkit/abi/testdata/edge-cases.yaml b/pkg/agentkit/abi/testdata/edge-cases.yaml new file mode 100644 index 0000000..2198b2e --- /dev/null +++ b/pkg/agentkit/abi/testdata/edge-cases.yaml @@ -0,0 +1,45 @@ +abiVersion: "v0" +metadata: + name: "acme-support" +model: + provider: "openai-compatible" + baseURL: "https://api.openai.com/v1" + name: "gpt-4o-mini" +instructions: "instructions NEL:\u0085LS:\u2028PS:\u2029end" +tools: [] +brokeredTools: +- name: "unicode-numeric-tool" + description: "description NEL:\u0085LS:\u2028PS:\u2029end" + brokeredClass: "read" + parameters: + "description": "schema NEL:\u0085LS:\u2028PS:\u2029end" + "properties": + ".inf": + "default": ".inf" + "type": "string" + "12:34:56": + "default": "2001-12-14 21:59:43.10 -5" + "type": "string" + "<<": + "default": "<<" + "type": "string" + "=": + "default": "=" + "type": "string" + "? ask": + "default": "before\tafter" + "type": "string" + "binary": + "default": "SGVsbG8=" + "type": "string" + "line:\u2028break": + "description": "property NEL:\u0085LS:\u2028PS:\u2029end" + "minimum": -0.0 + "type": "number" + "required": + - "line:\u2028break" + "type": "object" + schemaDigest: "sha256:ce77aaf228491b5007ed2ee703e57180acec8def6214c84d4324719b7f4f1fb6" +expose: + openai: true + port: 8080 diff --git a/pkg/agentkit/config/config_test.go b/pkg/agentkit/config/config_test.go index 047ed26..a4fb7e8 100644 --- a/pkg/agentkit/config/config_test.go +++ b/pkg/agentkit/config/config_test.go @@ -13,6 +13,7 @@ const ( jsonSchemaPropertiesKey = "properties" brokeredSiteField = "site" brokeredSafeDescription = "safe schema" + benignNearMissValue = "near miss" ) // TestKindProbeRejectsKindlessFile is the regression guard for the AIKit @@ -130,6 +131,65 @@ expose: } } +func TestValidateRejectsControlPlaneMetadataLabels(t *testing.T) { + reserved := []string{ + nativeImageLabelNamespace, + ImageLabelNativeRuntime, + ImageLabelNativeName, + ImageLabelNativeABI, + nativeImageLabelNamespace + ".future-control", + portableImageLabelNamespace, + ImageLabelPortableABI, + ImageLabelPortableRuntime, + ImageLabelPortableProtocols, + ImageLabelPortableCapabilities, + portableImageLabelNamespace + ".future-control", + orkaImageLabelNamespace, + ImageLabelOrkaHarnessVersion, + orkaImageLabelNamespace + ".future-control", + ImageLabelOCITitle, + } + + for _, label := range reserved { + t.Run(label, func(t *testing.T) { + cfg, err := NewFromBytes(agentBaseYAML("")) + if err != nil { + t.Fatalf("parse error: %v", err) + } + cfg.Metadata.Labels = map[string]string{label: "spoofed"} + + verr := cfg.Validate() + if verr == nil { + t.Fatalf("expected reserved metadata label %q to be rejected", label) + } + for _, want := range []string{"metadata.labels", label, "reserved"} { + if !strings.Contains(verr.Error(), want) { + t.Errorf("validation error missing %q; full: %s", want, verr) + } + } + }) + } +} + +func TestValidateAllowsUnrelatedMetadataLabels(t *testing.T) { + cfg, err := NewFromBytes(agentBaseYAML("")) + if err != nil { + t.Fatalf("parse error: %v", err) + } + cfg.Metadata.Labels = map[string]string{ + "com.example/team": "agentkit", + "org.opencontainers.image.description": "helpful agent", + nativeImageLabelNamespace + "-runtime": benignNearMissValue, + portableImageLabelNamespace + "-custom": benignNearMissValue, + orkaImageLabelNamespace + "-tools": benignNearMissValue, + "ai.example.agentkit.runtime": "unrelated namespace", + } + + if verr := cfg.Validate(); verr != nil { + t.Fatalf("unrelated metadata labels should be allowed: %v", verr) + } +} + func TestValidateRejectsSecretLiteralInApiKeyEnv(t *testing.T) { in := []byte(`apiVersion: v1alpha1 kind: Agent @@ -1404,8 +1464,24 @@ func TestValidateAcceptsBrokeredToolsWithMatchingDigest(t *testing.T) { } } +func TestValidateAcceptsHarmlessBrokeredDescriptions(t *testing.T) { + for _, description := range []string{"Read basic telemetry", "Count model tokens", "Count model tokens 100", "Count model tokens: 1,000."} { + cfg := validMinimalConfig() + cfg.BrokeredTools = []BrokeredTool{{ + Name: safeLookupToolName, + Description: description, + BrokeredClass: BrokeredClassRead, + Parameters: map[string]any{jsonSchemaTypeKey: jsonSchemaTypeObject}, + }} + + if err := cfg.Validate(); err != nil { + t.Errorf("description %q should be accepted: %v", description, err) + } + } +} + func TestValidateRejectsUnsafeBrokeredDescriptions(t *testing.T) { - for _, description := range []string{"contains sk-secret", "execution at https://tool.default", "Bearer token required", "execution at tool.default.svc.cluster.local"} { + for _, description := range []string{"contains sk-secret", "execution at https://tool.default", "Bearer token required", "Basic Zm9vOmJhcg==", "Basic=dXNlcjpwYXNz", "Basic user:pass", "Basic alice: pass1", "alice:pass is the Basic credential", "Use Basic credential abc123", "Use Basic `dXNlcjpwYXNz`", "Use HTTP Basic value dXNlcjpwYXNz", "Use HTTP Basic value=dXNlcjpwYXNz", "Use HTTP-Basic value dXNlcjpwYXNz", "token=abc123", `{"token":"abc123"}`, `'access_token'=abc123`, "token.value=abc123", "token[value]=abc123", "Use --token abc123", "access token abc123", "token abc123", "token 123", "token abc", "token ABCDEF", `"token" "abc123"`, "model token abc123", `input token "abc123"`, `token "abcdef"`, `"token" "ABCDEF"`, `token "abc"`, `input token '123'`, "token.value abc123", "token[value] abc123", "request.token abc123", "request/token abc123", "request-token abc123", "requesttoken abc123", `token used: 'abc'`, "Use abc123 as model [token]", "Count requests. Use model token 123", "Finished counting. Model token 123", `Finished "counting." Model token 123`, "Count model tokens 123456789012345678901234567890", "Read {auth}", "Read (header)", "Read [REDACTED_AUTH_HEADER]", "execution at tool.default.svc.cluster.local"} { cfg := validMinimalConfig() cfg.BrokeredTools = []BrokeredTool{{ Name: safeLookupToolName, @@ -1458,7 +1534,7 @@ func TestValidateRejectsPrivateKeyBrokeredNamesAndStrings(t *testing.T) { } func TestValidateRejectsUnsafeBrokeredSchemaStringValues(t *testing.T) { - for _, value := range []string{"see https://internal-tool", "Bearer abc", "Bearer: abc", "Bearer=abc", "authorization header", "tool.default.svc.cluster.local", "example ghp_not_real", "AWS key AKIAEXAMPLE"} { + for _, value := range []string{"see https://internal-tool", "Bearer abc", "Bearer: abc", "Bearer=abc", "authorization header", "token=abc123", "tool.default.svc.cluster.local", "example ghp_not_real", "AWS key AKIAEXAMPLE"} { cfg := validMinimalConfig() cfg.BrokeredTools = []BrokeredTool{{ Name: safeLookupToolName, @@ -1479,7 +1555,7 @@ func TestValidateRejectsUnsafeBrokeredSchemaStringValues(t *testing.T) { } func TestValidateRejectsCommonCredentialBrokeredParameterNames(t *testing.T) { - for _, field := range []string{"authentication", "authConfig", "clientSecret", "dbPassword", "passphrase", "pwd", "apiKey", credentialHeaderAPIKey, "baseUrl", "callbackURL", "apiEndpoint", "sessionCookie", "cookies"} { + for _, field := range []string{brokeredAuthenticationWord, "authConfig", "clientSecret", "dbPassword", "passphrase", "pwd", "apiKey", credentialHeaderAPIKey, "baseUrl", "callbackURL", "apiEndpoint", "sessionCookie", "cookies"} { cfg := validMinimalConfig() cfg.BrokeredTools = []BrokeredTool{{ Name: safeLookupToolName, @@ -1500,7 +1576,7 @@ func TestValidateRejectsCommonCredentialBrokeredParameterNames(t *testing.T) { } func TestValidateRejectsCommonCredentialBrokeredRequiredNames(t *testing.T) { - for _, field := range []string{"authentication", "authConfig", "clientSecret", "dbPassword", "passphrase", "pwd", "apiKey", credentialHeaderAPIKey, "baseUrl", "callbackURL", "apiEndpoint", "sessionCookie", "cookies"} { + for _, field := range []string{brokeredAuthenticationWord, "authConfig", "clientSecret", "dbPassword", "passphrase", "pwd", "apiKey", credentialHeaderAPIKey, "baseUrl", "callbackURL", "apiEndpoint", "sessionCookie", "cookies"} { cfg := validMinimalConfig() cfg.BrokeredTools = []BrokeredTool{{ Name: safeLookupToolName, @@ -1539,7 +1615,7 @@ func TestValidateAcceptsHarmlessBrokeredAuthorField(t *testing.T) { } func TestValidateRejectsUnsafeBrokeredToolSchema(t *testing.T) { - for _, unsafeName := range []string{"token", "authHeader", "authorizationHeader", "httpHeaders", "accessKey", "clientSecretValue", "tokenValue", "apiSecretKey", brokeredUnsafeCookieKey, "subscriptionKey", "xFunctionsKey"} { + for _, unsafeName := range []string{brokeredTokenWord, "access_token", "authHeader", "authorizationHeader", "httpHeaders", "accessKey", "clientSecretValue", "tokenValue", "apiSecretKey", brokeredUnsafeCookieKey, "subscriptionKey", "xFunctionsKey"} { cfg := validMinimalConfig() cfg.BrokeredTools = []BrokeredTool{{ Name: safeLookupToolName, diff --git a/pkg/agentkit/config/labels.go b/pkg/agentkit/config/labels.go new file mode 100644 index 0000000..2b05110 --- /dev/null +++ b/pkg/agentkit/config/labels.go @@ -0,0 +1,63 @@ +package config + +import ( + "strings" + + "github.com/sozercan/agentkit/pkg/utils" +) + +const ( + nativeImageLabelNamespace = utils.LabelPrefix + portableImageLabelNamespace = "ai.agentkit" + orkaImageLabelNamespace = "ai.orka" + + // ImageLabelNativeRuntime identifies the canonical AgentKit runtime. + ImageLabelNativeRuntime = nativeImageLabelNamespace + ".runtime" + // ImageLabelNativeName identifies the authored AgentKit agent name. + ImageLabelNativeName = nativeImageLabelNamespace + ".name" + // ImageLabelNativeABI identifies the baked Agent YAML ABI version. + ImageLabelNativeABI = nativeImageLabelNamespace + ".abi" + + // ImageLabelPortableABI is the cross-orchestrator Agent YAML ABI label. + ImageLabelPortableABI = portableImageLabelNamespace + ".abi" + // ImageLabelPortableRuntime is the cross-orchestrator runtime identity label. + ImageLabelPortableRuntime = portableImageLabelNamespace + ".runtime" + // ImageLabelPortableProtocols lists the protocols exposed by AgentKit images. + ImageLabelPortableProtocols = portableImageLabelNamespace + ".protocols" + // ImageLabelPortableCapabilities lists the selected runtime's capabilities. + ImageLabelPortableCapabilities = portableImageLabelNamespace + ".capabilities" + + // ImageLabelOrkaHarnessVersion identifies the supported Orka harness contract. + ImageLabelOrkaHarnessVersion = orkaImageLabelNamespace + ".harness.version" + + // ImageLabelOCITitle is the standard OCI title generated from metadata.name. + ImageLabelOCITitle = "org.opencontainers.image.title" +) + +var reservedMetadataLabelNamespaces = [...]string{ + nativeImageLabelNamespace, + portableImageLabelNamespace, + orkaImageLabelNamespace, +} + +var reservedMetadataLabelKeys = [...]string{ + ImageLabelOCITitle, +} + +func reservedMetadataLabelNamespace(key string) (string, bool) { + for _, namespace := range reservedMetadataLabelNamespaces { + if key == namespace || strings.HasPrefix(key, namespace+".") { + return namespace, true + } + } + return "", false +} + +func isReservedMetadataLabelKey(key string) bool { + for _, reserved := range reservedMetadataLabelKeys { + if key == reserved { + return true + } + } + return false +} diff --git a/pkg/agentkit/config/validate.go b/pkg/agentkit/config/validate.go index 1de618a..0fbd33e 100644 --- a/pkg/agentkit/config/validate.go +++ b/pkg/agentkit/config/validate.go @@ -3,6 +3,7 @@ package config import ( "bytes" "crypto/sha256" + "encoding/base64" "encoding/hex" "encoding/json" "errors" @@ -11,6 +12,7 @@ import ( "math/big" pathpkg "path" "reflect" + "regexp" "sort" "strconv" "strings" @@ -55,11 +57,19 @@ const ( jsonSchemaMaximumKey = "maximum" brokeredDigestDescriptionKey = "description" brokeredDigestNumberKey = "\u0000agentkit_json_number" + brokeredAuthenticationWord = "authentication" + brokeredSensitiveWord = "credential" + brokeredSensitivePluralWord = "credentials" + brokeredTokenWord = "token" + brokeredTokensWord = "tokens" brokeredUnsafeCookieKey = "cookie" + authorizationKey = "authorization" credentialHeaderAPIKey = "api-key" maxExactJSONFloatInteger = float64(1<<53 - 1) ) +var brokeredBasicValuePattern = regexp.MustCompile(`(?i)(?:^|[^A-Za-z0-9_])basic[^A-Za-z0-9_]+?([A-Za-z0-9+/]+={0,2})`) + // Validate reports every problem with the config at once via errors.Join (plan // §16.2 #3 — one report-all validator, not scattered first-error-wins funcs). // @@ -86,6 +96,20 @@ func (c *AgentConfig) Validate() error { if c.Metadata.Name == "" { add("metadata.name is required") } + metadataLabelKeys := make([]string, 0, len(c.Metadata.Labels)) + for key := range c.Metadata.Labels { + metadataLabelKeys = append(metadataLabelKeys, key) + } + sort.Strings(metadataLabelKeys) + for _, key := range metadataLabelKeys { + if namespace, reserved := reservedMetadataLabelNamespace(key); reserved { + add("metadata.labels[%q] uses reserved AgentKit/Orka control-plane label namespace %q", key, namespace) + continue + } + if isReservedMetadataLabelKey(key) { + add("metadata.labels[%q] is reserved for AgentKit-generated image identity metadata", key) + } + } // --- runtime ----------------------------------------------------------- runtimeName := c.Runtime @@ -223,7 +247,7 @@ func validateBrokeredTools(add func(string, ...any), tools []BrokeredTool, owned } if tool.Description == "" { add("%s.description is required", path) - } else if hasUnsafeBrokeredText(tool.Description) { + } else if hasUnsafeBrokeredDescription(tool.Description) { add("%s.description must not contain URLs or secret-like material", path) } switch tool.BrokeredClass { @@ -776,9 +800,229 @@ func isSchemaDigest(value string) bool { } func hasUnsafeBrokeredText(value string) bool { + lowered := strings.ToLower(value) + return hasUnsafeBrokeredDescription(value) || containsBrokeredWord(lowered, "basic") || strings.Contains(lowered, brokeredTokenWord) +} + +func hasUnsafeBrokeredDescription(value string) bool { lowered := strings.ToLower(value) normalized := normalizeKey(lowered) - return containsSecretPrefix(value) || strings.Contains(value, "://") || containsBrokeredWord(lowered, "bearer") || containsBrokeredWord(lowered, "basic") || strings.Contains(lowered, "authorization") || strings.Contains(lowered, "secret") || strings.Contains(lowered, "token") || strings.Contains(lowered, "password") || strings.Contains(lowered, "passphrase") || strings.Contains(lowered, "pwd") || strings.Contains(lowered, "api key") || strings.Contains(lowered, "apikey") || strings.Contains(normalized, "apikey") || strings.Contains(normalized, "xapikey") || strings.Contains(normalized, "subscriptionkey") || strings.Contains(normalized, "xfunctionskey") || strings.Contains(lowered, brokeredUnsafeCookieKey) || strings.Contains(lowered, "set-cookie") || strings.Contains(lowered, "x-api-key") || strings.Contains(lowered, credentialHeaderAPIKey) || strings.Contains(lowered, "subscription-key") || strings.Contains(lowered, "x-functions-key") || strings.Contains(lowered, "ocp-apim-subscription-key") || strings.Contains(lowered, "private key") || strings.Contains(lowered, "privatekey") || strings.Contains(lowered, "key material") || strings.Contains(lowered, ".svc") || strings.Contains(lowered, "cluster.local") + return containsSecretPrefix(value) || strings.Contains(value, "://") || containsBrokeredWord(lowered, "bearer") || containsBrokeredWord(lowered, brokeredSensitiveWord) || containsBrokeredWord(lowered, brokeredSensitivePluralWord) || containsBrokeredBasicAuthReference(value) || containsBrokeredCredentialAssignment(value) || containsBrokeredCredentialReference(value) || strings.Contains(lowered, authorizationKey) || strings.Contains(lowered, "secret") || strings.Contains(lowered, "password") || strings.Contains(lowered, "passphrase") || strings.Contains(lowered, "pwd") || strings.Contains(lowered, "api key") || strings.Contains(lowered, "apikey") || strings.Contains(normalized, "apikey") || strings.Contains(normalized, "xapikey") || strings.Contains(normalized, "subscriptionkey") || strings.Contains(normalized, "xfunctionskey") || strings.Contains(lowered, brokeredUnsafeCookieKey) || strings.Contains(lowered, "set-cookie") || strings.Contains(lowered, "x-api-key") || strings.Contains(lowered, credentialHeaderAPIKey) || strings.Contains(lowered, "subscription-key") || strings.Contains(lowered, "x-functions-key") || strings.Contains(lowered, "ocp-apim-subscription-key") || strings.Contains(lowered, "private key") || strings.Contains(lowered, "privatekey") || strings.Contains(lowered, "key material") || strings.Contains(lowered, ".svc") || strings.Contains(lowered, "cluster.local") +} + +func containsBrokeredBasicAuthReference(value string) bool { + lowered := strings.ToLower(value) + if !containsBrokeredWord(lowered, "basic") { + return false + } + if containsBrokeredWord(lowered, "auth") || containsBrokeredWord(lowered, brokeredAuthenticationWord) || containsBrokeredWord(lowered, authorizationKey) { + return true + } + + for _, match := range brokeredBasicValuePattern.FindAllStringSubmatch(value, -1) { + if isBrokeredBasicValue(match[1]) { + return true + } + } + fields := strings.Fields(value) + for _, field := range fields { + if isBrokeredBasicValue(field) { + return true + } + } + for i, field := range fields { + if !containsBrokeredWord(strings.ToLower(field), "basic") { + continue + } + for j, candidate := range fields[i+1:] { + if isBrokeredBasicValue(candidate) { + return true + } + plain := strings.Trim(candidate, "\"'`()[]{}<>,;.") + if strings.HasSuffix(plain, ":") && j+1 < len(fields[i+1:]) { + return true + } + if strings.ContainsAny(candidate, ".!?;") { + break + } + } + } + return false +} + +func isBrokeredBasicValue(value string) bool { + plain := strings.Trim(value, "\"'`()[]{}<>,;.") + if separator := strings.IndexByte(plain, ':'); separator > 0 && separator+1 < len(plain) { + return true + } + if decodesBrokeredBasicValue(value) { + return true + } + for separator, r := range value { + if (r == ':' || r == '=') && decodesBrokeredBasicValue(value[separator+1:]) { + return true + } + } + return false +} + +func decodesBrokeredBasicValue(value string) bool { + value = strings.TrimFunc(value, func(r rune) bool { return !isBrokeredBase64Rune(r) }) + if value == "" { + return false + } + decoded, err := base64.StdEncoding.DecodeString(value) + if err != nil { + decoded, err = base64.RawStdEncoding.DecodeString(value) + } + return err == nil && bytes.Contains(decoded, []byte(":")) +} + +func containsBrokeredCredentialAssignment(value string) bool { + for separator := 0; separator < len(value); separator++ { + if value[separator] != ':' && value[separator] != '=' { + continue + } + if isHarmlessBrokeredTokenCountAssignment(value, separator) { + continue + } + fields := strings.Fields(value[:separator]) + if len(fields) == 0 { + continue + } + if isUnsafeBrokeredKey(fields[len(fields)-1]) { + return true + } + } + return false +} + +func containsBrokeredCredentialReference(value string) bool { + fields := strings.Fields(value) + previous := "" + for i, field := range fields { + normalized := normalizeKey(field) + if normalized == "" { + continue + } + if normalized == brokeredTokenWord || normalized == brokeredTokensWord { + if !isHarmlessBrokeredTokenUse(fields, i, normalizeKey(previous)) { + return true + } + previous = field + continue + } + if strings.Contains(normalized, brokeredTokenWord) && !isHarmlessBrokeredTokenWord(normalized) { + return true + } + if hasBrokeredStructuredKeyShape(field) && isUnsafeBrokeredKey(field) { + return true + } + switch normalized { + case "accesstoken", "apitoken", "authtoken", "authenticationtoken", "authorizationtoken", "bearertoken", "credentialtoken", "identitytoken", "oauthtoken", "refreshtoken", "secrettoken", "sessiontoken": + return true + } + previous = field + } + return false +} + +func isHarmlessBrokeredTokenWord(value string) bool { + switch value { + case "tokenization", "tokenize", "tokenized", "tokenizer", "tokenizers", "tokenizing": + return true + default: + return false + } +} + +func isHarmlessBrokeredTokenUse(fields []string, index int, previous string) bool { + if hasBrokeredStructuredKeyShape(fields[index]) || !isHarmlessBrokeredTokenQualifier(previous) || !hasAdjacentBrokeredTokenCountIntent(fields, index) { + return false + } + if index+1 == len(fields) { + return true + } + return index+2 == len(fields) && isBrokeredNumericCount(fields[index+1]) +} + +func isHarmlessBrokeredTokenQualifier(value string) bool { + switch value { + case "completion", "context", "count", "counting", "counts", "input", "model", "output", "prompt", "usage": + return true + default: + return false + } +} + +func hasAdjacentBrokeredTokenCountIntent(fields []string, tokenIndex int) bool { + if tokenIndex > 0 && !endsBrokeredSentence(fields[tokenIndex-1]) && isHarmlessBrokeredTokenCountIntent(normalizeKey(fields[tokenIndex-1])) { + return true + } + return tokenIndex > 1 && !endsBrokeredSentence(fields[tokenIndex-2]) && isHarmlessBrokeredTokenCountIntent(normalizeKey(fields[tokenIndex-2])) +} + +func endsBrokeredSentence(value string) bool { + value = strings.TrimRight(value, "\"'`)]}>,") + return value != "" && strings.ContainsRune(".!?;", rune(value[len(value)-1])) +} + +func isHarmlessBrokeredTokenCountIntent(value string) bool { + switch value { + case "count", "counting", "counts", "measure", "measures", "measuring", "report", "reporting", "reports", "track", "tracking": + return true + default: + return false + } +} + +func isHarmlessBrokeredTokenCountAssignment(value string, separator int) bool { + if value[separator] != ':' { + return false + } + leftFields := strings.Fields(value[:separator]) + rightFields := strings.Fields(value[separator+1:]) + if len(leftFields) < 2 || len(rightFields) != 1 { + return false + } + key := normalizeKey(leftFields[len(leftFields)-1]) + if key != brokeredTokenWord && key != brokeredTokensWord { + return false + } + qualifier := normalizeKey(leftFields[len(leftFields)-2]) + return isHarmlessBrokeredTokenQualifier(qualifier) && hasAdjacentBrokeredTokenCountIntent(leftFields, len(leftFields)-1) && isBrokeredNumericCount(rightFields[0]) +} + +func isBrokeredNumericCount(value string) bool { + wrapped := strings.TrimLeft(value, "([{<") + if wrapped != "" && strings.ContainsRune("\"'`", rune(wrapped[0])) { + return false + } + value = strings.Trim(value, "\"'`()[]{}<>,;:.") + if value == "" { + return false + } + parts := strings.Split(value, ",") + digits := 0 + for i, part := range parts { + if part == "" || i == 0 && len(parts) > 1 && len(part) > 3 || i > 0 && len(part) != 3 { + return false + } + for _, r := range part { + if r < '0' || r > '9' { + return false + } + digits++ + if digits > 20 { + return false + } + } + } + return digits > 0 +} + +func hasBrokeredStructuredKeyShape(value string) bool { + return strings.HasPrefix(strings.TrimLeft(value, "\"'`([{<"), "-") || strings.ContainsAny(value, "_/.[]{}()<>\"'`") || value != strings.ToLower(value) } func containsBrokeredWord(value string, word string) bool { @@ -803,14 +1047,18 @@ func isBrokeredWordByte(value byte) bool { return (value >= 'a' && value <= 'z') || (value >= 'A' && value <= 'Z') || (value >= '0' && value <= '9') || value == '_' } +func isBrokeredBase64Rune(value rune) bool { + return (value >= 'a' && value <= 'z') || (value >= 'A' && value <= 'Z') || (value >= '0' && value <= '9') || value == '+' || value == '/' || value == '=' +} + func isUnsafeBrokeredKey(value string) bool { normalized := normalizeKey(value) switch normalized { - case "auth", "authorization", "apikey", "bearer", brokeredUnsafeCookieKey, "credential", "credentials", "endpoint", "endpoints", "executionendpoint", "executionurl", "header", "headers", "ocpapimsubscriptionkey", "password", "proxyauthorization", "secret", "secretref", "setcookie", "subscriptionkey", "token", "tokens", "url", "urls", "xapikey", "xfunctionskey": + case "auth", authorizationKey, "apikey", "bearer", brokeredUnsafeCookieKey, brokeredSensitiveWord, brokeredSensitivePluralWord, "endpoint", "endpoints", "executionendpoint", "executionurl", "header", "headers", "ocpapimsubscriptionkey", "password", "proxyauthorization", "secret", "secretref", "setcookie", "subscriptionkey", brokeredTokenWord, brokeredTokensWord, "url", "urls", "xapikey", "xfunctionskey": return true } authLike := (strings.HasPrefix(normalized, "auth") && !strings.HasPrefix(normalized, "author")) || strings.HasSuffix(normalized, "auth") - return authLike || strings.Contains(normalized, "authorization") || strings.Contains(normalized, "header") || strings.Contains(normalized, "url") || strings.Contains(normalized, "endpoint") || strings.Contains(normalized, brokeredUnsafeCookieKey) || strings.Contains(normalized, "secret") || strings.Contains(normalized, "token") || strings.Contains(normalized, "password") || strings.Contains(normalized, "passphrase") || strings.Contains(normalized, "pwd") || strings.Contains(normalized, "apikey") || strings.Contains(normalized, "accesskey") || strings.Contains(normalized, "privatekey") || strings.Contains(normalized, "keymaterial") || strings.Contains(normalized, "credential") || strings.Contains(normalized, "executionurl") || strings.Contains(normalized, "executionendpoint") + return authLike || strings.Contains(normalized, authorizationKey) || strings.Contains(normalized, "header") || strings.Contains(normalized, "url") || strings.Contains(normalized, "endpoint") || strings.Contains(normalized, brokeredUnsafeCookieKey) || strings.Contains(normalized, "secret") || strings.Contains(normalized, brokeredTokenWord) || strings.Contains(normalized, "password") || strings.Contains(normalized, "passphrase") || strings.Contains(normalized, "pwd") || strings.Contains(normalized, "apikey") || strings.Contains(normalized, "accesskey") || strings.Contains(normalized, "privatekey") || strings.Contains(normalized, "keymaterial") || strings.Contains(normalized, brokeredSensitiveWord) || strings.Contains(normalized, "executionurl") || strings.Contains(normalized, "executionendpoint") } func normalizeKey(value string) string { @@ -1202,7 +1450,7 @@ func validateToolHeaders(add func(string, ...any), i int, t Tool) { if values != 1 { add("%s must set exactly one of value or valueEnv", path) } - if t.Auth != nil && strings.EqualFold(h.Name, "authorization") { + if t.Auth != nil && strings.EqualFold(h.Name, authorizationKey) { add("%s must not set Authorization when auth is also configured; use one auth path", path) } if h.Value != "" && isCredentialHeaderName(h.Name) { @@ -1337,7 +1585,7 @@ func isHTTPHeaderName(v string) bool { func isCredentialHeaderName(name string) bool { switch strings.ToLower(name) { - case "authorization", "proxy-authorization", brokeredUnsafeCookieKey, "set-cookie", "x-api-key", credentialHeaderAPIKey, "ocp-apim-subscription-key", "subscription-key", "x-functions-key": + case authorizationKey, "proxy-authorization", brokeredUnsafeCookieKey, "set-cookie", "x-api-key", credentialHeaderAPIKey, "ocp-apim-subscription-key", "subscription-key", "x-functions-key": return true default: return false diff --git a/pkg/agentkit/render/orka.go b/pkg/agentkit/render/orka.go index b515ddc..51219dc 100644 --- a/pkg/agentkit/render/orka.go +++ b/pkg/agentkit/render/orka.go @@ -8,7 +8,9 @@ import ( "flag" "fmt" "io" + "net/url" "strings" + "unicode" "github.com/goccy/go-yaml" ) @@ -66,6 +68,22 @@ type agentRuntimeCaps struct { SupportsRuntimeSessions bool `yaml:"supportsRuntimeSessions"` } +const invalidExternalEndpointMessage = "--external-endpoint must be an absolute http(s) URL with a host and no userinfo, query, fragment, or whitespace" + +func validateExternalEndpoint(endpoint string) error { + if (!strings.HasPrefix(endpoint, "http://") && !strings.HasPrefix(endpoint, "https://")) || + strings.IndexFunc(endpoint, unicode.IsSpace) >= 0 || strings.ContainsAny(endpoint, "@?#") { + return errors.New(invalidExternalEndpointMessage) + } + parsed, err := url.Parse(endpoint) + if err != nil || !parsed.IsAbs() || parsed.Opaque != "" || parsed.Hostname() == "" || parsed.User != nil || + parsed.RawQuery != "" || parsed.ForceQuery || parsed.Fragment != "" || parsed.RawFragment != "" || + (parsed.Scheme != "http" && parsed.Scheme != "https") { + return errors.New(invalidExternalEndpointMessage) + } + return nil +} + // OrkaAgentRuntime renders a core.orka.ai AgentRuntime manifest for an // AgentKit image that exposes the observed-mode orka.harness.v1 protocol. func OrkaAgentRuntime(opts OrkaAgentRuntimeOptions) ([]byte, error) { @@ -76,10 +94,13 @@ func OrkaAgentRuntime(opts OrkaAgentRuntimeOptions) ([]byte, error) { if strings.TrimSpace(opts.Image) != "" { return nil, errors.New("the current Orka AgentRuntime CRD supports external endpoints only; deploy the image with AGENTKIT_PROTOCOL=orka and AGENTKIT_AUTH_TOKEN from the bearer token Secret, then pass --external-endpoint") } - endpoint := strings.TrimSpace(opts.ExternalEndpoint) + endpoint := opts.ExternalEndpoint if endpoint == "" { return nil, errors.New("--external-endpoint is required for the current Orka AgentRuntime CRD") } + if err := validateExternalEndpoint(endpoint); err != nil { + return nil, err + } authSecretName := strings.TrimSpace(opts.AuthSecretName) if authSecretName == "" { authSecretName = name + "-harness-token" diff --git a/pkg/agentkit/render/orka_test.go b/pkg/agentkit/render/orka_test.go index 7023b15..5db4d4e 100644 --- a/pkg/agentkit/render/orka_test.go +++ b/pkg/agentkit/render/orka_test.go @@ -6,7 +6,12 @@ import ( "testing" ) -const testRuntimeName = "fibey" +const ( + testRuntimeName = "fibey" + testTargetFlag = "--target" + testNameFlag = "--name" + testExternalEndpointFlag = "--external-endpoint" +) func TestOrkaAgentRuntimeExternalEndpoint(t *testing.T) { got, err := OrkaAgentRuntime(OrkaAgentRuntimeOptions{ @@ -43,9 +48,9 @@ spec: func TestRunCLIRendersOrkaAgentRuntime(t *testing.T) { var stdout, stderr bytes.Buffer code := RunCLI([]string{ - "--target", TargetOrkaAgentRuntime, - "--external-endpoint", "http://fibey-agentkit.default.svc.cluster.local:8080", - "--name", "fibey-agentkit", + testTargetFlag, TargetOrkaAgentRuntime, + testExternalEndpointFlag, "http://fibey-agentkit.default.svc.cluster.local:8080", + testNameFlag, "fibey-agentkit", "--auth-secret-name", "fibey-auth", }, &stdout, &stderr) if code != 0 { @@ -65,8 +70,68 @@ func TestOrkaAgentRuntimeValidation(t *testing.T) { t.Fatalf("expected image unsupported error, got %v", err) } var stdout, stderr bytes.Buffer - code := RunCLI([]string{"--target", "other", "--name", testRuntimeName, "--external-endpoint", "http://example.invalid"}, &stdout, &stderr) + code := RunCLI([]string{testTargetFlag, "other", testNameFlag, testRuntimeName, testExternalEndpointFlag, "http://example.invalid"}, &stdout, &stderr) if code == 0 || !strings.Contains(stderr.String(), "unsupported --target") { t.Fatalf("RunCLI() code=%d stderr=%q", code, stderr.String()) } } + +func TestOrkaAgentRuntimeRejectsExternalEndpointsOutsideCRDPattern(t *testing.T) { + const redactionMarker = "redaction-marker" + invalid := []string{ + "example.com:8080", + "ftp://example.com", + "HTTP://example.com", + "http:///missing-host", + "http://:8080", + "http://" + "user:" + redactionMarker + "@example.com", + "http://example.com/path?q=" + redactionMarker, + "http://example.com/path?", + "http://example.com/path#" + redactionMarker, + "http://example.com/path#", + "http://example.com/path@segment", + " http://example.com", + "http://example.com ", + "http://exa mple.com", + "http://example.com\nnext", + } + for _, endpoint := range invalid { + t.Run(endpoint, func(t *testing.T) { + _, err := OrkaAgentRuntime(OrkaAgentRuntimeOptions{ + Name: testRuntimeName, + ExternalEndpoint: endpoint, + }) + if err == nil { + t.Fatalf("OrkaAgentRuntime() accepted invalid endpoint %q", endpoint) + } + if strings.Contains(err.Error(), redactionMarker) { + t.Fatalf("validation error leaked endpoint material %q: %v", redactionMarker, err) + } + }) + } +} + +func TestRunCLIExternalEndpointValidationRedactsCredentialBearingInput(t *testing.T) { + markers := []string{"user-marker", "userinfo-marker", "query-marker", "anchor-marker"} + endpoint := "https://" + markers[0] + ":" + markers[1] + "@example.com/path?q=" + markers[2] + "#" + markers[3] + var stdout, stderr bytes.Buffer + code := RunCLI([]string{ + testTargetFlag, TargetOrkaAgentRuntime, + testNameFlag, testRuntimeName, + testExternalEndpointFlag, endpoint, + }, &stdout, &stderr) + if code == 0 { + t.Fatalf("RunCLI() accepted credential-bearing endpoint; stdout=%q", stdout.String()) + } + if stdout.Len() != 0 { + t.Fatalf("RunCLI() wrote stdout for invalid endpoint: %q", stdout.String()) + } + for _, marker := range append([]string{endpoint}, markers...) { + if strings.Contains(stderr.String(), marker) { + t.Fatalf("RunCLI() leaked endpoint material %q in stderr: %q", marker, stderr.String()) + } + } + if !strings.Contains(stderr.String(), testExternalEndpointFlag) { + t.Fatalf("RunCLI() stderr lacks generic endpoint guidance: %q", stderr.String()) + } +} diff --git a/pkg/agentkit2llb/agent/image.go b/pkg/agentkit2llb/agent/image.go index bc7641f..18d88d7 100644 --- a/pkg/agentkit2llb/agent/image.go +++ b/pkg/agentkit2llb/agent/image.go @@ -7,20 +7,23 @@ import ( "github.com/moby/buildkit/util/system" specs "github.com/opencontainers/image-spec/specs-go/v1" "github.com/sozercan/agentkit/pkg/agentkit/abi" + "github.com/sozercan/agentkit/pkg/agentkit/config" "github.com/sozercan/agentkit/pkg/agentkit/effective" "github.com/sozercan/agentkit/pkg/agentkit/runtimes" "github.com/sozercan/agentkit/pkg/utils" ) +const ( + imageProtocols = "openai,foundry,orka" + orkaHarnessVersion = "orka.harness.v1" +) + // NewImageConfig builds the OCI image config for the agent image. It deliberately // does NOT inherit AIKit's root user: per plan §10 the agent runs non-root, binds // loopback by default, and exposes the serve port. func NewImageConfig(agent effective.Agent, platform *specs.Platform) *specs.Image { img := &specs.Image{ - Platform: specs.Platform{ - Architecture: platform.Architecture, - OS: utils.PlatformLinux, - }, + Platform: *platform, } img.RootFS.Type = "layers" @@ -50,22 +53,29 @@ func NewImageConfig(agent effective.Agent, platform *specs.Platform) *specs.Imag capabilities = strings.Join(runtimeSpec.Capabilities, ",") } - img.Config.Labels = map[string]string{ + generatedLabels := map[string]string{ // Current AgentKit label namespace. - utils.LabelPrefix + ".runtime": agent.Runtime, - utils.LabelPrefix + ".name": agent.Metadata.Name, - utils.LabelPrefix + ".abi": abi.Version, + config.ImageLabelNativeRuntime: agent.Runtime, + config.ImageLabelNativeName: agent.Metadata.Name, + config.ImageLabelNativeABI: abi.Version, // Cross-orchestrator metadata consumed by Orka and other registries. - "ai.agentkit.abi": abi.Version, - "ai.agentkit.runtime": agent.Runtime, - "ai.agentkit.protocols": "openai,foundry,orka", - "ai.agentkit.capabilities": capabilities, - "ai.orka.harness.version": "orka.harness.v1", - "org.opencontainers.image.title": agent.Metadata.Name, + config.ImageLabelPortableABI: abi.Version, + config.ImageLabelPortableRuntime: agent.Runtime, + config.ImageLabelPortableProtocols: imageProtocols, + config.ImageLabelPortableCapabilities: capabilities, + config.ImageLabelOrkaHarnessVersion: orkaHarnessVersion, + config.ImageLabelOCITitle: agent.Metadata.Name, } + + img.Config.Labels = make(map[string]string, len(agent.Metadata.Labels)+len(generatedLabels)) for k, v := range agent.Metadata.Labels { img.Config.Labels[k] = v } + // Generated identity and capability labels are applied last so they remain + // authoritative even if an invalid effective.Agent bypasses config validation. + for k, v := range generatedLabels { + img.Config.Labels[k] = v + } return img } diff --git a/pkg/agentkit2llb/agent/image_test.go b/pkg/agentkit2llb/agent/image_test.go index 6f9d314..b6eb66c 100644 --- a/pkg/agentkit2llb/agent/image_test.go +++ b/pkg/agentkit2llb/agent/image_test.go @@ -2,6 +2,7 @@ package agent import ( "fmt" + "reflect" "strings" "testing" @@ -14,6 +15,7 @@ import ( ) const ( + testAgentName = "acme-support" testTeamLabel = "com.example/team" testTeamValue = "agentkit" ) @@ -21,7 +23,7 @@ const ( func imageAgent(runtime string, port int) effective.Agent { cfg := &config.AgentConfig{ Metadata: config.Metadata{ - Name: "acme-support", + Name: testAgentName, Labels: map[string]string{testTeamLabel: testTeamValue}, }, Runtime: runtime, @@ -52,25 +54,25 @@ func TestNewImageConfigUsesEffectiveAgentContract(t *testing.T) { if _, ok := img.Config.ExposedPorts[fmt.Sprintf("%d/tcp", utils.DefaultFoundryPort)]; !ok { t.Fatalf("ExposedPorts = %#v, want Foundry default port %d", img.Config.ExposedPorts, utils.DefaultFoundryPort) } - if got := img.Config.Labels[utils.LabelPrefix+".runtime"]; got != runtimes.MAF { + if got := img.Config.Labels[config.ImageLabelNativeRuntime]; got != runtimes.MAF { t.Fatalf("runtime label = %q, want canonical %q", got, runtimes.MAF) } - if got := img.Config.Labels[utils.LabelPrefix+".abi"]; got != abi.Version { + if got := img.Config.Labels[config.ImageLabelNativeABI]; got != abi.Version { t.Fatalf("abi label = %q, want %q", got, abi.Version) } - if got := img.Config.Labels["ai.agentkit.abi"]; got != abi.Version { + if got := img.Config.Labels[config.ImageLabelPortableABI]; got != abi.Version { t.Fatalf("portable abi label = %q, want %q", got, abi.Version) } - if got := img.Config.Labels["ai.agentkit.runtime"]; got != runtimes.MAF { + if got := img.Config.Labels[config.ImageLabelPortableRuntime]; got != runtimes.MAF { t.Fatalf("portable runtime label = %q, want %q", got, runtimes.MAF) } - if got := img.Config.Labels["ai.agentkit.protocols"]; got != "openai,foundry,orka" { + if got := img.Config.Labels[config.ImageLabelPortableProtocols]; got != imageProtocols { t.Fatalf("protocols label = %q", got) } - if got := img.Config.Labels["ai.orka.harness.version"]; got != "orka.harness.v1" { + if got := img.Config.Labels[config.ImageLabelOrkaHarnessVersion]; got != orkaHarnessVersion { t.Fatalf("orka harness label = %q", got) } - if got := img.Config.Labels["ai.agentkit.capabilities"]; !strings.Contains(got, runtimes.CapabilityOrkaHarnessV1) { + if got := img.Config.Labels[config.ImageLabelPortableCapabilities]; !strings.Contains(got, runtimes.CapabilityOrkaHarnessV1) { t.Fatalf("capabilities label = %q, want %s", got, runtimes.CapabilityOrkaHarnessV1) } if got := img.Config.Labels[testTeamLabel]; got != testTeamValue { @@ -78,6 +80,56 @@ func TestNewImageConfigUsesEffectiveAgentContract(t *testing.T) { } } +func TestNewImageConfigPreservesTargetPlatformIdentity(t *testing.T) { + platform := &specs.Platform{ + Architecture: "arm", + OS: "windows", + OSVersion: "10.0.20348.2113", + OSFeatures: []string{"win32k"}, + Variant: "v7", + } + + img := NewImageConfig(imageAgent(runtimes.MAF, 0), platform) + + if !reflect.DeepEqual(img.Platform, *platform) { + t.Fatalf("Platform = %#v, want full target platform %#v", img.Platform, *platform) + } +} + +func TestNewImageConfigGeneratedLabelsOverrideMetadataLabels(t *testing.T) { + agent := imageAgent(runtimes.MAFAlias, 0) + runtimeSpec, ok := runtimes.RuntimeByName(agent.Runtime) + if !ok { + t.Fatalf("runtime %q is not registered", agent.Runtime) + } + expected := map[string]string{ + config.ImageLabelNativeRuntime: runtimes.MAF, + config.ImageLabelNativeName: testAgentName, + config.ImageLabelNativeABI: abi.Version, + config.ImageLabelPortableABI: abi.Version, + config.ImageLabelPortableRuntime: runtimes.MAF, + config.ImageLabelPortableProtocols: imageProtocols, + config.ImageLabelPortableCapabilities: strings.Join(runtimeSpec.Capabilities, ","), + config.ImageLabelOrkaHarnessVersion: orkaHarnessVersion, + config.ImageLabelOCITitle: testAgentName, + } + agent.Metadata.Labels = map[string]string{testTeamLabel: testTeamValue} + for key := range expected { + agent.Metadata.Labels[key] = "user-controlled" + } + + img := NewImageConfig(agent, &specs.Platform{Architecture: utils.PlatformAMD64}) + + for key, want := range expected { + if got := img.Config.Labels[key]; got != want { + t.Errorf("generated label %q = %q, want %q", key, got, want) + } + } + if got := img.Config.Labels[testTeamLabel]; got != testTeamValue { + t.Fatalf("unrelated user label = %q, want %q", got, testTeamValue) + } +} + func TestNewImageConfigPreservesEffectivePort(t *testing.T) { agent := imageAgent("", 9090) img := NewImageConfig(agent, &specs.Platform{Architecture: utils.PlatformAMD64}) diff --git a/pkg/build/build.go b/pkg/build/build.go index ac180ff..463a7f1 100644 --- a/pkg/build/build.go +++ b/pkg/build/build.go @@ -7,9 +7,7 @@ import ( "github.com/containerd/platforms" controlapi "github.com/moby/buildkit/api/services/control" - "github.com/moby/buildkit/client/llb" "github.com/moby/buildkit/exporter/containerimage/exptypes" - "github.com/moby/buildkit/frontend/dockerui" "github.com/moby/buildkit/frontend/gateway/client" specs "github.com/opencontainers/image-spec/specs-go/v1" "github.com/pkg/errors" @@ -36,22 +34,33 @@ const ( func Build(ctx context.Context, c client.Client) (*client.Result, error) { opts := c.BuildOpts().Opts - cfg, err := getAgentkitfileConfig(ctx, c) + loaded, err := loadAgentkitfile(ctx, c) if err != nil { return nil, errors.Wrap(err, "getting agentkitfile") } + cfg := loaded.config if err := validateAgentConfig(cfg); err != nil { return nil, errors.Wrap(err, "validating agentkitfile") } target := opts[keyTarget] - matched, route, rc, ok := lookupRoute(target, cfg.Runtime) + effectiveRuntime := canonicalEffectiveRuntime(cfg.Runtime) + matched, route, rc, ok := lookupRoute(target, effectiveRuntime) if !ok { - return nil, errors.Errorf("no route for target %q with runtime %q", target, cfg.Runtime) + if targetRuntime, namesRuntime := targetRuntimeSegment(target); namesRuntime && targetRuntime != effectiveRuntime { + return nil, errors.Errorf( + "no route for target %q: target runtime %q does not match effective runtime %q", + target, + targetRuntime, + effectiveRuntime, + ) + } + return nil, errors.Errorf("no route for target %q with effective runtime %q", target, effectiveRuntime) + } + if handler := contextRouteHandlers[matched]; handler != nil { + return handler(ctx, c, cfg, rc, loaded.instructions) } - _ = matched - return route.Handler(ctx, c, cfg, rc) } @@ -64,6 +73,10 @@ func validateAgentConfig(cfg *config.AgentConfig) error { // then solves the agent image for every target platform in parallel — mirroring // AIKit's buildInference multi-platform errgroup. func HandleAgent(ctx context.Context, c client.Client, cfg *config.AgentConfig, rc *RuntimeConfig) (*client.Result, error) { + return handleAgent(ctx, c, cfg, rc, localContextReader{client: c}) +} + +func handleAgent(ctx context.Context, c client.Client, cfg *config.AgentConfig, rc *RuntimeConfig, reader contextFileReader) (*client.Result, error) { opts := c.BuildOpts().Opts cacheImports, err := parseCacheOptions(opts) @@ -73,7 +86,7 @@ func HandleAgent(ctx context.Context, c client.Client, cfg *config.AgentConfig, // Resolve instructions (inline → as-is; file → read from build context) BEFORE // converting, so the baked agent.yaml carries a fully-resolved scalar (ABI). - instructions, err := resolveInstructions(ctx, c, cfg) + instructions, err := resolveInstructionSource(ctx, reader, cfg.Instructions) if err != nil { return nil, errors.Wrap(err, "resolving instructions") } @@ -193,74 +206,6 @@ func buildImage(ctx context.Context, c client.Client, agentSpec effective.Agent, return &result, nil } -// getAgentkitfileConfig resolves the agentkitfile from the build context -// (git/http/local — copied verbatim from AIKit's getAikitfileConfig) and parses -// it with the strict, kind-probed loader. -func getAgentkitfileConfig(ctx context.Context, c client.Client) (*config.AgentConfig, error) { - opts := c.BuildOpts().Opts - filename := opts[keyFilename] - if filename == "" { - filename = defaultAgentkitfileName - } - - name := "load agentkitfile" - if filename != defaultAgentkitfileName { - name += " from " + filename - } - - contextName := opts[localNameContext] - - var st *llb.State - var ok bool - keepGit := true - switch { - case strings.HasPrefix(contextName, "git"): - st, ok, _ = dockerui.DetectGitContext(contextName, &keepGit) - if !ok { - return nil, errors.Errorf("invalid git context %s", contextName) - } - case strings.HasPrefix(contextName, "http") || strings.HasPrefix(contextName, "https"): - st, ok, _ = dockerui.DetectGitContext(contextName, &keepGit) - if !ok { - st, filename, _ = dockerui.DetectHTTPContext(contextName) - } - default: - localSt := llb.Local(localNameDockerfile, - llb.IncludePatterns([]string{filename}), - llb.SessionID(c.BuildOpts().SessionID), - llb.SharedKeyHint(defaultAgentkitfileName), - dockerui.WithInternalName(name), - ) - st = &localSt - } - - def, err := st.Marshal(ctx) - if err != nil { - return nil, errors.Wrap(err, "failed to marshal local source") - } - res, err := c.Solve(ctx, client.SolveRequest{Definition: def.ToPB()}) - if err != nil { - return nil, errors.Wrap(err, "failed to resolve agentkitfile") - } - ref, err := res.SingleRef() - if err != nil { - return nil, err - } - dt, err := ref.ReadFile(ctx, client.ReadRequest{Filename: filename}) - if err != nil { - return nil, errors.Wrap(err, "failed to read agentkitfile") - } - - cfg, err := config.NewFromBytes(dt) - if err != nil { - return nil, errors.Wrap(err, "getting config") - } - if err := parseBuildArgs(opts, cfg); err != nil { - return nil, errors.Wrap(err, "parsing build args") - } - return cfg, nil -} - // getBuildArg returns the value of build-arg:, or "". func getBuildArg(opts map[string]string, k string) string { if opts != nil { @@ -305,7 +250,10 @@ func parseCacheOptions(opts map[string]string) ([]client.CacheOptionsEntry, erro if err := json.Unmarshal([]byte(cacheImportsStr), &cacheImportsUM); err != nil { return nil, errors.Wrapf(err, "failed to unmarshal %s", keyCacheImports) } - for _, um := range cacheImportsUM { + for i, um := range cacheImportsUM { + if um == nil { + return nil, errors.Errorf("%s entry %d is null", keyCacheImports, i) + } cacheImports = append(cacheImports, client.CacheOptionsEntry{Type: um.Type, Attrs: um.Attrs}) } } diff --git a/pkg/build/build_test.go b/pkg/build/build_test.go new file mode 100644 index 0000000..55f5748 --- /dev/null +++ b/pkg/build/build_test.go @@ -0,0 +1,801 @@ +package build + +import ( + "archive/tar" + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "net/url" + "strings" + "sync" + "testing" + + "github.com/moby/buildkit/client/llb" + "github.com/moby/buildkit/frontend/gateway/client" + "github.com/moby/buildkit/solver/pb" + fstypes "github.com/tonistiigi/fsutil/types" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +const ( + remoteContextURL = "https://example.com/agent.git#main" + localDockerfileSourcePrefix = "local://dockerfile" + dockerImageSourcePrefix = "docker-image://" + downloadedHTTPContextFilename = "context" + redactionTestPrefix = "redaction" +) + +type memoryReference struct { + mu sync.Mutex + files map[string][]byte + reads []string + readErr error +} + +func (r *memoryReference) ToState() (llb.State, error) { + return llb.Scratch(), nil +} + +func (r *memoryReference) Evaluate(context.Context) error { + return nil +} + +func (r *memoryReference) ReadFile(_ context.Context, req client.ReadRequest) ([]byte, error) { + r.mu.Lock() + defer r.mu.Unlock() + + name := strings.TrimPrefix(req.Filename, "/") + r.reads = append(r.reads, name) + if r.readErr != nil { + return nil, r.readErr + } + dt, ok := r.files[name] + if !ok { + return nil, fmt.Errorf("file %q not found", req.Filename) + } + if req.Range == nil { + return append([]byte(nil), dt...), nil + } + start := req.Range.Offset + if start < 0 || start > len(dt) { + return nil, fmt.Errorf("invalid offset %d for %q", start, req.Filename) + } + end := len(dt) + if req.Range.Length > 0 && start+req.Range.Length < end { + end = start + req.Range.Length + } + return append([]byte(nil), dt[start:end]...), nil +} + +func (r *memoryReference) StatFile(context.Context, client.StatRequest) (*fstypes.Stat, error) { + return nil, fmt.Errorf("StatFile not implemented") +} + +func (r *memoryReference) ReadDir(context.Context, client.ReadDirRequest) ([]*fstypes.Stat, error) { + return nil, fmt.Errorf("ReadDir not implemented") +} + +func (r *memoryReference) readPaths() []string { + r.mu.Lock() + defer r.mu.Unlock() + return append([]string(nil), r.reads...) +} + +type expectedSolve struct { + sourcePrefix string + requireAttemptUnpack bool + ref client.Reference + err error +} + +type fakeBuildClient struct { + client.Client + + opts client.BuildOpts + + mu sync.Mutex + expected []expectedSolve + definitions []*pb.Definition + sources [][]string + operationNames [][]string +} + +func (c *fakeBuildClient) BuildOpts() client.BuildOpts { + return c.opts +} + +func (c *fakeBuildClient) Solve(_ context.Context, req client.SolveRequest) (*client.Result, error) { + sources, err := definitionSources(req.Definition) + if err != nil { + return nil, err + } + names := definitionOperationNames(req.Definition) + + c.mu.Lock() + defer c.mu.Unlock() + c.definitions = append(c.definitions, req.Definition) + c.sources = append(c.sources, sources) + c.operationNames = append(c.operationNames, names) + index := len(c.sources) - 1 + if index >= len(c.expected) { + return nil, fmt.Errorf("unexpected solve %d with sources %v", index+1, sources) + } + expected := c.expected[index] + if !hasSourcePrefix(sources, expected.sourcePrefix) { + return nil, fmt.Errorf("solve %d sources %v do not include %q", index+1, sources, expected.sourcePrefix) + } + if expected.requireAttemptUnpack && !definitionAttemptsUnpack(req.Definition) { + return nil, fmt.Errorf("solve %d does not unpack the HTTP archive", index+1) + } + if expected.err != nil { + return nil, expected.err + } + result := client.NewResult() + result.SetRef(expected.ref) + return result, nil +} + +func (c *fakeBuildClient) allOperationNames() []string { + c.mu.Lock() + defer c.mu.Unlock() + var out []string + for _, names := range c.operationNames { + out = append(out, names...) + } + return out +} + +func (c *fakeBuildClient) solveSources() [][]string { + c.mu.Lock() + defer c.mu.Unlock() + out := make([][]string, len(c.sources)) + for i := range c.sources { + out[i] = append([]string(nil), c.sources[i]...) + } + return out +} + +func (c *fakeBuildClient) solveDefinition(index int) *pb.Definition { + c.mu.Lock() + defer c.mu.Unlock() + if index < 0 || index >= len(c.definitions) { + return nil + } + return c.definitions[index] +} + +func definitionSources(def *pb.Definition) ([]string, error) { + if def == nil { + return nil, nil + } + var sources []string + for _, dt := range def.Def { + var op pb.Op + if err := op.Unmarshal(dt); err != nil { + return nil, fmt.Errorf("unmarshal solve op: %w", err) + } + if source := op.GetSource(); source != nil { + sources = append(sources, source.Identifier) + } + } + return sources, nil +} + +func definitionOperationNames(def *pb.Definition) []string { + if def == nil { + return nil + } + var names []string + for _, metadata := range def.Metadata { + if metadata.Description == nil { + continue + } + if name := metadata.Description["llb.customname"]; name != "" { + names = append(names, name) + } + } + return names +} + +func definitionAttemptsUnpack(def *pb.Definition) bool { + if def == nil { + return false + } + for _, dt := range def.Def { + var op pb.Op + if err := op.Unmarshal(dt); err != nil { + continue + } + file := op.GetFile() + if file == nil { + continue + } + for _, action := range file.Actions { + if copyAction := action.GetCopy(); copyAction != nil && copyAction.AttemptUnpackDockerCompatibility { + return true + } + } + } + return false +} + +func hasSourcePrefix(sources []string, prefix string) bool { + for _, source := range sources { + if strings.HasPrefix(source, prefix) { + return true + } + } + return false +} + +func assertLocalSourceFollowsPath(t *testing.T, def *pb.Definition, sourcePrefix, wantPath string) { + t.Helper() + if def == nil { + t.Fatal("solve definition is nil") + } + for _, dt := range def.Def { + var op pb.Op + if err := op.Unmarshal(dt); err != nil { + t.Fatalf("unmarshal solve op: %v", err) + } + source := op.GetSource() + if source == nil || !strings.HasPrefix(source.Identifier, sourcePrefix) { + continue + } + + var followPaths []string + if err := json.Unmarshal([]byte(source.Attrs[pb.AttrFollowPaths]), &followPaths); err != nil { + t.Fatalf("decode %s for %s: %v", pb.AttrFollowPaths, source.Identifier, err) + } + if len(followPaths) != 1 || followPaths[0] != wantPath { + t.Fatalf("%s = %v, want [%q]", pb.AttrFollowPaths, followPaths, wantPath) + } + if include := source.Attrs[pb.AttrIncludePatterns]; include != "" { + t.Fatalf("%s = %q, want empty when using BuildKit follow-path filtering", pb.AttrIncludePatterns, include) + } + return + } + t.Fatalf("definition has no source with prefix %q", sourcePrefix) +} + +func fileBackedAgentkitfile(path string) []byte { + return []byte(fmt.Sprintf(`apiVersion: v1alpha1 +kind: Agent +metadata: + name: reliability-test +model: + provider: openai-compatible + baseURL: https://api.openai.com/v1 + name: gpt-4o-mini +instructions: + file: %s +expose: + openai: true +`, path)) +} + +func inlineAgentkitfile() []byte { + return []byte(`apiVersion: v1alpha1 +kind: Agent +metadata: + name: reliability-test +model: + provider: openai-compatible + baseURL: https://api.openai.com/v1 + name: gpt-4o-mini +instructions: Be reliable. +expose: + openai: true +`) +} + +func inlineAgentkitfileWithRuntime(runtime string) []byte { + return []byte(strings.Replace( + string(inlineAgentkitfile()), + "kind: Agent\n", + "kind: Agent\nruntime: "+runtime+"\n", + 1, + )) +} + +func TestLoadLocalAgentkitfileFollowsSymlinkPath(t *testing.T) { + const filename = "configs/agentkitfile-link.yaml" + configRef := &memoryReference{files: map[string][]byte{filename: inlineAgentkitfile()}} + c := &fakeBuildClient{ + opts: client.BuildOpts{Opts: map[string]string{}}, + expected: []expectedSolve{ + {sourcePrefix: localDockerfileSourcePrefix, ref: configRef}, + }, + } + + if _, err := loadLocalAgentkitfile(context.Background(), c, c.opts.Opts, filename); err != nil { + t.Fatalf("loadLocalAgentkitfile() error = %v", err) + } + assertLocalSourceFollowsPath(t, c.solveDefinition(0), localDockerfileSourcePrefix, filename) +} + +func TestLocalContextReaderFollowsInstructionSymlinkPath(t *testing.T) { + const filename = "prompts/current.md" + contextRef := &memoryReference{files: map[string][]byte{filename: []byte(filePrompt)}} + c := &fakeBuildClient{ + opts: client.BuildOpts{Opts: map[string]string{}}, + expected: []expectedSolve{ + {sourcePrefix: "local://context", ref: contextRef}, + }, + } + + dt, err := (localContextReader{client: c}).ReadFile(context.Background(), filename) + if err != nil { + t.Fatalf("ReadFile() error = %v", err) + } + if got := string(dt); got != filePrompt { + t.Fatalf("ReadFile() = %q, want %q", got, filePrompt) + } + assertLocalSourceFollowsPath(t, c.solveDefinition(0), "local://context", filename) +} + +func TestBuildRejectsTargetRuntimeMismatchBeforeImageSolve(t *testing.T) { + tests := []struct { + name string + agentkitfile []byte + target string + buildArgValue string + wantTargetRuntime string + }{ + { + name: "config runtime", + agentkitfile: inlineAgentkitfileWithRuntime(runtimePydca), + target: wantLangGraphRoute, + wantTargetRuntime: runtimeLangGraph, + }, + { + name: "build arg overrides config runtime", + agentkitfile: inlineAgentkitfileWithRuntime(runtimeLangGraph), + target: wantLangGraphRoute, + buildArgValue: runtimePydca, + wantTargetRuntime: runtimeLangGraph, + }, + { + name: "default runtime rejects alias-prefixed route", + agentkitfile: inlineAgentkitfile(), + target: runtimeMAFAls + "/image/debug", + wantTargetRuntime: runtimeMAFName, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + configRef := &memoryReference{files: map[string][]byte{defaultAgentkitfileName: tt.agentkitfile}} + opts := map[string]string{keyTarget: tt.target} + if tt.buildArgValue != "" { + opts["build-arg:runtime"] = tt.buildArgValue + } + c := &fakeBuildClient{ + opts: client.BuildOpts{Opts: opts}, + expected: []expectedSolve{ + {sourcePrefix: localDockerfileSourcePrefix, ref: configRef}, + }, + } + + _, err := Build(context.Background(), c) + if err == nil { + t.Fatal("Build() error = nil, want runtime mismatch") + } + for _, want := range []string{ + "no route", + fmt.Sprintf("target runtime %q", tt.wantTargetRuntime), + `effective runtime "pydantic-ai"`, + } { + if !strings.Contains(err.Error(), want) { + t.Fatalf("Build() error = %q, want substring %q", err, want) + } + } + if got := len(c.solveSources()); got != 1 { + t.Fatalf("Build() solve count = %d, want 1 config solve and no image solve", got) + } + }) + } +} + +func tarContext(t *testing.T, files map[string][]byte) []byte { + t.Helper() + var buf bytes.Buffer + tw := tar.NewWriter(&buf) + for name, data := range files { + if err := tw.WriteHeader(&tar.Header{Name: name, Mode: 0o644, Size: int64(len(data))}); err != nil { + t.Fatalf("write tar header: %v", err) + } + if _, err := tw.Write(data); err != nil { + t.Fatalf("write tar file: %v", err) + } + } + if err := tw.Close(); err != nil { + t.Fatalf("close tar: %v", err) + } + return buf.Bytes() +} + +type redactionTestMaterial struct { + username string + userInfo string + queryValue string + fragmentValue string +} + +func newRedactionTestMaterial() redactionTestMaterial { + return redactionTestMaterial{ + username: strings.Join([]string{redactionTestPrefix, "user"}, "-"), + userInfo: strings.Join([]string{redactionTestPrefix, "userinfo"}, "-"), + queryValue: strings.Join([]string{redactionTestPrefix, "query"}, "-"), + fragmentValue: strings.Join([]string{redactionTestPrefix, "fragment"}, "-"), + } +} + +func redactionTestURL(t *testing.T, baseURL, queryKey string, includeFragment bool) string { + t.Helper() + parsed, err := url.Parse(baseURL) + if err != nil { + t.Fatalf("parse redaction test URL: %v", err) + } + material := newRedactionTestMaterial() + parsed.User = url.UserPassword(material.username, material.userInfo) + if queryKey != "" { + query := parsed.Query() + query.Set(queryKey, material.queryValue) + parsed.RawQuery = query.Encode() + } + if includeFragment { + parsed.Fragment = material.fragmentValue + } + return parsed.String() +} + +func TestBuildGitContextInstructionsFileUsesSameRemoteContext(t *testing.T) { + remote := &memoryReference{files: map[string][]byte{ + defaultAgentkitfileName: fileBackedAgentkitfile(filePath), + filePath: []byte(filePrompt), + }} + final := &memoryReference{} + c := &fakeBuildClient{ + opts: client.BuildOpts{Opts: map[string]string{localNameContext: remoteContextURL}}, + expected: []expectedSolve{ + {sourcePrefix: "git://example.com/agent.git", ref: remote}, + {sourcePrefix: dockerImageSourcePrefix, ref: final}, + }, + } + + if _, err := Build(context.Background(), c); err != nil { + t.Fatalf("Build() error = %v", err) + } + if got := remote.readPaths(); !containsString(got, filePath) { + t.Fatalf("remote context reads = %v, want instructions path %q", got, filePath) + } + for _, solveSources := range c.solveSources() { + if hasSourcePrefix(solveSources, "local://context") { + t.Fatalf("remote build unexpectedly created local context source: %v", solveSources) + } + } +} + +func TestBuildHTTPArchiveInstructionsFileUsesSameRemoteContext(t *testing.T) { + archive := tarContext(t, map[string][]byte{ + defaultAgentkitfileName: fileBackedAgentkitfile(filePath), + filePath: []byte(filePrompt), + }) + rawHTTP := &memoryReference{files: map[string][]byte{downloadedHTTPContextFilename: archive}} + resolvedHTTP := &memoryReference{files: map[string][]byte{ + defaultAgentkitfileName: fileBackedAgentkitfile(filePath), + filePath: []byte(filePrompt), + }} + final := &memoryReference{} + contextURL := "https://example.com/agent-context.tar" + c := &fakeBuildClient{ + opts: client.BuildOpts{Opts: map[string]string{localNameContext: contextURL}}, + expected: []expectedSolve{ + {sourcePrefix: contextURL, ref: rawHTTP}, + {sourcePrefix: contextURL, requireAttemptUnpack: true, ref: resolvedHTTP}, + {sourcePrefix: dockerImageSourcePrefix, ref: final}, + }, + } + + if _, err := Build(context.Background(), c); err != nil { + t.Fatalf("Build() error = %v", err) + } + if got := resolvedHTTP.readPaths(); !containsString(got, filePath) { + t.Fatalf("resolved HTTP context reads = %v, want instructions path %q", got, filePath) + } + for _, solveSources := range c.solveSources() { + if hasSourcePrefix(solveSources, "local://context") { + t.Fatalf("remote build unexpectedly created local context source: %v", solveSources) + } + } +} + +func TestBuildHTTPArchiveRedactsContextURLFromOperationNames(t *testing.T) { + contextURL := redactionTestURL(t, "https://example.com/contexts/agent.tar", "sig", true) + archive := tarContext(t, map[string][]byte{ + defaultAgentkitfileName: fileBackedAgentkitfile(filePath), + filePath: []byte(filePrompt), + }) + rawHTTP := &memoryReference{files: map[string][]byte{downloadedHTTPContextFilename: archive}} + resolvedHTTP := &memoryReference{files: map[string][]byte{ + defaultAgentkitfileName: fileBackedAgentkitfile(filePath), + filePath: []byte(filePrompt), + }} + final := &memoryReference{} + c := &fakeBuildClient{ + opts: client.BuildOpts{Opts: map[string]string{localNameContext: contextURL}}, + expected: []expectedSolve{ + {sourcePrefix: contextURL, ref: rawHTTP}, + {sourcePrefix: contextURL, requireAttemptUnpack: true, ref: resolvedHTTP}, + {sourcePrefix: dockerImageSourcePrefix, ref: final}, + }, + } + + if _, err := Build(context.Background(), c); err != nil { + t.Fatalf("Build() error = %v", err) + } + assertSafeContextText(t, strings.Join(c.allOperationNames(), "\n"), "example.com/contexts/agent.tar") + if sources := fmt.Sprint(c.solveSources()); !strings.Contains(sources, contextURL) { + t.Fatalf("source resolution = %s, want raw context URL preserved", sources) + } +} + +func TestBuildGitContextRedactsContextURLFromOperationNames(t *testing.T) { + contextURL := redactionTestURL(t, "https://example.com/repos/agent.git", "ref", false) + remote := &memoryReference{files: map[string][]byte{ + defaultAgentkitfileName: fileBackedAgentkitfile(filePath), + filePath: []byte(filePrompt), + }} + final := &memoryReference{} + c := &fakeBuildClient{ + opts: client.BuildOpts{Opts: map[string]string{localNameContext: contextURL}}, + expected: []expectedSolve{ + {sourcePrefix: "git://example.com/repos/agent.git#" + newRedactionTestMaterial().queryValue, ref: remote}, + {sourcePrefix: dockerImageSourcePrefix, ref: final}, + }, + } + + if _, err := Build(context.Background(), c); err != nil { + t.Fatalf("Build() error = %v", err) + } + assertSafeContextText(t, strings.Join(c.allOperationNames(), "\n"), "example.com/repos/agent.git") +} + +func TestBuildHTTPContextRedactsContextURLFromErrors(t *testing.T) { + contextURL := redactionTestURL(t, "https://example.com/contexts/agent.yaml", "sig", true) + c := &fakeBuildClient{ + opts: client.BuildOpts{Opts: map[string]string{localNameContext: contextURL}}, + expected: []expectedSolve{ + { + sourcePrefix: contextURL, + err: fmt.Errorf("backend failed to fetch %s", contextURL), + }, + }, + } + + _, err := Build(context.Background(), c) + if err == nil { + t.Fatal("Build() error = nil, want context resolution error") + } + assertSafeContextText(t, err.Error(), "example.com/contexts/agent.yaml") +} + +func assertSafeContextText(t *testing.T, text, usefulContext string) { + t.Helper() + if !strings.Contains(text, usefulContext) { + t.Fatalf("text = %q, want useful context %q", text, usefulContext) + } + material := newRedactionTestMaterial() + for _, forbidden := range []string{ + material.username, + material.userInfo, + "?sig=", + "?ref=", + material.queryValue, + "#" + material.fragmentValue, + material.fragmentValue, + } { + if strings.Contains(text, forbidden) { + t.Fatalf("text = %q, leaked %q", text, forbidden) + } + } +} + +func TestBuildHostlessHTTPContextRedactsQueryAndFragmentFromErrors(t *testing.T) { + material := newRedactionTestMaterial() + query := url.Values{"sig": []string{material.queryValue}} + contextURL := "https://?" + query.Encode() + "#" + url.PathEscape(material.fragmentValue) + c := &fakeBuildClient{opts: client.BuildOpts{Opts: map[string]string{localNameContext: contextURL}}} + + _, err := Build(context.Background(), c) + if err == nil { + t.Fatal("Build() error = nil, want invalid HTTP context error") + } + assertSafeContextText(t, err.Error(), "HTTP(S) remote context") +} + +func TestBuildNetworkPathContextRedactsCredentialsFromErrors(t *testing.T) { + contextName := redactionTestURL(t, "//example.com/context", "sig", true) + c := &fakeBuildClient{opts: client.BuildOpts{Opts: map[string]string{localNameContext: contextName}}} + + _, err := Build(context.Background(), c) + if err == nil { + t.Fatal("Build() error = nil, want unsupported context error") + } + assertSafeContextText(t, err.Error(), "example.com/context") +} + +func TestBuildMalformedContextUsesGenericDisplayInsteadOfRawCredentials(t *testing.T) { + material := newRedactionTestMaterial() + for _, contextName := range []string{ + fmt.Sprintf("///%s:%s@example.com/context", material.username, material.userInfo), + fmt.Sprintf("%s:%s@example.com/context", material.username, material.userInfo), + } { + t.Run(contextName, func(t *testing.T) { + c := &fakeBuildClient{opts: client.BuildOpts{Opts: map[string]string{localNameContext: contextName}}} + + _, err := Build(context.Background(), c) + if err == nil { + t.Fatal("Build() error = nil, want unsupported context error") + } + assertSafeContextText(t, err.Error(), "remote context") + }) + } +} + +func TestBuildRemoteSolveErrorPreservesCancellationWithoutURLLeak(t *testing.T) { + contextURL := redactionTestURL(t, "https://example.com/contexts/agent.yaml", "sig", true) + c := &fakeBuildClient{ + opts: client.BuildOpts{Opts: map[string]string{localNameContext: contextURL}}, + expected: []expectedSolve{ + {sourcePrefix: contextURL, err: context.Canceled}, + }, + } + + _, err := Build(context.Background(), c) + if !errors.Is(err, context.Canceled) { + t.Fatalf("Build() error = %v, want context.Canceled", err) + } + assertSafeContextText(t, err.Error(), "example.com/contexts/agent.yaml") +} + +func TestBuildRemoteSolveErrorPreservesGRPCStatusWithoutURLLeak(t *testing.T) { + contextURL := redactionTestURL(t, "https://example.com/contexts/agent.yaml", "sig", true) + c := &fakeBuildClient{ + opts: client.BuildOpts{Opts: map[string]string{localNameContext: contextURL}}, + expected: []expectedSolve{ + { + sourcePrefix: contextURL, + err: status.Error(codes.Unauthenticated, "backend rejected "+contextURL), + }, + }, + } + + _, err := Build(context.Background(), c) + if got := status.Code(err); got != codes.Unauthenticated { + t.Fatalf("Build() status code = %s, want %s (error: %v)", got, codes.Unauthenticated, err) + } + assertSafeContextText(t, err.Error(), "example.com/contexts/agent.yaml") +} + +func TestBuildRemoteReadErrorPreservesDeadlineWithoutURLLeak(t *testing.T) { + contextURL := redactionTestURL(t, "https://example.com/repos/agent.git", "ref", false) + remote := &memoryReference{readErr: context.DeadlineExceeded} + c := &fakeBuildClient{ + opts: client.BuildOpts{Opts: map[string]string{localNameContext: contextURL}}, + expected: []expectedSolve{ + {sourcePrefix: "git://example.com/repos/agent.git#" + newRedactionTestMaterial().queryValue, ref: remote}, + }, + } + + _, err := Build(context.Background(), c) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("Build() error = %v, want context.DeadlineExceeded", err) + } + assertSafeContextText(t, err.Error(), "example.com/repos/agent.git") +} + +func TestBuildLocalContextInstructionsFileUsesLocalContext(t *testing.T) { + configRef := &memoryReference{files: map[string][]byte{ + defaultAgentkitfileName: fileBackedAgentkitfile(filePath), + }} + contextRef := &memoryReference{files: map[string][]byte{filePath: []byte(filePrompt)}} + final := &memoryReference{} + c := &fakeBuildClient{ + opts: client.BuildOpts{Opts: map[string]string{}}, + expected: []expectedSolve{ + {sourcePrefix: localDockerfileSourcePrefix, ref: configRef}, + {sourcePrefix: "local://context", ref: contextRef}, + {sourcePrefix: dockerImageSourcePrefix, ref: final}, + }, + } + + if _, err := Build(context.Background(), c); err != nil { + t.Fatalf("Build() error = %v", err) + } + if got := contextRef.readPaths(); !containsString(got, filePath) { + t.Fatalf("local context reads = %v, want instructions path %q", got, filePath) + } +} + +func TestBuildSingleFileHTTPContextRejectsInstructionsFileClearly(t *testing.T) { + contextURL := "https://example.com/agentkitfile.yaml" + rawHTTP := &memoryReference{files: map[string][]byte{downloadedHTTPContextFilename: fileBackedAgentkitfile(filePath)}} + c := &fakeBuildClient{ + opts: client.BuildOpts{Opts: map[string]string{localNameContext: contextURL}}, + expected: []expectedSolve{ + {sourcePrefix: contextURL, ref: rawHTTP}, + }, + } + + _, err := Build(context.Background(), c) + if err == nil { + t.Fatal("Build() error = nil, want unsupported instructions.file error") + } + for _, want := range []string{contextURL, "single-file HTTP context", "instructions.file"} { + if !strings.Contains(err.Error(), want) { + t.Fatalf("Build() error = %q, want substring %q", err, want) + } + } +} + +func TestBuildUnsupportedRemoteContextReturnsError(t *testing.T) { + contextName := "ftp://example.com/agent-context.tar" + c := &fakeBuildClient{opts: client.BuildOpts{Opts: map[string]string{localNameContext: contextName}}} + + _, err := Build(context.Background(), c) + if err == nil { + t.Fatal("Build() error = nil, want unsupported context error") + } + for _, want := range []string{contextName, "unsupported build context", "Git", "HTTP"} { + if !strings.Contains(err.Error(), want) { + t.Fatalf("Build() error = %q, want substring %q", err, want) + } + } +} + +func TestBuildMalformedGitContextReturnsError(t *testing.T) { + contextName := redactionTestURL(t, "https://example.com/agent.git", "ref", true) + c := &fakeBuildClient{opts: client.BuildOpts{Opts: map[string]string{localNameContext: contextName}}} + + _, err := Build(context.Background(), c) + if err == nil { + t.Fatal("Build() error = nil, want malformed context error") + } + if !strings.Contains(err.Error(), "ref conflicts") { + t.Fatalf("Build() error = %q, want ref-conflict context", err) + } + assertSafeContextText(t, err.Error(), "example.com/agent.git") +} + +func TestBuildNullCacheImportReturnsError(t *testing.T) { + configRef := &memoryReference{files: map[string][]byte{defaultAgentkitfileName: inlineAgentkitfile()}} + c := &fakeBuildClient{ + opts: client.BuildOpts{Opts: map[string]string{keyCacheImports: "[null]"}}, + expected: []expectedSolve{ + {sourcePrefix: localDockerfileSourcePrefix, ref: configRef}, + }, + } + + _, err := Build(context.Background(), c) + if err == nil { + t.Fatal("Build() error = nil, want invalid cache import error") + } + for _, want := range []string{keyCacheImports, "entry 0", "null"} { + if !strings.Contains(err.Error(), want) { + t.Fatalf("Build() error = %q, want substring %q", err, want) + } + } +} + +func containsString(values []string, want string) bool { + for _, value := range values { + if value == want { + return true + } + } + return false +} diff --git a/pkg/build/commands_test.go b/pkg/build/commands_test.go new file mode 100644 index 0000000..c1a7e74 --- /dev/null +++ b/pkg/build/commands_test.go @@ -0,0 +1,233 @@ +package build + +import ( + "os" + "os/exec" + "path/filepath" + "strings" + "testing" +) + +func TestRuntimeAdapterBuildTargetsHonorPlatform(t *testing.T) { + tests := []struct { + target string + dockerfile string + image string + }{ + {target: "build-serve", dockerfile: "runtimes/pydantic-ai/Dockerfile", image: "agentkit-serve:platform-test"}, + {target: "build-serve-maf", dockerfile: "runtimes/microsoft-agent-framework/Dockerfile", image: "agentkit-serve-maf:platform-test"}, + {target: "build-serve-langgraph", dockerfile: "runtimes/langgraph/Dockerfile", image: "agentkit-serve-langgraph:platform-test"}, + } + + for _, tt := range tests { + t.Run(tt.target, func(t *testing.T) { + cmd := makeAdapterDryRunCommand(tt.target) + cmd.Dir = filepath.Join("..", "..") + out, err := cmd.CombinedOutput() + if err != nil { + t.Fatalf("make dry run failed: %v\n%s", err, out) + } + command := string(out) + for _, want := range []string{ + "docker buildx build", + "-f " + tt.dockerfile, + "-t " + tt.image, + "--platform linux/arm64", + "--load", + } { + if !strings.Contains(command, want) { + t.Fatalf("%s command = %q, want substring %q", tt.target, command, want) + } + } + }) + } +} + +func makeAdapterDryRunCommand(target string) *exec.Cmd { + switch target { + case "build-serve": + return exec.Command("make", "--no-print-directory", "-n", "build-serve", "PLATFORM=linux/arm64", "TAG=platform-test") + case "build-serve-maf": + return exec.Command("make", "--no-print-directory", "-n", "build-serve-maf", "PLATFORM=linux/arm64", "TAG=platform-test") + case "build-serve-langgraph": + return exec.Command("make", "--no-print-directory", "-n", "build-serve-langgraph", "PLATFORM=linux/arm64", "TAG=platform-test") + default: + panic("unsupported adapter build target: " + target) + } +} + +func TestRunTestAgentIsHostReachableAndCapturesCurlToken(t *testing.T) { + repoRoot, err := filepath.Abs(filepath.Join("..", "..")) + if err != nil { + t.Fatalf("resolve repository root: %v", err) + } + tempDir := t.TempDir() + binDir := filepath.Join(tempDir, "bin") + if err := os.Mkdir(binDir, 0o755); err != nil { + t.Fatalf("create fake bin directory: %v", err) + } + commandLog := filepath.Join(tempDir, "commands.log") + + writeCommandStub(t, binDir, "docker", ` +{ + printf 'docker' + for arg in "$@"; do printf '\t%s' "$arg"; done + printf '\n' +} >>"${COMMAND_LOG}" +`) + + const modelKey = "model-key-must-not-appear" + const localToken = "command-capture-token" + cmd := exec.Command( + "make", + "--no-print-directory", + "run-test-agent", + "PLATFORM=linux/arm64", + "TAG=command-capture", + "LOCAL_AUTH_TOKEN="+localToken, + ) + cmd.Dir = repoRoot + cmd.Env = replaceEnvironment(os.Environ(), map[string]string{ + "COMMAND_LOG": commandLog, + "MAKEFLAGS": "", + "OPENAI_API_KEY": modelKey, + "PATH": binDir + string(os.PathListSeparator) + os.Getenv("PATH"), + }) + out, err := cmd.CombinedOutput() + if err != nil { + t.Fatalf("run-test-agent command capture failed: %v\n%s", err, out) + } + logBytes, err := os.ReadFile(commandLog) + if err != nil { + t.Fatalf("read command log: %v", err) + } + command := string(logBytes) + for _, want := range []string{ + "\t-p\t127.0.0.1:8080:8080\t", + "\t-e\tAGENTKIT_BIND=0.0.0.0\t", + "\t-e\tAGENTKIT_AUTH_TOKEN=" + localToken + "\t", + "\t-e\tOPENAI_API_KEY\t", + "\thello-agent:command-capture\n", + } { + if !strings.Contains(command, want) { + t.Fatalf("captured docker command = %q, want substring %q", command, want) + } + } + combined := string(out) + command + if strings.Contains(combined, modelKey) { + t.Fatalf("run-test-agent output leaked model key: %q", combined) + } + for _, want := range []string{ + "Authorization: Bearer " + localToken, + "http://127.0.0.1:8080/v1/models", + } { + if !strings.Contains(string(out), want) { + t.Fatalf("run-test-agent output = %q, want substring %q", out, want) + } + } +} + +func TestLiveCopilotScriptForwardsDetectedPlatformToAdapterBuild(t *testing.T) { + repoRoot, err := filepath.Abs(filepath.Join("..", "..")) + if err != nil { + t.Fatalf("resolve repository root: %v", err) + } + tempDir := t.TempDir() + binDir := filepath.Join(tempDir, "bin") + if err := os.Mkdir(binDir, 0o755); err != nil { + t.Fatalf("create fake bin directory: %v", err) + } + cacheDir := filepath.Join(tempDir, "vekil-cache") + if err := os.Mkdir(cacheDir, 0o755); err != nil { + t.Fatalf("create fake Vekil cache: %v", err) + } + commandLog := filepath.Join(tempDir, "commands.log") + + writeCommandStub(t, binDir, "docker", ` +{ + printf 'docker' + for arg in "$@"; do printf '\t%s' "$arg"; done + printf '\n' +} >>"${COMMAND_LOG}" +if [ "${1:-}" = info ]; then + printf 'arm64\n' +fi +`) + writeCommandStub(t, binDir, "make", ` +{ + printf 'make' + for arg in "$@"; do printf '\t%s' "$arg"; done + printf '\n' +} >>"${COMMAND_LOG}" +`) + writeCommandStub(t, binDir, "curl", ` +case "$*" in + */v1/models*) + printf '{"data":[{"id":"claude-haiku-4.5"}]}' + ;; + */v1/chat/completions*) + printf '{"model":"claude-haiku-4.5","choices":[{"message":{"content":"DONE42"}}]}' + ;; +esac +`) + writeCommandStub(t, binDir, "jq", ` +case "$*" in + *'.data[].id'*) printf 'claude-haiku-4.5\n' ;; + *'{model,'*) printf '{"model":"claude-haiku-4.5","content":"DONE42"}\n' ;; +esac +`) + writeCommandStub(t, binDir, "go", "") + + cmd := exec.Command("bash", "scripts/live-copilot-agent-e2e.sh") + cmd.Dir = repoRoot + cmd.Env = replaceEnvironment(os.Environ(), map[string]string{ + "COMMAND_LOG": commandLog, + "COPILOT_GITHUB_TOKEN": "", + "PATH": binDir + string(os.PathListSeparator) + os.Getenv("PATH"), + "PLATFORM": "", + "RUNNER_TEMP": tempDir, + "TAG": "command-capture", + "VEKIL_CACHE_DIR": cacheDir, + }) + out, err := cmd.CombinedOutput() + if err != nil { + t.Fatalf("live script command capture failed: %v\n%s", err, out) + } + logBytes, err := os.ReadFile(commandLog) + if err != nil { + t.Fatalf("read command log: %v", err) + } + want := "make\tbuild-serve-maf\tTAG=command-capture\tPLATFORM=linux/arm64\n" + if !strings.Contains(string(logBytes), want) { + t.Fatalf("captured commands = %q, want %q", logBytes, want) + } +} + +func writeCommandStub(t *testing.T, dir, name, body string) { + t.Helper() + path := filepath.Join(dir, name) + contents := "#!/bin/sh\nset -eu\n" + body + if err := os.WriteFile(path, []byte(contents), 0o600); err != nil { + t.Fatalf("write %s stub: %v", name, err) + } + if err := os.Chmod(path, 0o700); err != nil { + t.Fatalf("make %s stub executable: %v", name, err) + } +} + +func replaceEnvironment(base []string, replacements map[string]string) []string { + out := make([]string, 0, len(base)+len(replacements)) + for _, entry := range base { + key, _, ok := strings.Cut(entry, "=") + if ok { + if _, replace := replacements[key]; replace { + continue + } + } + out = append(out, entry) + } + for key, value := range replacements { + out = append(out, key+"="+value) + } + return out +} diff --git a/pkg/build/context.go b/pkg/build/context.go new file mode 100644 index 0000000..5406185 --- /dev/null +++ b/pkg/build/context.go @@ -0,0 +1,378 @@ +package build + +import ( + "archive/tar" + "bytes" + "context" + "fmt" + "net/url" + "path" + "strings" + + "github.com/moby/buildkit/client/llb" + "github.com/moby/buildkit/frontend/dockerfile/dfgitutil" + "github.com/moby/buildkit/frontend/dockerui" + "github.com/moby/buildkit/frontend/gateway/client" + "github.com/pkg/errors" + "github.com/sozercan/agentkit/pkg/agentkit/config" + "google.golang.org/grpc/status" +) + +const httpContextProbeSize = 1024 + +type loadedAgentkitfile struct { + config *config.AgentConfig + instructions contextFileReader +} + +type remoteContextIdentity struct { + rawURL string + displayURL string +} + +func newRemoteContextIdentity(rawURL string) remoteContextIdentity { + return remoteContextIdentity{ + rawURL: rawURL, + displayURL: safeBuildContextDisplay(rawURL), + } +} + +func (i remoteContextIdentity) description(kind string) string { + return kind + " " + i.displayURL +} + +// loadAgentkitfile resolves the authored file and binds instruction reads to the +// same build context. Remote contexts retain one resolved reference for both the +// Agentkitfile and its relative files; local builds keep BuildKit's distinct +// dockerfile and context session inputs. +func loadAgentkitfile(ctx context.Context, c client.Client) (*loadedAgentkitfile, error) { + opts := c.BuildOpts().Opts + filename := opts[keyFilename] + if filename == "" { + filename = defaultAgentkitfileName + } + + contextName := opts[localNameContext] + if contextName == "" { + return loadLocalAgentkitfile(ctx, c, opts, filename) + } + identity := newRemoteContextIdentity(contextName) + + keepGit := true + gitState, isGit, gitErr := detectGitContext(identity, &keepGit) + if isGit { + if gitErr != nil { + return nil, invalidGitContextError(identity, gitErr) + } + if gitState == nil { + return nil, errors.Errorf("invalid Git build context %q: context detection returned no state", identity.displayURL) + } + return loadRemoteAgentkitfile(ctx, c, opts, filename, gitState, identity, "Git build context") + } + + httpState, downloadedFilename, isHTTP := dockerui.DetectHTTPContext(identity.rawURL) + if isHTTP { + if err := validateHTTPContextURL(identity); err != nil { + return nil, err + } + if httpState == nil || downloadedFilename == "" { + return nil, errors.Errorf("invalid HTTP build context %q: context detection returned no source", identity.displayURL) + } + return loadHTTPAgentkitfile(ctx, c, opts, filename, identity, httpState, downloadedFilename) + } + + return nil, errors.Errorf("unsupported build context %q: expected a Git or HTTP(S) context", identity.displayURL) +} + +func loadLocalAgentkitfile(ctx context.Context, c client.Client, opts map[string]string, filename string) (*loadedAgentkitfile, error) { + name := "load agentkitfile" + if filename != defaultAgentkitfileName { + name += " from " + filename + } + state := llb.Local(localNameDockerfile, + llb.FollowPaths([]string{filename}), + llb.SessionID(c.BuildOpts().SessionID), + llb.SharedKeyHint(defaultAgentkitfileName), + dockerui.WithInternalName(name), + ) + ref, err := solveStateReference(ctx, c, &state, "local agentkitfile source") + if err != nil { + return nil, err + } + dt, err := readReferenceFile(ctx, ref, filename, "local agentkitfile source") + if err != nil { + return nil, errors.Wrap(err, "failed to read agentkitfile") + } + return parseLoadedAgentkitfile(dt, opts, localContextReader{client: c}) +} + +func loadRemoteAgentkitfile(ctx context.Context, c client.Client, opts map[string]string, filename string, state *llb.State, identity remoteContextIdentity, kind string) (*loadedAgentkitfile, error) { + description := identity.description(kind) + ref, err := solveRemoteStateReference(ctx, c, state, description) + if err != nil { + return nil, err + } + reader := referenceContextReader{ref: ref, description: description} + dt, err := reader.ReadFile(ctx, filename) + if err != nil { + return nil, errors.Wrap(err, "failed to read agentkitfile") + } + return parseLoadedAgentkitfile(dt, opts, reader) +} + +func loadHTTPAgentkitfile(ctx context.Context, c client.Client, opts map[string]string, filename string, identity remoteContextIdentity, state *llb.State, downloadedFilename string) (*loadedAgentkitfile, error) { + description := identity.description("HTTP build context") + rawRef, err := solveRemoteStateReference(ctx, c, state, description) + if err != nil { + return nil, err + } + + header, err := readRemoteReferenceFile(ctx, rawRef, client.ReadRequest{ + Filename: downloadedFilename, + Range: &client.FileRange{Length: httpContextProbeSize}, + }, description) + if err != nil { + return nil, errors.Wrap(err, "failed to inspect HTTP build context") + } + + if !isArchiveHeader(header) { + dt, err := readRemoteReferenceFile(ctx, rawRef, client.ReadRequest{Filename: downloadedFilename}, description) + if err != nil { + return nil, errors.Wrap(err, "failed to read agentkitfile") + } + reader := unsupportedContextReader{err: errors.Errorf( + "instructions.file is not supported for single-file HTTP context %q; use an HTTP archive, Git context, local context, or inline instructions", + identity.displayURL, + )} + return parseLoadedAgentkitfile(dt, opts, reader) + } + + unpacked := llb.Scratch().File( + llb.Copy(*state, path.Join("/", downloadedFilename), "/", &llb.CopyInfo{AttemptUnpack: true}), + dockerui.WithInternalName("unpack "+description), + ) + return loadRemoteAgentkitfile(ctx, c, opts, filename, &unpacked, identity, "HTTP build context") +} + +func parseLoadedAgentkitfile(dt []byte, opts map[string]string, reader contextFileReader) (*loadedAgentkitfile, error) { + cfg, err := config.NewFromBytes(dt) + if err != nil { + return nil, errors.Wrap(err, "getting config") + } + if err := parseBuildArgs(opts, cfg); err != nil { + return nil, errors.Wrap(err, "parsing build args") + } + return &loadedAgentkitfile{config: cfg, instructions: reader}, nil +} + +func solveStateReference(ctx context.Context, c client.Client, state *llb.State, description string) (client.Reference, error) { + if state == nil { + return nil, errors.Errorf("failed to resolve %s: context detection returned no state", description) + } + def, err := state.Marshal(ctx) + if err != nil { + return nil, errors.Wrapf(err, "failed to marshal %s", description) + } + result, err := c.Solve(ctx, client.SolveRequest{Definition: def.ToPB()}) + if err != nil { + return nil, errors.Wrapf(err, "failed to resolve %s", description) + } + if result == nil { + return nil, errors.Errorf("failed to resolve %s: solve returned no result", description) + } + ref, err := result.SingleRef() + if err != nil { + return nil, errors.Wrapf(err, "failed to resolve %s reference", description) + } + if ref == nil { + return nil, errors.Errorf("failed to resolve %s: solve returned no reference", description) + } + return ref, nil +} + +// solveRemoteStateReference deliberately does not expose marshal/solve error +// causes. BuildKit or a transport may echo the raw source URL in those causes; +// callers instead receive the safe host/path display URL. +func solveRemoteStateReference(ctx context.Context, c client.Client, state *llb.State, description string) (client.Reference, error) { + if state == nil { + return nil, errors.Errorf("failed to resolve %s: context detection returned no state", description) + } + def, err := state.Marshal(ctx) + if err != nil { + return nil, errors.Errorf("failed to marshal %s", description) + } + + result, err := c.Solve(ctx, client.SolveRequest{Definition: def.ToPB()}) + if err != nil { + return nil, redactedRemoteError(err, "failed to resolve %s", description) + } + if result == nil { + return nil, errors.Errorf("failed to resolve %s: solve returned no result", description) + } + ref, err := result.SingleRef() + if err != nil { + return nil, errors.Errorf("failed to resolve %s reference", description) + } + if ref == nil { + return nil, errors.Errorf("failed to resolve %s: solve returned no reference", description) + } + return ref, nil +} + +func readReferenceFile(ctx context.Context, ref client.Reference, filename, description string) ([]byte, error) { + if ref == nil { + return nil, errors.Errorf("failed to read %q from %s: context reference is nil", filename, description) + } + dt, err := ref.ReadFile(ctx, client.ReadRequest{Filename: filename}) + if err != nil { + return nil, errors.Wrapf(err, "failed to read %q from %s", filename, description) + } + return dt, nil +} + +// readRemoteReferenceFile omits the underlying cause because gateway read +// errors may include the credential-bearing source URL. +func readRemoteReferenceFile(ctx context.Context, ref client.Reference, request client.ReadRequest, description string) ([]byte, error) { + if ref == nil { + return nil, errors.Errorf("failed to read %q from %s: context reference is nil", request.Filename, description) + } + dt, err := ref.ReadFile(ctx, request) + if err != nil { + return nil, redactedRemoteError(err, "failed to read %q from %s", request.Filename, description) + } + return dt, nil +} + +func redactedRemoteError(err error, format string, args ...any) error { + message := fmt.Sprintf(format, args...) + switch { + case errors.Is(err, context.Canceled): + return errors.Wrap(context.Canceled, message) + case errors.Is(err, context.DeadlineExceeded): + return errors.Wrap(context.DeadlineExceeded, message) + } + if grpcStatus, ok := status.FromError(err); ok { + return status.Error(grpcStatus.Code(), message) + } + return errors.New(message) +} + +func validateHTTPContextURL(identity remoteContextIdentity) error { + u, err := url.Parse(identity.rawURL) + if err != nil { + return errors.Errorf("invalid HTTP build context %q: malformed URL", identity.displayURL) + } + if u.Host == "" || (u.Scheme != "http" && u.Scheme != "https") { + return errors.Errorf("invalid HTTP build context %q: expected an absolute HTTP(S) URL", identity.displayURL) + } + return nil +} + +func detectGitContext(identity remoteContextIdentity, keepGit *bool) (*llb.State, bool, error) { + gitRef, isGit, err := dfgitutil.ParseGitRef(identity.rawURL) + if err != nil { + return nil, isGit, err + } + + gitOpts := []llb.GitOption{ + llb.GitRef(gitRef.Ref), + dockerui.WithInternalName("load git source " + identity.displayURL), + } + if gitRef.KeepGitDir != nil && *gitRef.KeepGitDir { + gitOpts = append(gitOpts, llb.KeepGitDir()) + } + if keepGit != nil && *keepGit { + gitOpts = append(gitOpts, llb.KeepGitDir()) + } + if gitRef.SubDir != "" { + gitOpts = append(gitOpts, llb.GitSubDir(gitRef.SubDir)) + } + if gitRef.Checksum != "" { + gitOpts = append(gitOpts, llb.GitChecksum(gitRef.Checksum)) + } + if gitRef.Submodules != nil && !*gitRef.Submodules { + gitOpts = append(gitOpts, llb.GitSkipSubmodules()) + } + if gitRef.MTime != "" { + gitOpts = append(gitOpts, llb.GitMTime(gitRef.MTime)) + } + if gitRef.FetchByCommit { + gitOpts = append(gitOpts, llb.GitFetchByCommit()) + } + + state := llb.Git(gitRef.Remote, "", gitOpts...) + return &state, true, nil +} + +func invalidGitContextError(identity remoteContextIdentity, err error) error { + return errors.Errorf("invalid Git build context %q: %s", identity.displayURL, safeGitContextErrorSummary(err)) +} + +func safeGitContextErrorSummary(err error) string { + if err == nil { + return "malformed Git context" + } + message := err.Error() + switch { + case strings.Contains(message, "ref conflicts"): + return "ref conflicts" + case strings.Contains(message, "subdir conflicts"): + return "subdir conflicts" + case strings.Contains(message, "branch conflicts with tag"): + return "branch conflicts with tag" + case strings.Contains(message, "multiple values"): + return "invalid Git context query" + case strings.Contains(message, "invalid keep-git-dir value"): + return "invalid keep-git-dir value" + case strings.Contains(message, "invalid submodules value"): + return "invalid submodules value" + case strings.Contains(message, "invalid fetch-by-commit value"): + return "invalid fetch-by-commit value" + case strings.Contains(message, "invalid mtime value"): + return "invalid mtime value" + default: + return "malformed Git context options" + } +} + +func safeBuildContextDisplay(raw string) string { + u, err := url.Parse(raw) + if err != nil { + if strings.HasPrefix(strings.ToLower(raw), "http://") || strings.HasPrefix(strings.ToLower(raw), "https://") { + return "HTTP(S) remote context" + } + return "remote context" + } + + scheme := strings.ToLower(u.Scheme) + if u.Host != "" { + u.User = nil + u.RawQuery = "" + u.ForceQuery = false + u.Fragment = "" + u.RawFragment = "" + return u.String() + } + if scheme == "http" || scheme == "https" { + return "HTTP(S) remote context" + } + return "remote context" +} + +// isArchiveHeader mirrors BuildKit's HTTP-context archive probe so compressed +// tarballs and plain tar streams are unpacked before the Agentkitfile is read. +func isArchiveHeader(header []byte) bool { + for _, magic := range [][]byte{ + {0x42, 0x5A, 0x68}, // bzip2 + {0x1F, 0x8B, 0x08}, // gzip + {0xFD, 0x37, 0x7A, 0x58, 0x5A, 0x00}, // xz + } { + if len(header) >= len(magic) && bytes.Equal(magic, header[:len(magic)]) { + return true + } + } + + tr := tar.NewReader(bytes.NewReader(header)) + _, err := tr.Next() + return err == nil +} diff --git a/pkg/build/instructions.go b/pkg/build/instructions.go index 0ab6298..090da23 100644 --- a/pkg/build/instructions.go +++ b/pkg/build/instructions.go @@ -11,53 +11,63 @@ import ( ) // contextFileReader is the small seam between authored instruction sources and -// the BuildKit context. Tests use an in-memory Adapter; production uses the -// BuildKit gateway client Adapter below. +// the BuildKit context. Tests use an in-memory adapter; production binds it to +// either a resolved remote reference or the local context session input. type contextFileReader interface { ReadFile(ctx context.Context, path string) ([]byte, error) } -// buildkitContextReader reads files from the build context through BuildKit. -type buildkitContextReader struct { +// localContextReader reads files from BuildKit's local context session input. +type localContextReader struct { client client.Client } -func (r buildkitContextReader) ReadFile(ctx context.Context, path string) ([]byte, error) { - localSt := llb.Local(localNameContext, - llb.IncludePatterns([]string{path}), +func (r localContextReader) ReadFile(ctx context.Context, path string) ([]byte, error) { + state := llb.Local(localNameContext, + llb.FollowPaths([]string{path}), llb.SessionID(r.client.BuildOpts().SessionID), llb.SharedKeyHint("agentkit-instructions"), dockerui.WithInternalName("load instructions "+path), ) - def, err := localSt.Marshal(ctx) - if err != nil { - return nil, errors.Wrap(err, "failed to marshal instructions source") - } - res, err := r.client.Solve(ctx, client.SolveRequest{Definition: def.ToPB()}) - if err != nil { - return nil, errors.Wrap(err, "failed to resolve instructions source") - } - ref, err := res.SingleRef() + ref, err := solveStateReference(ctx, r.client, &state, "local instructions source") if err != nil { return nil, err } - dt, err := ref.ReadFile(ctx, client.ReadRequest{Filename: path}) - if err != nil { - return nil, errors.Wrap(err, "failed to read context file") - } - return dt, nil + return readReferenceFile(ctx, ref, path, "local build context") +} + +// referenceContextReader reuses one solved remote context reference, ensuring +// relative instruction files come from the same Git or HTTP archive snapshot as +// the Agentkitfile. +type referenceContextReader struct { + ref client.Reference + description string +} + +func (r referenceContextReader) ReadFile(ctx context.Context, path string) ([]byte, error) { + return readRemoteReferenceFile(ctx, r.ref, client.ReadRequest{Filename: path}, r.description) } -// resolveInstructions returns the fully-resolved system prompt: inline content -// as-is, or file contents read from the build context. -func resolveInstructions(ctx context.Context, c client.Client, cfg *config.AgentConfig) (string, error) { - return resolveInstructionSource(ctx, buildkitContextReader{client: c}, cfg.Instructions) +type unsupportedContextReader struct { + err error } +func (r unsupportedContextReader) ReadFile(context.Context, string) ([]byte, error) { + if r.err == nil { + return nil, errors.New("instructions.file is not supported by this build context") + } + return nil, r.err +} + +// resolveInstructionSource returns the fully-resolved system prompt: inline +// content as-is, or file contents read from the supplied build-context reader. func resolveInstructionSource(ctx context.Context, reader contextFileReader, source config.Source) (string, error) { if source.File == "" { return source.Inline, nil } + if reader == nil { + return "", errors.Errorf("failed to read instructions file %s: build context reader is nil", source.File) + } dt, err := reader.ReadFile(ctx, source.File) if err != nil { diff --git a/pkg/build/route_compat_test.go b/pkg/build/route_compat_test.go new file mode 100644 index 0000000..1c1559c --- /dev/null +++ b/pkg/build/route_compat_test.go @@ -0,0 +1,17 @@ +package build_test + +import ( + "testing" + + build "github.com/sozercan/agentkit/pkg/build" +) + +// Downstream users historically constructed Route with an unkeyed one-field +// literal. Keep that source-compatible layout even though Build has an internal +// context-aware dispatch path. +func TestRouteUnkeyedLiteralSourceCompatibility(t *testing.T) { + route := build.Route{nil} + if route.Handler != nil { + t.Fatal("zero handler unexpectedly became non-nil") + } +} diff --git a/pkg/build/router.go b/pkg/build/router.go index 8eca5e5..b228085 100644 --- a/pkg/build/router.go +++ b/pkg/build/router.go @@ -14,9 +14,10 @@ import ( // RouteHandler builds a result for a resolved / route. type RouteHandler func(ctx context.Context, c client.Client, cfg *config.AgentConfig, rc *RuntimeConfig) (*client.Result, error) -// Route is one entry in the flat router (plan §7.1). Output kinds are -// re-packagings of the one agent layer, so v0 has a single handler (image); -// adding agentpack/compose later is a new route, not a Build() rewrite. +type contextRouteHandler func(ctx context.Context, c client.Client, cfg *config.AgentConfig, rc *RuntimeConfig, reader contextFileReader) (*client.Result, error) + +// Route is one entry in the flat router (plan §7.1). Keep this as the original +// single-field public layout: downstream packages may use unkeyed Route literals. type Route struct { Handler RouteHandler } @@ -46,14 +47,19 @@ func (rc *RuntimeConfig) AdapterRef(opts map[string]string) string { // runtime name to its config. Both are DERIVED in init() from runtimes.Runtimes // (the single source of truth) — dispatch as data. var ( - routes = map[string]Route{} - runtimeConfigs = map[string]*RuntimeConfig{} + routes = map[string]Route{} + // Keep request-scoped context dispatch outside the exported Route value so + // its historical one-field unkeyed literals remain source-compatible. + contextRouteHandlers = map[string]contextRouteHandler{} + runtimeConfigs = map[string]*RuntimeConfig{} ) // registerRuntime wires a runtime adapter and its image output route. func registerRuntime(rc *RuntimeConfig) { runtimeConfigs[rc.Name] = rc - routes[rc.Name+"/"+utils.OutputKindImage] = Route{Handler: HandleAgent} + key := rc.Name + "/" + utils.OutputKindImage + routes[key] = Route{Handler: HandleAgent} + contextRouteHandlers[key] = handleAgent } func init() { @@ -85,6 +91,23 @@ func defaultRuntime() string { return runtimes.DefaultRuntime() } +func canonicalEffectiveRuntime(runtime string) string { + if runtime == "" { + runtime = defaultRuntime() + } + return runtimes.CanonicalRuntime(runtime) +} + +func targetRuntimeSegment(target string) (string, bool) { + if target == "" { + return "", false + } + segment, _, _ := strings.Cut(target, "/") + runtime := runtimes.CanonicalRuntime(segment) + _, registered := runtimeConfigs[runtime] + return runtime, registered +} + // lookupRoute resolves a build target plus the effective runtime to a route and // its RuntimeConfig. // @@ -93,12 +116,9 @@ func defaultRuntime() string { // target resolves to "/image". A bare runtime ("pydantic-ai") resolves // to "/image". Otherwise exact match, then longest-prefix match. func lookupRoute(target, runtime string) (matched string, route Route, rc *RuntimeConfig, ok bool) { - if runtime == "" { - runtime = defaultRuntime() - } - // Resolve a user-written alias (e.g. "maf") to its canonical name so the - // registry lookup and the "/image" route key always agree. - runtime = runtimes.CanonicalRuntime(runtime) + // Resolve an omitted or aliased runtime to the canonical adapter identity so + // the registry lookup and the "/image" route key always agree. + runtime = canonicalEffectiveRuntime(runtime) rc, rcOK := runtimeConfigs[runtime] if !rcOK { return "", Route{}, nil, false @@ -110,6 +130,14 @@ func lookupRoute(target, runtime string) (matched string, route Route, rc *Runti // an alias target would miss every branch and fail to route. target = canonicalizeTargetRuntime(target) + // A target that names a registered runtime must agree with the effective + // runtime. Otherwise route and RuntimeConfig would describe different + // adapters (for example, langgraph/image with pydantic-ai config). + targetRuntime, namesRuntime := targetRuntimeSegment(target) + if namesRuntime && targetRuntime != runtime { + return "", Route{}, nil, false + } + // Empty or bare-runtime target → the runtime's image route. if target == "" || target == runtime { r, rOK := routes[runtime+"/"+utils.OutputKindImage] diff --git a/pkg/build/router_test.go b/pkg/build/router_test.go index a1e9298..82c4d75 100644 --- a/pkg/build/router_test.go +++ b/pkg/build/router_test.go @@ -45,6 +45,13 @@ func TestLookupRouteUnknownRuntime(t *testing.T) { } } +func TestLookupRouteRejectsConflictingRuntimeTarget(t *testing.T) { + matched, _, rc, ok := lookupRoute(wantLangGraphRoute, runtimePydca) + if ok { + t.Fatalf("conflicting target resolved: matched=%q rc=%+v, want no route", matched, rc) + } +} + // TestLookupRouteLangGraph proves the LangGraph runtime resolves through the // same data-derived router as pydantic-ai and MAF. func TestLookupRouteLangGraph(t *testing.T) { @@ -120,19 +127,44 @@ func TestLookupRouteAliasAsTarget(t *testing.T) { } } -// TestLookupRouteAliasTargetEmptyRuntime documents the contract seam: when the -// target names a runtime by alias but the runtime arg is empty (→ defaults to -// pydantic-ai), the ROUTE still canonicalizes to maf/image, but the returned -// RuntimeConfig follows the runtime arg (pydantic-ai). build.go always passes the -// authoritative cfg.Runtime, so this mixed case does not arise in practice; the -// test pins the behavior so a future refactor notices if it changes. -func TestLookupRouteAliasTargetEmptyRuntime(t *testing.T) { - matched, _, rc, ok := lookupRoute(runtimeMAFAls+"/image", "") - if !ok || matched != wantMAFRoute { - t.Fatalf("matched=%q ok=%v, want %s", matched, ok, wantMAFRoute) +func TestLookupRouteRejectsConflictingAliasAndPrefixedTargets(t *testing.T) { + cases := []struct { + target string + runtime string + }{ + {runtimeMAFAls, ""}, + {runtimeMAFAls + "/image", runtimePydca}, + {runtimeMAFAls + "/image/debug", runtimePydca}, + {runtimeMAFName + "/image/debug", runtimeLangGraph}, } - if rc == nil || rc.Name != runtimePydca { - t.Fatalf("rc=%+v, want pydantic-ai (rc follows the runtime arg, not the target)", rc) + for _, tc := range cases { + matched, _, rc, ok := lookupRoute(tc.target, tc.runtime) + if ok { + t.Errorf("lookupRoute(%q, %q) resolved: matched=%q rc=%+v, want no route", tc.target, tc.runtime, matched, rc) + } + } +} + +func TestLookupRouteMatchingRuntimePreservesOutputSuffix(t *testing.T) { + cases := []struct { + target string + runtime string + wantRoute string + wantRuntime string + }{ + {wantImageRoute + "/debug", runtimePydca, wantImageRoute, runtimePydca}, + {runtimeMAFAls + "/image/debug", runtimeMAFName, wantMAFRoute, runtimeMAFName}, + {wantMAFRoute + "/debug/more", runtimeMAFAls, wantMAFRoute, runtimeMAFName}, + {wantLangGraphRoute + "/debug", runtimeLangGraph, wantLangGraphRoute, runtimeLangGraph}, + } + for _, tc := range cases { + matched, _, rc, ok := lookupRoute(tc.target, tc.runtime) + if !ok || matched != tc.wantRoute { + t.Errorf("lookupRoute(%q, %q): matched=%q ok=%v, want %q", tc.target, tc.runtime, matched, ok, tc.wantRoute) + } + if rc == nil || rc.Name != tc.wantRuntime { + t.Errorf("lookupRoute(%q, %q): rc=%+v, want runtime %q", tc.target, tc.runtime, rc, tc.wantRuntime) + } } } diff --git a/runtimes/common/agentkit_serve_common/adapter_support.py b/runtimes/common/agentkit_serve_common/adapter_support.py index 767ea07..d932fc7 100644 --- a/runtimes/common/agentkit_serve_common/adapter_support.py +++ b/runtimes/common/agentkit_serve_common/adapter_support.py @@ -2,9 +2,10 @@ The adapter modules should spend their complexity budget on translating the frozen ``agent.yaml`` ABI into their framework's concrete agent/client/tool -objects. Cross-runtime invariants live here instead: model API-key resolution, -secret-safe tool env projection, MCP timeout parsing, command validation, and -normalizing framework/model failures into the common run error. +objects. Cross-runtime invariants live here instead: cancellation-safe lifecycle +cleanup, model API-key resolution, secret-safe tool env projection, MCP timeout +parsing, command validation, and normalizing framework/model failures into the +common run error. """ from __future__ import annotations @@ -14,7 +15,10 @@ import re import shlex import subprocess -from typing import Mapping +from collections.abc import Awaitable, Callable +from contextlib import AsyncExitStack +from types import TracebackType +from typing import Mapping, TypeVar from .config import AgentSpec, ToolSpec from .conversation import FORWARDED_ROLES @@ -30,12 +34,217 @@ WORKLOAD_TOKEN_ENV = "AGENTKIT_WORKLOAD_IDENTITY_TOKEN" WORKLOAD_TOKEN_COMMAND_ENV = "AGENTKIT_WORKLOAD_IDENTITY_TOKEN_COMMAND" _BRACED_ENV_REF_RE = re.compile(r"\$\{([^}]+)\}") +_T = TypeVar("_T") class AgentBuildError(Exception): """Raised when an adapter cannot construct its concrete agent.""" +def _attach_secondary_error( + primary: BaseException, + secondary: BaseException, + *, + label: str, +) -> None: + if secondary is primary: + return + primary.add_note(f"{label} with {secondary.__class__.__name__}") + for note in getattr(secondary, "__notes__", ()): + primary.add_note(note) + if primary.__cause__ is None: + primary.__cause__ = secondary + + +async def _wait_for_owner_task( + task: asyncio.Task[_T], + *, + preserve: BaseException | None = None, +) -> _T | None: + """Wait for a lifecycle-owner task without forwarding caller cancellation.""" + cancellation: asyncio.CancelledError | None = None + while True: + try: + result = await asyncio.shield(task) + except asyncio.CancelledError as exc: + cancellation = cancellation or exc + if not task.done(): + continue + try: + result = task.result() + except BaseException as owner_error: + if preserve is not None: + _attach_secondary_error( + preserve, + owner_error, + label="lifecycle owner also failed", + ) + return None + _attach_secondary_error( + cancellation, + owner_error, + label="cleanup also failed", + ) + raise cancellation from owner_error + if preserve is not None: + preserve.add_note("cleanup wait was also cancelled") + return None + raise cancellation + except BaseException as owner_error: + if preserve is not None: + _attach_secondary_error( + preserve, + owner_error, + label="lifecycle owner also failed", + ) + return None + if cancellation is not None: + _attach_secondary_error( + cancellation, + owner_error, + label="cleanup also failed", + ) + raise cancellation from owner_error + raise + else: + if cancellation is not None: + if preserve is not None: + preserve.add_note("cleanup wait was also cancelled") + return None + raise cancellation + return result + + +async def _close_exit_stack_inline( + stack: AsyncExitStack, + exc_type: type[BaseException] | None = None, + exc: BaseException | None = None, + tb: TracebackType | None = None, + *, + preserve: BaseException | None = None, +) -> bool | None: + """Close a stack in its owner task without replacing a primary exception.""" + try: + return await stack.__aexit__(exc_type, exc, tb) + except BaseException as cleanup_error: + if preserve is None: + raise + _attach_secondary_error( + preserve, + cleanup_error, + label="cleanup also failed", + ) + return None + + +class AsyncExitStackLifecycle: + """Keep async resource entry and exit on one cancellation-isolated task.""" + + def __init__(self, stack: AsyncExitStack) -> None: + self.stack = stack + self._owner_task: asyncio.Task[bool | None] | None = None + self._close_request: asyncio.Future[ + tuple[ + type[BaseException] | None, + BaseException | None, + TracebackType | None, + ] + ] | None = None + self._closing = False + + async def enter(self, start: Callable[[], Awaitable[_T]]) -> _T: + """Run startup in the owner task and return its initialized runtime.""" + if self._owner_task is not None: + raise RuntimeError("async lifecycle is already entered") + + loop = asyncio.get_running_loop() + ready: asyncio.Future[_T] = loop.create_future() + close_request = loop.create_future() + self._close_request = close_request + owner_task = asyncio.create_task(self._run(start, ready, close_request)) + self._owner_task = owner_task + + try: + return await asyncio.shield(ready) + except BaseException as exc: + if not owner_task.done() and not self._closing: + owner_task.cancel() + await _wait_for_owner_task(owner_task, preserve=exc) + if ready.done() and not ready.cancelled(): + owner_error = ready.exception() + if owner_error is not None: + _attach_secondary_error( + exc, + owner_error, + label="lifecycle owner also failed", + ) + self._clear_if_done() + raise + + async def exit( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> bool | None: + """Request owner-task cleanup and wait for every exit callback to finish.""" + owner_task = self._owner_task + close_request = self._close_request + if owner_task is None or close_request is None: + return None + if not close_request.done(): + close_request.set_result((exc_type, exc, tb)) + try: + return await _wait_for_owner_task(owner_task, preserve=exc) + finally: + self._clear_if_done() + + async def _run( + self, + start: Callable[[], Awaitable[_T]], + ready: asyncio.Future[_T], + close_request: asyncio.Future[ + tuple[ + type[BaseException] | None, + BaseException | None, + TracebackType | None, + ] + ], + ) -> bool | None: + try: + result = await start() + ready.set_result(result) + exc_type, exc, tb = await close_request + except BaseException as error: + self._closing = True + await _close_exit_stack_inline( + self.stack, + type(error), + error, + error.__traceback__, + preserve=error, + ) + if not ready.done(): + ready.set_exception(error) + return None + raise + + self._closing = True + return await _close_exit_stack_inline( + self.stack, + exc_type, + exc, + tb, + preserve=exc, + ) + + def _clear_if_done(self) -> None: + if self._owner_task is not None and self._owner_task.done(): + self._owner_task = None + self._close_request = None + self._closing = False + + def _env_get(name: str, env: Mapping[str, str] | None = None) -> str | None: """Resolve one env var with per-run values taking precedence over process env.""" if env is not None and name in env: @@ -203,13 +412,34 @@ def resolve_workload_identity_token(audience: str, env: Mapping[str, str] | None f"{WORKLOAD_TOKEN_COMMAND_ENV}, set {WORKLOAD_TOKEN_ENV}, or install azure-identity" ) from exc - credential = DefaultAzureCredential() + identity_client = DefaultAzureCredential() + try: + value = identity_client.get_token(audience).token + except BaseException as exc: + error: BaseException + if isinstance(exc, Exception): + error = AgentBuildError( + f"workload identity token acquisition failed for audience {audience!r}: {exc}" + ) + else: + error = exc + try: + identity_client.close() + except BaseException as cleanup_error: + error.add_note( + f"credential cleanup also failed with {cleanup_error.__class__.__name__}" + ) + if error is exc: + raise + raise error from exc + try: - return credential.get_token(audience).token + identity_client.close() except Exception as exc: # noqa: BLE001 - normalize provider failures. raise AgentBuildError( - f"workload identity token acquisition failed for audience {audience!r}: {exc}" + f"workload identity credential cleanup failed for audience {audience!r}: {exc}" ) from exc + return value def same_origin_mcp_httpx_client_factory(tool: ToolSpec, url: str, *, timeout: float | int | None): diff --git a/runtimes/common/agentkit_serve_common/brokered.py b/runtimes/common/agentkit_serve_common/brokered.py index e3a4fee..cf1a0b1 100644 --- a/runtimes/common/agentkit_serve_common/brokered.py +++ b/runtimes/common/agentkit_serve_common/brokered.py @@ -28,6 +28,10 @@ def _load_orka_tool_crd_documents(raw: str) -> list[Any]: return [doc for doc in safe_load_all_lossless(raw) if doc is not None] +_ORKA_TOOL_API_VERSION = "core.orka.ai/v1alpha1" +_ORKA_TOOL_KIND = "Tool" +_ORKA_BROKERED_TOOL_CLASSES = ("read", "write", "coordination") + def brokered_tool_definitions(spec: AgentSpec) -> list[BrokeredToolDefinition]: """Return runtime-safe brokered tool definitions from an AgentSpec.""" diff --git a/runtimes/common/agentkit_serve_common/config.py b/runtimes/common/agentkit_serve_common/config.py index 15db8c4..6cd7a92 100644 --- a/runtimes/common/agentkit_serve_common/config.py +++ b/runtimes/common/agentkit_serve_common/config.py @@ -8,6 +8,8 @@ from __future__ import annotations +import base64 +import binascii import hashlib import json import math @@ -55,6 +57,62 @@ } _BROKERED_CLASSES = {"read", "write", "coordination"} _BROKERED_SCHEMA_DIGEST_RE = re.compile(r"^sha256:[0-9a-f]{64}$") +_BROKERED_BASIC_VALUE_RE = re.compile( + r"(?:^|[^A-Za-z0-9_ſK])" + r"[bB][aA][sSſ][iI][cC]" + r"[^A-Za-z0-9_ſK]+?([A-Za-z0-9+/ſK]+={0,2})" +) +_BROKERED_BASE64_CHARS = frozenset( + "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=" +) +_BROKERED_GO_WHITESPACE = frozenset( + "\t\n\v\f\r \u0085\u00a0\u1680" + "\u2000\u2001\u2002\u2003\u2004\u2005\u2006\u2007\u2008\u2009\u200a" + "\u2028\u2029\u202f\u205f\u3000" +) +# Go 1.26's unicode.Version is 15.0. Python 3.14 uses Unicode 16 and lowercases +# these newly assigned capitals, while Go deliberately leaves them unchanged. +_BROKERED_GO15_LOWER_IDENTITY_OVERRIDES = frozenset( + {"\u1c89", "\ua7cb", "\ua7cc", "\ua7da", "\ua7dc"} + | {chr(codepoint) for codepoint in range(0x10D50, 0x10D66)} +) +_HARMLESS_BROKERED_TOKEN_COUNT_INTENTS = frozenset( + { + "count", + "counting", + "counts", + "measure", + "measures", + "measuring", + "report", + "reporting", + "reports", + "track", + "tracking", + } +) +_HARMLESS_BROKERED_TOKEN_QUALIFIERS = frozenset( + {"completion", "context", "count", "counting", "counts", "input", "model", "output", "prompt", "usage"} +) +_HARMLESS_BROKERED_TOKEN_WORDS = frozenset( + {"tokenization", "tokenize", "tokenized", "tokenizer", "tokenizers", "tokenizing"} +) +_BROKERED_CREDENTIAL_TOKEN_WORDS = frozenset( + { + "accesstoken", + "apitoken", + "authtoken", + "authenticationtoken", + "authorizationtoken", + "bearertoken", + "credentialtoken", + "identitytoken", + "oauthtoken", + "refreshtoken", + "secrettoken", + "sessiontoken", + } +) _MAX_BROKERED_SCHEMA_BYTES = 64 * 1024 _MAX_EXACT_JSON_FLOAT_INTEGER = (1 << 53) - 1 _UNSAFE_BROKERED_FIELD_NAMES = { @@ -113,17 +171,284 @@ def _contains_secret_prefix(value: str) -> bool: return any(prefix in value for prefix in _SECRET_VALUE_PREFIXES) -def _unsafe_brokered_text(value: str) -> bool: - lowered = value.lower() - normalized = re.sub(r"[^a-z0-9]", "", lowered) +def _normalize_brokered_key(value: str) -> str: + return "".join( + char.lower() + for char in value + if "a" <= char <= "z" or "A" <= char <= "Z" or "0" <= char <= "9" + ) + + +def _go_lower(value: str) -> str: + """Match Go's rune-by-rune strings.ToLower behavior for validator parity.""" + + lowered: list[str] = [] + for char in value: + if char == "\u0130": + # Python expands this to ``i`` + COMBINING DOT; Go's simple rune + # mapping produces one ASCII ``i``. + lowered.append("i") + elif char in _BROKERED_GO15_LOWER_IDENTITY_OVERRIDES: + lowered.append(char) + else: + lowered.append(char.lower()) + return "".join(lowered) + + +def _is_brokered_word_char(value: str) -> bool: + return ( + "a" <= value <= "z" + or "A" <= value <= "Z" + or "0" <= value <= "9" + or value == "_" + ) + + +def _contains_brokered_word(value: str, word: str) -> bool: + start = 0 + while True: + index = value.find(word, start) + if index < 0: + return False + after = index + len(word) + before_ok = index == 0 or not _is_brokered_word_char(value[index - 1]) + after_ok = after == len(value) or not _is_brokered_word_char(value[after]) + if before_ok and after_ok: + return True + start = after + + +def _split_brokered_fields(value: str) -> list[str]: + """Match Go strings.Fields for cross-language validator parity.""" + + fields: list[str] = [] + field_start: int | None = None + for index, char in enumerate(value): + if char in _BROKERED_GO_WHITESPACE: + if field_start is not None: + fields.append(value[field_start:index]) + field_start = None + elif field_start is None: + field_start = index + if field_start is not None: + fields.append(value[field_start:]) + return fields + + +def _decode_brokered_basic_value(value: str) -> bool: + start = 0 + while start < len(value) and value[start] not in _BROKERED_BASE64_CHARS: + start += 1 + end = len(value) + while end > start and value[end - 1] not in _BROKERED_BASE64_CHARS: + end -= 1 + encoded = value[start:end] + if not encoded: + return False + + candidates = [encoded] + if "=" not in encoded: + candidates.append(encoded + "=" * (-len(encoded) % 4)) + for candidate in candidates: + try: + decoded = base64.b64decode(candidate, validate=True) + except (ValueError, binascii.Error): + continue + if b":" in decoded: + return True + return False + + +def _is_brokered_basic_value(value: str) -> bool: + plain = value.strip("\"'`()[]{}<>,;.") + separator = plain.find(":") + if 0 < separator < len(plain) - 1: + return True + if _decode_brokered_basic_value(value): + return True + return any( + char in ":=" and _decode_brokered_basic_value(value[index + 1 :]) + for index, char in enumerate(value) + ) + + +def _contains_brokered_basic_auth_reference(value: str) -> bool: + lowered = _go_lower(value) + if not _contains_brokered_word(lowered, "basic"): + return False + if any( + _contains_brokered_word(lowered, word) + for word in ("auth", "authentication", "authorization") + ): + return True + if any(_is_brokered_basic_value(match.group(1)) for match in _BROKERED_BASIC_VALUE_RE.finditer(value)): + return True + + fields = _split_brokered_fields(value) + if any(_is_brokered_basic_value(field) for field in fields): + return True + for index, field in enumerate(fields): + if not _contains_brokered_word(_go_lower(field), "basic"): + continue + candidates = fields[index + 1 :] + for candidate_index, candidate in enumerate(candidates): + if _is_brokered_basic_value(candidate): + return True + plain = candidate.strip("\"'`()[]{}<>,;.") + if plain.endswith(":") and candidate_index + 1 < len(candidates): + return True + if any(char in candidate for char in ".!?;"): + break + return False + + +def _has_brokered_structured_key_shape(value: str) -> bool: + return ( + value.lstrip("\"'`([{<").startswith("-") + or any(char in value for char in "_/.[]{}()<>\"'`") + or value != _go_lower(value) + ) + + +def _ends_brokered_sentence(value: str) -> bool: + value = value.rstrip("\"'`)]}>,") + return bool(value) and value[-1] in ".!?;" + + +def _is_harmless_brokered_token_count_intent(value: str) -> bool: + return value in _HARMLESS_BROKERED_TOKEN_COUNT_INTENTS + + +def _has_adjacent_brokered_token_count_intent(fields: list[str], token_index: int) -> bool: + if ( + token_index > 0 + and not _ends_brokered_sentence(fields[token_index - 1]) + and _is_harmless_brokered_token_count_intent(_normalize_brokered_key(fields[token_index - 1])) + ): + return True + return ( + token_index > 1 + and not _ends_brokered_sentence(fields[token_index - 2]) + and _is_harmless_brokered_token_count_intent(_normalize_brokered_key(fields[token_index - 2])) + ) + + +def _is_harmless_brokered_token_qualifier(value: str) -> bool: + return value in _HARMLESS_BROKERED_TOKEN_QUALIFIERS + + +def _is_brokered_numeric_count(value: str) -> bool: + wrapped = value.lstrip("([{<") + if wrapped and wrapped[0] in "\"'`": + return False + value = value.strip("\"'`()[]{}<>,;:.") + if not value: + return False + parts = value.split(",") + digits = 0 + for index, part in enumerate(parts): + if ( + not part + or index == 0 + and len(parts) > 1 + and len(part) > 3 + or index > 0 + and len(part) != 3 + ): + return False + for char in part: + if not "0" <= char <= "9": + return False + digits += 1 + if digits > 20: + return False + return digits > 0 + + +def _is_harmless_brokered_token_use(fields: list[str], index: int, previous: str) -> bool: + if ( + _has_brokered_structured_key_shape(fields[index]) + or not _is_harmless_brokered_token_qualifier(previous) + or not _has_adjacent_brokered_token_count_intent(fields, index) + ): + return False + if index + 1 == len(fields): + return True + return index + 2 == len(fields) and _is_brokered_numeric_count(fields[index + 1]) + + +def _is_harmless_brokered_token_count_assignment(value: str, separator: int) -> bool: + if value[separator] != ":": + return False + left_fields = _split_brokered_fields(value[:separator]) + right_fields = _split_brokered_fields(value[separator + 1 :]) + if len(left_fields) < 2 or len(right_fields) != 1: + return False + key = _normalize_brokered_key(left_fields[-1]) + if key not in {"token", "tokens"}: + return False + qualifier = _normalize_brokered_key(left_fields[-2]) return ( - (_looks_like_secret_literal(value) or _contains_secret_prefix(value)) + _is_harmless_brokered_token_qualifier(qualifier) + and _has_adjacent_brokered_token_count_intent(left_fields, len(left_fields) - 1) + and _is_brokered_numeric_count(right_fields[0]) + ) + + +def _contains_brokered_credential_assignment(value: str) -> bool: + for separator, char in enumerate(value): + if char not in ":=": + continue + if _is_harmless_brokered_token_count_assignment(value, separator): + continue + fields = _split_brokered_fields(value[:separator]) + if fields and _unsafe_brokered_key(fields[-1]) is not None: + return True + return False + + +def _is_harmless_brokered_token_word(value: str) -> bool: + return value in _HARMLESS_BROKERED_TOKEN_WORDS + + +def _contains_brokered_credential_reference(value: str) -> bool: + fields = _split_brokered_fields(value) + previous = "" + for index, field in enumerate(fields): + normalized = _normalize_brokered_key(field) + if not normalized: + continue + if normalized in {"token", "tokens"}: + if not _is_harmless_brokered_token_use(fields, index, _normalize_brokered_key(previous)): + return True + previous = field + continue + if "token" in normalized and not _is_harmless_brokered_token_word(normalized): + return True + if _has_brokered_structured_key_shape(field) and _unsafe_brokered_key(field) is not None: + return True + if normalized in _BROKERED_CREDENTIAL_TOKEN_WORDS: + return True + previous = field + return False + + +def _unsafe_brokered_description(value: str) -> bool: + # Keep this in lockstep with Go's hasUnsafeBrokeredDescription. Descriptions + # allow harmless prose about basic telemetry and token counts, but not values. + lowered = _go_lower(value) + normalized = _normalize_brokered_key(lowered) + return ( + _contains_secret_prefix(value) or "://" in value - or re.search(r"\bbearer\b", lowered) is not None - or re.search(r"\bbasic\b", lowered) is not None + or _contains_brokered_word(lowered, "bearer") + or _contains_brokered_word(lowered, "credential") + or _contains_brokered_word(lowered, "credentials") + or _contains_brokered_basic_auth_reference(value) + or _contains_brokered_credential_assignment(value) + or _contains_brokered_credential_reference(value) or "authorization" in lowered or "secret" in lowered - or "token" in lowered or "password" in lowered or "passphrase" in lowered or "pwd" in lowered @@ -148,6 +473,18 @@ def _unsafe_brokered_text(value: str) -> bool: ) +def _unsafe_brokered_text(value: str) -> bool: + # Arbitrary schema strings stay stricter, matching Go's hasUnsafeBrokeredText. + # The delegated description check uses _contains_secret_prefix, which + # subsumes the older startswith-only _looks_like_secret_literal guard. + lowered = _go_lower(value) + return ( + _unsafe_brokered_description(value) + or _contains_brokered_word(lowered, "basic") + or "token" in lowered + ) + + def _canonical_number(value: int | float) -> str: if isinstance(value, bool): raise TypeError("boolean is not a JSON number") @@ -164,6 +501,18 @@ def _canonical_number(value: int | float) -> str: return out +def _canonical_json_string(value: str) -> str: + # The Go digest writer deliberately restores these JSON-compatible line + # separators after encoding, matching ensure_ascii=False here. + return json.dumps(value, ensure_ascii=False, separators=(",", ":")) + + +def _parse_canonical_json_int(token: str) -> int | float: + # json.loads normally collapses JSON -0 to integer 0, losing the sign that + # participates in the Go writer's canonical digest. + return -0.0 if token == "-0" else int(token) + + def _canonical_json(value: Any) -> str: """Return a deterministic JSON representation for digests and drift checks.""" @@ -174,7 +523,7 @@ def _canonical_json(value: Any) -> str: if isinstance(value, (int, float)) and not isinstance(value, bool): return _canonical_number(value) if isinstance(value, str): - return json.dumps(value, ensure_ascii=False, separators=(",", ":")) + return _canonical_json_string(value) if isinstance(value, list): return "[" + ",".join(_canonical_json(item) for item in value) + "]" if isinstance(value, Mapping): @@ -182,17 +531,11 @@ def _canonical_json(value: Any) -> str: for key in sorted(value): if not isinstance(key, str): raise TypeError("JSON object keys must be strings") - parts.append(json.dumps(key, ensure_ascii=False, separators=(",", ":")) + ":" + _canonical_json(value[key])) + parts.append(_canonical_json_string(key) + ":" + _canonical_json(value[key])) return "{" + ",".join(parts) + "}" raise TypeError(f"unsupported JSON value {type(value).__name__}") -def _parse_canonical_int(value: str) -> int | float: - if value == "-0": - return -0.0 - return int(value) - - def brokered_tool_schema_digest( *, name: str, @@ -212,7 +555,7 @@ def brokered_tool_schema_digest( def _unsafe_brokered_key(value: str) -> str | None: - normalized = re.sub(r"[^A-Za-z0-9]", "", value).lower() + normalized = _normalize_brokered_key(value) if normalized in _UNSAFE_BROKERED_FIELD_NAMES: return value auth_like = (normalized.startswith("auth") and not normalized.startswith("author")) or normalized.endswith("auth") @@ -703,7 +1046,7 @@ def _reject_unsafe_top_level_fields(cls, data: Any) -> Any: @field_validator("description") @classmethod def _safe_description(cls, value: str) -> str: - if _unsafe_brokered_text(value): + if _unsafe_brokered_description(value): raise ValueError("brokered tool description must not contain URLs or secret-like material") return value @@ -730,7 +1073,10 @@ def _valid_json_schema(cls, value: dict[str, Any]) -> dict[str, Any]: raise ValueError("brokered tool parameters must be JSON serializable UTF-8") from exc if len(encoded_bytes) > _MAX_BROKERED_SCHEMA_BYTES: raise ValueError("brokered tool parameters schema is too large") - cloned = json.loads(encoded, parse_int=_parse_canonical_int) + cloned = json.loads( + encoded, + parse_int=_parse_canonical_json_int, + ) if cloned.get("type") != "object": raise ValueError("brokered tool parameters schema must set type: object") _validate_json_schema_subset(cloned, path="brokeredTools[].parameters") @@ -831,6 +1177,20 @@ class AgentSpec(_Strict): model_config = ConfigDict(extra="forbid", populate_by_name=True) + @field_validator("tools") + @classmethod + def _unique_tool_names(cls, value: list[ToolSpec]) -> list[ToolSpec]: + seen: set[str] = set() + duplicates: list[str] = [] + for entry in value: + if entry.name in seen: + duplicates.append(entry.name) + seen.add(entry.name) + if duplicates: + names = ", ".join(sorted(set(duplicates))) + raise ValueError(f"duplicate tool names: {names}") + return value + @field_validator("env") @classmethod def _unique_env_names(cls, value: list[EnvVarSpec]) -> list[EnvVarSpec]: diff --git a/runtimes/common/agentkit_serve_common/conformance.py b/runtimes/common/agentkit_serve_common/conformance.py index 4827ce0..28224c9 100644 --- a/runtimes/common/agentkit_serve_common/conformance.py +++ b/runtimes/common/agentkit_serve_common/conformance.py @@ -27,6 +27,8 @@ __all__ = [ "test_healthz_open", "test_models_listing", + "test_structurally_invalid_request_body_uses_openai_error", + "test_malformed_json_uses_secret_safe_openai_error", "test_stream_true_rejected", "test_caller_tools_rejected", "test_caller_tool_choice_required_rejected", @@ -72,6 +74,57 @@ def test_models_listing(make_client, model_name): assert body["data"][0]["id"] == model_name +def test_structurally_invalid_request_body_uses_openai_error(make_client): + with make_client() as c: + r = c.post( + "/v1/chat/completions", + json={"model": "x", "messages": {"role": "user", "content": "private-value"}}, + ) + + assert r.status_code == 400 + assert r.json() == { + "error": { + "message": "request body is invalid", + "type": "invalid_request_error", + "code": "invalid_request", + } + } + + +def test_malformed_json_uses_secret_safe_openai_error(make_client): + malformed_bodies = [ + ( + "sk-private-request-value", + '{"model":"x","messages":[{"role":"user","content":"sk-private-request-value"}', + ), + ( + "sk-another-private-value", + '{"model":"x" "private":"sk-another-private-value"}', + ), + ] + with make_client() as c: + responses = [ + c.post( + "/v1/chat/completions", + content=body, + headers={"Content-Type": "application/json"}, + ) + for _, body in malformed_bodies + ] + + expected = { + "error": { + "message": "request body must be valid JSON", + "type": "invalid_request_error", + "code": "invalid_json", + } + } + for response, (secret_marker, _) in zip(responses, malformed_bodies, strict=True): + assert response.status_code == 400 + assert response.json() == expected + assert secret_marker not in response.text + + def test_stream_true_rejected(make_client): with make_client() as c: r = c.post( @@ -143,7 +196,15 @@ def test_final_message_must_be_user(make_client): "/v1/chat/completions", json={"model": "x", "messages": [{"role": "assistant", "content": "hi"}]}, ) - assert r.status_code == 400 + + assert r.status_code == 400 + assert r.json() == { + "error": { + "message": "the final message must have role 'user'", + "type": "invalid_request_error", + "code": None, + } + } def test_multi_turn_history_accepted(make_client): diff --git a/runtimes/common/agentkit_serve_common/foundry.py b/runtimes/common/agentkit_serve_common/foundry.py index 1c8b132..bbf80ae 100644 --- a/runtimes/common/agentkit_serve_common/foundry.py +++ b/runtimes/common/agentkit_serve_common/foundry.py @@ -16,14 +16,17 @@ from __future__ import annotations import asyncio +import hashlib import hmac import json import logging import math import os import re +import threading import time import uuid +from copy import deepcopy from decimal import Decimal, InvalidOperation from fractions import Fraction from pathlib import Path @@ -52,8 +55,10 @@ _DEFAULT_MAX_PENDING_RESPONSES = 128 _DEFAULT_MAX_ARGUMENT_BYTES = 8192 _DEFAULT_MAX_OUTPUT_BYTES = 64 * 1024 +_DEFAULT_MAX_RESPONSE_STATE_BYTES = 4 * 1024 * 1024 _DEFAULT_MAX_REQUEST_BODY_BYTES = 1024 * 1024 _DEFAULT_MAX_MODEL_MESSAGES_BYTES = 1024 * 1024 +_REQUEST_BODY_OVERHEAD_BYTES = 16 * 1024 _MAX_SYNTHETIC_ARRAY_ITEMS = 32 _MAX_SYNTHETIC_STRING_LENGTH = 4096 _MAX_SYNTHETIC_VALUES = 256 @@ -62,7 +67,9 @@ _MAX_PENDING_ENV = "AGENTKIT_FOUNDRY_RESPONSE_STATE_MAX_PENDING" _MAX_ARGUMENT_BYTES_ENV = "AGENTKIT_FOUNDRY_BROKERED_MAX_ARGUMENT_BYTES" _MAX_OUTPUT_BYTES_ENV = "AGENTKIT_FOUNDRY_BROKERED_MAX_OUTPUT_BYTES" +_MAX_RESPONSE_STATE_BYTES_ENV = "AGENTKIT_FOUNDRY_RESPONSE_STATE_MAX_BYTES" _MAX_REQUEST_BODY_BYTES_ENV = "AGENTKIT_FOUNDRY_MAX_REQUEST_BODY_BYTES" +_LEGACY_MAX_REQUEST_BODY_BYTES_ENV = "AGENTKIT_FOUNDRY_REQUEST_BODY_MAX_BYTES" _MAX_MODEL_MESSAGES_BYTES_ENV = "AGENTKIT_FOUNDRY_BROKERED_MAX_MODEL_MESSAGES_BYTES" _CONTINUATION_PROOF_ENV = "AGENTKIT_FOUNDRY_BROKERED_CONTINUATION_PROOF" _CONTINUATION_PROOF_HEADER = "x-agentkit-brokered-continuation-proof" @@ -70,6 +77,7 @@ _MODEL_LOOP_ENV = "AGENTKIT_FOUNDRY_BROKERED_MODEL_LOOP" _STATE_FILE_ENV = "AGENTKIT_FOUNDRY_RESPONSE_STATE_FILE" _FOUNDRY_SESSION_ENV = "FOUNDRY_AGENT_SESSION_ID" +_TERMINAL_STATE_FULL = "state_full" def _new_response_id(previous_response_id: str | None = None) -> str: @@ -144,6 +152,67 @@ def _error(message: str, status: int = 400, code: str | None = None) -> JSONResp ) +def _state_storage_error() -> JSONResponse: + return _error( + "brokered response state storage unavailable", + status=503, + code="brokered_response_state_storage_error", + ) + + +def _state_too_large_error() -> JSONResponse: + return _error( + "brokered response state exceeds the configured byte limit", + status=413, + code="brokered_response_state_too_large", + ) + + +def _state_full_error() -> JSONResponse: + return _error( + "too many pending brokered responses", + status=429, + code="brokered_response_state_full", + ) + + +def _request_too_large_error() -> JSONResponse: + return _error( + "Foundry request body is too large", + status=413, + code="request_too_large", + ) + + +def _model_response_too_large_error() -> JSONResponse: + return _error( + "model response is too large to retain safely", + status=502, + code="ModelResponseTooLarge", + ) + + +def _non_brokered_agent_run_error(exc: AgentRunError) -> JSONResponse: + if exc.status < 500: + return _error(str(exc), status=exc.status, code=exc.code) + logger.warning("non-brokered Foundry runtime request failed: %s", exc, exc_info=True) + status = exc.status if 500 <= exc.status <= 599 else 502 + return _error( + "agent runtime request failed", + status=status, + code="RuntimeFailure", + ) + + +def _non_brokered_unexpected_runtime_error() -> JSONResponse: + logger.exception("non-brokered Foundry runtime request failed unexpectedly") + return _error( + "agent runtime request failed", + status=502, + code="RuntimeFailure", + ) + + def _message_to_prompt(message: Any) -> str: if isinstance(message, str): return message @@ -151,9 +220,11 @@ def _message_to_prompt(message: Any) -> str: def _session_id_from_request(request: Request) -> str | None: - # Foundry hosted agents may pass the session as a query parameter to the - # container and expose it as x-agent-session-id externally. The AgentKit - # header keeps local standalone validation provider-neutral. + # The hosted sandbox identity is authoritative when present. Query and + # header carriers remain ordered compatibility fallbacks for local use. + value = os.environ.get(_FOUNDRY_SESSION_ENV) + if value and value.strip(): + return value.strip() for name in ("agent_session_id", "session_id"): value = request.query_params.get(name) if value and value.strip(): @@ -162,25 +233,60 @@ def _session_id_from_request(request: Request) -> str | None: value = request.headers.get(name) if value and value.strip(): return value.strip() - value = os.environ.get(_FOUNDRY_SESSION_ENV) - if value and value.strip(): + return None + + +class _SessionIdentityConflict(ValueError): + pass + + +def _clean_session_id(value: Any) -> str | None: + if isinstance(value, str) and value.strip(): return value.strip() return None -def _effective_responses_session_id(request: Request, data: Mapping[str, Any]) -> str | None: +def _effective_responses_session_id( + request: Request, + data: Mapping[str, Any], + *, + enforce_trusted_precedence: bool = False, +) -> str | None: # The hosted platform identity describes the sandbox that actually received # the request, so prefer it over caller-controlled routing fields. Body # fields remain useful for direct/local protocol fidelity when the hosted # runtime environment is unavailable. - hosted = os.environ.get(_FOUNDRY_SESSION_ENV) - if hosted and hosted.strip(): - return hosted.strip() - for name in ("agent_session_id", "session_id"): - value = data.get(name) - if isinstance(value, str) and value.strip(): - return value.strip() - return _session_id_from_request(request) + hosted = _clean_session_id(os.environ.get(_FOUNDRY_SESSION_ENV)) + if not enforce_trusted_precedence: + if hosted: + return hosted + for name in ("agent_session_id", "session_id"): + value = _clean_session_id(data.get(name)) + if value: + return value + return _session_id_from_request(request) + + # Brokered continuations persist a session binding. The hosted ingress + # contract strips/replaces x-agent-session-id, so that gateway-owned header + # and the sandbox environment must not be silently replaced by local + # body/query compatibility fields. Matching duplicates remain valid. + gateway = _clean_session_id(request.headers.get("x-agent-session-id")) + compatibility = [ + _clean_session_id(data.get("agent_session_id")), + _clean_session_id(data.get("session_id")), + _clean_session_id(request.query_params.get("agent_session_id")), + _clean_session_id(request.query_params.get("session_id")), + _clean_session_id(request.headers.get("x-agentkit-session-id")), + ] + trusted = hosted or gateway + if trusted is not None and any( + value is not None and value != trusted + for value in (hosted, gateway, *compatibility) + ): + raise _SessionIdentityConflict("conflicting Foundry session identities") + if trusted is not None: + return trusted + return next((value for value in compatibility if value), None) def _continuation_proof_matches( @@ -305,10 +411,13 @@ class _HostedResponseState: pending_calls: dict[str, _PendingCall] expires_at: float status: str = "pending" - accepted_outputs: dict[str, str] = field(default_factory=dict) + accepted_output_digests: dict[str, str] = field(default_factory=dict) + accepted_output_sizes: dict[str, int] = field(default_factory=dict) final_payload: dict[str, Any] | None = None + terminal_error: str | None = None model_messages: list[dict[str, Any]] | None = None initial_usage: dict[str, int] = field(default_factory=dict) + final_persistence_pending: bool = False class _StateExpired(KeyError): @@ -321,6 +430,26 @@ class _StateStoreFull(Exception): pass +class _StatePersistenceError(Exception): + pass + + +class _StateSizeLimitExceeded(Exception): + pass + + +class _SerializedPayloadTooLarge(ValueError): + pass + + +class _InvalidUnicodeValue(ValueError): + pass + + +class _RequestBodyTooLarge(ValueError): + pass + + def _tool_to_state_payload(tool: BrokeredToolDefinition) -> dict[str, Any]: return { "name": tool.name, @@ -382,31 +511,55 @@ def _persistence_json_bytes(value: Any) -> bytes: def _state_to_payload(state: _HostedResponseState) -> dict[str, Any]: - return { + payload = { "responseID": state.response_id, "sessionID": state.session_id, "pendingCalls": {call_id: _pending_call_to_state_payload(call) for call_id, call in state.pending_calls.items()}, "expiresAt": state.expires_at, "status": state.status, - "acceptedOutputs": dict(state.accepted_outputs), + "acceptedOutputDigests": dict(state.accepted_output_digests), + "acceptedOutputSizes": dict(state.accepted_output_sizes), "finalPayload": state.final_payload, "modelMessages": state.model_messages, "initialUsage": dict(state.initial_usage), } + if state.terminal_error is not None: + payload["terminalError"] = state.terminal_error + return payload def _state_from_payload(data: Mapping[str, Any]) -> _HostedResponseState: pending_calls_raw = data.get("pendingCalls", {}) if not isinstance(pending_calls_raw, Mapping): raise ValueError("stored pendingCalls must be an object") - accepted_outputs = data.get("acceptedOutputs", {}) - if not isinstance(accepted_outputs, Mapping): - raise ValueError("stored acceptedOutputs must be an object") + accepted_output_digests = data.get("acceptedOutputDigests") + accepted_output_sizes = data.get("acceptedOutputSizes") + if accepted_output_digests is None: + accepted_outputs = data.get("acceptedOutputs", {}) + if not isinstance(accepted_outputs, Mapping): + raise ValueError("stored acceptedOutputs must be an object") + accepted_output_digests = { + str(key): _output_digest(str(value)) for key, value in accepted_outputs.items() + } + accepted_output_sizes = { + str(key): len(str(value).encode("utf-8")) for key, value in accepted_outputs.items() + } + elif not isinstance(accepted_output_digests, Mapping): + raise ValueError("stored acceptedOutputDigests must be an object") + if accepted_output_sizes is None: + accepted_output_sizes = {} + if not isinstance(accepted_output_sizes, Mapping): + raise ValueError("stored acceptedOutputSizes must be an object") final_payload = data.get("finalPayload") + terminal_error = data.get("terminalError") model_messages = data.get("modelMessages") initial_usage = data.get("initialUsage", {}) 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): + raise ValueError("stored terminalError must be a string") + if terminal_error not in {None, _TERMINAL_STATE_FULL}: + raise ValueError("stored terminalError is invalid") if model_messages is not None and not isinstance(model_messages, list): raise ValueError("stored modelMessages must be an array") if not isinstance(initial_usage, Mapping): @@ -415,164 +568,545 @@ def _state_from_payload(data: Mapping[str, Any]) -> _HostedResponseState: expires_at = float(data.get("expiresAt") or 0) if not math.isfinite(expires_at): raise ValueError("stored expiresAt must be finite") - accepted = {str(key): str(value) for key, value in accepted_outputs.items()} - if final_payload is None and accepted: - # A persisted accepted output without a final payload means the process + accepted = {str(key): str(value) for key, value in accepted_output_digests.items()} + accepted_sizes: dict[str, int] = {} + for key, value in accepted_output_sizes.items(): + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise ValueError("stored acceptedOutputSizes values must be non-negative integers") + accepted_sizes[str(key)] = value + if any(key not in accepted for key in accepted_sizes): + raise ValueError("stored acceptedOutputSizes contains an unknown call id") + if final_payload is None and terminal_error is None and (accepted or status == "resuming"): + # Persisted in-progress state without a final payload means the process # stopped between accepting the continuation and completing the resume. - # Clear it on load so Orka can retry the same continuation instead of - # being stuck behind duplicate_continuation_in_progress until TTL expiry. + # 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 = {} return _HostedResponseState( response_id=str(data.get("responseID") or ""), session_id=str(data["sessionID"]) if data.get("sessionID") is not None else None, pending_calls={str(call_id): _pending_call_from_state_payload(call) for call_id, call in pending_calls_raw.items() if isinstance(call, Mapping)}, expires_at=expires_at, status=status, - accepted_outputs=accepted, + accepted_output_digests=accepted, + accepted_output_sizes=accepted_sizes, final_payload=final_payload, + terminal_error=terminal_error, model_messages=model_messages, initial_usage={str(key): int(value or 0) for key, value in initial_usage.items()}, ) +def _state_has_replay(state: _HostedResponseState) -> bool: + return state.final_payload is not None or state.terminal_error is not None + + class _FoundryResponseStateStore: """Continuation store for hosted Responses brokered calls. Without a file path this is in-memory only. With a file path, the store persists pending/final response state using atomic JSON writes so a restarted - single-replica/sticky deployment can resume known response IDs. + single-replica/sticky deployment can resume known response IDs. Initial model + work reserves capacity in memory; a reservation may provisionally claim a + completed replay entry, but eviction is committed only when pending state is + installed in its place. """ - def __init__(self, ttl_seconds: float, max_entries: int, state_file: str | Path | None = None) -> None: + def __init__( + self, + ttl_seconds: float, + max_entries: int, + max_bytes: int, + state_file: str | Path | None = None, + ) -> None: self.ttl_seconds = ttl_seconds self.max_entries = max_entries + self.max_bytes = max_bytes self.state_file = Path(state_file) if state_file else None self._states: dict[str, _HostedResponseState] = {} - self._reservations: set[str] = set() + self._entry_sizes: dict[str, int] = {} + self._active_resumes: dict[str, asyncio.Task[Any]] = {} + self._reservations: dict[str, str | None] = {} + self._lock = threading.RLock() self._load() @property def backend_name(self) -> str: return "file" if self.state_file else "memory" - def reserve(self, response_id: str) -> None: - self.purge_expired() - if response_id in self._states or response_id in self._reservations: - return - non_evictable = sum( + def max_accepted_output_bytes(self) -> int: + with self._lock: + now = time.time() + active_resume_ids = self._active_resume_ids() + return max( + ( + size + for response_id, state in self._states.items() + if state.expires_at > now + or (state.status == "resuming" and response_id in active_resume_ids) + for size in state.accepted_output_sizes.values() + ), + default=0, + ) + + def _active_resume_ids(self) -> set[str]: + finished = [response_id for response_id, task in self._active_resumes.items() if task.done()] + for response_id in finished: + self._active_resumes.pop(response_id, None) + return set(self._active_resumes) + + def _claimed_response_ids( + self, + states: Mapping[str, _HostedResponseState], + *, + exclude_reservation_id: str | None = None, + ) -> set[str]: + return { + claimed_response_id + for reservation_id, claimed_response_id in self._reservations.items() + if reservation_id != exclude_reservation_id + and claimed_response_id is not None + and claimed_response_id in states + } + + def _reservation_slots( + self, + states: Mapping[str, _HostedResponseState], + *, + exclude_reservation_id: str | None = None, + ) -> int: + return sum( 1 - for state in self._states.values() - if state.status != "completed" or state.final_payload is None + for reservation_id, claimed_response_id in self._reservations.items() + if reservation_id != exclude_reservation_id + and (claimed_response_id is None or claimed_response_id not in states) + ) + + @staticmethod + def _purge_expired_from( + states: dict[str, _HostedResponseState], + *, + now: float, + active_resume_ids: set[str], + ) -> bool: + expired = [ + response_id + for response_id, state in states.items() + if state.expires_at <= now and not (state.status == "resuming" and response_id in active_resume_ids) + ] + for response_id in expired: + states.pop(response_id, None) + return bool(expired) + + @staticmethod + def _evict_completed_from( + states: dict[str, _HostedResponseState], + *, + target: int, + excluded_response_ids: set[str] | None = None, + ) -> bool: + if len(states) <= target: + return False + excluded = excluded_response_ids or set() + completed = sorted( + ( + entry + for entry in states.values() + if entry.status == "completed" + and _state_has_replay(entry) + and not entry.final_persistence_pending + and entry.response_id not in excluded + ), + key=lambda entry: entry.expires_at, + ) + changed = False + for entry in completed: + states.pop(entry.response_id, None) + changed = True + if len(states) <= target: + break + return changed + + @staticmethod + def _completed_reservation_candidate( + states: Mapping[str, _HostedResponseState], + *, + excluded_response_ids: set[str], + ) -> str | None: + candidates = sorted( + ( + entry + for entry in states.values() + if entry.status == "completed" + and _state_has_replay(entry) + and not entry.final_persistence_pending + and entry.response_id not in excluded_response_ids + ), + key=lambda entry: entry.expires_at, ) - if non_evictable + len(self._reservations) >= self.max_entries: + return candidates[0].response_id if candidates else None + + def _commit( + self, + states: dict[str, _HostedResponseState], + *, + data: bytes | None = None, + entry_sizes: Mapping[str, int] | None = None, + ) -> None: + if data is None: + data = self._serialize(states) + if entry_sizes is None: + entry_sizes = self._entry_sizes_for(states) + self._persist(data) + installed = states + 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 pending_finals: + installed = dict(states) + for response_id in pending_finals: + durable = deepcopy(states[response_id]) + durable.final_persistence_pending = False + installed[response_id] = durable + self._states = installed + self._entry_sizes = {response_id: entry_sizes[response_id] for response_id in installed} + + def _serialized_state_entry_size(self, response_id: str, state: _HostedResponseState) -> int: + try: + key_data = _bounded_json_bytes(response_id, max_bytes=self.max_bytes) + state_data = _bounded_json_bytes(_state_to_payload(state), max_bytes=self.max_bytes) + except _SerializedPayloadTooLarge as exc: + raise _StateSizeLimitExceeded("Foundry response state exceeds configured byte limit") from exc + except (TypeError, ValueError) as exc: + raise _StatePersistenceError("Foundry response state serialization failed") from exc + return len(key_data) + 1 + len(state_data) + + def _entry_sizes_for( + self, + states: Mapping[str, _HostedResponseState], + *, + recompute_response_id: str | None = None, + ) -> dict[str, int]: + return { + response_id: ( + self._serialized_state_entry_size(response_id, state) + if response_id == recompute_response_id or response_id not in self._entry_sizes + else self._entry_sizes[response_id] + ) + for response_id, state in states.items() + } + + def _serialize_candidate( + self, + states: dict[str, _HostedResponseState], + *, + response_id: str, + excluded_response_ids: set[str], + ) -> tuple[bytes, dict[str, int]]: + entry_sizes = self._entry_sizes_for(states, recompute_response_id=response_id) + + envelope_bytes = len(b'{"states":{}}') + candidate_bytes = envelope_bytes + entry_sizes[response_id] + if candidate_bytes > self.max_bytes: + raise _StateSizeLimitExceeded("Foundry response state exceeds configured byte limit") + + state_count = len(states) + total_bytes = envelope_bytes + sum(entry_sizes.values()) + max(state_count - 1, 0) + excluded = {*excluded_response_ids, response_id} + evictable = sorted( + ( + state + for state in states.values() + if state.status == "completed" + and _state_has_replay(state) + and not state.final_persistence_pending + and state.response_id not in excluded + ), + key=lambda state: state.expires_at, + ) + for completed in evictable: + if total_bytes <= self.max_bytes: + break + states.pop(completed.response_id, None) + total_bytes -= entry_sizes[completed.response_id] + if state_count > 1: + total_bytes -= 1 + state_count -= 1 + if total_bytes > self.max_bytes: + raise _StateStoreFull("brokered response state byte capacity is full") + return self._serialize(states), {current_response_id: entry_sizes[current_response_id] for current_response_id in states} + + def _add_locked(self, state: _HostedResponseState, *, reservation_id: str | None = None) -> None: + if reservation_id is not None: + if reservation_id != state.response_id or reservation_id not in self._reservations: + raise RuntimeError("Foundry response state reservation is missing") + states = dict(self._states) + self._purge_expired_from(states, now=time.time(), active_resume_ids=self._active_resume_ids()) + is_new_state = state.response_id not in states + reservation_slots = self._reservation_slots( + states, + exclude_reservation_id=reservation_id, + ) + if ( + is_new_state + and len(states) + reservation_slots >= self.max_entries + and any( + entry.final_persistence_pending and entry.status == "completed" and _state_has_replay(entry) + for entry in states.values() + ) + ): + self._commit(states) + states = dict(self._states) + + claimed_response_id = self._reservations.get(reservation_id) if reservation_id is not None else None + if claimed_response_id is not None: + states.pop(claimed_response_id, None) + reservation_slots = self._reservation_slots( + states, + exclude_reservation_id=reservation_id, + ) + protected_response_ids = self._claimed_response_ids( + states, + exclude_reservation_id=reservation_id, + ) + if len(states) + reservation_slots >= self.max_entries and is_new_state: + self._evict_completed_from( + states, + target=max(self.max_entries - reservation_slots - 1, 0), + excluded_response_ids=protected_response_ids, + ) + if len(states) + reservation_slots >= self.max_entries and is_new_state: raise _StateStoreFull("too many pending brokered responses") - self._reservations.add(response_id) + states[state.response_id] = state + data, entry_sizes = self._serialize_candidate( + states, + response_id=state.response_id, + excluded_response_ids=protected_response_ids, + ) + states[state.response_id] = deepcopy(state) + self._commit(states, data=data, entry_sizes=entry_sizes) + if reservation_id is not None: + self._reservations.pop(reservation_id, None) + + def reserve(self, response_id: str) -> None: + with self._lock: + if response_id in self._states or response_id in self._reservations: + raise RuntimeError("Foundry response state ID is already in use") + states = dict(self._states) + changed = self._purge_expired_from( + states, + now=time.time(), + active_resume_ids=self._active_resume_ids(), + ) + reservation_slots = self._reservation_slots(states) + if ( + len(states) + reservation_slots >= self.max_entries + and any( + entry.final_persistence_pending and entry.status == "completed" and _state_has_replay(entry) + for entry in states.values() + ) + ): + self._commit(states) + states = dict(self._states) + changed = False + reservation_slots = self._reservation_slots(states) + + if len(states) + reservation_slots < self.max_entries: + if changed: + self._commit(states) + self._reservations[response_id] = None + return + + claimed_response_id = self._completed_reservation_candidate( + states, + excluded_response_ids=self._claimed_response_ids(states), + ) + if claimed_response_id is None: + raise _StateStoreFull("too many pending brokered responses") + if changed: + self._commit(states) + self._reservations[response_id] = claimed_response_id def release_reservation(self, response_id: str) -> None: - self._reservations.discard(response_id) + with self._lock: + self._reservations.pop(response_id, None) def add(self, state: _HostedResponseState) -> None: - self.purge_expired() - if state.response_id in self._states: - self._reservations.discard(state.response_id) - else: - reserved = state.response_id in self._reservations - self._reservations.discard(state.response_id) - self.evict_completed_to_capacity(reserve_slots=len(self._reservations) + 1) - if len(self._states) + len(self._reservations) >= self.max_entries: - if reserved: - self._reservations.add(state.response_id) - raise _StateStoreFull("too many pending brokered responses") - self._states[state.response_id] = state - self._persist() + with self._lock: + self._add_locked(state) + + def add_reserved(self, state: _HostedResponseState) -> None: + with self._lock: + self._add_locked(state, reservation_id=state.response_id) def save(self, state: _HostedResponseState) -> None: - if state.response_id in self._states: - self._states[state.response_id] = state - self._persist() + with self._lock: + if state.response_id not in self._states: + return + states = dict(self._states) + states[state.response_id] = state + data, entry_sizes = self._serialize_candidate( + states, + response_id=state.response_id, + excluded_response_ids=self._claimed_response_ids(states), + ) + states[state.response_id] = deepcopy(state) + self._commit(states, data=data, entry_sizes=entry_sizes) + + def cache_in_memory(self, state: _HostedResponseState) -> bool: + """Cache one entry after a durable transition could not be written. + + Persisted unfinalized continuations are normalized back to pending by + ``_state_from_payload`` on restart. In-process caching also preserves a + successfully computed final payload until an identical retry can persist it. + """ + + with self._lock: + if state.response_id not in self._states: + return False + states = dict(self._states) + states[state.response_id] = state + try: + _, entry_sizes = self._serialize_candidate( + states, + response_id=state.response_id, + excluded_response_ids=self._claimed_response_ids(states), + ) + except (_StateSizeLimitExceeded, _StateStoreFull, _StatePersistenceError): + return False + states[state.response_id] = deepcopy(state) + self._states = states + self._entry_sizes = entry_sizes + return True + + def mark_resume_active(self, response_id: str) -> None: + with self._lock: + task = asyncio.current_task() + if task is not None: + self._active_resumes[response_id] = task + + def mark_resume_inactive(self, response_id: str) -> None: + with self._lock: + self._active_resumes.pop(response_id, None) def evict_completed_to_capacity(self, *, reserve_slots: int = 0) -> None: - target = max(self.max_entries - reserve_slots, 0) - if len(self._states) <= target: - return - completed = sorted( - (entry for entry in self._states.values() if entry.status == "completed" and entry.final_payload is not None), - key=lambda entry: entry.expires_at, - ) - changed = False - for entry in completed: - self._states.pop(entry.response_id, None) - changed = True - if len(self._states) <= target: - break - if changed: - self._persist() + with self._lock: + states = dict(self._states) + reservation_slots = self._reservation_slots(states) + target = max(self.max_entries - reservation_slots - reserve_slots, 0) + if self._evict_completed_from( + states, + target=target, + excluded_response_ids=self._claimed_response_ids(states), + ): + self._commit(states) def get(self, response_id: str) -> _HostedResponseState: - state = self._states.get(response_id) - if state is None: - raise KeyError(response_id) - if state.expires_at <= time.time(): - self._states.pop(response_id, None) - self._persist() - raise _StateExpired(response_id, state) - self.purge_expired() - return state + with self._lock: + state = self._states.get(response_id) + if state is None: + raise KeyError(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 + ) + 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) + if current is None: + raise KeyError(response_id) + return deepcopy(current) def purge_expired(self) -> None: - now = time.time() - expired = [response_id for response_id, state in self._states.items() if state.expires_at <= now] - for response_id in expired: - self._states.pop(response_id, None) - if expired: - self._persist() + with self._lock: + states = dict(self._states) + if self._purge_expired_from(states, now=time.time(), active_resume_ids=self._active_resume_ids()): + self._commit(states) def _load(self) -> None: if self.state_file is None or not self.state_file.exists(): return try: - data = json.loads(self.state_file.read_text(encoding="utf-8")) + if self.state_file.stat().st_size > self.max_bytes: + raise _StateSizeLimitExceeded("Foundry response state file exceeds configured byte limit") + with self.state_file.open("rb") as handle: + raw = handle.read(self.max_bytes + 1) + if len(raw) > self.max_bytes: + raise _StateSizeLimitExceeded("Foundry response state file exceeds configured byte limit") + data = json.loads(raw) if not isinstance(data, Mapping): raise ValueError("Foundry response state file root must be an object") states = data.get("states", {}) if not isinstance(states, Mapping): raise ValueError("Foundry response state file states must be an object") + if len(states) > self.max_entries: + raise ValueError("Foundry response state file exceeds configured entry limit") if not all(isinstance(state, Mapping) for state in states.values()): raise ValueError("Foundry response state entries must be objects") - self._states = {str(response_id): _state_from_payload(state) for response_id, state in states.items()} + loaded_states = { + str(response_id): _state_from_payload(state) + for response_id, state in states.items() + } + self._serialize(loaded_states) + self._states = loaded_states + self._entry_sizes = {} self.purge_expired() - except (OSError, json.JSONDecodeError, ValueError) as exc: + except ( + OSError, + _StatePersistenceError, + TypeError, + UnicodeDecodeError, + json.JSONDecodeError, + RecursionError, + ValueError, + _StateSizeLimitExceeded, + ) as exc: raise RuntimeError(f"invalid Foundry response state file {self.state_file}: {exc}") from exc - def _persist(self) -> None: + def _serialize(self, states: Mapping[str, _HostedResponseState]) -> bytes: + payload = {"states": {response_id: _state_to_payload(state) for response_id, state in states.items()}} + try: + return _bounded_json_bytes(payload, max_bytes=self.max_bytes) + except _SerializedPayloadTooLarge as exc: + raise _StateSizeLimitExceeded("Foundry response state exceeds configured byte limit") from exc + except (TypeError, ValueError) as exc: + raise _StatePersistenceError("Foundry response state serialization failed") from exc + + def _persist(self, data: bytes) -> None: if self.state_file is None: return - self.state_file.parent.mkdir(parents=True, exist_ok=True) - payload = {"states": {response_id: _state_to_payload(state) for response_id, state in self._states.items()}} tmp = self.state_file.with_name(f".{self.state_file.name}.tmp") - data = _persistence_json_bytes(payload) - try: - tmp.unlink(missing_ok=True) - except OSError: - pass - flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL - if hasattr(os, "O_NOFOLLOW"): - flags |= os.O_NOFOLLOW - fd = os.open(tmp, flags, 0o600) try: + self.state_file.parent.mkdir(parents=True, exist_ok=True) + try: + tmp.unlink(missing_ok=True) + except OSError: + pass + flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL + if hasattr(os, "O_NOFOLLOW"): + flags |= os.O_NOFOLLOW + fd = os.open(tmp, flags, 0o600) with os.fdopen(fd, "wb") as handle: handle.write(data) handle.flush() os.fsync(handle.fileno()) - except Exception: + os.chmod(tmp, 0o600) + tmp.replace(self.state_file) + except OSError as exc: try: tmp.unlink(missing_ok=True) - finally: - raise - os.chmod(tmp, 0o600) - tmp.replace(self.state_file) - os.chmod(self.state_file, 0o600) + except OSError: + pass + raise _StatePersistenceError("Foundry response state persistence failed") from exc def _state_ttl_seconds(value: float | None = None) -> float: @@ -614,10 +1148,24 @@ def _max_output_bytes(value: int | None = None) -> int: return _positive_int_setting(value, env_name=_MAX_OUTPUT_BYTES_ENV, default=_DEFAULT_MAX_OUTPUT_BYTES) -def _max_request_body_bytes(value: int | None = None) -> int: +def _max_response_state_bytes(value: int | None = None) -> int: return _positive_int_setting( value, - env_name=_MAX_REQUEST_BODY_BYTES_ENV, + env_name=_MAX_RESPONSE_STATE_BYTES_ENV, + default=_DEFAULT_MAX_RESPONSE_STATE_BYTES, + ) + + +def _max_request_body_bytes(value: int | None = None) -> int: + if value is not None or os.environ.get(_MAX_REQUEST_BODY_BYTES_ENV): + return _positive_int_setting( + value, + env_name=_MAX_REQUEST_BODY_BYTES_ENV, + default=_DEFAULT_MAX_REQUEST_BODY_BYTES, + ) + return _positive_int_setting( + None, + env_name=_LEGACY_MAX_REQUEST_BODY_BYTES_ENV, default=_DEFAULT_MAX_REQUEST_BODY_BYTES, ) @@ -630,21 +1178,6 @@ def _max_model_messages_bytes(value: int | None = None) -> int: ) -class _RequestBodyTooLarge(Exception): - pass - - -async def _bounded_request_body(request: Request, *, max_bytes: int) -> bytes: - chunks: list[bytes] = [] - total = 0 - async for chunk in request.stream(): - total += len(chunk) - if total > max_bytes: - raise _RequestBodyTooLarge - chunks.append(chunk) - return b"".join(chunks) - - def _brokered_model_loop_enabled(value: bool | None = None) -> bool: if value is not None: return value @@ -673,20 +1206,33 @@ def _function_call_outputs_from_input(input_value: Any) -> list[dict[str, Any]]: return outputs +def _mapping_value_paths(value: Mapping[Any, Any], parent_path: str): + for key, child in value.items(): + yield child, f"{parent_path}.{key}" + + +def _list_value_paths(value: list[Any], parent_path: str): + for idx, child in enumerate(value): + yield child, f"{parent_path}[{idx}]" + + def _reject_nonfinite_json_values(value: Any, *, path: str = "value") -> None: - pending: list[tuple[Any, str, int]] = [(value, path, 0)] + pending: list[tuple[Any, int]] = [(iter(((value, path),)), 0)] while pending: - current, current_path, depth = pending.pop() + iterator, depth = pending[-1] + try: + current, current_path = next(iterator) + except StopIteration: + pending.pop() + continue if depth > _MAX_OUTPUT_DEPTH: raise ValueError(f"{path} is nested too deeply") if isinstance(current, float) and not math.isfinite(current): raise ValueError(f"{current_path} must be finite") if isinstance(current, Mapping): - for key, child in current.items(): - pending.append((child, f"{current_path}.{key}", depth + 1)) + pending.append((_mapping_value_paths(current, current_path), depth + 1)) elif isinstance(current, list): - for idx, child in enumerate(current): - pending.append((child, f"{current_path}[{idx}]", depth + 1)) + pending.append((_list_value_paths(current, current_path), depth + 1)) def _parse_output_float(raw: str) -> float: @@ -743,7 +1289,102 @@ def _json_object_from_output(output: Any) -> dict[str, Any]: if "error" not in parsed or not isinstance(parsed["error"], dict): raise ValueError("denied function_call_output.output.error must be an object") _reject_nonfinite_json_values(parsed, path="function_call_output.output") - return json.loads(json.dumps(parsed, allow_nan=False, separators=(",", ":"), sort_keys=True)) + return parsed + + +def _utf8_exceeds_limit(value: str, max_bytes: int) -> bool: + if len(value) > max_bytes: + return True + try: + return len(value.encode("utf-8")) > max_bytes + except UnicodeEncodeError as exc: + raise _InvalidUnicodeValue("function_call_output.output must contain valid Unicode") from exc + + +def _mapping_children(value: Mapping[Any, Any]): + for key, child in value.items(): + yield key + yield child + + +def _json_string_encoded_size(value: str, *, max_bytes: int) -> int: + size = 2 + for char in value: + codepoint = ord(char) + if 0xD800 <= codepoint <= 0xDFFF: + raise _InvalidUnicodeValue("JSON strings must contain valid Unicode") + if char in {'"', "\\", "\b", "\f", "\n", "\r", "\t"}: + size += 2 + elif codepoint < 0x20 or codepoint == 0x7F or 0x80 <= codepoint <= 0xFFFF: + size += 6 + elif codepoint > 0xFFFF: + size += 12 + else: + size += 1 + if size > max_bytes: + raise _SerializedPayloadTooLarge + return size + + +def _encode_json_bounded(value: Any, *, max_bytes: int, sort_keys: bool, collect: bool) -> bytes: + encoded = bytearray() if collect else None + total = 0 + encoder = json.JSONEncoder(allow_nan=False, separators=(",", ":"), sort_keys=sort_keys) + try: + for chunk in encoder.iterencode(value): + remaining = max_bytes - total + if len(chunk) > remaining: + raise _SerializedPayloadTooLarge + chunk_bytes = chunk.encode("utf-8") + if len(chunk_bytes) > remaining: + raise _SerializedPayloadTooLarge + total += len(chunk_bytes) + if encoded is not None: + encoded.extend(chunk_bytes) + except RecursionError as exc: + raise _SerializedPayloadTooLarge from exc + return bytes(encoded or b"") + + +async def _read_json_request_bounded(request: Request, *, max_bytes: int) -> tuple[Any, int]: + content_length = request.headers.get("content-length") + if content_length is not None: + try: + declared_length = int(content_length) + except ValueError: + declared_length = -1 + if declared_length > max_bytes: + raise _RequestBodyTooLarge + + body = bytearray() + async for chunk in request.stream(): + if len(body) + len(chunk) > max_bytes: + raise _RequestBodyTooLarge + body.extend(chunk) + return json.loads(body), len(body) + + +def _bounded_json_bytes(value: Any, *, max_bytes: int) -> bytes: + pending = [iter((value,))] + visited = 0 + while pending: + try: + current = next(pending[-1]) + except StopIteration: + pending.pop() + continue + visited += 1 + if visited > max_bytes: + raise _SerializedPayloadTooLarge + if isinstance(current, str): + _json_string_encoded_size(current, max_bytes=max_bytes) + elif isinstance(current, Mapping): + pending.append(iter(_mapping_children(current))) + elif isinstance(current, (list, tuple)): + pending.append(iter(current)) + + _encode_json_bounded(value, max_bytes=max_bytes, sort_keys=False, collect=False) + return _encode_json_bounded(value, max_bytes=max_bytes, sort_keys=True, collect=True) def _canonical_output_json(output: dict[str, Any]) -> str: @@ -751,6 +1392,10 @@ def _canonical_output_json(output: dict[str, Any]) -> str: return json.dumps(output, allow_nan=False, separators=(",", ":"), sort_keys=True) +def _output_digest(output_json: str) -> str: + return "sha256:" + hashlib.sha256(output_json.encode("utf-8")).hexdigest() + + def _first_typed_value(schema: Mapping[str, Any], key: str, expected_type: type) -> Any: if key in schema and isinstance(schema[key], expected_type): return schema[key] @@ -1514,6 +2159,62 @@ def _final_text_from_tool_output(call: _PendingCall, output: dict[str, Any]) -> return f"Brokered tool {call.tool.name} completed with output: {_canonical_output_json(tool_output)}" +def _reset_unfinalized_continuation( + store: _FoundryResponseStateStore, + state: _HostedResponseState, + *, + call_id: str, +) -> bool: + state.accepted_output_digests.pop(call_id, None) + state.accepted_output_sizes.pop(call_id, None) + state.status = "pending" + state.final_payload = None + state.terminal_error = None + state.final_persistence_pending = False + state.expires_at = time.time() + store.ttl_seconds + try: + store.save(state) + except _StatePersistenceError as exc: + logger.warning("failed to persist Foundry continuation rollback: %s", exc) + store.cache_in_memory(state) + return False + return True + + +def _complete_with_state_full( + store: _FoundryResponseStateStore, + state: _HostedResponseState, + *, + call_id: str, + resume_model_messages: list[dict[str, Any]] | None, + resume_initial_usage: Mapping[str, int], +) -> JSONResponse: + state.status = "completed" + state.final_payload = None + state.terminal_error = _TERMINAL_STATE_FULL + state.model_messages = None + state.initial_usage = {} + state.final_persistence_pending = False + state.expires_at = time.time() + store.ttl_seconds + try: + store.save(state) + except (_StateSizeLimitExceeded, _StateStoreFull): + state.model_messages = resume_model_messages + state.initial_usage = dict(resume_initial_usage) + if not _reset_unfinalized_continuation(store, state, call_id=call_id): + return _state_storage_error() + return _state_full_error() + except _StatePersistenceError as exc: + logger.warning("failed to persist terminal Foundry brokered capacity state: %s", exc) + state.final_persistence_pending = True + try: + store.cache_in_memory(state) + except (_StateSizeLimitExceeded, _StateStoreFull): + return _state_storage_error() + return _state_storage_error() + return _state_full_error() + + async def _handle_brokered_continuation( *, spec: AgentSpec, @@ -1571,6 +2272,9 @@ async def _handle_brokered_continuation( ) try: state = store.get(str(previous_response_id)) + except _StatePersistenceError as exc: + logger.warning("failed to access Foundry brokered response state: %s", exc) + return _state_storage_error() except _StateExpired: return _error("previous_response_id state has expired", status=410, code="response_state_expired") except KeyError: @@ -1592,34 +2296,74 @@ async def _handle_brokered_continuation( if call is None: return _error("unknown function_call_output call_id", status=400, code="unknown_call_id") raw_output = item.get("output") - if isinstance(raw_output, str): - try: - if len(raw_output) > max_output_bytes or len(raw_output.encode("utf-8")) > max_output_bytes: - return _error( + if not isinstance(raw_output, str): + return _error( + "function_call_output.output must be a JSON object string", + status=400, + code="invalid_function_call_output", + ) + output_exceeds_current_limit = False + existing_output_digest = state.accepted_output_digests.get(call_id) + persisted_output_size = state.accepted_output_sizes.get(call_id, 0) + replay_parse_limit = ( + max_output_bytes + if existing_output_digest is None + else max(store.max_bytes, max_output_bytes, persisted_output_size) + ) + raw_output_size = 0 + try: + if _utf8_exceeds_limit(raw_output, replay_parse_limit): + return _error( "brokered function_call_output is too large", status=413, code="brokered_output_too_large", ) - except UnicodeEncodeError as exc: - return _error(str(exc), status=400, code="invalid_function_call_output") + raw_output_size = len(raw_output.encode("utf-8")) + output_exceeds_current_limit = raw_output_size > max_output_bytes + except _InvalidUnicodeValue as exc: + return _error(str(exc), status=400, code="invalid_function_call_output") try: parsed_output = _json_object_from_output(raw_output) except ValueError as exc: return _error(str(exc), status=400, code="invalid_function_call_output") - output_json = _canonical_output_json(parsed_output) - output_size = len(output_json.encode("utf-8")) - if output_size > max_output_bytes: + try: + output_bytes = _bounded_json_bytes(parsed_output, max_bytes=replay_parse_limit) + except _SerializedPayloadTooLarge: return _error( "brokered function_call_output is too large", status=413, code="brokered_output_too_large", ) - - existing_output = state.accepted_outputs.get(call_id) - if existing_output is not None: - if existing_output == output_json and state.final_payload is not None: + except _InvalidUnicodeValue as exc: + return _error(str(exc), status=400, code="invalid_function_call_output") + if len(output_bytes) > max_output_bytes: + output_exceeds_current_limit = True + accepted_output_size = max(raw_output_size, len(output_bytes)) + output_json = output_bytes.decode("utf-8") + output_digest = _output_digest(output_json) + + if existing_output_digest is not None: + if existing_output_digest == output_digest and _state_has_replay(state): + if state.final_persistence_pending: + state.final_persistence_pending = False + try: + store.save(state) + except _StateStoreFull: + state.final_persistence_pending = True + return _state_full_error() + except _StateSizeLimitExceeded: + state.final_persistence_pending = True + return _state_too_large_error() + except _StatePersistenceError as exc: + logger.warning("failed to persist cached Foundry brokered completion: %s", exc) + state.final_persistence_pending = True + store.cache_in_memory(state) + return _state_storage_error() + if state.terminal_error == _TERMINAL_STATE_FULL: + return _state_full_error() + assert state.final_payload is not None return JSONResponse(state.final_payload) - if existing_output == output_json: + if existing_output_digest == output_digest: return _error( "matching function_call_output is already being processed", status=409, @@ -1630,6 +2374,12 @@ async def _handle_brokered_continuation( status=409, code="conflicting_duplicate_continuation", ) + if output_exceeds_current_limit: + return _error( + "brokered function_call_output is too large", + status=413, + code="brokered_output_too_large", + ) if state.model_messages is not None and model_loop is None: return _error( "brokered model-loop continuation is unavailable for this pending response", @@ -1639,50 +2389,103 @@ 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") - state.accepted_outputs[call_id] = output_json - state.expires_at = time.time() + store.ttl_seconds - store.save(state) if state.model_messages is not None and model_loop is not None: - state.status = "resuming" - store.save(state) try: - model_result = await model_loop.resume(state.model_messages, call_id=call_id, output=output_json) + model_loop.validate_static_credentials() except AgentRunError as exc: - state.accepted_outputs.pop(call_id, None) - state.status = "pending" - store.save(state) - 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) - except asyncio.CancelledError: - state.accepted_outputs.pop(call_id, None) - state.status = "pending" - store.save(state) - raise - except Exception as exc: # noqa: BLE001 - reset continuation state before surfacing unexpected model failures. - logger.exception("brokered model-loop resume failed unexpectedly") - state.accepted_outputs.pop(call_id, None) - state.status = "pending" - store.save(state) - return _error("model resume failed", status=502, code="ModelResumeError") - if not isinstance(model_result, ModelLoopFinal): - state.accepted_outputs.pop(call_id, None) - state.status = "pending" + state.accepted_output_digests[call_id] = output_digest + state.accepted_output_sizes[call_id] = accepted_output_size + state.status = "resuming" + state.expires_at = time.time() + store.ttl_seconds + try: store.save(state) - return _error( - "model requested another brokered tool after resume", - status=400, - code="tool_loop_limit_exceeded", - ) - result = RunResult(text=model_result.text, usage=_combine_usage(state.initial_usage, model_result.usage)) + except _StateStoreFull: + return _state_full_error() + except _StateSizeLimitExceeded: + return _state_too_large_error() + except _StatePersistenceError as exc: + logger.warning("failed to persist Foundry brokered continuation state: %s", exc) + return _state_storage_error() + 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) + 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) + 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") + 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", + ) + result = RunResult(text=model_result.text, usage=_combine_usage(state.initial_usage, model_result.usage)) + finally: + store.mark_resume_inactive(state.response_id) else: result = RunResult(text=_final_text_from_tool_output(call, parsed_output)) + resume_model_messages = state.model_messages + 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) + state.accepted_output_digests[call_id] = output_digest + state.accepted_output_sizes[call_id] = accepted_output_size state.status = "completed" state.final_payload = final_payload - store.save(state) - store.evict_completed_to_capacity() + state.terminal_error = None + if has_resume_transcript: + state.model_messages = None + state.initial_usage = {} + state.final_persistence_pending = False + state.expires_at = time.time() + store.ttl_seconds + try: + store.save(state) + except _StateStoreFull: + return _complete_with_state_full( + store, + state, + call_id=call_id, + resume_model_messages=resume_model_messages, + resume_initial_usage=resume_initial_usage, + ) + except _StateSizeLimitExceeded: + if has_resume_transcript: + state.model_messages = resume_model_messages + state.initial_usage = resume_initial_usage + if not _reset_unfinalized_continuation(store, state, call_id=call_id): + return _state_storage_error() + if used_model_resume: + return _model_response_too_large_error() + return _state_too_large_error() + except _StatePersistenceError as exc: + logger.warning("failed to persist completed Foundry brokered response state: %s", exc) + state.final_persistence_pending = True + store.cache_in_memory(state) + return _state_storage_error() + try: + store.evict_completed_to_capacity() + except _StatePersistenceError as exc: + logger.warning("failed to evict completed Foundry brokered response state: %s", exc) + return _state_storage_error() return JSONResponse(final_payload) @@ -1696,6 +2499,7 @@ def create_foundry_app( max_pending_responses: int | None = None, max_brokered_argument_bytes: int | None = None, max_brokered_output_bytes: int | None = None, + max_response_state_bytes: int | None = None, max_request_body_bytes: int | None = None, max_model_messages_bytes: int | None = None, brokered_model_loop_enabled: bool | None = None, @@ -1713,8 +2517,18 @@ def create_foundry_app( response_states = _FoundryResponseStateStore( ttl_seconds=_state_ttl_seconds(state_ttl_seconds), max_entries=_max_pending_responses(max_pending_responses), + max_bytes=_max_response_state_bytes(max_response_state_bytes), state_file=_response_state_file(response_state_file) if brokered_tools else None, ) + + def max_brokered_request_body_bytes() -> int: + replay_ceiling = max( + response_states.max_bytes, + max_output_bytes, + response_states.max_accepted_output_bytes(), + ) + return max(6 * replay_ceiling, max_argument_bytes) + _REQUEST_BODY_OVERHEAD_BYTES + model_loop = ( BrokeredChatModelLoop( spec, @@ -1722,6 +2536,7 @@ def create_foundry_app( http_client=brokered_model_http_client, max_argument_bytes=max_argument_bytes, max_output_bytes=max_output_bytes, + max_response_bytes=response_states.max_bytes, ) if brokered_tools and _brokered_model_loop_enabled(brokered_model_loop_enabled) else None @@ -1755,6 +2570,13 @@ async def readiness(): } if not continuation_proof: body["ready"] = False + if model_loop is not None: + try: + await model_loop.validate_credentials() + except AgentRunError: + body["ready"] = False + body["foundryResponses"]["modelAuth"] = "missing" + if not body["ready"]: return JSONResponse(body, status_code=503) return body @@ -1767,12 +2589,10 @@ async def invocations(request: Request): code="invocations_disabled_in_brokered_mode", ) try: - raw_body = await _bounded_request_body(request, max_bytes=request_body_limit) + data, _ = await _read_json_request_bounded(request, max_bytes=request_body_limit) except _RequestBodyTooLarge: - return Response("Request body is too large", status_code=413) - try: - data = json.loads(raw_body) - except (json.JSONDecodeError, UnicodeDecodeError, RecursionError): + return _request_too_large_error() + except (UnicodeDecodeError, RecursionError, ValueError): return Response("Request body must be JSON", status_code=400) if not isinstance(data, dict): @@ -1786,21 +2606,28 @@ async def invocations(request: Request): RunRequest(prompt=prompt, session_id=_session_id_from_request(request)) ) except AgentRunError as exc: - return _error(str(exc), status=exc.status, code=exc.code) - except Exception as exc: # noqa: BLE001 - deterministic protocol envelope. - return _error(str(exc), status=502, code=exc.__class__.__name__) + return _non_brokered_agent_run_error(exc) + except Exception: # noqa: BLE001 - deterministic protocol envelope. + return _non_brokered_unexpected_runtime_error() return JSONResponse({"response": result.text, "usage": _usage(result)}) @app.post("/responses", dependencies=[auth]) async def responses(request: Request): try: - raw_body = await _bounded_request_body(request, max_bytes=request_body_limit) + data, request_body_size = await _read_json_request_bounded( + request, + max_bytes=max_brokered_request_body_bytes() if brokered_tools else request_body_limit, + ) except _RequestBodyTooLarge: - return _error("Request body is too large", status=413, code="request_body_too_large") - try: - data = json.loads(raw_body) - except (json.JSONDecodeError, UnicodeDecodeError, RecursionError): + if not brokered_tools: + return _request_too_large_error() + return _error( + "brokered Responses request body is too large", + status=413, + code="brokered_request_too_large", + ) + except (UnicodeDecodeError, RecursionError, ValueError): return _error("Request body must be JSON", status=400, code="invalid_json") if not isinstance(data, dict): @@ -1824,9 +2651,26 @@ async def responses(request: Request): if "input" not in data: return _error("Missing 'input' in request", status=400, code="missing_input") - session_id = _effective_responses_session_id(request, data) + try: + session_id = _effective_responses_session_id( + request, + data, + enforce_trusted_precedence=bool(brokered_tools), + ) + except _SessionIdentityConflict: + return _error( + "request contains conflicting Foundry session identities", + status=409, + code="response_session_mismatch", + ) 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: + return _error( + "Request body is too large", + status=413, + code="request_body_too_large", + ) if brokered_tools and function_outputs: return await _handle_brokered_continuation( spec=spec, @@ -1843,6 +2687,9 @@ async def responses(request: Request): if brokered_tools and isinstance(previous_response_id, str) and previous_response_id: try: previous_state = response_states.get(previous_response_id) + except _StatePersistenceError as exc: + 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"}: return _error("previous_response_id state has expired", status=410, code="response_state_expired") @@ -1873,6 +2720,10 @@ async def responses(request: Request): ) previous_response_id_for_output = previous_response_id if isinstance(previous_response_id, str) and previous_response_id else None if model_loop is not None: + try: + model_loop.validate_static_credentials() + except AgentRunError as exc: + return _error(str(exc), status=exc.status, code=exc.code) response_id = _new_response_id(previous_response_id_for_output) call_id = f"call_{response_id}_1" try: @@ -1883,13 +2734,24 @@ async def responses(request: Request): status=429, code="brokered_response_state_full", ) + except _StatePersistenceError as exc: + logger.warning("failed to persist Foundry brokered response state: %s", exc) + return _state_storage_error() try: try: 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) 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)) + return JSONResponse( + _responses_payload( + spec, + RunResult(text=model_result.text, usage=model_result.usage), + previous_response_id=previous_response_id_for_output, + ) + ) 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") @@ -1905,21 +2767,19 @@ async def responses(request: Request): code="brokered_arguments_too_large", ) try: - model_messages_bytes = len( - _persistence_json_bytes(model_result.messages) + _bounded_json_bytes(model_result.messages, max_bytes=model_messages_limit) + except _SerializedPayloadTooLarge: + return _error( + "model loop messages are too large for pending state", + status=413, + code="brokered_model_messages_too_large", ) - except (TypeError, ValueError, RecursionError, UnicodeEncodeError) as exc: + except (_InvalidUnicodeValue, TypeError, ValueError) as exc: return _error( f"model loop messages are invalid: {exc}", status=502, code="InvalidModelResponse", ) - if model_messages_bytes > model_messages_limit: - return _error( - "model loop messages are too large for pending state", - status=413, - code="brokered_model_messages_too_large", - ) call = _PendingCall( call_id=call_id, item_id=_new_function_call_id(response_id), @@ -1934,7 +2794,19 @@ async def responses(request: Request): model_messages=model_result.messages, initial_usage=dict(model_result.usage), ) - response_states.add(state) + try: + response_states.add_reserved(state) + except _StateStoreFull: + return _error( + "too many pending brokered responses", + status=429, + code="brokered_response_state_full", + ) + except _StateSizeLimitExceeded: + return _state_too_large_error() + except _StatePersistenceError as exc: + logger.warning("failed to persist Foundry brokered response state: %s", exc) + return _state_storage_error() return JSONResponse( _function_call_response_payload( spec, @@ -1987,6 +2859,11 @@ async def responses(request: Request): status=429, code="brokered_response_state_full", ) + except _StateSizeLimitExceeded: + return _state_too_large_error() + except _StatePersistenceError as exc: + logger.warning("failed to persist Foundry brokered response state: %s", exc) + return _state_storage_error() return JSONResponse( _function_call_response_payload( spec, @@ -1999,9 +2876,9 @@ async def responses(request: Request): try: result = await request.app.state.runtime.run(run_request) except AgentRunError as exc: - return _error(str(exc), status=exc.status, code=exc.code) - except Exception as exc: # noqa: BLE001 - deterministic protocol envelope. - return _error(str(exc), status=502, code=exc.__class__.__name__) + 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)) diff --git a/runtimes/common/agentkit_serve_common/foundry_model_loop.py b/runtimes/common/agentkit_serve_common/foundry_model_loop.py index 82c94dc..7cbba5b 100644 --- a/runtimes/common/agentkit_serve_common/foundry_model_loop.py +++ b/runtimes/common/agentkit_serve_common/foundry_model_loop.py @@ -19,7 +19,7 @@ import httpx -from .adapter_support import AgentBuildError, resolve_api_key, resolve_workload_identity_token +from .adapter_support import AgentBuildError, NO_AUTH_API_KEY, resolve_api_key, resolve_workload_identity_token from .config import AgentSpec from .conversation import FORWARDED_ROLES, RunRequest from .runtime import AgentRunError, BrokeredToolDefinition @@ -52,12 +52,14 @@ def __init__( http_client: httpx.AsyncClient | None = None, max_argument_bytes: int = 8192, max_output_bytes: int = 64 * 1024, + max_response_bytes: int = 4 * 1024 * 1024, ) -> None: self.spec = spec self.tools = list(tools) self.http_client = http_client self.max_argument_bytes = max_argument_bytes self.max_output_bytes = max_output_bytes + self.max_response_bytes = max_response_bytes self.tools_by_name = {tool.name: tool for tool in self.tools} async def start(self, request: RunRequest, *, call_id: str) -> ModelLoopFinal | ModelLoopToolRequest: @@ -67,7 +69,7 @@ async def start(self, request: RunRequest, *, call_id: str) -> ModelLoopFinal | usage = _usage(data) tool_calls = message.get("tool_calls") if not tool_calls: - return ModelLoopFinal(text=_message_text(message), usage=usage) + return ModelLoopFinal(text=_message_text(message, max_bytes=self.max_output_bytes), usage=usage) 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", @@ -95,7 +97,7 @@ async def start(self, request: RunRequest, *, call_id: str) -> ModelLoopFinal | except UnicodeEncodeError as exc: raise AgentRunError("model tool arguments must contain valid Unicode", status=400, code="InvalidToolArguments") from exc arguments = _parse_arguments(raw_arguments) - _validate_json_unicode(arguments) + _validate_argument_unicode(arguments) argument_text = json.dumps(arguments, separators=(",", ":"), sort_keys=True) assistant_message = { "role": "assistant", @@ -111,7 +113,17 @@ 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: - if len(output.encode("utf-8")) > self.max_output_bytes: + 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: + output_bytes = output.encode("utf-8") + except UnicodeEncodeError as exc: + raise AgentRunError( + "brokered tool output must contain valid Unicode", + status=400, + code="InvalidToolOutput", + ) from exc + if len(output_bytes) > self.max_output_bytes: 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}) @@ -119,7 +131,18 @@ async def resume(self, messages: Sequence[Mapping[str, Any]], *, call_id: str, o 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), usage=_usage(data)) + return ModelLoopFinal(text=_message_text(message, max_bytes=self.max_output_bytes), usage=_usage(data)) + + async def validate_credentials(self) -> None: + await self._auth_headers() + + def validate_static_credentials(self) -> None: + if self.spec.model.auth is not None: + return + try: + resolve_api_key(self.spec) + except AgentBuildError as exc: + raise _model_auth_missing_error() from exc def _initial_messages(self, request: RunRequest) -> list[dict[str, Any]]: messages: list[dict[str, Any]] = [] @@ -152,39 +175,93 @@ async def _chat(self, messages: Sequence[Mapping[str, Any]], *, tools: Sequence[ if tools: payload["tools"] = list(tools) payload["tool_choice"] = "auto" + headers = await self._auth_headers() client = self.http_client close_client = False if client is None: - headers: dict[str, str] = {} - try: - auth = self.spec.model.auth - if auth is not None and auth.type == "workload-identity-token": - token = os.environ.get("AGENTKIT_MODEL_WORKLOAD_IDENTITY_TOKEN") - if not token: - token = await asyncio.to_thread(resolve_workload_identity_token, auth.audience or "") - headers["Authorization"] = f"Bearer {token}" - else: - api_key = resolve_api_key(self.spec) - headers["Authorization"] = f"Bearer {api_key}" - except AgentBuildError as exc: - raise AgentRunError(str(exc), status=400, code="ModelAuthMissing") from exc client = httpx.AsyncClient(headers=headers, timeout=60) close_client = True try: - response = await client.post(_chat_completions_url(self.spec.model.base_url), json=payload) - response.raise_for_status() - data = response.json() + 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, + ) except httpx.HTTPStatusError as exc: - raise AgentRunError(str(exc), status=exc.response.status_code, code="ModelHTTPError") from exc - except Exception as exc: # noqa: BLE001 - normalize transport/model failures. - raise AgentRunError(str(exc), status=502, code=exc.__class__.__name__) from exc + raise _normalized_model_http_error(exc.response.status_code) from exc + except AgentRunError: + raise + except Exception as exc: # noqa: BLE001 - normalize transport/model failures without leaking request URLs. + raise AgentRunError( + "model service request failed", + status=502, + code="ModelUpstreamError", + ) from exc finally: if close_client: await client.aclose() + try: + data = json.loads(response_body) + except Exception as exc: # noqa: BLE001 - normalize decoder failures without exposing response internals. + raise AgentRunError( + "model service returned an invalid JSON response", + status=502, + code="InvalidModelResponse", + ) from exc if not isinstance(data, dict): raise AgentRunError("model response must be a JSON object", status=502, code="InvalidModelResponse") + _validate_model_response_unicode(data) return data + async def _auth_headers(self) -> dict[str, str]: + headers: dict[str, str] = {} + try: + auth = self.spec.model.auth + if auth is not None and auth.type == "workload-identity-token": + token = os.environ.get("AGENTKIT_MODEL_WORKLOAD_IDENTITY_TOKEN") + if not token: + token = await asyncio.to_thread(resolve_workload_identity_token, auth.audience or "") + headers["Authorization"] = f"Bearer {token}" + else: + api_key = resolve_api_key(self.spec) + if api_key != NO_AUTH_API_KEY: + headers["Authorization"] = f"Bearer {api_key}" + except AgentBuildError as exc: + raise _model_auth_missing_error() from exc + return headers + + +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: + try: + declared_length = int(content_length) + except ValueError: + declared_length = -1 + if declared_length > max_bytes: + raise _model_response_too_large_error() + + body = bytearray() + async for chunk in response.aiter_bytes(): + if len(body) + len(chunk) > max_bytes: + raise _model_response_too_large_error() + body.extend(chunk) + return body + + +def _model_response_too_large_error() -> AgentRunError: + return AgentRunError( + "model response is too large to retain safely", + status=502, + code="ModelResponseTooLarge", + ) + def _chat_completions_url(base_url: str) -> str: root = base_url.rstrip("/") @@ -193,6 +270,55 @@ def _chat_completions_url(base_url: str) -> str: return f"{root}/chat/completions" +def _normalized_model_http_error(status_code: int) -> AgentRunError: + if status_code in {401, 403}: + return AgentRunError( + "model service rejected configured credentials", + status=503, + code="ModelAuthRejected", + ) + if status_code == 429 or status_code >= 500: + return AgentRunError( + "model service is unavailable", + status=503, + code="ModelUnavailable", + ) + return AgentRunError( + "model service request failed", + status=502, + code="ModelUpstreamError", + ) + + +def _validate_model_response_unicode(value: Any) -> None: + pending = [iter((value,))] + while pending: + try: + current = next(pending[-1]) + except StopIteration: + pending.pop() + continue + if isinstance(current, str): + if any(0xD800 <= ord(char) <= 0xDFFF for char in current): + raise AgentRunError( + "model service returned an invalid JSON response", + status=502, + code="InvalidModelResponse", + ) + elif isinstance(current, Mapping): + pending.append(iter(item for pair in current.items() for item in pair)) + elif isinstance(current, (list, tuple)): + pending.append(iter(current)) + + +def _model_auth_missing_error() -> AgentRunError: + return AgentRunError( + "model authentication is not configured", + status=503, + code="ModelAuthMissing", + ) + + def _choice_message(data: Mapping[str, Any]) -> Mapping[str, Any]: choices = data.get("choices") if not isinstance(choices, list) or not choices: @@ -206,7 +332,7 @@ def _choice_message(data: Mapping[str, Any]) -> Mapping[str, Any]: return message -def _message_text(message: Mapping[str, Any]) -> str: +def _message_text(message: Mapping[str, Any], *, max_bytes: int) -> str: content = message.get("content") if isinstance(content, str): text = content @@ -219,14 +345,18 @@ def _message_text(message: Mapping[str, Any]) -> str: code="InvalidModelResponse", ) text = refusal + if len(text) > max_bytes: + raise _model_response_too_large_error() try: - text.encode("utf-8") + text_bytes = text.encode("utf-8") except UnicodeEncodeError as exc: raise AgentRunError( "model response final assistant content must contain valid Unicode", status=502, code="InvalidModelResponse", ) from exc + if len(text_bytes) > max_bytes: + raise _model_response_too_large_error() return text @@ -258,7 +388,7 @@ def _reject_duplicate_argument_keys(pairs: list[tuple[str, Any]]) -> dict[str, A return out -def _validate_json_unicode(value: Any, *, path: str = "arguments") -> None: +def _validate_argument_unicode(value: Any, *, path: str = "arguments") -> None: pending: list[tuple[Any, str, int]] = [(value, path, 0)] while pending: current, current_path, depth = pending.pop() diff --git a/runtimes/common/agentkit_serve_common/orka.py b/runtimes/common/agentkit_serve_common/orka.py index 123f8d7..09b1c11 100644 --- a/runtimes/common/agentkit_serve_common/orka.py +++ b/runtimes/common/agentkit_serve_common/orka.py @@ -10,6 +10,7 @@ from __future__ import annotations import asyncio +import hashlib import json import os import re @@ -17,6 +18,7 @@ from dataclasses import dataclass, field from datetime import UTC, datetime from typing import Any, Awaitable, Callable, Mapping +from urllib.parse import quote from fastapi import Depends, FastAPI, HTTPException, Query, Request from fastapi.responses import StreamingResponse @@ -49,12 +51,28 @@ _ENABLE_BROKERED_COORDINATION_ENV = "AGENTKIT_ORKA_ENABLE_BROKERED_COORDINATION" _MAX_TOOL_SCHEMA_BYTES = 65536 _TERMINAL_TYPES = frozenset({"TurnCompleted", "TurnFailed", "TurnCancelled"}) +# Orka's canonical Go client scans one SSE line with a 1 MiB token ceiling. The +# advertised payload ceiling deliberately leaves roughly half the line for the +# native frame envelope and JSON escaping. Runtime text counts encoded UTF-8 +# bytes; brokered values count their compact JSON UTF-8 representation. +_MAX_OUTPUT_BYTES = 512 * 1024 +_ORKA_CLIENT_MAX_SSE_TOKEN_BYTES = 1 << 20 +_SSE_DATA_PREFIX = "data: " +_OUTPUT_LIMIT_CODE = "MaxOutputBytesExceeded" _DEFAULT_MAX_TERMINAL_TURNS = 256 _DEFAULT_MAX_RUNTIME_SESSIONS = 64 _MAX_TERMINAL_TURNS_ENV = "AGENTKIT_ORKA_MAX_TERMINAL_TURNS" _MAX_RUNTIME_SESSIONS_ENV = "AGENTKIT_ORKA_MAX_RUNTIME_SESSIONS" -_TURN_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,127}$") _ENV_NAME_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") +_JSON_VALUE_ABSENT = object() +# Go net/url.PathEscape leaves these reserved bytes unescaped in a path segment. +_ORKA_PATH_SEGMENT_SAFE = "$&+-.:=@_~" +# Mirrors Go strings.TrimSpace/unicode.IsSpace (which excludes U+001C-U+001F). +_ORKA_TRIM_SPACE_CHARS = ( + "\t\n\v\f\r \x85\xa0\u1680" + "\u2000\u2001\u2002\u2003\u2004\u2005\u2006\u2007\u2008\u2009\u200a" + "\u2028\u2029\u202f\u205f\u3000" +) @dataclass @@ -95,7 +113,7 @@ class TurnEvent: created_at: str = field(default_factory=_now_iso) severity: str = "info" summary: str = "" - content: Mapping[str, Any] | None = None + content: Any = _JSON_VALUE_ABSENT content_text: str = "" completed: Mapping[str, Any] | None = None failed: Mapping[str, Any] | None = None @@ -110,7 +128,7 @@ def terminal(self) -> bool: return self.type in _TERMINAL_TYPES def as_frame(self) -> dict[str, Any]: - return { + frame = { "version": ORKA_HARNESS_VERSION, "type": self.type, "runtimeSessionID": self.runtime_session_id, @@ -120,7 +138,6 @@ def as_frame(self) -> dict[str, Any]: "createdAt": self.created_at, "severity": self.severity, "summary": self.summary, - "content": dict(self.content or {}), "contentText": self.content_text, "toolName": self.tool_name, "toolCallID": self.tool_call_id, @@ -130,13 +147,77 @@ def as_frame(self) -> dict[str, Any]: "error": dict(self.error) if self.error is not None else None, "metadata": dict(self.metadata), } + if self.content is not _JSON_VALUE_ABSENT: + frame["content"] = self.content + return frame + + +class _SSEFrameTooLargeError(ValueError): + pass + + +class _SSEFrameEncodingError(ValueError): + pass + + +def _frame_json(event: TurnEvent) -> str: + return json.dumps(event.as_frame(), ensure_ascii=False, separators=(",", ":"), sort_keys=True) + + +def _compact_json_bytes(value: Any) -> bytes: + return json.dumps( + value, + allow_nan=False, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ).encode() + + +def _json_snapshot(value: Any) -> Any: + return json.loads(_compact_json_bytes(value)) + + +def _sse_data_line_bytes(event: TurnEvent) -> int: + return len(_SSE_DATA_PREFIX.encode()) + len(_frame_json(event).encode()) + + +def _ensure_sse_frame_fits(event: TurnEvent) -> None: + try: + line_bytes = _sse_data_line_bytes(event) + except UnicodeEncodeError as exc: + raise _SSEFrameEncodingError(f"{event.type} contains text that is not valid UTF-8") from exc + if line_bytes >= _ORKA_CLIENT_MAX_SSE_TOKEN_BYTES: + raise _SSEFrameTooLargeError( + f"{event.type} SSE data line is {line_bytes} UTF-8 bytes; " + f"Orka client limit is {_ORKA_CLIENT_MAX_SSE_TOKEN_BYTES - 1}" + ) + + +def _ensure_terminal_frame_fits(state: "TurnState") -> None: + # Reserve ample room for deterministic failure detail even when identity + # fields are unusually long. + message = "x" * 1024 + _ensure_sse_frame_fits( + TurnEvent( + seq=9_223_372_036_854_775_807, + type="TurnFailed", + runtime_session_id=state.runtime_session_id, + turn_id=state.turn_id, + correlation_id=state.correlation_id, + severity="error", + summary="turn output rejected", + failed={"reason": _OUTPUT_LIMIT_CODE, "message": message, "retryable": False}, + error={"code": _OUTPUT_LIMIT_CODE, "message": message, "retryable": False}, + metadata={}, + ) + ) @dataclass class PendingBrokeredTool: call: BrokeredToolCall future: asyncio.Future[BrokeredToolResult] - accepted_result: BrokeredToolResult | None = None class TurnState: @@ -168,14 +249,18 @@ def __init__( self.condition = asyncio.Condition() self.terminal_event: TurnEvent | None = None self.pending_tools: dict[str, PendingBrokeredTool] = {} + # Preserve idempotent /continue replay without retaining a second copy of + # each potentially large accepted tool result. + self.accepted_tool_result_digests: dict[str, str] = {} + self.rejected_tool_result_digests: dict[str, tuple[str, str, str]] = {} - async def append( + def _append_locked( self, event_type: str, *, severity: str = "info", summary: str = "", - content: Mapping[str, Any] | None = None, + content: Any = _JSON_VALUE_ABSENT, content_text: str = "", completed: Mapping[str, Any] | None = None, failed: Mapping[str, Any] | None = None, @@ -185,19 +270,76 @@ async def append( approval_id: str = "", metadata: Mapping[str, str] | None = None, ) -> tuple[TurnEvent, bool]: - async with self.condition: - if event_type in _TERMINAL_TYPES and self.terminal_event is not None: + if self.terminal_event is not None: + if event_type in _TERMINAL_TYPES: return self.terminal_event, False - seq = len(self.events) + 1 - if event_type == "TurnCompleted": - completed = dict(completed or {}) - completed.setdefault("finalEventSeq", seq) + raise AgentRunError("turn is already terminal", status=409, code="TurnTerminal") + seq = len(self.events) + 1 + if event_type == "TurnCompleted": + completed = dict(completed or {}) + completed.setdefault("finalEventSeq", seq) + event = TurnEvent( + seq=seq, + type=event_type, + runtime_session_id=self.runtime_session_id, + turn_id=self.turn_id, + correlation_id=self.correlation_id, + severity=severity, + summary=summary, + content=content, + content_text=content_text, + completed=completed, + failed=failed, + error=error, + tool_name=tool_name, + tool_call_id=tool_call_id, + approval_id=approval_id, + metadata=self.metadata if metadata is None else metadata, + ) + try: + _ensure_sse_frame_fits(event) + except (_SSEFrameTooLargeError, _SSEFrameEncodingError): + if event_type != "TurnFailed": + raise + message = "terminal failure details could not be emitted safely" event = TurnEvent( seq=seq, - type=event_type, + type="TurnFailed", runtime_session_id=self.runtime_session_id, turn_id=self.turn_id, correlation_id=self.correlation_id, + severity="error", + summary="turn failed", + failed={"reason": "TerminalFrameRejected", "message": message, "retryable": False}, + error={"code": "TerminalFrameRejected", "message": message, "retryable": False}, + metadata={}, + ) + _ensure_sse_frame_fits(event) + self.events.append(event) + if event.terminal: + self.terminal_event = event + self.condition.notify_all() + return event, True + + async def append( + self, + event_type: str, + *, + severity: str = "info", + summary: str = "", + content: Any = _JSON_VALUE_ABSENT, + content_text: str = "", + completed: Mapping[str, Any] | None = None, + failed: Mapping[str, Any] | None = None, + error: Mapping[str, Any] | None = None, + tool_name: str = "", + tool_call_id: str = "", + approval_id: str = "", + metadata: Mapping[str, str] | None = None, + ) -> tuple[TurnEvent, bool]: + async with self.condition: + return self._append_locked( + event_type, severity=severity, summary=summary, content=content, @@ -208,13 +350,8 @@ async def append( tool_name=tool_name, tool_call_id=tool_call_id, approval_id=approval_id, - metadata=metadata or self.metadata, + metadata=metadata, ) - self.events.append(event) - if event.terminal: - self.terminal_event = event - self.condition.notify_all() - return event, True async def events_after(self, seq: int) -> list[TurnEvent]: async with self.condition: @@ -319,6 +456,7 @@ async def _append_terminal_if_missing( completed=completed, failed=failed, error=error, + metadata={}, ) if created: _record_terminal_turn(state.turn_id, terminal_order, turns, max_terminal_turns) @@ -364,7 +502,7 @@ def _ensure_terminal_on_task_done( def _sse_frame(event: TurnEvent) -> str: - data = json.dumps(event.as_frame(), separators=(",", ":"), sort_keys=True) + data = _frame_json(event) return f"id: {event.seq}\nevent: {event.type}\ndata: {data}\n\n" @@ -383,13 +521,18 @@ def _required_string(data: Mapping[str, Any], field_name: str) -> str: def _turn_id_from_payload(data: Mapping[str, Any]) -> str: - turn_id = _required_string(data, "turnID") - if not _TURN_ID_RE.fullmatch(turn_id): + value = data.get("turnID") + if not isinstance(value, str): + raise HTTPException(status_code=400, detail="turnID is required") + trimmed = value.strip(_ORKA_TRIM_SPACE_CHARS) + if not trimmed: + raise HTTPException(status_code=400, detail="turnID is required") + if value != trimmed or value in {".", ".."} or "/" in value or "\\" in value: raise HTTPException( status_code=400, - detail="turnID must be URL-safe: letters, numbers, '.', '_', or '-'", + detail="turnID must be a single safe path segment without surrounding whitespace", ) - return turn_id + return value def _mapping_of_strings(data: Any, *, field_name: str) -> dict[str, str]: @@ -626,6 +769,59 @@ def _usage_payload(result: RunResult) -> dict[str, int]: return {key: int(usage.get(key, 0) or 0) for key in sorted(usage)} +def _utf8_bytes(value: str) -> int: + return len(value.encode()) + + +def _output_limit_message(output_kind: str, actual_bytes: int) -> str: + return f"{output_kind} is {actual_bytes} UTF-8 bytes; maxOutputBytes is {_MAX_OUTPUT_BYTES}" + + +def _brokered_result_digest(result: BrokeredToolResult) -> str: + value: dict[str, Any] = { + "approved": result.approved, + "outputPresent": result.output_present, + } + if result.output_present: + value["output"] = result.output + if result.error is not None: + value["error"] = dict(result.error) + return hashlib.sha256(_compact_json_bytes(value)).hexdigest() + + +async def _append_output_failure( + state: TurnState, + terminal_order: list[str], + turns: dict[str, TurnState], + max_terminal_turns: int, + *, + message: str, + code: str = _OUTPUT_LIMIT_CODE, +) -> None: + await _append_terminal_if_missing( + state, + terminal_order, + turns, + max_terminal_turns, + "TurnFailed", + summary="turn output rejected", + failed={"reason": code, "message": message, "retryable": False}, + error={"code": code, "message": message, "retryable": False}, + ) + + +def _append_output_failure_locked(state: TurnState, message: str, code: str) -> bool: + _, created = state._append_locked( + "TurnFailed", + severity="error", + summary="turn output rejected", + failed={"reason": code, "message": message, "retryable": False}, + error={"code": code, "message": message, "retryable": False}, + metadata={}, + ) + return created + + def health_response(spec: AgentSpec) -> dict[str, Any]: return { "version": ORKA_HARNESS_VERSION, @@ -650,6 +846,7 @@ def capabilities_response(spec: AgentSpec, *, brokered_classes: set[str] | None "supportsSuspend": False, "supportsWorkspaceSnapshot": False, "maxConcurrentTurns": 1, + "maxOutputBytes": _MAX_OUTPUT_BYTES, "metadata": { "agentName": spec.metadata.name, "model": spec.model.name, @@ -675,7 +872,7 @@ def start_turn_response(turn: TurnState) -> dict[str, Any]: "runtimeSessionID": turn.runtime_session_id, "turnID": turn.turn_id, "correlationID": turn.correlation_id, - "eventStreamPath": f"/v1/turns/{turn.turn_id}/events", + "eventStreamPath": f"/v1/turns/{quote(turn.turn_id, safe=_ORKA_PATH_SEGMENT_SAFE)}/events", } @@ -729,31 +926,42 @@ async def request_tool(self, call: BrokeredToolCall) -> BrokeredToolResult: if call.tool_call_id in self.state.pending_tools: raise AgentRunError(f"duplicate brokered tool call id {call.tool_call_id!r}", status=400, code="DuplicateToolCallID") future: asyncio.Future[BrokeredToolResult] = asyncio.get_running_loop().create_future() - self.state.pending_tools[call.tool_call_id] = PendingBrokeredTool(call=call, future=future) - await self.state.append( - "ToolCallRequested", - summary="brokered tool requested", - content=arguments, - tool_name=call.name, - tool_call_id=call.tool_call_id, - ) - if self.state.deadline is None: - result = await future - else: - seconds = (self.state.deadline - datetime.now(UTC)).total_seconds() - if seconds <= 0: - raise TimeoutError("turn deadline exceeded while waiting for brokered tool result") - async with asyncio.timeout(seconds): + pending = PendingBrokeredTool(call=call, future=future) + self.state.pending_tools[call.tool_call_id] = pending + try: + await self.state.append( + "ToolCallRequested", + summary="brokered tool requested", + content=arguments, + tool_name=call.name, + tool_call_id=call.tool_call_id, + metadata={}, + ) + if self.state.deadline is None: result = await future - await self.state.append( - "ToolResultReceived", - summary="tool result received", - content=dict(result.output or {}), - error=dict(result.error) if result.error is not None else None, - tool_name=call.name, - tool_call_id=call.tool_call_id, - ) - return result + else: + seconds = (self.state.deadline - datetime.now(UTC)).total_seconds() + if seconds <= 0: + raise TimeoutError("turn deadline exceeded while waiting for brokered tool result") + async with asyncio.timeout(seconds): + result = await future + await self.state.append( + "ToolResultReceived", + summary="tool result received", + content=_json_snapshot(result.output) if result.output_present else _JSON_VALUE_ABSENT, + error=_json_snapshot(dict(result.error)) if result.error is not None else None, + tool_name=call.name, + tool_call_id=call.tool_call_id, + metadata={}, + ) + return result + finally: + async with self.state.condition: + if self.state.pending_tools.get(call.tool_call_id) is pending: + self.state.pending_tools.pop(call.tool_call_id, None) + if not future.done(): + future.cancel() + self.state.condition.notify_all() async def _run_turn( @@ -833,22 +1041,56 @@ async def _run_with_runtime() -> RunResult: ) return - if result.text: - await state.append( - "RuntimeOutput", - summary="runtime output", - content={"message": result.text, "usage": _usage_payload(result)}, - content_text=result.text, + try: + output_bytes = _utf8_bytes(result.text) + except UnicodeEncodeError: + await _append_output_failure( + state, + terminal_order, + turns, + max_terminal_turns, + message="runtime output is not valid UTF-8", + code="InvalidOutputEncoding", + ) + return + if output_bytes > _MAX_OUTPUT_BYTES: + await _append_output_failure( + state, + terminal_order, + turns, + max_terminal_turns, + message=_output_limit_message("runtime output", output_bytes), + ) + return + try: + if result.text: + await state.append( + "RuntimeOutput", + summary="runtime output", + # contentText is the native text channel; keep only structured + # usage in content so one SSE frame does not duplicate the text. + content={"usage": _usage_payload(result)}, + content_text=result.text, + metadata={}, + ) + await _append_terminal_if_missing( + state, + terminal_order, + turns, + max_terminal_turns, + "TurnCompleted", + summary="turn completed", + completed={"result": result.text}, + ) + except _SSEFrameTooLargeError as exc: + await _append_output_failure( + state, + terminal_order, + turns, + max_terminal_turns, + message=str(exc), + code="HarnessFrameTooLarge", ) - await _append_terminal_if_missing( - state, - terminal_order, - turns, - max_terminal_turns, - "TurnCompleted", - summary="turn completed", - completed={"result": result.text}, - ) def create_orka_app( @@ -882,6 +1124,120 @@ def create_orka_app( active_runtimes: dict[str, ActiveRuntime] = {} runtime_order: list[str] = [] background_tasks: set[asyncio.Task[None]] = set() + runtime_close_tasks: set[asyncio.Task[BaseException | None]] = set() + runtime_close_tasks_by_session: dict[str, asyncio.Task[BaseException | None]] = {} + runtime_close_failed = False + + async def run_runtime_close(active: ActiveRuntime) -> BaseException | None: + try: + await active.context.__aexit__(None, None, None) + except BaseException as exc: + return exc + return None + + def runtime_close_error(close_task: asyncio.Task[BaseException | None]) -> BaseException | None: + try: + return close_task.result() + except BaseException as exc: + return exc + + def record_runtime_close_failure(close_error: BaseException | None) -> None: + nonlocal runtime_close_failed + if close_error is not None: + # Do not retain exception text or tracebacks across turns: adapter + # cleanup errors may contain prior-turn URLs, headers, or env data. + runtime_close_failed = True + + def new_runtime_close_failure() -> AgentRunError | None: + if not runtime_close_failed: + return None + return AgentRunError( + "runtime cleanup failed; restart required before opening another runtime session", + status=500, + code="RuntimeCloseFailed", + ) + + def runtime_close_done(runtime_session_id: str, close_task: asyncio.Task[BaseException | None]) -> None: + runtime_close_tasks.discard(close_task) + if runtime_close_tasks_by_session.get(runtime_session_id) is close_task: + runtime_close_tasks_by_session.pop(runtime_session_id, None) + # Retrieve every result even when the awaiting caller was cancelled, + # and fail closed after any cleanup error because framework resources + # may still occupy the configured runtime-session capacity. + record_runtime_close_failure(runtime_close_error(close_task)) + + async def close_runtime(runtime_session_id: str, active: ActiveRuntime) -> None: + close_task = asyncio.create_task(run_runtime_close(active), name="agentkit-orka-runtime-close") + runtime_close_tasks.add(close_task) + runtime_close_tasks_by_session[runtime_session_id] = close_task + close_task.add_done_callback(lambda task: runtime_close_done(runtime_session_id, task)) + try: + # asyncio.wait does not forward cancellation to the supplied task. + # Unlike an abandoned shield future, it also leaves no wrapper that + # can report a late close failure before shutdown retrieves it. + await asyncio.wait((close_task,)) + except asyncio.CancelledError: + # The runtime has already been removed from active ownership. Keep + # the close task rooted so turn cancellation cannot orphan cleanup; + # lifespan shutdown will await it if the caller cannot. + raise + close_error = runtime_close_error(close_task) + record_runtime_close_failure(close_error) + if close_error is not None: + close_failure = new_runtime_close_failure() + if close_failure is None: # pragma: no cover - record above establishes the invariant. + raise RuntimeError("runtime cleanup failed") from None + raise close_failure from None + + async def drain_runtime_close_tasks() -> list[BaseException]: + while runtime_close_tasks: + closing = list(runtime_close_tasks) + await asyncio.gather(*closing, return_exceptions=True) + runtime_close_tasks.difference_update(closing) + for close_task in closing: + record_runtime_close_failure(runtime_close_error(close_task)) + runtime_close_tasks_by_session.clear() + close_failure = new_runtime_close_failure() + return [close_failure] if close_failure is not None else [] + + async def wait_for_runtime_session_close(runtime_session_id: str) -> None: + while close_task := runtime_close_tasks_by_session.get(runtime_session_id): + await asyncio.wait((close_task,)) + runtime_close_tasks.discard(close_task) + if runtime_close_tasks_by_session.get(runtime_session_id) is close_task: + runtime_close_tasks_by_session.pop(runtime_session_id, None) + record_runtime_close_failure(runtime_close_error(close_task)) + + async def wait_for_runtime_close_progress() -> None: + closing = tuple(runtime_close_tasks) + if not closing: + return + done, _ = await asyncio.wait(closing, return_when=asyncio.FIRST_COMPLETED) + runtime_close_tasks.difference_update(done) + for close_task in done: + record_runtime_close_failure(runtime_close_error(close_task)) + + async def reserve_runtime_slot() -> None: + # create_turn enforces the advertised maxConcurrentTurns=1, and a turn + # publishes its terminal event only after it has stopped opening a + # runtime. Runtime admission is therefore serialized at this seam. + while True: + if close_failure := new_runtime_close_failure(): + # A failed close may still own framework resources. Keep the + # cache fail-closed instead of treating that slot as reusable. + raise close_failure + if len(active_runtimes) + len(runtime_close_tasks) < runtime_session_limit: + return + if runtime_close_tasks: + # A removed runtime still owns framework resources until its + # close task finishes. Count it against cache capacity so + # repeated cancelled turns cannot build an unbounded backlog. + await wait_for_runtime_close_progress() + continue + evict_id = runtime_order.pop(0) + evicted = active_runtimes.pop(evict_id, None) + if evicted is not None: + await close_runtime(evict_id, evicted) async def get_runtime(run_request: RunRequest) -> Any: runtime_session_id = run_request.session_id or "" @@ -892,10 +1248,14 @@ async def get_runtime(run_request: RunRequest) -> Any: runtime_order.remove(runtime_session_id) runtime_order.append(runtime_session_id) return active.session + if close_failure := new_runtime_close_failure(): + raise close_failure active_runtimes.pop(runtime_session_id, None) if runtime_session_id in runtime_order: runtime_order.remove(runtime_session_id) - await asyncio.shield(active.context.__aexit__(None, None, None)) + await close_runtime(runtime_session_id, active) + await wait_for_runtime_session_close(runtime_session_id) + await reserve_runtime_slot() # Runtime factories read the process environment today. Keep this scoped # section on the event loop thread so cancellation cannot leave a worker # thread running with turn credentials in process-global os.environ. @@ -904,11 +1264,6 @@ async def get_runtime(run_request: RunRequest) -> Any: session = await context.__aenter__() active_runtimes[runtime_session_id] = ActiveRuntime(context=context, session=session, env=dict(run_request.env)) runtime_order.append(runtime_session_id) - while len(runtime_order) > runtime_session_limit: - evict_id = runtime_order.pop(0) - evicted = active_runtimes.pop(evict_id, None) - if evicted is not None: - await evicted.context.__aexit__(None, None, None) return session @asynccontextmanager @@ -918,18 +1273,26 @@ async def lifespan(app: FastAPI): try: yield finally: - for state in turns.values(): - if state.task is not None and not state.task.done(): - state.task.cancel() - for task in list(background_tasks): + active_turn_tasks = [state.task for state in turns.values() if state.task is not None and not state.task.done()] + for task in active_turn_tasks: task.cancel() - if background_tasks: - await asyncio.gather(*background_tasks, return_exceptions=True) + if active_turn_tasks: + await asyncio.gather(*active_turn_tasks, return_exceptions=True) + while background_tasks: + terminal_tasks = list(background_tasks) + await asyncio.gather(*terminal_tasks, return_exceptions=True) + background_tasks.difference_update(terminal_tasks) background_tasks.clear() + close_errors = await drain_runtime_close_tasks() for active in reversed(list(active_runtimes.values())): - await active.context.__aexit__(None, None, None) + try: + await active.context.__aexit__(None, None, None) + except BaseException as exc: + close_errors.append(exc) active_runtimes.clear() runtime_order.clear() + if close_errors: + raise close_errors[0] app = FastAPI(title="agentkit-serve-orka", lifespan=lifespan) auth = Depends(make_auth_dependency(auth_token)) @@ -977,8 +1340,14 @@ async def create_turn(request: Request): deadline=run_request.deadline, metadata=run_request.metadata, ) + try: + await state.append("TurnStarted", summary="turn started") + _ensure_terminal_frame_fits(state) + except _SSEFrameTooLargeError as exc: + raise HTTPException(status_code=413, detail=str(exc)) from exc + except _SSEFrameEncodingError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc turns[turn_id] = state - await state.append("TurnStarted", summary="turn started") state.task = asyncio.create_task( _run_turn( get_runtime, @@ -1062,7 +1431,7 @@ async def continue_turn(turn_id: str, request: Request): raw_results = data.get("toolResults") if not isinstance(raw_results, list) or not raw_results: raise HTTPException(status_code=400, detail="toolResults must be a non-empty array") - results: list[BrokeredToolResult] = [] + results: list[tuple[BrokeredToolResult, int, str]] = [] for idx, raw_result in enumerate(raw_results): if not isinstance(raw_result, dict): raise HTTPException(status_code=400, detail=f"toolResults[{idx}] must be an object") @@ -1080,54 +1449,155 @@ async def continue_turn(turn_id: str, request: Request): approved = raw_result.get("approved", False) if not isinstance(approved, bool): raise HTTPException(status_code=400, detail=f"toolResults[{idx}].approved must be a boolean") + output_present = "output" in raw_result output_value = raw_result.get("output") error_value = raw_result.get("error") - if not approved and output_value is not None: + if not approved and output_present: raise HTTPException(status_code=400, detail=f"toolResults[{idx}].output is not allowed when approved is false") - if output_value is None and error_value is None: + if not output_present and error_value is None: if approved: raise HTTPException(status_code=400, detail=f"toolResults[{idx}] output or error is required") error_value = {"code": "ToolCallDenied", "message": "tool call was not approved", "retryable": False} - if output_value is not None and not isinstance(output_value, dict): - raise HTTPException(status_code=400, detail=f"toolResults[{idx}].output must be an object") if error_value is not None and not isinstance(error_value, dict): raise HTTPException(status_code=400, detail=f"toolResults[{idx}].error must be an object") - results.append( - BrokeredToolResult( - tool_call_id=tool_call_id, - approved=approved, - output=dict(output_value) if output_value is not None else None, - error=dict(error_value) if error_value is not None else None, - ) - ) + try: + output_bytes = len(_compact_json_bytes(output_value)) if output_present else 0 + except (TypeError, ValueError) as exc: + raise HTTPException(status_code=400, detail=f"toolResults[{idx}].output must be valid JSON") from exc + result_kwargs: dict[str, Any] = { + "tool_call_id": tool_call_id, + "approved": approved, + "error": dict(error_value) if error_value is not None else None, + } + if output_present: + result_kwargs["output"] = output_value + result = BrokeredToolResult(**result_kwargs) + try: + digest = _brokered_result_digest(result) + except (TypeError, ValueError) as exc: + raise HTTPException(status_code=400, detail=f"toolResults[{idx}] must be valid JSON") from exc + results.append((result, output_bytes, digest)) + + seen_results: set[str] = set() + for result, _, _ in results: + if result.tool_call_id in seen_results: + raise HTTPException(status_code=400, detail=f"duplicate toolCallID {result.tool_call_id!r}") + seen_results.add(result.tool_call_id) + + rejection_message: str | None = None + rejection_code = _OUTPUT_LIMIT_CODE + cancel_runtime_task = False async with state.condition: if state.terminal_event is not None: - known_done = all( - result.tool_call_id in state.pending_tools - and state.pending_tools[result.tool_call_id].accepted_result == result - for result in results - ) - if known_done: + rejected_outcomes: set[tuple[str, str]] = set() + for result, _, digest in results: + accepted_digest = state.accepted_tool_result_digests.get(result.tool_call_id) + if accepted_digest is not None: + if accepted_digest != digest: + raise HTTPException( + status_code=409, + detail=f"conflicting tool result for toolCallID {result.tool_call_id!r}", + ) + continue + rejected = state.rejected_tool_result_digests.get(result.tool_call_id) + if rejected is None: + raise HTTPException(status_code=409, detail="turn is already terminal") + rejected_digest, message, code = rejected + if rejected_digest != digest: + raise HTTPException( + status_code=409, + detail=f"conflicting tool result for toolCallID {result.tool_call_id!r}", + ) + rejected_outcomes.add((message, code)) + if not rejected_outcomes: return continue_turn_response(state) - raise HTTPException(status_code=409, detail="turn is already terminal") - seen_results: set[str] = set() - pending_results: list[tuple[PendingBrokeredTool, BrokeredToolResult]] = [] - for result in results: - if result.tool_call_id in seen_results: - raise HTTPException(status_code=400, detail=f"duplicate toolCallID {result.tool_call_id!r}") - seen_results.add(result.tool_call_id) - pending = state.pending_tools.get(result.tool_call_id) - if pending is None: - raise HTTPException(status_code=400, detail=f"unknown toolCallID {result.tool_call_id!r}") - if pending.accepted_result is not None and pending.accepted_result != result: - raise HTTPException(status_code=409, detail=f"conflicting tool result for toolCallID {result.tool_call_id!r}") - pending_results.append((pending, result)) - for pending, result in pending_results: - if pending.accepted_result is None: - pending.accepted_result = result - if not pending.future.done(): - pending.future.set_result(result) - state.condition.notify_all() + if len(rejected_outcomes) != 1: + raise HTTPException(status_code=409, detail="tool results do not match one rejected continuation") + rejection_message, rejection_code = rejected_outcomes.pop() + + pending_results: list[tuple[PendingBrokeredTool, BrokeredToolResult, str]] = [] + if state.terminal_event is None: + fresh_results: list[tuple[PendingBrokeredTool, BrokeredToolResult, int, str]] = [] + # Validate every identity and accepted-result digest before any + # oversized member can terminalize the whole continuation batch. + for result, output_bytes, digest in results: + accepted_digest = state.accepted_tool_result_digests.get(result.tool_call_id) + if accepted_digest is not None: + if accepted_digest != digest: + raise HTTPException( + status_code=409, + detail=f"conflicting tool result for toolCallID {result.tool_call_id!r}", + ) + continue + pending = state.pending_tools.get(result.tool_call_id) + if pending is None: + raise HTTPException(status_code=400, detail=f"unknown toolCallID {result.tool_call_id!r}") + fresh_results.append((pending, result, output_bytes, digest)) + + rejection_candidates: list[tuple[str, str, str]] = [] + for pending, result, output_bytes, digest in fresh_results: + if output_bytes > _MAX_OUTPUT_BYTES: + rejection_candidates.append( + ( + result.tool_call_id, + _output_limit_message("brokered tool output", output_bytes), + _OUTPUT_LIMIT_CODE, + ) + ) + continue + candidate = TurnEvent( + # Preflight with the widest native int64 sequence so an + # in-flight event cannot add a digit after HTTP 202. + seq=9_223_372_036_854_775_807, + type="ToolResultReceived", + runtime_session_id=state.runtime_session_id, + turn_id=state.turn_id, + correlation_id=state.correlation_id, + summary="tool result received", + content=result.output if result.output_present else _JSON_VALUE_ABSENT, + error=dict(result.error) if result.error is not None else None, + tool_name=pending.call.name, + tool_call_id=result.tool_call_id, + metadata={}, + ) + try: + _ensure_sse_frame_fits(candidate) + except _SSEFrameTooLargeError as exc: + rejection_candidates.append((result.tool_call_id, str(exc), "HarnessFrameTooLarge")) + continue + pending_results.append((pending, result, digest)) + if rejection_candidates: + _, rejection_message, rejection_code = min(rejection_candidates) + + if rejection_message is None and state.terminal_event is None: + for pending, result, digest in pending_results: + state.accepted_tool_result_digests[result.tool_call_id] = digest + if not pending.future.done(): + pending.future.set_result(result) + state.condition.notify_all() + elif state.terminal_event is None: + for result, _, digest in results: + if result.tool_call_id not in state.accepted_tool_result_digests: + state.rejected_tool_result_digests[result.tool_call_id] = ( + digest, + rejection_message, + rejection_code, + ) + created = _append_output_failure_locked(state, rejection_message, rejection_code) + if created: + _record_terminal_turn(state.turn_id, terminal_order, turns, retention_limit) + pending = list(state.pending_tools.values()) + state.pending_tools.clear() + for item in pending: + if not item.future.done(): + item.future.cancel() + state.condition.notify_all() + cancel_runtime_task = state.task is not None and not state.task.done() + + if rejection_message is not None: + if cancel_runtime_task and state.task is not None and state.task is not asyncio.current_task(): + state.task.cancel() + raise HTTPException(status_code=413, detail=rejection_message) return continue_turn_response(state) @@ -1146,8 +1616,9 @@ async def cancel_turn(turn_id: str, request: Request): raise HTTPException(status_code=400, detail="cancel turnID must match route turnID") runtime_session_id = _required_string(data, "runtimeSessionID") correlation_id = _required_string(data, "correlationID") - for field_name in ("namespace", "taskName", "sessionName"): - _required_string(data, field_name) + namespace = _required_string(data, "namespace") + task_name = _required_string(data, "taskName") + session_name = _required_string(data, "sessionName") state = turns.get(turn_id) if state is None: @@ -1156,6 +1627,8 @@ async def cancel_turn(turn_id: str, request: Request): raise HTTPException(status_code=400, detail="cancel runtimeSessionID must match turn runtimeSessionID") if correlation_id != state.correlation_id: raise HTTPException(status_code=400, detail="cancel correlationID must match turn correlationID") + if namespace != state.namespace or task_name != state.task_name or session_name != state.session_name: + raise HTTPException(status_code=400, detail="cancel namespace/taskName/sessionName must match turn") if state.terminal_event is None and state.task is not None and not state.task.done(): state.task.cancel() elif state.terminal_event is None: diff --git a/runtimes/common/agentkit_serve_common/runtime.py b/runtimes/common/agentkit_serve_common/runtime.py index c906061..77ffc13 100644 --- a/runtimes/common/agentkit_serve_common/runtime.py +++ b/runtimes/common/agentkit_serve_common/runtime.py @@ -20,6 +20,7 @@ OFFLINE_ORKA_ECHO_ENV = "AGENTKIT_ORKA_OFFLINE_ECHO" OFFLINE_ORKA_DELEGATE_AGENT_ENV = "AGENTKIT_ORKA_OFFLINE_DELEGATE_AGENT" +_BROKERED_TOOL_OUTPUT_ABSENT = object() def offline_orka_echo_enabled() -> bool: @@ -91,14 +92,33 @@ class BrokeredToolCall: brokered_class: Literal["read", "write", "coordination"] -@dataclass(frozen=True) +@dataclass(frozen=True, init=False) class BrokeredToolResult: - """Result returned by Orka after policy, approval, and tool execution.""" + """Result returned by Orka after policy, approval, and tool execution. + + ``output_present`` mirrors Go ``json.RawMessage`` presence semantics so an + omitted output remains distinct from a present JSON ``null`` value. + """ tool_call_id: str approved: bool - output: Mapping[str, Any] | None = None - error: Mapping[str, Any] | None = None + output: Any + error: Mapping[str, Any] | None + output_present: bool + + def __init__( + self, + tool_call_id: str, + approved: bool, + output: Any = _BROKERED_TOOL_OUTPUT_ABSENT, + error: Mapping[str, Any] | None = None, + ) -> None: + output_present = output is not _BROKERED_TOOL_OUTPUT_ABSENT + object.__setattr__(self, "tool_call_id", tool_call_id) + object.__setattr__(self, "approved", approved) + object.__setattr__(self, "output", None if not output_present else output) + object.__setattr__(self, "error", error) + object.__setattr__(self, "output_present", output_present) @runtime_checkable diff --git a/runtimes/common/agentkit_serve_common/server.py b/runtimes/common/agentkit_serve_common/server.py index 6ca068a..27e65fc 100644 --- a/runtimes/common/agentkit_serve_common/server.py +++ b/runtimes/common/agentkit_serve_common/server.py @@ -31,6 +31,7 @@ from typing import Any from fastapi import Depends, FastAPI, Header, HTTPException, Request +from fastapi.exceptions import RequestValidationError from fastapi.responses import JSONResponse from pydantic import BaseModel, ConfigDict @@ -224,6 +225,24 @@ async def chat_completions( return _completion_response(model_name, result) + @app.exception_handler(RequestValidationError) + async def _request_validation_exc_handler(request: Request, exc: RequestValidationError): + # Pydantic validation details can embed request inputs. Inspect only the + # machine-readable error type and return fixed, secret-safe messages. + if any(error.get("type") == "json_invalid" for error in exc.errors()): + return _error_response( + 400, + "request body must be valid JSON", + "invalid_request_error", + "invalid_json", + ) + return _error_response( + 400, + "request body is invalid", + "invalid_request_error", + "invalid_request", + ) + # Map HTTPExceptions raised in helpers to the OpenAI error envelope too. @app.exception_handler(HTTPException) async def _http_exc_handler(request: Request, exc: HTTPException): diff --git a/runtimes/common/tests/_brokered_description_cases.py b/runtimes/common/tests/_brokered_description_cases.py new file mode 100644 index 0000000..9312ca3 --- /dev/null +++ b/runtimes/common/tests/_brokered_description_cases.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +HARMLESS_BROKERED_DESCRIPTIONS = ( + "Read basic telemetry", + "Count model tokens", + "Count model tokens 100", + "Count model tokens: 1,000.", + "Count input tokens\U00010d50", +) + +UNSAFE_BROKERED_DESCRIPTIONS = ( + "contains sk-secret", + "execution at https://tool.default", + "Bearer token required", + "Basic Zm9vOmJhcg==", + "Basİc dXNlcjpwYXNz", + "Basic=dXNlcjpwYXNz", + "Basic user:pass", + "Basic alice: pass1", + "basic mode. baſic-@dXNlcjpwYXNz", + "BASIC baſic=Basic[baſic]Zm9vOmJhcg==", + "alice:pass is the Basic credential", + "Use Basic credential abc123", + "Use Basic `dXNlcjpwYXNz`", + "Use HTTP Basic value dXNlcjpwYXNz", + "Use HTTP Basic value=dXNlcjpwYXNz", + "Use HTTP-Basic value dXNlcjpwYXNz", + "token=abc123", + '{"token":"abc123"}', + "'access_token'=abc123", + "token.value=abc123", + "token[value]=abc123", + "Use --token abc123", + "access token abc123", + "token abc123", + "token 123", + "token abc", + "token ABCDEF", + '"token" "abc123"', + "model token abc123", + 'input token "abc123"', + 'token "abcdef"', + '"token" "ABCDEF"', + 'token "abc"', + "input token '123'", + "token.value abc123", + "token[value] abc123", + "request.token abc123", + "request/token abc123", + "request-token abc123", + "requesttoken abc123", + "report input\u001ctoken: 123456", + "token used: 'abc'", + "Use abc123 as model [token]", + "Count requests. Use model token 123", + "Finished counting. Model token 123", + 'Finished "counting." Model token 123', + "Count model tokens 123456789012345678901234567890", + "Read {auth}", + "Read (header)", + "Read [REDACTED_AUTH_HEADER]", + "execution at tool.default.svc.cluster.local", +) diff --git a/runtimes/common/tests/test_abi_contract.py b/runtimes/common/tests/test_abi_contract.py index 6d8620a..921e34e 100644 --- a/runtimes/common/tests/test_abi_contract.py +++ b/runtimes/common/tests/test_abi_contract.py @@ -1,7 +1,13 @@ from __future__ import annotations +import json +import math +import shutil +import subprocess +import textwrap from pathlib import Path +from _brokered_description_cases import HARMLESS_BROKERED_DESCRIPTIONS, UNSAFE_BROKERED_DESCRIPTIONS from agentkit_serve_common.config import load @@ -24,3 +30,128 @@ def test_go_rendered_agent_yaml_golden_loads_in_python_reader(): assert spec.env == [] assert spec.expose.openai is True assert spec.expose.port == 8080 + + +def test_go_rendered_edge_case_agent_yaml_loads_exactly_in_python_reader(): + repo = Path(__file__).resolve().parents[3] + golden = repo / "pkg" / "agentkit" / "abi" / "testdata" / "edge-cases.yaml" + line_break_text = "NEL:\u0085LS:\u2028PS:\u2029end" + property_name = "line:\u2028break" + + spec = load(golden) + + assert spec.instructions == "instructions " + line_break_text + assert len(spec.brokered_tools) == 1 + tool = spec.brokered_tools[0] + assert tool.description == "description " + line_break_text + assert tool.parameters["description"] == "schema " + line_break_text + assert tool.parameters["properties"][property_name]["description"] == "property " + line_break_text + binary_default = tool.parameters["properties"]["binary"]["default"] + assert isinstance(binary_default, str) + assert binary_default == "SGVsbG8=" + yaml_sensitive_default = tool.parameters["properties"]["? ask"]["default"] + assert yaml_sensitive_default == "before\tafter" + assert tool.parameters["properties"]["<<"]["default"] == "<<" + assert tool.parameters["properties"]["="]["default"] == "=" + assert tool.parameters["properties"][".inf"]["default"] == ".inf" + assert tool.parameters["properties"]["12:34:56"]["default"] == "2001-12-14 21:59:43.10 -5" + minimum = tool.parameters["properties"][property_name]["minimum"] + assert isinstance(minimum, float) + assert minimum == 0.0 + assert math.copysign(1.0, minimum) == -1.0 + assert tool.schema_digest == "sha256:ce77aaf228491b5007ed2ee703e57180acec8def6214c84d4324719b7f4f1fb6" + + +def test_current_go_validation_and_renderer_match_python_brokered_description_contract(tmp_path): + repo = Path(__file__).resolve().parents[3] + go = shutil.which("go") + assert go is not None, "Go is required for the cross-language ABI contract test" + descriptions = [*HARMLESS_BROKERED_DESCRIPTIONS, *UNSAFE_BROKERED_DESCRIPTIONS] + expected_validity = [True] * len(HARMLESS_BROKERED_DESCRIPTIONS) + [False] * len( + UNSAFE_BROKERED_DESCRIPTIONS + ) + + source = tmp_path / "render_brokered_descriptions.go" + source.write_text( + textwrap.dedent( + """ + package main + + import ( + "encoding/json" + "fmt" + "log" + "os" + + "github.com/sozercan/agentkit/pkg/agentkit/abi" + "github.com/sozercan/agentkit/pkg/agentkit/config" + "github.com/sozercan/agentkit/pkg/agentkit/effective" + ) + + type result struct { + Valid bool `json:"valid"` + Rendered string `json:"rendered,omitempty"` + Error string `json:"error,omitempty"` + } + + func main() { + var descriptions []string + if err := json.NewDecoder(os.Stdin).Decode(&descriptions); err != nil { + log.Fatal(err) + } + results := make([]result, 0, len(descriptions)) + for index, description := range descriptions { + cfg := &config.AgentConfig{ + APIVersion: "v1alpha1", + Kind: "Agent", + Metadata: config.Metadata{Name: fmt.Sprintf("description-regression-%d", index)}, + Model: config.Model{ + Provider: "openai-compatible", + BaseURL: "https://model.example/v1", + Name: "test-model", + }, + Instructions: config.Source{Inline: "Test brokered descriptions."}, + BrokeredTools: []config.BrokeredTool{{ + Name: "safe_lookup", + Description: description, + BrokeredClass: config.BrokeredClassRead, + Parameters: map[string]any{"type": "object"}, + }}, + Expose: config.Expose{OpenAI: true, Port: 8080}, + } + if err := cfg.Validate(); err != nil { + results = append(results, result{Valid: false, Error: err.Error()}) + continue + } + rendered, err := abi.Render(effective.FromConfig(cfg, cfg.Instructions.Inline)) + if err != nil { + log.Fatal(err) + } + results = append(results, result{Valid: true, Rendered: string(rendered)}) + } + if err := json.NewEncoder(os.Stdout).Encode(results); err != nil { + log.Fatal(err) + } + } + """ + ), + encoding="utf-8", + ) + + rendered = subprocess.run( + [go, "run", str(source)], + cwd=repo, + check=False, + capture_output=True, + input=json.dumps(descriptions), + text=True, + ) + assert rendered.returncode == 0, rendered.stderr + results = json.loads(rendered.stdout) + assert [result["valid"] for result in results] == expected_validity + + for index, description in enumerate(HARMLESS_BROKERED_DESCRIPTIONS): + golden = tmp_path / f"go-rendered-agent-{index}.yaml" + golden.write_text(results[index]["rendered"], encoding="utf-8") + spec = load(golden) + assert [tool.description for tool in spec.brokered_tools] == [description] diff --git a/runtimes/common/tests/test_adapter_support.py b/runtimes/common/tests/test_adapter_support.py index 981a251..f1c83dc 100644 --- a/runtimes/common/tests/test_adapter_support.py +++ b/runtimes/common/tests/test_adapter_support.py @@ -1,6 +1,8 @@ from __future__ import annotations import os +import sys +from types import ModuleType, SimpleNamespace from unittest import mock import pytest @@ -279,3 +281,68 @@ def test_resolve_tool_headers_prefers_per_run_env_and_keeps_errors_secret_free() msg = str(exc.value) assert "TOOLBOX_TOKEN" in msg assert "do-not-mention" not in msg + + +def _fake_azure_identity_module(factory_type: type) -> dict[str, ModuleType]: + azure = ModuleType("azure") + azure.__path__ = [] # type: ignore[attr-defined] + identity = ModuleType("azure.identity") + setattr(identity, "DefaultAzureCredential", factory_type) + azure.identity = identity # type: ignore[attr-defined] + return {"azure": azure, "azure.identity": identity} + + +def test_default_azure_credential_fallback_closes_after_success(): + instances = [] + + class FakeCredential: + def __init__(self) -> None: + self.closed = False + instances.append(self) + + def get_token(self, audience: str): + assert audience == "https://ai.azure.com/.default" + result = SimpleNamespace() + setattr(result, "token", "azure-token") + return result + + def close(self) -> None: + self.closed = True + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch.dict(sys.modules, _fake_azure_identity_module(FakeCredential)), + ): + resolved = support.resolve_workload_identity_token("https://ai.azure.com/.default") + + assert resolved == "azure-token" + assert len(instances) == 1 + assert instances[0].closed is True + + +def test_default_azure_credential_fallback_closes_after_failure_without_masking_token_error(): + instances = [] + acquisition_error = RuntimeError("token unavailable") + + class FakeCredential: + def __init__(self) -> None: + self.closed = False + instances.append(self) + + def get_token(self, audience: str): + raise acquisition_error + + def close(self) -> None: + self.closed = True + raise RuntimeError("close failed") + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch.dict(sys.modules, _fake_azure_identity_module(FakeCredential)), + ): + with pytest.raises(support.AgentBuildError, match="token unavailable") as exc_info: + support.resolve_workload_identity_token("https://ai.azure.com/.default") + + assert exc_info.value.__cause__ is acquisition_error + assert len(instances) == 1 + assert instances[0].closed is True diff --git a/runtimes/common/tests/test_config_validation.py b/runtimes/common/tests/test_config_validation.py index 47fc4cf..d7da11d 100644 --- a/runtimes/common/tests/test_config_validation.py +++ b/runtimes/common/tests/test_config_validation.py @@ -6,6 +6,7 @@ import pytest import yaml +from _brokered_description_cases import HARMLESS_BROKERED_DESCRIPTIONS, UNSAFE_BROKERED_DESCRIPTIONS from agentkit_serve_common.config import ConfigError, brokered_tool_schema_digest, load, load_or_exit, validate_required_env @@ -81,6 +82,18 @@ def test_load_rejects_invalid_tool_env_name(tmp_path): assert "[A-Z0-9_]+" in msg +def test_load_rejects_duplicate_direct_tool_names(tmp_path): + msg = _invalid_message( + tmp_path, + lambda spec: spec["tools"].append( + {"name": "fetch", "command": ["uvx", "another-mcp-server"]} + ), + ) + assert "tools" in msg + assert "duplicate tool name" in msg + assert "fetch" in msg + + def test_load_rejects_expose_openai_false(tmp_path): msg = _invalid_message(tmp_path, lambda spec: spec["expose"].update(openai=False)) assert "expose.openai" in msg @@ -467,7 +480,33 @@ def test_load_accepts_static_brokered_tools_with_matching_digest(tmp_path): assert spec.brokered_tools[0].schema_digest == digest -@pytest.mark.parametrize("description", ["contains sk-secret", "execution at https://tool.default", "Bearer token required"]) +@pytest.mark.parametrize( + "description", + HARMLESS_BROKERED_DESCRIPTIONS, +) +def test_load_accepts_harmless_brokered_tool_descriptions(tmp_path, description: str): + spec_dict = deepcopy(_BASE_SPEC) + spec_dict.update( + tools=[], + brokeredTools=[ + { + "name": "safe_lookup", + "description": description, + "brokeredClass": "read", + "parameters": {"type": "object"}, + } + ], + ) + + spec = load(_write_spec(tmp_path, spec_dict)) + + assert spec.brokered_tools[0].description == description + + +@pytest.mark.parametrize( + "description", + UNSAFE_BROKERED_DESCRIPTIONS, +) def test_load_rejects_unsafe_brokered_tool_descriptions(tmp_path, description: str): msg = _invalid_message( tmp_path, @@ -526,7 +565,21 @@ def test_load_rejects_unsafe_brokered_parameter_names(tmp_path, unsafe_name: str assert "not safe" in msg -@pytest.mark.parametrize("value", ["see https://internal-tool", "Bearer abc", "authorization header", "Cookie: session=abc", "x-api-key: abc", "api_key=abc", "x_api_key: abc", "format password", "enter passphrase"]) +@pytest.mark.parametrize( + "value", + [ + "see https://internal-tool", + "Bearer abc", + "authorization header", + "token=abc123", + "Cookie: session=abc", + "x-api-key: abc", + "api_key=abc", + "x_api_key: abc", + "format password", + "enter passphrase", + ], +) def test_load_rejects_unsafe_brokered_schema_string_values(tmp_path, value: str): msg = _invalid_message( tmp_path, @@ -751,14 +804,34 @@ def test_load_rejects_explicit_null_brokered_json_schema_keywords(tmp_path, bad_ assert "brokeredTools.0.parameters" in msg +@pytest.mark.parametrize("bad_child", [{"type": 123}, {"type": "strnig"}, {"items": "bad"}, {"items": [{"type": "number"}]}, {"multipleOf": 2}, {"uniqueItems": True}, {"minLength": -1}]) +def test_load_rejects_malformed_nested_brokered_json_schema(tmp_path, bad_child: dict): + msg = _invalid_message( + tmp_path, + lambda spec: spec.update( + tools=[], + brokeredTools=[ + { + "name": "safe_lookup", + "description": "safe schema", + "brokeredClass": "read", + "parameters": {"type": "object", "properties": {"site": bad_child}}, + } + ], + ), + ) + assert "brokeredTools.0.parameters" in msg + + @pytest.mark.parametrize( - "bad_parameters", + "schema", [ - {"type": "object", "enum": []}, - {"type": "object", "properties": {"site": {"type": "string", "enum": []}}}, + {"type": "object", "properties": {"site": {"type": None}}}, + {"type": "object", "properties": {"n": {"type": "integer", "default": "1"}}}, + {"type": "object", "properties": {"site": {"type": "string", "enum": [0, "ok"]}}}, ], ) -def test_load_rejects_empty_brokered_json_schema_enums(tmp_path, bad_parameters: dict): +def test_load_rejects_invalid_brokered_schema_type_values_and_defaults(tmp_path, schema: dict): msg = _invalid_message( tmp_path, lambda spec: spec.update( @@ -768,16 +841,50 @@ def test_load_rejects_empty_brokered_json_schema_enums(tmp_path, bad_parameters: "name": "safe_lookup", "description": "safe schema", "brokeredClass": "read", - "parameters": bad_parameters, + "parameters": schema, } ], ), ) - assert "enum must contain at least one value" in msg + assert "brokeredTools.0.parameters" in msg -@pytest.mark.parametrize("bad_child", [{"type": 123}, {"type": "strnig"}, {"items": "bad"}, {"items": [{"type": "number"}]}, {"multipleOf": 2}, {"uniqueItems": True}, {"minLength": -1}]) -def test_load_rejects_malformed_nested_brokered_json_schema(tmp_path, bad_child: dict): +def test_load_rejects_unknown_brokered_class_and_malformed_schema(tmp_path): + unknown_class = _invalid_message( + tmp_path, + lambda spec: spec.update( + tools=[], + brokeredTools=[ + { + "name": "safe_lookup", + "description": "safe schema", + "brokeredClass": "admin", + "parameters": {"type": "object"}, + } + ], + ), + ) + bad_schema = _invalid_message( + tmp_path, + lambda spec: spec.update( + tools=[], + brokeredTools=[ + { + "name": "safe_lookup", + "description": "safe schema", + "brokeredClass": "read", + "parameters": {"type": "array"}, + } + ], + ), + ) + + assert "brokeredTools.0.brokeredClass" in unknown_class + assert "brokeredTools.0.parameters" in bad_schema + assert "type: object" in bad_schema + + +def test_load_rejects_brokered_schema_digest_mismatch(tmp_path): msg = _invalid_message( tmp_path, lambda spec: spec.update( @@ -787,23 +894,158 @@ def test_load_rejects_malformed_nested_brokered_json_schema(tmp_path, bad_child: "name": "safe_lookup", "description": "safe schema", "brokeredClass": "read", - "parameters": {"type": "object", "properties": {"site": bad_child}}, + "parameters": {"type": "object"}, + "schemaDigest": "sha256:" + "0" * 64, } ], ), ) - assert "brokeredTools.0.parameters" in msg + assert "brokeredTools.0" in msg + assert "schemaDigest does not match" in msg + + +def test_load_rejects_owned_and_brokered_tool_name_overlap(tmp_path): + msg = _invalid_message( + tmp_path, + lambda spec: spec.update( + brokeredTools=[ + { + "name": "fetch", + "description": "safe schema", + "brokeredClass": "read", + "parameters": {"type": "object"}, + } + ] + ), + ) + assert "tools and brokeredTools cannot be mixed" in msg + + +def test_load_preserves_negative_zero_for_brokered_schema_digest(tmp_path): + from agentkit_serve_common.config import brokered_tool_schema_digest + + parameters = { + "type": "object", + "properties": {"n": {"type": "number", "minimum": -0.0}}, + } + digest = brokered_tool_schema_digest( + name="negative-zero-tool", + description="Negative zero schema.", + brokered_class="read", + parameters=parameters, + ) + spec_dict = deepcopy(_BASE_SPEC) + spec_dict.update( + tools=[], + brokeredTools=[ + { + "name": "negative-zero-tool", + "description": "Negative zero schema.", + "brokeredClass": "read", + "parameters": parameters, + "schemaDigest": digest, + } + ], + ) + + spec = load(_write_spec(tmp_path, spec_dict)) + + minimum = spec.brokered_tools[0].parameters["properties"]["n"]["minimum"] + assert isinstance(minimum, float) + assert minimum == 0.0 + assert math.copysign(1.0, minimum) == -1.0 + assert spec.brokered_tools[0].schema_digest == digest + + +def test_brokered_tool_schema_digest_uses_utf8_canonical_json(): + from agentkit_serve_common.config import brokered_tool_schema_digest + + parameters = { + "type": "object", + "properties": {"site": {"type": "string", "description": "São & R&D "}}, + "required": ["site"], + } + + assert brokered_tool_schema_digest( + name="check-network-telemetry", + description="São & R&D ", + brokered_class="read", + parameters=parameters, + ) == "sha256:7066de4e62dd1a6550701772aad901e1efcf5eb81a4f639252f82e6c7be8d4c1" + + +def test_brokered_tool_schema_digest_normalizes_integer_valued_floats(): + from agentkit_serve_common.config import brokered_tool_schema_digest + + parameters = {"type": "object", "properties": {"retries": {"type": "number", "minimum": 1.0}}} + + assert brokered_tool_schema_digest( + name="numeric-tool", + description="Numeric constraints.", + brokered_class="read", + parameters=parameters, + ) == "sha256:7ad9d43791e157981bcd65fd8452c9e71a64064875cc1330ced42d4956bf7d75" + + +def test_brokered_tool_schema_digest_canonicalizes_numeric_constraints_cross_language(): + from agentkit_serve_common.config import brokered_tool_schema_digest + + parameters = { + "type": "object", + "properties": {"n": {"type": "number", "minimum": 1e-6}}, + "required": ["n"], + } + + assert brokered_tool_schema_digest( + name="num-tool", + description="Numeric schema", + brokered_class="read", + parameters=parameters, + ) == "sha256:83bf12180154a21f8ba19049687e24acee9ef430966af67a83721a10bf7eee50" + + +def test_brokered_tool_schema_digest_canonicalizes_positive_exponent_numbers(): + from agentkit_serve_common.config import brokered_tool_schema_digest + + parameters = {"type": "object", "properties": {"n": {"type": "number", "maximum": 1e20}}} + + assert brokered_tool_schema_digest( + name="large-num-tool", + description="Large numeric schema", + brokered_class="read", + parameters=parameters, + ) == "sha256:51ebe9cdbb967453ba3ae9fb737566028814968d8121ba0d32825f7c8ffb5639" +def test_brokered_tool_schema_digest_canonicalizes_exponent_number_spellings(): + from agentkit_serve_common.config import brokered_tool_schema_digest + + parameters = { + "type": "object", + "properties": { + "small": {"type": "number", "minimum": 1e-7}, + "large": {"type": "number", "maximum": 1e21}, + }, + } + + assert brokered_tool_schema_digest( + name="exponent-num-tool", + description="Exponent numeric schema", + brokered_class="read", + parameters=parameters, + ) == "sha256:bb4fac58c3f65a33ed3c4ebfa10e5da9c3dc5a9250a1316f64d25e40e0e645e0" + + +# Regression coverage retained from the merged brokered-continuation stack. + @pytest.mark.parametrize( - "schema", + "bad_parameters", [ - {"type": "object", "properties": {"site": {"type": None}}}, - {"type": "object", "properties": {"n": {"type": "integer", "default": "1"}}}, - {"type": "object", "properties": {"site": {"type": "string", "enum": [0, "ok"]}}}, + {"type": "object", "enum": []}, + {"type": "object", "properties": {"site": {"type": "string", "enum": []}}}, ], ) -def test_load_rejects_invalid_brokered_schema_type_values_and_defaults(tmp_path, schema: dict): +def test_load_rejects_empty_brokered_json_schema_enums(tmp_path, bad_parameters: dict): msg = _invalid_message( tmp_path, lambda spec: spec.update( @@ -813,13 +1055,12 @@ def test_load_rejects_invalid_brokered_schema_type_values_and_defaults(tmp_path, "name": "safe_lookup", "description": "safe schema", "brokeredClass": "read", - "parameters": schema, + "parameters": bad_parameters, } ], ), ) - assert "brokeredTools.0.parameters" in msg - + assert "enum must contain at least one value" in msg def test_load_preserves_high_precision_integral_brokered_yaml_number(tmp_path): path = tmp_path / "agent.yaml" @@ -856,7 +1097,6 @@ def test_load_preserves_high_precision_integral_brokered_yaml_number(tmp_path): assert default == 9007199254740993 assert isinstance(default, int) - def test_load_rejects_lossy_fractional_brokered_yaml_number(tmp_path): path = tmp_path / "agent.yaml" path.write_text( @@ -889,7 +1129,6 @@ def test_load_rejects_lossy_fractional_brokered_yaml_number(tmp_path): with pytest.raises(ConfigError, match="cannot be represented exactly"): load(path) - def test_load_normalizes_cyclic_brokered_schema_to_config_error(tmp_path): path = tmp_path / "agent.yaml" path.write_text( @@ -920,7 +1159,6 @@ def test_load_normalizes_cyclic_brokered_schema_to_config_error(tmp_path): with pytest.raises(ConfigError, match="acyclic JSON-serializable"): load(path) - def test_load_normalizes_surrogate_brokered_schema_to_config_error(tmp_path): path = tmp_path / "agent.yaml" path.write_text( @@ -953,7 +1191,6 @@ def test_load_normalizes_surrogate_brokered_schema_to_config_error(tmp_path): with pytest.raises(ConfigError, match="JSON serializable UTF-8"): load(path) - def test_load_accepts_integral_float_values_for_integer_brokered_schema(tmp_path): spec_dict = deepcopy(_BASE_SPEC) spec_dict.update( @@ -982,39 +1219,6 @@ def test_load_accepts_integral_float_values_for_integer_brokered_schema(tmp_path assert properties["withConst"]["const"] == 2.0 assert properties["withEnum"]["enum"] == [3.0] - -def test_load_preserves_negative_zero_for_brokered_schema_digest(tmp_path): - parameters = { - "type": "object", - "properties": {"offset": {"type": "number", "default": -0.0}}, - } - spec_dict = deepcopy(_BASE_SPEC) - spec_dict.update( - tools=[], - brokeredTools=[ - { - "name": "negative_zero", - "description": "preserve negative zero", - "brokeredClass": "read", - "parameters": parameters, - "schemaDigest": brokered_tool_schema_digest( - name="negative_zero", - description="preserve negative zero", - brokered_class="read", - parameters=parameters, - ), - } - ], - ) - - spec = load(_write_spec(tmp_path, spec_dict)) - - default = spec.brokered_tools[0].parameters["properties"]["offset"]["default"] - assert isinstance(default, float) - assert default == 0.0 - assert math.copysign(1.0, default) == -1.0 - - def test_load_accepts_integral_float_integer_schema_constraints(tmp_path): spec_dict = deepcopy(_BASE_SPEC) spec_dict.update( @@ -1043,154 +1247,3 @@ def test_load_accepts_integral_float_integer_schema_constraints(tmp_path): values = spec.brokered_tools[0].parameters["properties"]["values"] assert values["minItems"] == 1 assert values["maxItems"] == 2 - - -def test_load_rejects_unknown_brokered_class_and_malformed_schema(tmp_path): - unknown_class = _invalid_message( - tmp_path, - lambda spec: spec.update( - tools=[], - brokeredTools=[ - { - "name": "safe_lookup", - "description": "safe schema", - "brokeredClass": "admin", - "parameters": {"type": "object"}, - } - ], - ), - ) - bad_schema = _invalid_message( - tmp_path, - lambda spec: spec.update( - tools=[], - brokeredTools=[ - { - "name": "safe_lookup", - "description": "safe schema", - "brokeredClass": "read", - "parameters": {"type": "array"}, - } - ], - ), - ) - - assert "brokeredTools.0.brokeredClass" in unknown_class - assert "brokeredTools.0.parameters" in bad_schema - assert "type: object" in bad_schema - - -def test_load_rejects_brokered_schema_digest_mismatch(tmp_path): - msg = _invalid_message( - tmp_path, - lambda spec: spec.update( - tools=[], - brokeredTools=[ - { - "name": "safe_lookup", - "description": "safe schema", - "brokeredClass": "read", - "parameters": {"type": "object"}, - "schemaDigest": "sha256:" + "0" * 64, - } - ], - ), - ) - assert "brokeredTools.0" in msg - assert "schemaDigest does not match" in msg - - -def test_load_rejects_owned_and_brokered_tool_name_overlap(tmp_path): - msg = _invalid_message( - tmp_path, - lambda spec: spec.update( - brokeredTools=[ - { - "name": "fetch", - "description": "safe schema", - "brokeredClass": "read", - "parameters": {"type": "object"}, - } - ] - ), - ) - assert "tools and brokeredTools cannot be mixed" in msg - - -def test_brokered_tool_schema_digest_uses_utf8_canonical_json(): - from agentkit_serve_common.config import brokered_tool_schema_digest - - parameters = { - "type": "object", - "properties": {"site": {"type": "string", "description": "São & R&D "}}, - "required": ["site"], - } - - assert brokered_tool_schema_digest( - name="check-network-telemetry", - description="São & R&D ", - brokered_class="read", - parameters=parameters, - ) == "sha256:7066de4e62dd1a6550701772aad901e1efcf5eb81a4f639252f82e6c7be8d4c1" - - -def test_brokered_tool_schema_digest_normalizes_integer_valued_floats(): - from agentkit_serve_common.config import brokered_tool_schema_digest - - parameters = {"type": "object", "properties": {"retries": {"type": "number", "minimum": 1.0}}} - - assert brokered_tool_schema_digest( - name="numeric-tool", - description="Numeric constraints.", - brokered_class="read", - parameters=parameters, - ) == "sha256:7ad9d43791e157981bcd65fd8452c9e71a64064875cc1330ced42d4956bf7d75" - - -def test_brokered_tool_schema_digest_canonicalizes_numeric_constraints_cross_language(): - from agentkit_serve_common.config import brokered_tool_schema_digest - - parameters = { - "type": "object", - "properties": {"n": {"type": "number", "minimum": 1e-6}}, - "required": ["n"], - } - - assert brokered_tool_schema_digest( - name="num-tool", - description="Numeric schema", - brokered_class="read", - parameters=parameters, - ) == "sha256:83bf12180154a21f8ba19049687e24acee9ef430966af67a83721a10bf7eee50" - - -def test_brokered_tool_schema_digest_canonicalizes_positive_exponent_numbers(): - from agentkit_serve_common.config import brokered_tool_schema_digest - - parameters = {"type": "object", "properties": {"n": {"type": "number", "maximum": 1e20}}} - - assert brokered_tool_schema_digest( - name="large-num-tool", - description="Large numeric schema", - brokered_class="read", - parameters=parameters, - ) == "sha256:51ebe9cdbb967453ba3ae9fb737566028814968d8121ba0d32825f7c8ffb5639" - - -def test_brokered_tool_schema_digest_canonicalizes_exponent_number_spellings(): - from agentkit_serve_common.config import brokered_tool_schema_digest - - parameters = { - "type": "object", - "properties": { - "small": {"type": "number", "minimum": 1e-7}, - "large": {"type": "number", "maximum": 1e21}, - }, - } - - assert brokered_tool_schema_digest( - name="exponent-num-tool", - description="Exponent numeric schema", - brokered_class="read", - parameters=parameters, - ) == "sha256:bb4fac58c3f65a33ed3c4ebfa10e5da9c3dc5a9250a1316f64d25e40e0e645e0" diff --git a/runtimes/common/tests/test_foundry_brokered_conformance_script.py b/runtimes/common/tests/test_foundry_brokered_conformance_script.py new file mode 100644 index 0000000..c35b293 --- /dev/null +++ b/runtimes/common/tests/test_foundry_brokered_conformance_script.py @@ -0,0 +1,224 @@ +from __future__ import annotations + +import json +import os +import subprocess +import threading +from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Any, Iterator + + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_HELPER = _REPO_ROOT / "deploy" / "foundry" / "scripts" / "foundry_brokered_conformance.sh" + + +@contextmanager +def _mock_gateway( + *, + agent_session_id: str | None = None, + echo_continuation_proof: bool = False, + malformed_continuation_response: bool = False, + serialize_continuation_body: bool = False, +) -> Iterator[tuple[str, list[dict[str, Any]]]]: + requests: list[dict[str, Any]] = [] + initial_response_id = "caresp_mock_initial" + + class Handler(BaseHTTPRequestHandler): + def do_POST(self) -> None: # noqa: N802 - stdlib HTTP handler API. + content_length = int(self.headers.get("content-length", "0")) + body = json.loads(self.rfile.read(content_length)) + requests.append({"headers": dict(self.headers), "body": body}) + + if len(requests) == 1: + response: dict[str, Any] = { + "id": initial_response_id, + "status": "completed", + "output": [ + { + "type": "function_call", + "response_id": initial_response_id, + "call_id": "call_conformance_1", + "name": "conformance_read", + "arguments": '{"probe":true}', + "status": "completed", + } + ], + } + if agent_session_id is not None: + response["agent_session_id"] = agent_session_id + else: + final_text = 'conformance complete: {"approved": true, "output": {"success": true}}' + if serialize_continuation_body: + final_text = "request failed: " + json.dumps(body, separators=(",", ":")) + elif echo_continuation_proof: + final_text = str(body.get("brokered_continuation_proof", "")) + response = { + "id": "caresp_mock_final", + "previous_response_id": initial_response_id, + "status": "completed", + "output": [ + { + "type": "message", + "response_id": "caresp_mock_final", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": final_text}], + } + ], + } + + if len(requests) > 1 and malformed_continuation_response: + encoded = ("not-json: " + json.dumps(body, separators=(",", ":"))).encode() + else: + encoded = json.dumps(response, separators=(",", ":")).encode() + self.send_response(200) + self.send_header("content-type", "application/json") + self.send_header("content-length", str(len(encoded))) + self.end_headers() + self.wfile.write(encoded) + + def log_message(self, format: str, *args: object) -> None: + return + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + host, port = server.server_address + yield f"http://{host}:{port}/responses", requests + finally: + server.shutdown() + thread.join(timeout=5) + server.server_close() + + +def _run_helper(endpoint: str, transcript_dir: Path, **environment: str) -> subprocess.CompletedProcess[str]: + env = os.environ.copy() + for name in ( + "AGENTKIT_CONFORMANCE_OUTPUT", + "AGENTKIT_CONTINUATION_PROOF", + "AGENTKIT_CONTINUATION_PROOF_BODY", + "AGENTKIT_EXPECTED_ARGUMENTS", + "AGENTKIT_EXPECTED_CALL_ID", + "AGENTKIT_EXPECTED_CALL_ID_PREFIX", + "AGENTKIT_EXPECTED_TOOL_NAME", + ): + env.pop(name, None) + env.update( + { + "AGENT_RESPONSES_ENDPOINT": endpoint, + "AGENT_RESPONSES_BEARER_TOKEN": "mock-token", + **environment, + } + ) + return subprocess.run( + [str(_HELPER), "conformance_read", str(transcript_dir)], + cwd=_REPO_ROOT, + env=env, + check=False, + capture_output=True, + text=True, + timeout=30, + ) + + +def test_brokered_conformance_helper_propagates_returned_agent_session_id(tmp_path: Path): + transcript = tmp_path / "transcript" + with _mock_gateway(agent_session_id="gateway-session") as (endpoint, requests): + result = _run_helper(endpoint, transcript) + + assert result.returncode == 0, result.stderr + assert len(requests) == 2 + assert requests[1]["body"]["agent_session_id"] == "gateway-session" + archived = json.loads((transcript / "03-continuation-request.json").read_text(encoding="utf-8")) + assert archived["agent_session_id"] == "gateway-session" + + +def test_brokered_conformance_helper_sends_body_proof_without_archiving_it(tmp_path: Path): + transcript = tmp_path / "transcript" + proof = "body-proof-must-stay-out-of-transcript" + with _mock_gateway(agent_session_id="gateway-session") as (endpoint, requests): + result = _run_helper(endpoint, transcript, AGENTKIT_CONTINUATION_PROOF_BODY=proof) + + assert result.returncode == 0, result.stderr + assert len(requests) == 2 + assert requests[1]["body"]["agent_session_id"] == "gateway-session" + assert requests[1]["body"]["brokered_continuation_proof"] == proof + assert "x-agentkit-brokered-continuation-proof" not in requests[1]["headers"] + archived_text = "\n".join(path.read_text(encoding="utf-8") for path in transcript.iterdir() if path.is_file()) + assert proof not in archived_text + archived_request = json.loads((transcript / "03-continuation-request.json").read_text(encoding="utf-8")) + assert "brokered_continuation_proof" not in archived_request + + +def test_brokered_conformance_helper_preserves_header_proof_option(tmp_path: Path): + transcript = tmp_path / "transcript" + proof = "header-proof-must-stay-out-of-transcript" + with _mock_gateway() as (endpoint, requests): + result = _run_helper(endpoint, transcript, AGENTKIT_CONTINUATION_PROOF=proof) + + assert result.returncode == 0, result.stderr + assert len(requests) == 2 + assert requests[1]["headers"]["x-agentkit-brokered-continuation-proof"] == proof + assert "brokered_continuation_proof" not in requests[1]["body"] + archived_text = "\n".join(path.read_text(encoding="utf-8") for path in transcript.iterdir() if path.is_file()) + assert proof not in archived_text + + +def test_brokered_conformance_helper_refuses_to_archive_echoed_proof(tmp_path: Path): + transcript = tmp_path / "transcript" + proof = 'echoed-"proof"\\with\nnewline-must-not-be-archived' + with _mock_gateway(echo_continuation_proof=True) as (endpoint, requests): + result = _run_helper(endpoint, transcript, AGENTKIT_CONTINUATION_PROOF_BODY=proof) + + assert len(requests) == 2 + assert result.returncode != 0 + assert not (transcript / "04-continuation-response.json").exists() + archived_text = "\n".join(path.read_text(encoding="utf-8") for path in transcript.iterdir() if path.is_file()) + assert proof not in archived_text + + +def test_brokered_conformance_helper_refuses_nested_json_proof_echo(tmp_path: Path): + transcript = tmp_path / "transcript" + proof = 'nested-"proof"\\with\nnewline-must-not-be-archived' + with _mock_gateway(serialize_continuation_body=True) as (endpoint, requests): + result = _run_helper(endpoint, transcript, AGENTKIT_CONTINUATION_PROOF_BODY=proof) + + assert len(requests) == 2 + assert result.returncode != 0 + assert not (transcript / "04-continuation-response.json").exists() + archived_text = "\n".join(path.read_text(encoding="utf-8") for path in transcript.iterdir() if path.is_file()) + assert proof not in archived_text + + +def test_brokered_conformance_helper_refuses_encoded_proof_in_malformed_response(tmp_path: Path): + transcript = tmp_path / "transcript" + proof = 'malformed-"proof"\\with\nnewline-must-not-be-archived' + with _mock_gateway(malformed_continuation_response=True) as (endpoint, requests): + result = _run_helper(endpoint, transcript, AGENTKIT_CONTINUATION_PROOF_BODY=proof) + + assert len(requests) == 2 + assert result.returncode != 0 + assert not (transcript / "04-continuation-response.json").exists() + archived_text = "\n".join(path.read_text(encoding="utf-8") for path in transcript.iterdir() if path.is_file()) + assert proof not in archived_text + + +def test_brokered_conformance_helper_removes_stale_final_artifacts_before_refused_rerun(tmp_path: Path): + transcript = tmp_path / "transcript" + transcript.mkdir() + stale_response = transcript / "04-continuation-response.json" + stale_summary = transcript / "summary.json" + stale_response.write_text('{"stale":true}', encoding="utf-8") + stale_summary.write_text('{"stale":true}', encoding="utf-8") + proof = "rerun-proof-must-not-be-archived" + + with _mock_gateway(echo_continuation_proof=True) as (endpoint, requests): + result = _run_helper(endpoint, transcript, AGENTKIT_CONTINUATION_PROOF_BODY=proof) + + assert len(requests) == 2 + assert result.returncode != 0 + assert not stale_response.exists() + assert not stale_summary.exists() diff --git a/runtimes/common/tests/test_foundry_brokered_protocol.py b/runtimes/common/tests/test_foundry_brokered_protocol.py index 11a0fa1..f2acf92 100644 --- a/runtimes/common/tests/test_foundry_brokered_protocol.py +++ b/runtimes/common/tests/test_foundry_brokered_protocol.py @@ -1,9 +1,13 @@ from __future__ import annotations +import asyncio import json import math +import threading import time +from concurrent.futures import ThreadPoolExecutor from copy import deepcopy +from pathlib import Path from types import TracebackType from typing import Any @@ -11,13 +15,13 @@ import httpx import pytest -from agentkit_serve_common import foundry as foundry_module +import agentkit_serve_common.foundry as foundry_module from agentkit_serve_common import foundry_model_loop as foundry_model_loop_module -from agentkit_serve_common.adapter_support import NO_AUTH_API_KEY from agentkit_serve_common.config import AgentSpec from agentkit_serve_common.conversation import RunRequest from agentkit_serve_common.foundry import create_foundry_app -from agentkit_serve_common.runtime import RunResult, RuntimeSession +from agentkit_serve_common.foundry_model_loop import BrokeredChatModelLoop +from agentkit_serve_common.runtime import AgentRunError, RunResult, RuntimeSession CONTINUATION_PROOF = "test-orka-continuation-proof" @@ -182,29 +186,6 @@ def test_foundry_brokered_initial_response_emits_static_function_call_without_di assert factory.runtime.run_requests == [] -@pytest.mark.parametrize("ttl_seconds", [float("nan"), float("inf"), float("-inf")]) -def test_foundry_brokered_nonfinite_explicit_state_ttl_uses_default(ttl_seconds: float): - app = _app(state_ttl_seconds=ttl_seconds) - - with TestClient(app) as client: - readiness = client.get("/readiness") - - assert readiness.status_code == 200 - assert readiness.json()["foundryResponses"]["stateTtlSeconds"] == 900.0 - - -@pytest.mark.parametrize("raw_ttl", ["nan", "inf", "-inf"]) -def test_foundry_brokered_nonfinite_state_ttl_env_uses_default(monkeypatch, raw_ttl: str): - monkeypatch.setenv("AGENTKIT_FOUNDRY_RESPONSE_STATE_TTL_SECONDS", raw_ttl) - app = _app() - - with TestClient(app) as client: - readiness = client.get("/readiness") - - assert readiness.status_code == 200 - assert readiness.json()["foundryResponses"]["stateTtlSeconds"] == 900.0 - - def test_foundry_brokered_deterministic_arguments_reject_unsafe_prompt_text(): spec = _spec(tool_name="check-network-telemetry") spec.brokered_tools[0].parameters = { @@ -259,7 +240,171 @@ def test_foundry_brokered_rejects_normal_followup_while_previous_response_is_pen assert resp.json()["error"]["code"] == "response_pending_function_call_output" -def test_foundry_brokered_rejects_normal_followup_while_previous_response_is_resuming(monkeypatch): +def test_foundry_brokered_pending_state_store_is_bounded(): + app = _app(max_pending_responses=1) + + with TestClient(app) as client: + first = _start(client) + second = client.post("/responses", json={"input": "another pending request"}) + + assert _call(first) + assert second.status_code == 429 + assert second.json()["error"]["code"] == "brokered_response_state_full" + + +def test_foundry_brokered_failed_initial_state_persist_is_retryable_without_consuming_capacity(monkeypatch, tmp_path): + state_file = tmp_path / "responses-state.json" + original_replace = Path.replace + failed_once = False + + def fail_first_replace(path: Path, target: Path) -> Path: + nonlocal failed_once + if not failed_once and path.name == f".{state_file.name}.tmp": + failed_once = True + raise OSError("simulated state storage failure") + return original_replace(path, target) + + monkeypatch.setattr(Path, "replace", fail_first_replace) + app = _app(response_state_file=state_file, max_pending_responses=1) + + with TestClient(app, raise_server_exceptions=False) as client: + failed = client.post("/responses", json={"input": "please read telemetry"}) + retried = client.post("/responses", json={"input": "please read telemetry"}) + + assert failed.status_code == 503 + assert failed.json()["error"] == { + "message": "brokered response state storage unavailable", + "code": "brokered_response_state_storage_error", + } + assert retried.status_code == 200, retried.text + assert _call(retried.json()) + + +def test_foundry_brokered_state_transactions_do_not_deepcopy_existing_state_graph(monkeypatch): + app = _app(max_pending_responses=3) + original_deepcopy = foundry_module.deepcopy + + def reject_full_store_deepcopy(value: Any, memo: dict[int, Any] | None = None) -> Any: + if isinstance(value, dict) and value and all(type(entry).__name__ == "_HostedResponseState" for entry in value.values()): + raise AssertionError("state transactions must not deepcopy the entire state store") + return original_deepcopy(value, memo) if memo is not None else original_deepcopy(value) + + with TestClient(app, raise_server_exceptions=False) as client: + first = _start(client) + first_call = _call(first) + monkeypatch.setattr(foundry_module, "deepcopy", reject_full_store_deepcopy) + second = client.post("/responses", json={"input": "please read telemetry again"}) + completed = client.post( + "/responses", + headers=CONTINUATION_AUTH, + json=_continuation(first["id"], first_call["call_id"], {"approved": True, "output": {"ok": True}}), + ) + + assert second.status_code == 200, second.text + assert completed.status_code == 200, completed.text + + +def test_foundry_bounded_json_walk_rejects_wide_mapping_without_bulk_child_copy(): + class WideMapping(dict[str, int]): + def __init__(self) -> None: + super().__init__() + self.item_yields = 0 + + def keys(self): + raise AssertionError("bounded traversal must not bulk-copy mapping keys") + + def values(self): + raise AssertionError("bounded traversal must not bulk-copy mapping values") + + def items(self): + for index in range(10_000): + self.item_yields += 1 + yield str(index), index + + value = WideMapping() + try: + foundry_module._bounded_json_bytes(value, max_bytes=128) + except foundry_module._SerializedPayloadTooLarge: + pass + else: + raise AssertionError("wide mapping must exceed the bounded JSON budget") + + assert value.item_yields <= 65 + + +def test_foundry_bounded_json_rejects_escaped_string_before_encoder(monkeypatch): + def reject_encoder(*_args: Any, **_kwargs: Any): + raise AssertionError("oversized escaped strings must be rejected before JSONEncoder") + + monkeypatch.setattr(json.JSONEncoder, "iterencode", reject_encoder) + try: + foundry_module._bounded_json_bytes("é" * 20, max_bytes=64) + except foundry_module._SerializedPayloadTooLarge: + pass + else: + raise AssertionError("escaped JSON string must exceed the byte budget") + + +def test_foundry_state_byte_eviction_serializes_aggregate_once(): + tool = foundry_module.brokered_tool_definitions(_spec())[0] + + def state(response_id: str, *, completed: bool) -> Any: + call_id = f"call_{response_id}" + call = foundry_module._PendingCall( + call_id=call_id, + item_id=f"item_{response_id}", + tool=tool, + arguments={}, + ) + return foundry_module._HostedResponseState( + response_id=response_id, + session_id=None, + pending_calls={call_id: call}, + expires_at=time.time() + 60, + status="completed" if completed else "pending", + accepted_output_digests={call_id: "sha256:" + "0" * 64} if completed else {}, + final_payload={"result": "x" * 200} if completed else None, + model_messages=[{"role": "user", "content": "y" * 500}] if not completed else None, + ) + + candidate = state("candidate", completed=False) + states = { + "completed-1": state("completed-1", completed=True), + "completed-2": state("completed-2", completed=True), + "completed-3": state("completed-3", completed=True), + candidate.response_id: candidate, + } + store = foundry_module._FoundryResponseStateStore( + ttl_seconds=60, + max_entries=10, + max_bytes=1_000_000, + ) + store.max_bytes = max( + len(store._serialize({response_id: states[response_id], candidate.response_id: candidate})) + for response_id in ("completed-1", "completed-2", "completed-3") + ) + original_serialize = store._serialize + aggregate_serializations = 0 + + def count_aggregate(current_states: dict[str, Any]) -> bytes: + nonlocal aggregate_serializations + if len(current_states) > 1: + aggregate_serializations += 1 + return original_serialize(current_states) + + store._serialize = count_aggregate # type: ignore[method-assign] + store._serialize_candidate( + states, + response_id=candidate.response_id, + excluded_response_ids=set(), + ) + + assert aggregate_serializations <= 1 + assert candidate.response_id in states + assert len(states) == 2 + + +def test_foundry_state_cache_fallback_contains_capacity_failure(monkeypatch): stores: list[Any] = [] original_store = foundry_module._FoundryResponseStateStore @@ -270,28 +415,16 @@ def capture_store(*args: Any, **kwargs: Any) -> Any: monkeypatch.setattr(foundry_module, "_FoundryResponseStateStore", capture_store) app = _app() - with TestClient(app) as client: initial = _start(client) - state = stores[0].get(initial["id"]) - state.status = "resuming" - stores[0].save(state) - resp = client.post("/responses", json={"previous_response_id": initial["id"], "input": "next question"}) - - assert resp.status_code == 409 - assert resp.json()["error"]["code"] == "response_pending_function_call_output" - -def test_foundry_brokered_pending_state_store_is_bounded(): - app = _app(max_pending_responses=1) + state = stores[0].get(initial["id"]) - with TestClient(app) as client: - first = _start(client) - second = client.post("/responses", json={"input": "another pending request"}) + def reject_cache(*_args: Any, **_kwargs: Any): + raise foundry_module._StateStoreFull("simulated concurrent capacity pressure") - assert _call(first) - assert second.status_code == 429 - assert second.json()["error"]["code"] == "brokered_response_state_full" + monkeypatch.setattr(stores[0], "_serialize_candidate", reject_cache) + assert stores[0].cache_in_memory(state) is False def test_foundry_brokered_completed_state_is_evicted_before_rejecting_new_pending_state(): @@ -392,31 +525,6 @@ def test_foundry_brokered_integer_arguments_honor_float_bounds(): assert json.loads(_call(body.json())["arguments"]) == {"count": 1, "after": 1, "below": 0} -def test_foundry_brokered_integer_arguments_preserve_large_integer_bounds(): - huge = 10**100 - spec = _spec(tool_name="large-integer-bounds") - spec.brokered_tools[0].parameters = { - "type": "object", - "properties": { - "minimum": {"type": "integer", "minimum": huge}, - "exclusiveMinimum": {"type": "integer", "exclusiveMinimum": huge}, - "exclusiveMaximum": {"type": "integer", "exclusiveMaximum": -huge}, - }, - "required": ["minimum", "exclusiveMinimum", "exclusiveMaximum"], - } - app = _app(spec) - - with TestClient(app) as client: - response = client.post("/responses", json={"input": "large-integer-bounds"}) - - assert response.status_code == 200, response.text - assert json.loads(_call(response.json())["arguments"]) == { - "minimum": huge, - "exclusiveMinimum": huge + 1, - "exclusiveMaximum": -huge - 1, - } - - def test_foundry_brokered_synthesizes_arguments_that_honor_basic_constraints(): spec = _spec(tool_name="bounded-lookup") tool = spec.brokered_tools[0] @@ -529,7 +637,7 @@ def test_foundry_brokered_argument_synthesis_honors_dependent_required_and_min_p assert json.loads(_call(body)["arguments"]) == {"site": "check-network-telemetry", "region": "west", "extra": True} -def test_foundry_brokered_integer_synthesis_rejects_unsupported_multiple_of(): +def test_foundry_brokered_integer_synthesis_rejects_multiple_of(): spec = _spec(tool_name="check-network-telemetry") spec.brokered_tools[0].parameters = { "type": "object", @@ -545,6 +653,7 @@ def test_foundry_brokered_integer_synthesis_rejects_unsupported_multiple_of(): assert response.json()["error"]["code"] == "UnsupportedBrokeredSchema" + def test_foundry_brokered_number_synthesis_uses_midpoint_for_fractional_exclusive_range(): spec = _spec(tool_name="check-network-telemetry") spec.brokered_tools[0].parameters = { @@ -560,347 +669,105 @@ def test_foundry_brokered_number_synthesis_uses_midpoint_for_fractional_exclusiv assert json.loads(_call(body)["arguments"]) == {"ratio": 0.5} -def test_foundry_brokered_number_synthesis_handles_large_finite_bounds_without_overflow(): - spec = _spec(tool_name="large-number-bounds") +def test_foundry_brokered_rejects_unbounded_schema_synthesis_before_allocating(): + spec = _spec(tool_name="check-network-telemetry") spec.brokered_tools[0].parameters = { "type": "object", - "properties": { - "midpoint": {"type": "number", "minimum": 1e308, "maximum": 1.1e308}, - "above": {"type": "number", "exclusiveMinimum": 1e308}, - "below": {"type": "number", "exclusiveMaximum": -1e308}, - "exactInteger": {"type": "number", "minimum": 9007199254740993, "maximum": 9007199254740994}, - "roundedMidpoint": {"type": "number", "exclusiveMinimum": 0, "exclusiveMaximum": 0.9999999999999999}, - "narrowValue": {"type": "number", "exclusiveMinimum": 1.0, "maximum": 1.0000000000000002}, - }, - "required": ["midpoint", "above", "below", "exactInteger", "roundedMidpoint", "narrowValue"], + "properties": {"site": {"type": "string", "minLength": 1_000_000_000}}, + "required": ["site"], } app = _app(spec) with TestClient(app) as client: - response = client.post("/responses", json={"input": "large-number-bounds"}) + resp = client.post("/responses", json={"input": "check-network-telemetry"}) - assert response.status_code == 200, response.text - arguments = json.loads(_call(response.json())["arguments"]) - assert arguments["midpoint"] == 1e308 - assert arguments["above"] == math.nextafter(1e308, math.inf) - assert arguments["below"] == math.nextafter(-1e308, -math.inf) - assert arguments["exactInteger"] == 9007199254740993 - assert 0 < arguments["roundedMidpoint"] < 0.9999999999999999 - assert arguments["narrowValue"] == 1.0000000000000002 + assert resp.status_code == 413 + assert resp.json()["error"]["code"] == "brokered_arguments_too_large" -def test_foundry_brokered_number_synthesis_rejects_unsupported_multiple_of(): - spec = _spec(tool_name="number-multiple") - spec.brokered_tools[0].parameters = { - "type": "object", - "properties": {"value": {"type": "number", "minimum": 0, "multipleOf": 0.5}}, - "required": ["value"], - } - app = _app(spec) +def test_foundry_brokered_tool_selection_requires_token_boundary_match(): + data = _multi_tool_spec().model_dump(by_alias=True) + data["brokeredTools"] = [ + {"name": "read", "description": "read", "brokeredClass": "read", "parameters": {"type": "object"}}, + {"name": "read_telemetry", "description": "read telemetry", "brokeredClass": "read", "parameters": {"type": "object"}}, + ] + app = _app(AgentSpec.model_validate(data)) with TestClient(app) as client: - response = client.post("/responses", json={"input": "number-multiple"}) + telemetry = client.post("/responses", json={"input": "please call read_telemetry"}) + unrelated = client.post("/responses", json={"input": "already done"}) - assert response.status_code == 400 - assert response.json()["error"]["code"] == "UnsupportedBrokeredSchema" + assert telemetry.status_code == 200, telemetry.text + assert _call(telemetry.json())["name"] == "read_telemetry" + assert unrelated.status_code == 400 + assert unrelated.json()["error"]["code"] == "brokered_tool_selection_required" -def test_foundry_brokered_number_synthesis_combines_all_declared_bounds(): - spec = _spec(tool_name="combined-number-bounds") +def test_foundry_brokered_rejects_schema_that_would_generate_huge_arguments_before_allocating(): + spec = _spec(tool_name="huge-args") spec.brokered_tools[0].parameters = { "type": "object", "properties": { - "value": {"type": "number", "minimum": 0, "exclusiveMinimum": 0, "maximum": 1}, + "items": {"type": "array", "minItems": 1000000000, "items": {"type": "string"}}, + "label": {"type": "string", "minLength": 1000000000}, }, - "required": ["value"], + "required": ["items", "label"], } app = _app(spec) with TestClient(app) as client: - response = client.post("/responses", json={"input": "combined-number-bounds"}) + resp = client.post("/responses", json={"input": "huge-args"}) - assert response.status_code == 200, response.text - assert json.loads(_call(response.json())["arguments"]) == {"value": 1} + assert resp.status_code == 413 + assert resp.json()["error"]["code"] == "brokered_arguments_too_large" -@pytest.mark.parametrize("bound", [float("nan"), float("inf"), float("-inf")]) -def test_foundry_brokered_number_synthesis_rejects_nonfinite_bounds(bound: float): - spec = _spec(tool_name="nonfinite-number-bound") - spec.brokered_tools[0].parameters = { - "type": "object", - "properties": {"value": {"type": "number", "minimum": bound}}, - "required": ["value"], - } - app = _app(spec) +def test_foundry_brokered_rejects_arguments_that_exceed_state_budget(): + spec = _spec(tool_name="check-network-telemetry") + spec.brokered_tools[0].parameters["required"] = ["site"] + app = _app(spec, max_brokered_argument_bytes=8) with TestClient(app) as client: - response = client.post("/responses", json={"input": "nonfinite-number-bound"}) + resp = client.post("/responses", json={"input": "check-network-telemetry: a prompt that is too large for the site argument"}) - assert response.status_code == 400 - assert response.json()["error"]["code"] == "UnsupportedBrokeredSchema" + assert resp.status_code == 413 + assert resp.json()["error"]["code"] == "brokered_arguments_too_large" -def test_foundry_brokered_number_synthesis_prefers_compact_zero_within_wide_bounds(): - spec = _spec(tool_name="compact-number-bounds") - properties = { - f"value{index}": {"type": "number", "minimum": -1e308, "maximum": 1e308} - for index in range(27) - } - spec.brokered_tools[0].parameters = { - "type": "object", - "properties": properties, - "required": list(properties), - } - app = _app(spec) +def test_foundry_brokered_single_read_tool_requires_explicit_name_except_conformance(): + app = _app(_spec(tool_name="check-network-telemetry")) with TestClient(app) as client: - response = client.post("/responses", json={"input": "compact-number-bounds"}) + unrelated = client.post("/responses", json={"input": "my password should stay in chat"}) + explicit = client.post("/responses", json={"input": "please call check-network-telemetry"}) - assert response.status_code == 200, response.text - assert json.loads(_call(response.json())["arguments"]) == {name: 0 for name in properties} + assert unrelated.status_code == 400 + assert unrelated.json()["error"]["code"] == "brokered_tool_selection_required" + assert explicit.status_code == 200, explicit.text + assert _call(explicit.json())["name"] == "check-network-telemetry" -def test_foundry_brokered_number_synthesis_prefers_compact_float_boundary_over_large_integer(): - spec = _spec(tool_name="compact-large-number-bounds") - properties = { - f"value{index}": {"type": "number", "minimum": 1e308, "maximum": 1.1e308} - for index in range(27) - } - spec.brokered_tools[0].parameters = { - "type": "object", - "properties": properties, - "required": list(properties), - } - app = _app(spec) +def test_foundry_brokered_single_write_tool_requires_explicit_tool_name(): + app = _app(_spec(tool_name="dispatch-work-order", brokered_class="write")) with TestClient(app) as client: - response = client.post("/responses", json={"input": "compact-large-number-bounds"}) + resp = client.post("/responses", json={"input": "hello"}) + explicit = client.post("/responses", json={"input": "please call dispatch-work-order"}) - assert response.status_code == 200, response.text - assert json.loads(_call(response.json())["arguments"]) == {name: 1e308 for name in properties} + assert resp.status_code == 400 + assert resp.json()["error"]["code"] == "brokered_tool_selection_required" + assert explicit.status_code == 200, explicit.text + assert _call(explicit.json())["name"] == "dispatch-work-order" -def test_foundry_brokered_rejects_unbounded_schema_synthesis_before_allocating(): - spec = _spec(tool_name="check-network-telemetry") - spec.brokered_tools[0].parameters = { - "type": "object", - "properties": {"site": {"type": "string", "minLength": 1_000_000_000}}, - "required": ["site"], - } - app = _app(spec) +def test_foundry_brokered_selects_named_tool_when_multiple_schemas_are_configured(): + app = _app(_multi_tool_spec()) with TestClient(app) as client: - resp = client.post("/responses", json={"input": "check-network-telemetry"}) + resp = client.post("/responses", json={"input": "please call get-active-incidents"}) - assert resp.status_code == 413 - assert resp.json()["error"]["code"] == "brokered_arguments_too_large" - - -def test_foundry_brokered_bounds_nested_array_synthesis_before_recursive_allocation(): - spec = _spec(tool_name="nested-array") - spec.brokered_tools[0].parameters = { - "type": "object", - "properties": { - "values": { - "type": "array", - "minItems": 32, - "items": { - "type": "array", - "minItems": 32, - "items": {"type": "string"}, - }, - } - }, - "required": ["values"], - } - app = _app(spec) - - with TestClient(app) as client: - response = client.post("/responses", json={"input": "nested-array"}) - - assert response.status_code == 413 - assert response.json()["error"]["code"] == "brokered_arguments_too_large" - - -def test_foundry_brokered_counts_wide_object_members_against_synthesis_budget(): - item_properties = {f"field{index}": {"type": "string"} for index in range(32)} - spec = _spec(tool_name="wide-object-array") - spec.brokered_tools[0].parameters = { - "type": "object", - "properties": { - "values": { - "type": "array", - "minItems": 32, - "items": { - "type": "object", - "properties": item_properties, - "required": list(item_properties), - }, - } - }, - "required": ["values"], - } - app = _app(spec) - - with TestClient(app) as client: - response = client.post("/responses", json={"input": "wide-object-array"}) - - assert response.status_code == 413 - assert response.json()["error"]["code"] == "brokered_arguments_too_large" - - -def test_foundry_brokered_rejects_overly_deep_deterministic_synthesis(): - child: dict[str, Any] = {"type": "string", "default": "leaf"} - for _ in range(129): - child = { - "type": "object", - "properties": {"child": child}, - "required": ["child"], - } - spec = _spec(tool_name="deep-synthesis") - spec.brokered_tools[0].parameters = { - "type": "object", - "properties": {"root": child}, - "required": ["root"], - } - app = _app(spec) - - with TestClient(app) as client: - response = client.post("/responses", json={"input": "deep-synthesis"}) - - assert response.status_code == 413 - assert response.json()["error"]["code"] == "brokered_arguments_too_large" - - -def test_foundry_brokered_rejects_overly_deep_root_literal_synthesis(): - literal: dict[str, Any] = {} - for _ in range(129): - literal = {"nested": literal} - spec = _spec(tool_name="deep-root-literal") - spec.brokered_tools[0].parameters = { - "type": "object", - "const": literal, - } - app = _app(spec) - - with TestClient(app) as client: - response = client.post("/responses", json={"input": "deep-root-literal"}) - - assert response.status_code == 413 - assert response.json()["error"]["code"] == "brokered_arguments_too_large" - - -@pytest.mark.parametrize( - ("literal", "expected"), - [ - ({"const": 2.0}, 2.0), - ({"default": 1.0}, 1.0), - ({"enum": [3.0]}, 3.0), - ], -) -def test_foundry_brokered_honors_integral_float_integer_literals( - literal: dict[str, Any], expected: float -): - spec = _spec(tool_name="integer-literal") - spec.brokered_tools[0].parameters = { - "type": "object", - "properties": {"count": {"type": "integer", **literal}}, - "required": ["count"], - } - app = _app(spec) - - with TestClient(app) as client: - response = client.post("/responses", json={"input": "integer-literal"}) - - assert response.status_code == 200, response.text - count = json.loads(_call(response.json())["arguments"])["count"] - assert count == expected - assert isinstance(count, float) - - -def test_foundry_brokered_tool_selection_requires_token_boundary_match(): - data = _multi_tool_spec().model_dump(by_alias=True) - data["brokeredTools"] = [ - {"name": "read", "description": "read", "brokeredClass": "read", "parameters": {"type": "object"}}, - {"name": "read_telemetry", "description": "read telemetry", "brokeredClass": "read", "parameters": {"type": "object"}}, - ] - app = _app(AgentSpec.model_validate(data)) - - with TestClient(app) as client: - telemetry = client.post("/responses", json={"input": "please call read_telemetry"}) - unrelated = client.post("/responses", json={"input": "already done"}) - - assert telemetry.status_code == 200, telemetry.text - assert _call(telemetry.json())["name"] == "read_telemetry" - assert unrelated.status_code == 400 - assert unrelated.json()["error"]["code"] == "brokered_tool_selection_required" - - -def test_foundry_brokered_rejects_schema_that_would_generate_huge_arguments_before_allocating(): - spec = _spec(tool_name="huge-args") - spec.brokered_tools[0].parameters = { - "type": "object", - "properties": { - "items": {"type": "array", "minItems": 1000000000, "items": {"type": "string"}}, - "label": {"type": "string", "minLength": 1000000000}, - }, - "required": ["items", "label"], - } - app = _app(spec) - - with TestClient(app) as client: - resp = client.post("/responses", json={"input": "huge-args"}) - - assert resp.status_code == 413 - assert resp.json()["error"]["code"] == "brokered_arguments_too_large" - - -def test_foundry_brokered_rejects_arguments_that_exceed_state_budget(): - spec = _spec(tool_name="check-network-telemetry") - spec.brokered_tools[0].parameters["required"] = ["site"] - app = _app(spec, max_brokered_argument_bytes=8) - - with TestClient(app) as client: - resp = client.post("/responses", json={"input": "check-network-telemetry: a prompt that is too large for the site argument"}) - - assert resp.status_code == 413 - assert resp.json()["error"]["code"] == "brokered_arguments_too_large" - - -def test_foundry_brokered_single_read_tool_requires_explicit_name_except_conformance(): - app = _app(_spec(tool_name="check-network-telemetry")) - - with TestClient(app) as client: - unrelated = client.post("/responses", json={"input": "my password should stay in chat"}) - explicit = client.post("/responses", json={"input": "please call check-network-telemetry"}) - - assert unrelated.status_code == 400 - assert unrelated.json()["error"]["code"] == "brokered_tool_selection_required" - assert explicit.status_code == 200, explicit.text - assert _call(explicit.json())["name"] == "check-network-telemetry" - - -def test_foundry_brokered_single_write_tool_requires_explicit_tool_name(): - app = _app(_spec(tool_name="dispatch-work-order", brokered_class="write")) - - with TestClient(app) as client: - resp = client.post("/responses", json={"input": "hello"}) - explicit = client.post("/responses", json={"input": "please call dispatch-work-order"}) - - assert resp.status_code == 400 - assert resp.json()["error"]["code"] == "brokered_tool_selection_required" - assert explicit.status_code == 200, explicit.text - assert _call(explicit.json())["name"] == "dispatch-work-order" - - -def test_foundry_brokered_selects_named_tool_when_multiple_schemas_are_configured(): - app = _app(_multi_tool_spec()) - - with TestClient(app) as client: - resp = client.post("/responses", json={"input": "please call get-active-incidents"}) - - assert resp.status_code == 200, resp.text - assert _call(resp.json())["name"] == "get-active-incidents" + assert resp.status_code == 200, resp.text + assert _call(resp.json())["name"] == "get-active-incidents" def test_foundry_brokered_rejects_ambiguous_multi_tool_prompt(): @@ -938,87 +805,231 @@ def test_foundry_brokered_rejects_nonfinite_function_call_output_values(): assert response.json()["error"]["code"] == "invalid_function_call_output" -def test_foundry_brokered_rejects_lossy_function_call_output_float_before_state_change(tmp_path): +def test_foundry_nonfinite_validation_traverses_wide_lists_lazily(): + class ExplodingTailList(list): + def __iter__(self): + yield float("nan") + raise AssertionError("non-finite validation eagerly consumed the full list") + + with pytest.raises(ValueError, match="must be finite"): + foundry_module._reject_nonfinite_json_values( + {"result": ExplodingTailList([0, 1])}, + path="function_call_output.output", + ) + + +def test_foundry_nonfinite_validation_binds_sibling_paths(): + with pytest.raises(ValueError, match=r"function_call_output\.output\.items\[1\] must be finite"): + foundry_module._reject_nonfinite_json_values( + {"items": [0.0, float("inf")]}, + path="function_call_output.output", + ) + + +def test_foundry_request_ceiling_ignores_expired_replay_sizes(): + store = foundry_module._FoundryResponseStateStore( + ttl_seconds=60, + max_entries=2, + max_bytes=4 * 1024, + ) + store._states["expired"] = foundry_module._HostedResponseState( + response_id="expired", + session_id=None, + pending_calls={}, + expires_at=time.time() - 1, + status="completed", + accepted_output_digests={"call-1": "sha256:" + "0" * 64}, + accepted_output_sizes={"call-1": 500_000}, + final_payload={"status": "completed"}, + ) + + assert store.max_accepted_output_bytes() == 0 + + +def test_foundry_brokered_rejects_function_output_with_lone_surrogate(): + app = _app() + + with TestClient(app) as client: + initial = _start(client) + call = _call(initial) + payload = _continuation(initial["id"], call["call_id"], {"approved": True, "output": {}}) + payload["input"][0]["output"] = r'{"approved":true,"output":{"value":"\ud800"}}' + rejected = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + accepted = client.post( + "/responses", + headers=CONTINUATION_AUTH, + json=_continuation(initial["id"], call["call_id"], {"approved": True, "output": {"ok": True}}), + ) + + assert rejected.status_code == 400 + assert rejected.json()["error"]["code"] == "invalid_function_call_output" + assert accepted.status_code == 200, accepted.text + + +def test_foundry_brokered_rejects_oversized_function_call_output_before_state_change(tmp_path): state_file = tmp_path / "responses-state.json" - app = _app(response_state_file=state_file) + app = _app(response_state_file=state_file, max_brokered_output_bytes=128) with TestClient(app) as client: initial = _start(client) call = _call(initial) persisted_before = state_file.read_bytes() - lossy = client.post( + oversized = client.post( "/responses", headers=CONTINUATION_AUTH, - json={ - "previous_response_id": initial["id"], - "input": [ - { - "type": "function_call_output", - "call_id": call["call_id"], - "output": '{"approved":true,"output":{"id":9007199254740993.0}}', - } - ], - }, + json=_continuation( + initial["id"], + call["call_id"], + {"approved": True, "output": {"blob": "x" * 256}}, + ), ) persisted_after_rejection = state_file.read_bytes() - exact = client.post( + accepted = client.post( "/responses", headers=CONTINUATION_AUTH, - json={ - "previous_response_id": initial["id"], - "input": [ - { - "type": "function_call_output", - "call_id": call["call_id"], - "output": '{"approved":true,"output":{"id":9007199254740992.0}}', - } - ], - }, + json=_continuation(initial["id"], call["call_id"], {"approved": True, "output": {"ok": True}}), ) - assert lossy.status_code == 400 - assert lossy.json()["error"]["code"] == "invalid_function_call_output" + assert oversized.status_code == 413 + assert oversized.json()["error"] == { + "message": "brokered function_call_output is too large", + "code": "brokered_output_too_large", + } assert persisted_after_rejection == persisted_before - assert exact.status_code == 200, exact.text - assert _message_text(exact.json()).endswith('{"id":9007199254740992.0}') + assert accepted.status_code == 200, accepted.text -def test_foundry_brokered_rejects_duplicate_keys_in_function_call_output(tmp_path): +def test_foundry_brokered_replay_ceiling_is_independent_from_non_brokered_request_limit(): + app = _app( + max_request_body_bytes=64, + max_brokered_output_bytes=512, + max_response_state_bytes=4 * 1024, + ) + + with TestClient(app) as client: + initial = _start(client) + call = _call(initial) + continuation = _continuation( + initial["id"], + call["call_id"], + {"approved": True, "output": {"blob": "x" * 128}}, + ) + assert len(json.dumps(continuation).encode("utf-8")) > 64 + response = client.post("/responses", headers=CONTINUATION_AUTH, json=continuation) + + assert response.status_code == 200, response.text + + +def test_foundry_brokered_rejects_oversized_request_body_without_truncation_or_state_change(tmp_path): state_file = tmp_path / "responses-state.json" - app = _app(response_state_file=state_file) + app = _app( + response_state_file=state_file, + max_brokered_output_bytes=128, + max_response_state_bytes=4 * 1024, + ) with TestClient(app) as client: initial = _start(client) call = _call(initial) persisted_before = state_file.read_bytes() - response = client.post( + oversized_payload = _continuation( + initial["id"], + call["call_id"], + {"approved": True, "output": {"ok": True}}, + ) + oversized_payload["ignored_padding"] = "x" * 100_000 + rejected = client.post( + "/responses", + headers={**CONTINUATION_AUTH, "content-type": "application/json"}, + content=json.dumps(oversized_payload, separators=(",", ":")), + ) + persisted_after_rejection = state_file.read_bytes() + accepted = client.post( "/responses", headers=CONTINUATION_AUTH, - json={ - "previous_response_id": initial["id"], - "input": [ - { - "type": "function_call_output", - "call_id": call["call_id"], - "output": '{"approved":false,"approved":true,"output":{}}', - } - ], - }, + json=_continuation(initial["id"], call["call_id"], {"approved": True, "output": {"ok": True}}), + ) + + assert rejected.status_code == 413 + assert rejected.json()["error"] == { + "message": "brokered Responses request body is too large", + "code": "brokered_request_too_large", + } + assert persisted_after_rejection == persisted_before + assert accepted.status_code == 200, accepted.text + + +def test_foundry_brokered_bounded_json_value_error_returns_400(): + app = _app(max_response_state_bytes=32 * 1024) + oversized_integer = '{"input":' + '9' * 5_000 + '}' + + with TestClient(app, raise_server_exceptions=False) as client: + response = client.post( + "/responses", + headers={"content-type": "application/json"}, + content=oversized_integer, ) assert response.status_code == 400 - assert response.json()["error"]["code"] == "invalid_function_call_output" - assert state_file.read_bytes() == persisted_before + assert response.json()["error"] == { + "message": "Request body must be JSON", + "code": "invalid_json", + } -def test_foundry_brokered_rejects_object_valued_function_call_output_before_state_change(tmp_path): - state_file = tmp_path / "responses-state.json" - app = _app(response_state_file=state_file) +def test_foundry_brokered_request_limit_allows_valid_reescaped_model_output(): + spec = _spec(tool_name="check-network-telemetry") + fake = _FakeChatTransport( + [ + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "model-generated-call", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": "{}"}, + } + ], + } + ), + _chat_response({"role": "assistant", "content": "Large tool output accepted."}), + ] + ) + app = _model_loop_app( + spec, + fake, + max_brokered_output_bytes=64 * 1024, + max_response_state_bytes=96 * 1024, + ) + message = "\x7f" * 10_800 + raw_output = '{"approved":false,"error":{"message":"' + message + '"}}' + raw_output += "\n" * (64 * 1024 - len(raw_output) - 100) + parsed_output = json.loads(raw_output) + canonical_output = json.dumps(parsed_output, separators=(",", ":"), sort_keys=True) + assert len(raw_output.encode("utf-8")) < 64 * 1024 + assert len(canonical_output.encode("utf-8")) < 64 * 1024 + + with TestClient(app) as client: + initial = client.post("/responses", json={"input": "check-network-telemetry"}) + call = _call(initial.json()) + payload = _continuation(initial.json()["id"], call["call_id"], {"approved": False}) + payload["input"][0]["output"] = raw_output + encoded_request = json.dumps(payload, separators=(",", ":")) + assert len(encoded_request.encode("utf-8")) > 2 * 64 * 1024 + 16 * 1024 + final = client.post("/responses", headers=CONTINUATION_AUTH, content=encoded_request) + + assert final.status_code == 200, final.text + assert _message_text(final.json()) == "Large tool output accepted." + + +def test_foundry_brokered_rejects_object_valued_function_call_output(): + app = _app(max_brokered_output_bytes=64) with TestClient(app) as client: initial = _start(client) call = _call(initial) - persisted_before = state_file.read_bytes() response = client.post( "/responses", headers=CONTINUATION_AUTH, @@ -1028,7 +1039,7 @@ def test_foundry_brokered_rejects_object_valued_function_call_output_before_stat { "type": "function_call_output", "call_id": call["call_id"], - "output": {"approved": True, "output": {"id": 9007199254740993.0}}, + "output": {"approved": True, "output": {"blob": "x" * 256}}, } ], }, @@ -1036,24 +1047,28 @@ def test_foundry_brokered_rejects_object_valued_function_call_output_before_stat assert response.status_code == 400 assert response.json()["error"]["code"] == "invalid_function_call_output" - assert state_file.read_bytes() == persisted_before -def test_foundry_brokered_rejects_oversized_function_call_output_before_state_change(tmp_path): + +def test_foundry_brokered_rejects_continuation_that_would_exceed_total_state_bytes(tmp_path): state_file = tmp_path / "responses-state.json" - app = _app(response_state_file=state_file, max_brokered_output_bytes=128) + app = _app( + response_state_file=state_file, + max_brokered_output_bytes=8 * 1024, + max_response_state_bytes=4 * 1024, + ) with TestClient(app) as client: initial = _start(client) call = _call(initial) persisted_before = state_file.read_bytes() - oversized = client.post( + oversized_state = client.post( "/responses", headers=CONTINUATION_AUTH, json=_continuation( initial["id"], call["call_id"], - {"approved": True, "output": {"blob": "x" * 256}}, + {"approved": True, "output": {"blob": "x" * 5_000}}, ), ) persisted_after_rejection = state_file.read_bytes() @@ -1063,42 +1078,61 @@ def test_foundry_brokered_rejects_oversized_function_call_output_before_state_ch json=_continuation(initial["id"], call["call_id"], {"approved": True, "output": {"ok": True}}), ) - assert oversized.status_code == 413 - assert oversized.json()["error"] == { - "message": "brokered function_call_output is too large", - "code": "brokered_output_too_large", + assert oversized_state.status_code == 413 + assert oversized_state.json()["error"] == { + "message": "brokered response state exceeds the configured byte limit", + "code": "brokered_response_state_too_large", } assert persisted_after_rejection == persisted_before assert accepted.status_code == 200, accepted.text -def test_foundry_brokered_rejects_oversized_raw_output_before_json_parsing(monkeypatch): - app = _app(max_brokered_output_bytes=64) - - def fail_if_parsed(_output: Any) -> dict[str, Any]: - raise AssertionError("oversized raw output must be rejected before JSON parsing") +def test_foundry_brokered_model_loop_counts_output_limit_in_utf8_bytes(): + spec = _spec(tool_name="check-network-telemetry") + fake = _FakeChatTransport( + [ + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "model_generated_call_id", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": "{}"}, + } + ], + } + ), + _chat_response({"role": "assistant", "content": "Bounded resume worked."}), + ] + ) + app = _model_loop_app(spec, fake, max_brokered_output_bytes=80) - monkeypatch.setattr(foundry_module, "_json_object_from_output", fail_if_parsed) with TestClient(app) as client: - initial = _start(client) - call = _call(initial) - response = client.post( + initial = client.post("/responses", json={"input": "check-network-telemetry"}) + call = _call(initial.json()) + oversized_payload = _continuation(initial.json()["id"], call["call_id"], {"approved": True, "output": {}}) + oversized_output = json.dumps( + {"approved": True, "output": {"value": "é" * 21}}, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ) + assert len(oversized_output) < 80 < len(oversized_output.encode("utf-8")) + oversized_payload["input"][0]["output"] = oversized_output + rejected = client.post("/responses", headers=CONTINUATION_AUTH, json=oversized_payload) + accepted = client.post( "/responses", headers=CONTINUATION_AUTH, - json={ - "previous_response_id": initial["id"], - "input": [ - { - "type": "function_call_output", - "call_id": call["call_id"], - "output": '{"approved":true,"output":{"blob":"' + ("x" * 128) + '"}}', - } - ], - }, + json=_continuation(initial.json()["id"], call["call_id"], {"approved": True, "output": {"ok": True}}), ) - assert response.status_code == 413 - assert response.json()["error"]["code"] == "brokered_output_too_large" + assert rejected.status_code == 413 + assert rejected.json()["error"]["code"] == "brokered_output_too_large" + assert accepted.status_code == 200, accepted.text + assert _message_text(accepted.json()) == "Bounded resume worked." + assert len(fake.requests) == 2 def test_foundry_brokered_continuation_accepts_matching_tool_output_and_completes(): @@ -1121,34 +1155,6 @@ def test_foundry_brokered_continuation_accepts_matching_tool_output_and_complete assert _message_text(final) == 'Brokered tool conformance_read completed with output: {"success":true}' -def test_foundry_brokered_accepting_continuation_refreshes_state_expiry(monkeypatch): - stores: list[Any] = [] - original_store = foundry_module._FoundryResponseStateStore - - def capture_store(*args: Any, **kwargs: Any) -> Any: - store = original_store(*args, **kwargs) - stores.append(store) - return store - - monkeypatch.setattr(foundry_module, "_FoundryResponseStateStore", capture_store) - app = _app(state_ttl_seconds=60) - - with TestClient(app) as client: - initial = _start(client) - call = _call(initial) - state = stores[0].get(initial["id"]) - state.expires_at = time.time() + 1 - stores[0].save(state) - response = client.post( - "/responses", - headers=CONTINUATION_AUTH, - json=_continuation(initial["id"], call["call_id"], {"approved": True, "output": {"ok": True}}), - ) - - assert response.status_code == 200, response.text - assert state.expires_at > time.time() + 50 - - def test_foundry_brokered_accepts_body_continuation_proof_without_header(): app = _app() @@ -1307,25 +1313,30 @@ def test_foundry_brokered_duplicate_continuation_is_idempotent_but_conflicts_are assert conflict.json()["error"]["code"] == "conflicting_duplicate_continuation" -def test_foundry_brokered_rejects_oversized_persisted_replay_after_limit_is_lowered(tmp_path): - state_file = tmp_path / "responses-state.json" +def test_foundry_brokered_persists_one_output_copy_while_preserving_idempotency(tmp_path): + state_file = tmp_path / "foundry-state.json" + marker = "unique-brokered-output-marker" + app = _app(response_state_file=state_file) - with TestClient(_app(response_state_file=state_file, max_brokered_output_bytes=1024)) as client: + with TestClient(app) as client: initial = _start(client) call = _call(initial) payload = _continuation( initial["id"], call["call_id"], - {"approved": True, "output": {"blob": "x" * 256}}, + {"approved": True, "output": {"value": marker}}, ) - completed = client.post("/responses", json=payload, headers=CONTINUATION_AUTH) - - with TestClient(_app(response_state_file=state_file, max_brokered_output_bytes=128)) as client: - replay = client.post("/responses", json=payload, headers=CONTINUATION_AUTH) + completed = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + duplicate = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + conflicting = deepcopy(payload) + conflicting["input"][0]["output"] = '{"approved":true,"output":{"value":"different"}}' + conflict = client.post("/responses", headers=CONTINUATION_AUTH, json=conflicting) assert completed.status_code == 200, completed.text - assert replay.status_code == 413 - assert replay.json()["error"]["code"] == "brokered_output_too_large" + assert duplicate.json() == completed.json() + assert conflict.status_code == 409 + assert conflict.json()["error"]["code"] == "conflicting_duplicate_continuation" + assert state_file.read_text(encoding="utf-8").count(marker) == 1 def test_foundry_brokered_file_state_survives_restart_for_deterministic_continuation(tmp_path): @@ -1360,61 +1371,319 @@ def test_foundry_brokered_file_state_survives_restart_for_deterministic_continua assert duplicate.json() == final.json() -def test_foundry_brokered_file_state_tracks_session_and_rejects_cross_session_continuation(monkeypatch, tmp_path): - state_file = tmp_path / "foundry-session-state.json" - monkeypatch.delenv("FOUNDRY_AGENT_SESSION_ID", raising=False) +def test_foundry_brokered_loads_legacy_full_output_state_for_idempotent_replay(tmp_path): + state_file = tmp_path / "foundry-legacy-state.json" with TestClient(_app(response_state_file=state_file)) as client: - initial_response = client.post( - "/responses", - json={"input": "please read telemetry", "agent_session_id": "session-a"}, - ) - assert initial_response.status_code == 200, initial_response.text - initial = initial_response.json() + initial = _start(client) call = _call(initial) + payload = _continuation(initial["id"], call["call_id"], {"approved": True, "output": {"success": True}}) + completed = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) - stored = json.loads(state_file.read_text(encoding="utf-8"))["states"][initial["id"]] - assert stored["sessionID"] == "session-a" - - payload = _continuation(initial["id"], call["call_id"], {"approved": True, "output": {"success": True}}) - payload[CONTINUATION_PROOF_BODY_FIELD] = CONTINUATION_PROOF - with TestClient(_app(response_state_file=state_file)) as client: - missing = client.post("/responses", json=payload) - - assert missing.status_code == 409 - assert missing.json()["error"]["code"] == "response_session_mismatch" - - payload["agent_session_id"] = "session-b" - with TestClient(_app(response_state_file=state_file)) as client: - mismatch = client.post("/responses", json=payload) - - assert mismatch.status_code == 409 - assert mismatch.json()["error"]["code"] == "response_session_mismatch" + persisted = json.loads(state_file.read_text(encoding="utf-8")) + state = persisted["states"][initial["id"]] + state.pop("acceptedOutputDigests") + state["acceptedOutputs"] = {call["call_id"]: payload["input"][0]["output"]} + state_file.write_text(json.dumps(persisted, separators=(",", ":"), sort_keys=True), encoding="utf-8") - payload["agent_session_id"] = "session-a" with TestClient(_app(response_state_file=state_file)) as client: - completed = client.post("/responses", json=payload) + duplicate = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) assert completed.status_code == 200, completed.text + assert duplicate.status_code == 200, duplicate.text + assert duplicate.json() == completed.json() -def test_foundry_brokered_body_proof_is_not_persisted_or_echoed(tmp_path): - state_file = tmp_path / "foundry-proof-state.json" - app = _app(response_state_file=state_file) - - with TestClient(app) as client: - initial_response = client.post( - "/responses", - json={"input": "please read telemetry", CONTINUATION_PROOF_BODY_FIELD: CONTINUATION_PROOF}, - ) - assert initial_response.status_code == 200, initial_response.text - initial = initial_response.json() - call = _call(initial) - payload = _continuation(initial["id"], call["call_id"], {"approved": True, "output": {"success": True}}) - payload[CONTINUATION_PROOF_BODY_FIELD] = CONTINUATION_PROOF - final_response = client.post("/responses", json=payload) - - assert final_response.status_code == 200, final_response.text +def test_foundry_brokered_replays_legacy_output_larger_than_current_limit(tmp_path): + state_file = tmp_path / "foundry-legacy-large-output-state.json" + spec = _spec(tool_name="check-network-telemetry") + large_output = {"approved": True, "output": {"blob": "x" * 80_000}} + fake = _FakeChatTransport( + [ + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "model-generated-call", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": "{}"}, + } + ], + } + ), + _chat_response({"role": "assistant", "content": "Persisted large output replay."}), + ] + ) + with TestClient( + _model_loop_app( + spec, + fake, + response_state_file=state_file, + max_brokered_output_bytes=128 * 1024, + ) + ) as client: + initial = client.post("/responses", json={"input": "check-network-telemetry"}) + call = _call(initial.json()) + payload = _continuation(initial.json()["id"], call["call_id"], large_output) + completed = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + + persisted = json.loads(state_file.read_text(encoding="utf-8")) + state = persisted["states"][initial.json()["id"]] + state.pop("acceptedOutputDigests") + state["acceptedOutputs"] = {call["call_id"]: payload["input"][0]["output"]} + state_file.write_text(json.dumps(persisted, separators=(",", ":"), sort_keys=True), encoding="utf-8") + + with TestClient( + _app( + spec, + response_state_file=state_file, + max_brokered_output_bytes=64 * 1024, + ) + ) as client: + duplicate = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + + assert completed.status_code == 200, completed.text + assert duplicate.status_code == 200, duplicate.text + assert duplicate.json() == completed.json() + + +def test_foundry_brokered_replays_digest_output_larger_than_current_and_state_limits(tmp_path): + state_file = tmp_path / "foundry-digest-large-output-state.json" + spec = _spec(tool_name="check-network-telemetry") + message = "\x7f" * 10_800 + raw_output = '{"approved":false,"error":{"message":"' + message + '"}}' + raw_output += "\n" * (400 * 1024 - len(raw_output) - 100) + parsed_output = json.loads(raw_output) + canonical_output = json.dumps(parsed_output, separators=(",", ":"), sort_keys=True) + assert len(raw_output.encode("utf-8")) < 512 * 1024 + assert len(canonical_output.encode("utf-8")) < 128 * 1024 + fake = _FakeChatTransport( + [ + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "model-generated-call", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": "{}"}, + } + ], + } + ), + _chat_response({"role": "assistant", "content": "Persisted digest replay."}), + ] + ) + with TestClient( + _model_loop_app( + spec, + fake, + response_state_file=state_file, + max_brokered_output_bytes=512 * 1024, + max_response_state_bytes=128 * 1024, + ) + ) as client: + initial = client.post("/responses", json={"input": "check-network-telemetry"}) + call = _call(initial.json()) + payload = _continuation(initial.json()["id"], call["call_id"], {"approved": False}) + payload["input"][0]["output"] = raw_output + completed = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + + persisted = json.loads(state_file.read_text(encoding="utf-8")) + stored = persisted["states"][initial.json()["id"]] + assert stored["acceptedOutputSizes"][call["call_id"]] == len(raw_output.encode("utf-8")) + assert state_file.stat().st_size < 128 * 1024 + + with TestClient( + _app( + spec, + response_state_file=state_file, + max_brokered_output_bytes=64 * 1024, + max_response_state_bytes=128 * 1024, + ) + ) as client: + encoded_request = json.dumps(payload, separators=(",", ":")) + assert len(encoded_request.encode("utf-8")) > 6 * 128 * 1024 + 16 * 1024 + duplicate = client.post("/responses", headers=CONTINUATION_AUTH, content=encoded_request) + + assert completed.status_code == 200, completed.text + assert duplicate.status_code == 200, duplicate.text + assert duplicate.json() == completed.json() + + +def test_foundry_brokered_rejects_legacy_state_that_expands_past_byte_limit(tmp_path): + state_file = tmp_path / "foundry-legacy-state.json" + + with TestClient(_app(response_state_file=state_file)) as client: + initial = _start(client) + call = _call(initial) + payload = _continuation( + initial["id"], + call["call_id"], + {"approved": False, "error": {"code": "denied", "message": "denied"}}, + ) + completed = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + + assert completed.status_code == 200, completed.text + persisted = json.loads(state_file.read_text(encoding="utf-8")) + state = persisted["states"][initial["id"]] + state.pop("acceptedOutputDigests") + state["acceptedOutputs"] = {call["call_id"]: payload["input"][0]["output"]} + raw_state = json.dumps(persisted, separators=(",", ":"), sort_keys=True) + state_file.write_text(raw_state, encoding="utf-8") + legacy_bytes = state_file.stat().st_size + + with pytest.raises(RuntimeError, match="exceeds configured byte limit"): + _app(response_state_file=state_file, max_response_state_bytes=legacy_bytes) + + assert state_file.read_text(encoding="utf-8") == raw_state + + + +def test_foundry_brokered_rejects_caller_session_conflicts_with_trusted_gateway_header(monkeypatch): + monkeypatch.delenv("FOUNDRY_AGENT_SESSION_ID", raising=False) + app = _app() + + with TestClient(app) as client: + body_conflict = client.post( + "/responses", + headers={"x-agent-session-id": "trusted-session"}, + json={"input": "please read telemetry", "agent_session_id": "caller-session"}, + ) + query_conflict = client.post( + "/responses?session_id=caller-session", + headers={"x-agent-session-id": "trusted-session"}, + json={"input": "please read telemetry"}, + ) + + for response in (body_conflict, query_conflict): + assert response.status_code == 409 + assert response.json()["error"]["code"] == "response_session_mismatch" + + +def test_foundry_brokered_rejects_session_conflicts_with_hosted_environment(monkeypatch): + monkeypatch.setenv("FOUNDRY_AGENT_SESSION_ID", "hosted-session") + app = _app() + + with TestClient(app) as client: + gateway_conflict = client.post( + "/responses", + headers={"x-agent-session-id": "other-session"}, + json={"input": "please read telemetry"}, + ) + caller_conflict = client.post( + "/responses", + json={"input": "please read telemetry", "agent_session_id": "other-session"}, + ) + matching = client.post( + "/responses?session_id=hosted-session", + headers={"x-agent-session-id": "hosted-session"}, + json={"input": "please read telemetry", "agent_session_id": "hosted-session"}, + ) + + for response in (gateway_conflict, caller_conflict): + assert response.status_code == 409 + assert response.json()["error"]["code"] == "response_session_mismatch" + assert matching.status_code == 200, matching.text + + +def test_foundry_brokered_trusted_gateway_session_accepts_matching_local_compatibility(monkeypatch, tmp_path): + state_file = tmp_path / "foundry-trusted-session-state.json" + monkeypatch.delenv("FOUNDRY_AGENT_SESSION_ID", raising=False) + + with TestClient(_app(response_state_file=state_file)) as client: + response = client.post( + "/responses?session_id=gateway-session", + headers={ + "x-agent-session-id": "gateway-session", + "x-agentkit-session-id": "gateway-session", + }, + json={"input": "please read telemetry", "agent_session_id": "gateway-session"}, + ) + + assert response.status_code == 200, response.text + stored = json.loads(state_file.read_text(encoding="utf-8"))["states"][response.json()["id"]] + assert stored["sessionID"] == "gateway-session" + + +def test_foundry_brokered_preserves_local_session_fallback_precedence_without_trusted_identity(monkeypatch, tmp_path): + state_file = tmp_path / "foundry-local-session-state.json" + monkeypatch.delenv("FOUNDRY_AGENT_SESSION_ID", raising=False) + + with TestClient(_app(response_state_file=state_file)) as client: + response = client.post( + "/responses?agent_session_id=query-session", + headers={"x-agentkit-session-id": "legacy-header-session"}, + json={ + "input": "please read telemetry", + "agent_session_id": "body-session", + "session_id": "legacy-body-session", + }, + ) + + assert response.status_code == 200, response.text + stored = json.loads(state_file.read_text(encoding="utf-8"))["states"][response.json()["id"]] + assert stored["sessionID"] == "body-session" + + +def test_foundry_brokered_file_state_tracks_session_and_rejects_cross_session_continuation(monkeypatch, tmp_path): + state_file = tmp_path / "foundry-session-state.json" + monkeypatch.delenv("FOUNDRY_AGENT_SESSION_ID", raising=False) + + with TestClient(_app(response_state_file=state_file)) as client: + initial_response = client.post( + "/responses", + json={"input": "please read telemetry", "agent_session_id": "session-a"}, + ) + assert initial_response.status_code == 200, initial_response.text + initial = initial_response.json() + call = _call(initial) + + stored = json.loads(state_file.read_text(encoding="utf-8"))["states"][initial["id"]] + assert stored["sessionID"] == "session-a" + + payload = _continuation(initial["id"], call["call_id"], {"approved": True, "output": {"success": True}}) + payload[CONTINUATION_PROOF_BODY_FIELD] = CONTINUATION_PROOF + with TestClient(_app(response_state_file=state_file)) as client: + missing = client.post("/responses", json=payload) + + assert missing.status_code == 409 + assert missing.json()["error"]["code"] == "response_session_mismatch" + + payload["agent_session_id"] = "session-b" + with TestClient(_app(response_state_file=state_file)) as client: + mismatch = client.post("/responses", json=payload) + + assert mismatch.status_code == 409 + assert mismatch.json()["error"]["code"] == "response_session_mismatch" + + payload["agent_session_id"] = "session-a" + with TestClient(_app(response_state_file=state_file)) as client: + completed = client.post("/responses", json=payload) + + assert completed.status_code == 200, completed.text + + +def test_foundry_brokered_body_proof_is_not_persisted_or_echoed(tmp_path): + state_file = tmp_path / "foundry-proof-state.json" + app = _app(response_state_file=state_file) + + with TestClient(app) as client: + initial_response = client.post( + "/responses", + json={"input": "please read telemetry", CONTINUATION_PROOF_BODY_FIELD: CONTINUATION_PROOF}, + ) + assert initial_response.status_code == 200, initial_response.text + initial = initial_response.json() + call = _call(initial) + payload = _continuation(initial["id"], call["call_id"], {"approved": True, "output": {"success": True}}) + payload[CONTINUATION_PROOF_BODY_FIELD] = CONTINUATION_PROOF + final_response = client.post("/responses", json=payload) + + assert final_response.status_code == 200, final_response.text assert CONTINUATION_PROOF not in json.dumps(initial, sort_keys=True) assert CONTINUATION_PROOF not in json.dumps(final_response.json(), sort_keys=True) assert CONTINUATION_PROOF not in state_file.read_text(encoding="utf-8") @@ -1459,10 +1728,10 @@ def test_foundry_brokered_file_state_survives_restart_for_model_loop_continuatio assert second_fake.requests[0]["messages"][-1]["tool_call_id"] == call["call_id"] -def test_foundry_brokered_rejects_persisted_model_loop_continuation_when_loop_is_disabled(tmp_path): +def test_foundry_brokered_model_loop_continuation_requires_model_loop_after_restart(tmp_path): state_file = tmp_path / "foundry-model-state.json" spec = _spec(tool_name="check-network-telemetry") - fake = _FakeChatTransport( + first_fake = _FakeChatTransport( [ _chat_response( { @@ -1479,21 +1748,23 @@ def test_foundry_brokered_rejects_persisted_model_loop_continuation_when_loop_is ) ] ) - with TestClient(_model_loop_app(spec, fake, response_state_file=state_file)) as client: + + with TestClient(_model_loop_app(spec, first_fake, response_state_file=state_file)) as client: initial = client.post("/responses", json={"input": "check-network-telemetry"}).json() call = _call(initial) - persisted_before = state_file.read_bytes() + before = state_file.read_bytes() with TestClient(_app(spec, response_state_file=state_file, brokered_model_loop_enabled=False)) as client: response = client.post( "/responses", headers=CONTINUATION_AUTH, - json=_continuation(initial["id"], call["call_id"], {"approved": True, "output": {"ok": True}}), + json=_continuation(initial["id"], call["call_id"], {"approved": True, "output": {"status": "ok"}}), ) assert response.status_code == 503 assert response.json()["error"]["code"] == "brokered_model_loop_unavailable" - assert state_file.read_bytes() == persisted_before + assert state_file.read_bytes() == before + def test_foundry_brokered_rejects_expired_response_state(): @@ -1513,54 +1784,6 @@ def test_foundry_brokered_rejects_expired_response_state(): assert resp.json()["error"]["code"] == "response_state_expired" -def test_foundry_brokered_rejects_normal_followup_to_expired_pending_response(): - app = _app(state_ttl_seconds=0) - - with TestClient(app) as client: - initial = _start(client) - time.sleep(0.01) - response = client.post( - "/responses", - json={"previous_response_id": initial["id"], "input": "next question"}, - ) - - assert response.status_code == 410 - assert response.json()["error"]["code"] == "response_state_expired" - - -def test_foundry_brokered_allows_normal_followup_after_completed_state_expires(monkeypatch): - stores: list[Any] = [] - original_store = foundry_module._FoundryResponseStateStore - - def capture_store(*args: Any, **kwargs: Any) -> Any: - store = original_store(*args, **kwargs) - stores.append(store) - return store - - monkeypatch.setattr(foundry_module, "_FoundryResponseStateStore", capture_store) - app = _app() - - with TestClient(app) as client: - initial = _start(client) - call = _call(initial) - completed = client.post( - "/responses", - headers=CONTINUATION_AUTH, - json=_continuation(initial["id"], call["call_id"], {"approved": True, "output": {"ok": True}}), - ) - state = stores[0].get(initial["id"]) - state.expires_at = time.time() - 1 - stores[0].save(state) - followup = client.post( - "/responses", - json={"previous_response_id": initial["id"], "input": "next question"}, - ) - - assert completed.status_code == 200, completed.text - assert followup.status_code == 200, followup.text - assert _call(followup.json()) - - def test_foundry_brokered_refuses_to_synthesize_nonliteral_write_arguments(): spec = _spec(tool_name="dispatch-work-order", brokered_class="write") spec.brokered_tools[0].parameters = { @@ -1593,98 +1816,33 @@ def test_foundry_brokered_refuses_multi_value_enum_write_arguments(): assert resp.json()["error"]["code"] == "UnsupportedBrokeredSchema" -@pytest.mark.parametrize("brokered_class", ["write", "coordination"]) -def test_foundry_brokered_refuses_multi_value_root_enum_side_effecting_arguments(brokered_class: str): - spec = _spec(tool_name="dispatch-work-order", brokered_class=brokered_class) +def test_foundry_brokered_allows_single_value_enum_write_arguments(): + spec = _spec(tool_name="dispatch-work-order", brokered_class="write") spec.brokered_tools[0].parameters = { "type": "object", - "enum": [ - {"operation": "delete"}, - {"operation": "create"}, - ], + "properties": {"operation": {"type": "string", "enum": ["create"]}}, + "required": ["operation"], } app = _app(spec) with TestClient(app) as client: resp = client.post("/responses", json={"input": "dispatch-work-order"}) - assert resp.status_code == 400 - assert resp.json()["error"]["code"] == "UnsupportedBrokeredSchema" + assert resp.status_code == 200 + assert json.loads(_call(resp.json())["arguments"]) == {"operation": "create"} -def test_foundry_brokered_allows_single_value_root_enum_write_arguments(): +def test_foundry_brokered_allows_literal_write_arguments_for_conformance(): spec = _spec(tool_name="dispatch-work-order", brokered_class="write") spec.brokered_tools[0].parameters = { "type": "object", - "enum": [{"operation": "create"}], + "properties": {"incident": {"type": "string", "const": "INC-1"}}, + "required": ["incident"], } app = _app(spec) with TestClient(app) as client: - resp = client.post("/responses", json={"input": "dispatch-work-order"}) - - assert resp.status_code == 200 - assert json.loads(_call(resp.json())["arguments"]) == {"operation": "create"} - - -def test_foundry_brokered_allows_single_value_enum_write_arguments(): - spec = _spec(tool_name="dispatch-work-order", brokered_class="write") - spec.brokered_tools[0].parameters = { - "type": "object", - "properties": {"operation": {"type": "string", "enum": ["create"]}}, - "required": ["operation"], - } - app = _app(spec) - - with TestClient(app) as client: - resp = client.post("/responses", json={"input": "dispatch-work-order"}) - - assert resp.status_code == 200 - assert json.loads(_call(resp.json())["arguments"]) == {"operation": "create"} - - -@pytest.mark.parametrize("brokered_class", ["write", "coordination"]) -def test_foundry_brokered_refuses_nonliteral_optional_prompt_for_side_effecting_tools(brokered_class: str): - spec = _spec(tool_name="dispatch-work-order", brokered_class=brokered_class) - spec.brokered_tools[0].parameters = { - "type": "object", - "properties": {"prompt": {"type": "string"}}, - } - app = _app(spec) - - with TestClient(app) as client: - resp = client.post("/responses", json={"input": "dispatch-work-order with arbitrary payload"}) - - assert resp.status_code == 400 - assert resp.json()["error"]["code"] == "UnsupportedBrokeredSchema" - - -def test_foundry_brokered_allows_literal_optional_prompt_for_write_tools(): - spec = _spec(tool_name="dispatch-work-order", brokered_class="write") - spec.brokered_tools[0].parameters = { - "type": "object", - "properties": {"prompt": {"type": "string", "const": "fixed"}}, - } - app = _app(spec) - - with TestClient(app) as client: - resp = client.post("/responses", json={"input": "dispatch-work-order with arbitrary payload"}) - - assert resp.status_code == 200 - assert json.loads(_call(resp.json())["arguments"]) == {"prompt": "fixed"} - - -def test_foundry_brokered_allows_literal_write_arguments_for_conformance(): - spec = _spec(tool_name="dispatch-work-order", brokered_class="write") - spec.brokered_tools[0].parameters = { - "type": "object", - "properties": {"incident": {"type": "string", "const": "INC-1"}}, - "required": ["incident"], - } - app = _app(spec) - - with TestClient(app) as client: - body = client.post("/responses", json={"input": "dispatch-work-order"}) + body = client.post("/responses", json={"input": "dispatch-work-order"}) assert body.status_code == 200, body.text assert json.loads(_call(body.json())["arguments"]) == {"incident": "INC-1"} @@ -1731,26 +1889,6 @@ def test_foundry_brokered_rejects_multiple_tool_outputs_deterministically(): -def test_foundry_brokered_rejects_continuation_with_extra_non_output_item(): - app = _app() - - with TestClient(app) as client: - initial = _start(client) - call = _call(initial) - request = _continuation(initial["id"], call["call_id"], {"approved": True, "output": {"ok": True}}) - request["input"].append({"type": "message", "role": "user", "content": "ignored"}) - rejected = client.post("/responses", json=request, headers=CONTINUATION_AUTH) - accepted = client.post( - "/responses", - json=_continuation(initial["id"], call["call_id"], {"approved": True, "output": {"ok": True}}), - headers=CONTINUATION_AUTH, - ) - - assert rejected.status_code == 400 - assert rejected.json()["error"]["code"] == "multiple_tool_outputs_unsupported" - assert accepted.status_code == 200, accepted.text - - def test_foundry_brokered_disables_invocations_direct_runtime_bypass(): app = _app() @@ -1789,150 +1927,211 @@ def _model_loop_app(spec: AgentSpec, fake: _FakeChatTransport, **kwargs: Any): return _app(spec, brokered_model_loop_enabled=True, brokered_model_http_client=client, **kwargs) -def test_foundry_brokered_model_loop_sends_placeholder_api_key_when_auth_is_omitted(monkeypatch): - captured_headers: dict[str, str] = {} +def test_foundry_brokered_model_loop_bounds_streamed_upstream_body_before_json_materialization(): + class TrackingStream(httpx.AsyncByteStream): + def __init__(self) -> None: + self.chunks_read = 0 + self.tail_read = False + + async def __aiter__(self): + chunks = [ + b'{"choices":[{"message":{"role":"assistant","content":"', + b"x" * 300, + b'"}}],"usage":{}}', + ] + for index, chunk in enumerate(chunks): + self.chunks_read += 1 + if index == 2: + self.tail_read = True + yield chunk - class FakeClient: - def __init__(self, *, headers: dict[str, str], timeout: int) -> None: - assert timeout == 60 - captured_headers.update(headers) + async def aclose(self) -> None: + return None + + class Transport(httpx.AsyncBaseTransport): + def __init__(self) -> None: + self.stream = TrackingStream() + self.calls = 0 - async def post(self, url: str, *, json: dict[str, Any]) -> httpx.Response: - request = httpx.Request("POST", url, json=json) + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + self.calls += 1 + if self.calls == 1: + return httpx.Response(200, request=request, stream=self.stream) return httpx.Response( 200, request=request, - json=_chat_response({"role": "assistant", "content": "done"}), + json=_chat_response({"role": "assistant", "content": "Retry stayed available."}), ) - async def aclose(self) -> None: - return None + transport = Transport() + app = _app( + _spec(tool_name="check-network-telemetry"), + brokered_model_loop_enabled=True, + brokered_model_http_client=httpx.AsyncClient(transport=transport), + max_pending_responses=1, + max_response_state_bytes=256, + ) - monkeypatch.setattr(foundry_model_loop_module.httpx, "AsyncClient", FakeClient) - app = _app(_spec(tool_name="check-network-telemetry"), brokered_model_loop_enabled=True) + with TestClient(app, raise_server_exceptions=False) as client: + oversized = client.post("/responses", json={"input": "Say hello"}) + retry = client.post("/responses", json={"input": "Say hello again"}) - with TestClient(app) as client: - response = client.post("/responses", json={"input": "check-network-telemetry"}) + assert oversized.status_code == 502 + assert oversized.json()["error"] == { + "message": "model response is too large to retain safely", + "code": "ModelResponseTooLarge", + } + assert transport.stream.tail_read is False + assert retry.status_code == 200, retry.text + assert _message_text(retry.json()) == "Retry stayed available." - assert response.status_code == 200, response.text - assert captured_headers["Authorization"] == f"Bearer {NO_AUTH_API_KEY}" +def test_foundry_brokered_max_pending_reserves_capacity_before_initial_model_work(): + spec = _spec(tool_name="check-network-telemetry") -@pytest.mark.parametrize( - "usage", - [ - {"prompt_tokens": {"unexpected": 1}}, - {"completion_tokens": "not-a-number"}, - {"prompt_tokens": "12"}, - {"total_tokens": 1.5}, - ], -) -def test_foundry_brokered_model_loop_normalizes_malformed_usage(usage: dict[str, Any]): - fake = _FakeChatTransport( - [ - { - "choices": [{"message": {"role": "assistant", "content": "done"}}], - "usage": usage, - } - ] + class BlockingInitialTransport(httpx.AsyncBaseTransport): + def __init__(self) -> None: + self.requests: list[dict[str, Any]] = [] + self.first_started = threading.Event() + self.release_first = threading.Event() + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + self.requests.append(json.loads(request.content.decode("utf-8"))) + if len(self.requests) == 1: + self.first_started.set() + while not self.release_first.is_set(): + await asyncio.sleep(0.001) + return httpx.Response( + 200, + request=request, + json=_chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": f"model_call_{len(self.requests)}", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": "{}"}, + } + ], + } + ), + ) + + transport = BlockingInitialTransport() + app = _app( + spec, + brokered_model_loop_enabled=True, + brokered_model_http_client=httpx.AsyncClient(transport=transport), + max_pending_responses=1, ) - app = _model_loop_app(_spec(tool_name="check-network-telemetry"), fake) - with TestClient(app) as client: - response = client.post("/responses", json={"input": "check-network-telemetry"}) + with TestClient(app) as client, ThreadPoolExecutor(max_workers=1) as executor: + first_future = executor.submit(client.post, "/responses", json={"input": "check-network-telemetry"}) + try: + assert transport.first_started.wait(timeout=2) + excess = client.post("/responses", json={"input": "check-network-telemetry again"}) + finally: + transport.release_first.set() + first = first_future.result(timeout=2) - assert response.status_code == 502 - assert response.json()["error"]["code"] == "InvalidModelResponse" + assert first.status_code == 200, first.text + assert excess.status_code == 429 + assert excess.json()["error"]["code"] == "brokered_response_state_full" + assert len(transport.requests) == 1 -def test_foundry_brokered_model_loop_emits_model_requested_tool_and_resumes_to_final_answer(): +def test_foundry_brokered_initial_model_reservation_releases_for_final_and_error_paths(): spec = _spec(tool_name="check-network-telemetry") - spec.brokered_tools[0].parameters["required"] = ["site"] fake = _FakeChatTransport( [ + _chat_response({"role": "assistant", "content": "No tool needed."}), _chat_response( { "role": "assistant", "content": None, "tool_calls": [ { - "id": "model_generated_call_id", + "id": "unknown_call", "type": "function", - "function": {"name": "check-network-telemetry", "arguments": '{"site":"sfo"}'}, + "function": {"name": "unknown-tool", "arguments": "{}"}, } ], - }, - prompt_tokens=2, - completion_tokens=3, + } + ), + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "valid_call", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": "{}"}, + } + ], + } ), - _chat_response({"role": "assistant", "content": "Telemetry is healthy."}, prompt_tokens=5, completion_tokens=7), ] ) - app = _model_loop_app(spec, fake) + app = _model_loop_app(spec, fake, max_pending_responses=1) with TestClient(app) as client: - initial = client.post("/responses", json={"input": "Check SFO telemetry"}) - call = _call(initial.json()) - continuation = _continuation(initial.json()["id"], call["call_id"], {"approved": True, "output": {"status": "healthy"}}) - continuation[CONTINUATION_PROOF_BODY_FIELD] = CONTINUATION_PROOF - final = client.post( - "/responses", - json=continuation, - ) + final = client.post("/responses", json={"input": "say hello"}) + invalid = client.post("/responses", json={"input": "request an unknown tool"}) + pending = client.post("/responses", json={"input": "check-network-telemetry"}) - assert initial.status_code == 200, initial.text - assert call["name"] == "check-network-telemetry" - assert call["call_id"] == f"call_{initial.json()['id']}_1" - assert json.loads(call["arguments"]) == {"site": "sfo"} - assert initial.json()["usage"] == {"input_tokens": 2, "output_tokens": 3, "total_tokens": 5} assert final.status_code == 200, final.text - assert _message_text(final.json()) == "Telemetry is healthy." - assert final.json()["usage"] == {"input_tokens": 7, "output_tokens": 10, "total_tokens": 17} - assert fake.requests[0]["tools"][0]["function"]["name"] == "check-network-telemetry" - 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 CONTINUATION_PROOF not in json.dumps(fake.requests, sort_keys=True) + assert _message_text(final.json()) == "No tool needed." + assert invalid.status_code == 400 + assert invalid.json()["error"]["code"] == "unknown_brokered_tool" + assert pending.status_code == 200, pending.text + assert _call(pending.json())["name"] == "check-network-telemetry" + assert len(fake.requests) == 3 -def test_foundry_brokered_model_loop_preserves_total_only_usage_across_resume(): +def test_foundry_brokered_final_model_result_does_not_evict_completed_replay_state(): spec = _spec(tool_name="check-network-telemetry") - initial_model_response = _chat_response( - { - "role": "assistant", - "content": None, - "tool_calls": [ + fake = _FakeChatTransport( + [ + _chat_response( { - "id": "model_generated_call_id", - "type": "function", - "function": {"name": "check-network-telemetry", "arguments": "{}"}, + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "initial_tool_call", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": "{}"}, + } + ], } - ], - } + ), + _chat_response({"role": "assistant", "content": "Stored completion."}), + _chat_response({"role": "assistant", "content": "No pending state needed."}), + ] ) - initial_model_response["usage"] = {"total_tokens": 5} - final_model_response = _chat_response({"role": "assistant", "content": "done"}) - final_model_response["usage"] = {"total_tokens": 7} - app = _model_loop_app(spec, _FakeChatTransport([initial_model_response, final_model_response])) + app = _model_loop_app(spec, fake, max_pending_responses=1) with TestClient(app) as client: initial = client.post("/responses", json={"input": "check-network-telemetry"}) call = _call(initial.json()) - final = client.post( - "/responses", - headers=CONTINUATION_AUTH, - json=_continuation(initial.json()["id"], call["call_id"], {"approved": True, "output": {"ok": True}}), - ) + continuation = _continuation(initial.json()["id"], call["call_id"], {"approved": True, "output": {"ok": True}}) + completed = client.post("/responses", headers=CONTINUATION_AUTH, json=continuation) + final_only = client.post("/responses", json={"input": "say hello without a tool"}) + replay = client.post("/responses", headers=CONTINUATION_AUTH, json=continuation) - assert initial.json()["usage"] == {"input_tokens": 0, "output_tokens": 0, "total_tokens": 5} - assert final.status_code == 200, final.text - assert final.json()["usage"] == {"input_tokens": 0, "output_tokens": 0, "total_tokens": 12} + assert completed.status_code == 200, completed.text + assert final_only.status_code == 200, final_only.text + assert _message_text(final_only.json()) == "No pending state needed." + assert replay.status_code == 200, replay.text + assert replay.json() == completed.json() -def test_foundry_brokered_model_loop_rejects_noncanonical_denied_output_before_resume(tmp_path): - state_file = tmp_path / "responses-state.json" - spec = _spec(tool_name="dispatch-work-order", brokered_class="write") +def test_foundry_brokered_initial_model_reservation_releases_after_state_persist_failure(monkeypatch, tmp_path): + state_file = tmp_path / "model-loop-responses-state.json" + spec = _spec(tool_name="check-network-telemetry") fake = _FakeChatTransport( [ _chat_response( @@ -1941,180 +2140,212 @@ def test_foundry_brokered_model_loop_rejects_noncanonical_denied_output_before_r "content": None, "tool_calls": [ { - "id": "model_generated_call_id", + "id": "first_call", "type": "function", - "function": {"name": "dispatch-work-order", "arguments": "{}"}, + "function": {"name": "check-network-telemetry", "arguments": "{}"}, } ], } - ) + ), + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "retry_call", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": "{}"}, + } + ], + } + ), ] ) - app = _model_loop_app(spec, fake, response_state_file=state_file) - - with TestClient(app) as client: - initial = client.post("/responses", json={"input": "dispatch-work-order"}) - call = _call(initial.json()) - persisted_before = state_file.read_bytes() - denied = client.post( - "/responses", - headers=CONTINUATION_AUTH, - json=_continuation( - initial.json()["id"], - call["call_id"], - { - "approved": False, - "error": {"code": "approval_declined", "message": "denied"}, - "output": {"sensitiveDiagnostic": "must-not-reach-model"}, - }, - ), - ) - - assert denied.status_code == 400 - assert denied.json()["error"]["code"] == "invalid_function_call_output" - assert state_file.read_bytes() == persisted_before - assert len(fake.requests) == 1 - + app = _model_loop_app( + spec, + fake, + response_state_file=state_file, + max_pending_responses=1, + ) + original_replace = Path.replace + failed_once = False + + def fail_first_replace(path: Path, target: Path) -> Path: + nonlocal failed_once + if not failed_once and path.name == f".{state_file.name}.tmp": + failed_once = True + raise OSError("simulated model-loop state storage failure") + return original_replace(path, target) + + monkeypatch.setattr(Path, "replace", fail_first_replace) + with TestClient(app, raise_server_exceptions=False) as client: + failed = client.post("/responses", json={"input": "check-network-telemetry"}) + retried = client.post("/responses", json={"input": "check-network-telemetry"}) + + assert failed.status_code == 503 + assert failed.json()["error"]["code"] == "brokered_response_state_storage_error" + assert retried.status_code == 200, retried.text + assert _call(retried.json())["name"] == "check-network-telemetry" + assert len(fake.requests) == 2 -@pytest.mark.parametrize("payload", [{"approved": True}, {"approved": True, "output": None}]) -def test_foundry_brokered_rejects_approved_continuation_without_object_output(payload: dict[str, Any]): - app = _app() - with TestClient(app) as client: - initial = _start(client) - call = _call(initial) - response = client.post( - "/responses", - headers=CONTINUATION_AUTH, - json=_continuation(initial["id"], call["call_id"], payload), - ) +def test_foundry_brokered_cancelled_initial_model_work_releases_reserved_capacity(): + spec = _spec(tool_name="check-network-telemetry") - assert response.status_code == 400 - assert response.json()["error"]["code"] == "invalid_function_call_output" + class CancellingInitialTransport(httpx.AsyncBaseTransport): + def __init__(self) -> None: + self.calls = 0 + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + self.calls += 1 + if self.calls == 1: + raise asyncio.CancelledError + return httpx.Response( + 200, + request=request, + json=_chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "model_call_after_cancellation", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": "{}"}, + } + ], + } + ), + ) -@pytest.mark.parametrize("payload", [{"approved": False}, {"approved": False, "error": None}]) -def test_foundry_brokered_rejects_denied_continuation_without_error_object(payload: dict[str, Any]): - app = _app() + transport = CancellingInitialTransport() + app = _app( + spec, + brokered_model_loop_enabled=True, + brokered_model_http_client=httpx.AsyncClient(transport=transport), + max_pending_responses=1, + ) - with TestClient(app) as client: - initial = _start(client) - call = _call(initial) - response = client.post( - "/responses", - headers=CONTINUATION_AUTH, - json=_continuation(initial["id"], call["call_id"], payload), - ) + async def exercise() -> httpx.Response: + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), + base_url="http://testserver", + ) as client: + try: + await client.post("/responses", json={"input": "cancel this model call"}) + except asyncio.CancelledError: + pass + return await client.post("/responses", json={"input": "check-network-telemetry"}) - assert response.status_code == 400 - assert response.json()["error"]["code"] == "invalid_function_call_output" + response = asyncio.run(exercise()) + assert response.status_code == 200, response.text + assert _call(response.json())["name"] == "check-network-telemetry" + assert transport.calls == 2 -def test_foundry_brokered_rejects_deep_continuation_output_without_consuming_state(): - app = _app() - nested: dict[str, Any] = {} - for _ in range(200): - nested = {"nested": nested} - with TestClient(app) as client: - initial = _start(client) - call = _call(initial) - rejected = client.post( - "/responses", - headers=CONTINUATION_AUTH, - json=_continuation(initial["id"], call["call_id"], {"approved": True, "output": nested}), - ) - accepted = client.post( - "/responses", - headers=CONTINUATION_AUTH, - json=_continuation(initial["id"], call["call_id"], {"approved": True, "output": {"ok": True}}), - ) +def test_foundry_brokered_active_resume_survives_ttl_and_completed_state_retains_from_completion(): + spec = _spec(tool_name="check-network-telemetry") - assert rejected.status_code == 400 - assert rejected.json()["error"]["code"] == "invalid_function_call_output" - assert accepted.status_code == 200, accepted.text + class BlockingResumeTransport(httpx.AsyncBaseTransport): + def __init__(self) -> None: + self.requests: list[dict[str, Any]] = [] + self.resume_started = threading.Event() + self.release_resume = threading.Event() + self.initial_calls = 0 + 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: + self.initial_calls += 1 + return httpx.Response( + 200, + request=request, + json=_chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": f"model_call_{self.initial_calls}", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": "{}"}, + } + ], + } + ), + ) + self.resume_started.set() + while not self.release_resume.is_set(): + await asyncio.sleep(0.001) + return httpx.Response( + 200, + request=request, + json=_chat_response({"role": "assistant", "content": "Resume completed after the lease window."}), + ) -def test_foundry_brokered_bounds_request_body_before_model_call(): - fake = _FakeChatTransport([_chat_response({"role": "assistant", "content": "must not run"})]) - app = _model_loop_app( - _spec(tool_name="check-network-telemetry"), - fake, - max_request_body_bytes=64, + transport = BlockingResumeTransport() + model_client = httpx.AsyncClient(transport=transport) + app = _app( + spec, + brokered_model_loop_enabled=True, + brokered_model_http_client=model_client, + state_ttl_seconds=0.05, ) - with TestClient(app) as client: - response = client.post( - "/responses", - content=json.dumps({"input": "x" * 128}), - headers={"content-type": "application/json"}, - ) - - assert response.status_code == 413 - assert response.json()["error"]["code"] == "request_body_too_large" - assert fake.requests == [] - - -def test_foundry_brokered_model_message_size_matches_persistence_encoding(): - messages = [{"role": "user", "content": "😀" * 8}] + with TestClient(app) as client, ThreadPoolExecutor(max_workers=1) as executor: + initial = client.post("/responses", json={"input": "check-network-telemetry"}) + call = _call(initial.json()) + payload = _continuation(initial.json()["id"], call["call_id"], {"approved": True, "output": {"ok": True}}) + future = executor.submit(client.post, "/responses", headers=CONTINUATION_AUTH, json=payload) + try: + assert transport.resume_started.wait(timeout=2) + time.sleep(0.08) + another_pending = client.post("/responses", json={"input": "check-network-telemetry again"}) + active_duplicate = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + finally: + transport.release_resume.set() + completed = future.result(timeout=2) + immediate_duplicate = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + + assert another_pending.status_code == 200, another_pending.text + assert active_duplicate.status_code == 409 + assert active_duplicate.json()["error"]["code"] == "duplicate_continuation_in_progress" + assert completed.status_code == 200, completed.text + assert _message_text(completed.json()) == "Resume completed after the lease window." + assert immediate_duplicate.status_code == 200, immediate_duplicate.text + assert immediate_duplicate.json() == completed.json() - measured = foundry_module._persistence_json_bytes(messages) - persisted = json.dumps( - messages, - allow_nan=False, - separators=(",", ":"), - sort_keys=True, - ).encode("utf-8") - utf8_compact = json.dumps( - messages, - allow_nan=False, - ensure_ascii=False, - separators=(",", ":"), - sort_keys=True, - ).encode("utf-8") - assert measured == persisted - assert len(measured) > len(utf8_compact) +def test_foundry_brokered_abandoned_resuming_state_expires_and_releases_capacity(monkeypatch): + stores: list[Any] = [] + original_store = foundry_module._FoundryResponseStateStore + def capture_store(*args: Any, **kwargs: Any) -> Any: + store = original_store(*args, **kwargs) + stores.append(store) + return store -def test_foundry_brokered_bounds_persisted_model_messages_and_releases_reservation(): - tool_response = _chat_response( - { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_model", - "type": "function", - "function": {"name": "check-network-telemetry", "arguments": "{}"}, - } - ], - } - ) - fake = _FakeChatTransport([tool_response, deepcopy(tool_response)]) - app = _model_loop_app( - _spec(tool_name="check-network-telemetry"), - fake, - max_request_body_bytes=4096, - max_model_messages_bytes=512, - max_pending_responses=1, - ) + monkeypatch.setattr(foundry_module, "_FoundryResponseStateStore", capture_store) + app = _app(max_pending_responses=1) with TestClient(app) as client: - oversized = client.post( - "/responses", - json={"input": "check-network-telemetry " + ("x" * 1024)}, - ) - accepted = client.post("/responses", json={"input": "check-network-telemetry"}) + initial = _start(client) + abandoned = stores[0].get(initial["id"]) + abandoned.status = "resuming" + abandoned.expires_at = time.time() - 1 + stores[0].save(abandoned) + replacement = client.post("/responses", json={"input": "please read telemetry again"}) - assert oversized.status_code == 413 - 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 replacement.status_code == 200, replacement.text + assert _call(replacement.json()) -def test_foundry_brokered_model_loop_reserves_capacity_before_model_call(): +def test_foundry_brokered_model_loop_emits_model_requested_tool_and_resumes_to_final_answer(): + spec = _spec(tool_name="check-network-telemetry") + spec.brokered_tools[0].parameters["required"] = ["site"] fake = _FakeChatTransport( [ _chat_response( @@ -2123,82 +2354,131 @@ def test_foundry_brokered_model_loop_reserves_capacity_before_model_call(): "content": None, "tool_calls": [ { - "id": "call_model", + "id": "model_generated_call_id", "type": "function", - "function": {"name": "check-network-telemetry", "arguments": "{}"}, + "function": {"name": "check-network-telemetry", "arguments": '{"site":"sfo"}'}, } ], - } + }, + prompt_tokens=2, + completion_tokens=3, ), - _chat_response({"role": "assistant", "content": "must not be called"}), + _chat_response({"role": "assistant", "content": "Telemetry is healthy."}, prompt_tokens=5, completion_tokens=7), ] ) - app = _model_loop_app( - _spec(tool_name="check-network-telemetry"), - fake, - max_pending_responses=1, - ) + app = _model_loop_app(spec, fake) with TestClient(app) as client: - first = client.post("/responses", json={"input": "check-network-telemetry"}) - second = client.post("/responses", json={"input": "check-network-telemetry"}) + initial = client.post("/responses", json={"input": "Check SFO telemetry"}) + call = _call(initial.json()) + continuation = _continuation(initial.json()["id"], call["call_id"], {"approved": True, "output": {"status": "healthy"}}) + continuation[CONTINUATION_PROOF_BODY_FIELD] = CONTINUATION_PROOF + final = client.post( + "/responses", + json=continuation, + ) - assert first.status_code == 200, first.text - assert _call(first.json()) - assert second.status_code == 429 - assert second.json()["error"]["code"] == "brokered_response_state_full" - assert len(fake.requests) == 1 + assert initial.status_code == 200, initial.text + assert call["name"] == "check-network-telemetry" + assert call["call_id"] == f"call_{initial.json()['id']}_1" + assert json.loads(call["arguments"]) == {"site": "sfo"} + assert initial.json()["usage"] == {"input_tokens": 2, "output_tokens": 3, "total_tokens": 5} + assert final.status_code == 200, final.text + assert _message_text(final.json()) == "Telemetry is healthy." + assert final.json()["usage"] == {"input_tokens": 7, "output_tokens": 10, "total_tokens": 17} + assert fake.requests[0]["tools"][0]["function"]["name"] == "check-network-telemetry" + 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 CONTINUATION_PROOF not in json.dumps(fake.requests, sort_keys=True) -def test_foundry_brokered_model_loop_unused_reservation_preserves_completed_replay(): - fake = _FakeChatTransport( - [ - _chat_response( - { - "role": "assistant", - "content": None, - "tool_calls": [ +def test_foundry_brokered_model_loop_missing_api_key_is_not_ready_or_caller_error(monkeypatch): + spec_data = _spec(tool_name="check-network-telemetry").model_dump(by_alias=True) + spec_data["model"]["apiKeyEnv"] = "MISSING_FOUNDRY_MODEL_KEY" + spec = AgentSpec.model_validate(spec_data) + monkeypatch.delenv("MISSING_FOUNDRY_MODEL_KEY", raising=False) + fake = _FakeChatTransport([_chat_response({"role": "assistant", "content": "must not run"})]) + app = _model_loop_app(spec, fake) + + with TestClient(app) as client: + readiness = client.get("/readiness") + response = client.post("/responses", json={"input": "check-network-telemetry"}) + + assert readiness.status_code == 503 + assert readiness.json()["ready"] is False + assert readiness.json()["foundryResponses"]["modelAuth"] == "missing" + assert response.status_code == 503 + assert response.json()["error"] == { + "message": "model authentication is not configured", + "code": "ModelAuthMissing", + } + assert "MISSING_FOUNDRY_MODEL_KEY" not in readiness.text + response.text + assert fake.requests == [] + + +def test_foundry_brokered_model_loop_missing_key_precedes_capacity_errors(monkeypatch, tmp_path): + state_file = tmp_path / "responses-state.json" + spec_data = _spec(tool_name="check-network-telemetry").model_dump(by_alias=True) + spec_data["model"]["apiKeyEnv"] = "FOUNDRY_MODEL_KEY" + spec = AgentSpec.model_validate(spec_data) + monkeypatch.setenv("FOUNDRY_MODEL_KEY", "mock-token") + first_fake = _FakeChatTransport( + [ + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ { - "id": "call_model", + "id": "model_generated_call_id", "type": "function", "function": {"name": "check-network-telemetry", "arguments": "{}"}, } ], } - ), - _chat_response({"role": "assistant", "content": "First response completed."}), - _chat_response({"role": "assistant", "content": "No tool needed."}), + ) ] ) - app = _model_loop_app( - _spec(tool_name="check-network-telemetry"), - fake, - max_pending_responses=1, - ) - - with TestClient(app) as client: + with TestClient( + _model_loop_app( + spec, + first_fake, + response_state_file=state_file, + max_pending_responses=1, + ) + ) as client: initial = client.post("/responses", json={"input": "check-network-telemetry"}) - call = _call(initial.json()) - continuation = _continuation( - initial.json()["id"], - call["call_id"], - {"approved": True, "output": {"ok": True}}, + assert initial.status_code == 200, initial.text + + monkeypatch.delenv("FOUNDRY_MODEL_KEY") + second_fake = _FakeChatTransport([_chat_response({"role": "assistant", "content": "must not run"})]) + with TestClient( + _model_loop_app( + spec, + second_fake, + response_state_file=state_file, + max_pending_responses=1, ) - completed = client.post("/responses", headers=CONTINUATION_AUTH, json=continuation) - direct = client.post("/responses", json={"input": "answer without a tool"}) - replayed = client.post("/responses", headers=CONTINUATION_AUTH, json=continuation) + ) as client: + response = client.post("/responses", json={"input": "check-network-telemetry again"}) - assert completed.status_code == 200, completed.text - assert direct.status_code == 200, direct.text - assert replayed.status_code == 200, replayed.text - assert replayed.json() == completed.json() - assert len(fake.requests) == 3 + assert response.status_code == 503 + assert response.json()["error"] == { + "message": "model authentication is not configured", + "code": "ModelAuthMissing", + } + assert second_fake.requests == [] -def test_foundry_brokered_model_loop_rejects_oversized_output_before_resume_or_state_change(tmp_path): +def test_foundry_brokered_model_loop_missing_key_precedes_continuation_state_errors(monkeypatch, tmp_path): state_file = tmp_path / "responses-state.json" - spec = _spec(tool_name="check-network-telemetry") - fake = _FakeChatTransport( + spec_data = _spec(tool_name="check-network-telemetry").model_dump(by_alias=True) + spec_data["model"]["apiKeyEnv"] = "FOUNDRY_MODEL_KEY" + spec = AgentSpec.model_validate(spec_data) + monkeypatch.setenv("FOUNDRY_MODEL_KEY", "mock-token") + first_fake = _FakeChatTransport( [ _chat_response( { @@ -2206,7 +2486,7 @@ def test_foundry_brokered_model_loop_rejects_oversized_output_before_resume_or_s "content": None, "tool_calls": [ { - "id": "model_generated_call_id", + "id": "model-generated-call", "type": "function", "function": {"name": "check-network-telemetry", "arguments": "{}"}, } @@ -2215,36 +2495,282 @@ def test_foundry_brokered_model_loop_rejects_oversized_output_before_resume_or_s ) ] ) - app = _model_loop_app( - spec, - fake, - response_state_file=state_file, - max_brokered_output_bytes=128, - ) - - with TestClient(app) as client: + with TestClient(_model_loop_app(spec, first_fake, response_state_file=state_file)) as client: initial = client.post("/responses", json={"input": "check-network-telemetry"}) call = _call(initial.json()) - persisted_before = state_file.read_bytes() - oversized = client.post( + persisted_before = state_file.read_bytes() + + monkeypatch.delenv("FOUNDRY_MODEL_KEY") + second_fake = _FakeChatTransport([_chat_response({"role": "assistant", "content": "must not run"})]) + with TestClient(_model_loop_app(spec, second_fake, response_state_file=state_file)) as client: + response = client.post( "/responses", headers=CONTINUATION_AUTH, - json=_continuation( - initial.json()["id"], - call["call_id"], - {"approved": True, "output": {"blob": "x" * 256}}, - ), + json=_continuation(initial.json()["id"], call["call_id"], {"approved": True, "output": {"ok": True}}), ) - assert oversized.status_code == 413 - assert oversized.json()["error"]["code"] == "brokered_output_too_large" + assert response.status_code == 503 + assert response.json()["error"] == { + "message": "model authentication is not configured", + "code": "ModelAuthMissing", + } assert state_file.read_bytes() == persisted_before - assert len(fake.requests) == 1 + assert second_fake.requests == [] -@pytest.mark.parametrize("raw_state", ["{not valid json", "[]"]) -def test_foundry_brokered_invalid_file_state_fails_startup_without_overwriting(tmp_path, raw_state: str): +def test_foundry_brokered_model_loop_sanitizes_upstream_auth_failures(): + spec_data = _spec(tool_name="check-network-telemetry").model_dump(by_alias=True) + internal_url = "https://private-model.internal.example/v1" + spec_data["model"]["baseURL"] = internal_url + spec = AgentSpec.model_validate(spec_data) + + for upstream_status in (401, 403): + def reject(request: httpx.Request, *, status: int = upstream_status) -> httpx.Response: + return httpx.Response(status, request=request, json={"error": "credential rejected"}) + + model_client = httpx.AsyncClient(transport=httpx.MockTransport(reject)) + app = _app( + spec, + brokered_model_loop_enabled=True, + brokered_model_http_client=model_client, + ) + with TestClient(app) as client: + response = client.post("/responses", json={"input": "check-network-telemetry"}) + + assert response.status_code == 503 + assert response.json()["error"] == { + "message": "model service rejected configured credentials", + "code": "ModelAuthRejected", + } + assert internal_url not in response.text + assert "/chat/completions" not in response.text + + +def test_foundry_brokered_model_loop_sanitizes_transport_failure_urls(): + spec_data = _spec(tool_name="check-network-telemetry").model_dump(by_alias=True) + internal_url = "https://private-model.internal.example/v1" + spec_data["model"]["baseURL"] = internal_url + spec = AgentSpec.model_validate(spec_data) + + def fail(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError(f"cannot connect to {internal_url}", request=request) + + model_client = httpx.AsyncClient(transport=httpx.MockTransport(fail)) + app = _app( + spec, + brokered_model_loop_enabled=True, + brokered_model_http_client=model_client, + ) + + with TestClient(app) as client: + response = client.post("/responses", json={"input": "check-network-telemetry"}) + + assert response.status_code == 502 + assert response.json()["error"] == { + "message": "model service request failed", + "code": "ModelUpstreamError", + } + assert internal_url not in response.text + assert "/chat/completions" not in response.text + + +def test_foundry_brokered_model_loop_normalizes_non_object_json_response(): + non_object_json = b"[0]" + + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + request=request, + content=non_object_json, + headers={"content-type": "application/json"}, + ) + + app = _app( + _spec(tool_name="check-network-telemetry"), + brokered_model_loop_enabled=True, + brokered_model_http_client=httpx.AsyncClient(transport=httpx.MockTransport(respond)), + ) + + with TestClient(app, raise_server_exceptions=False) as client: + response = client.post("/responses", json={"input": "check-network-telemetry"}) + + assert response.status_code == 502 + assert response.json()["error"] == { + "message": "model response must be a JSON object", + "code": "InvalidModelResponse", + } + + +def test_foundry_model_loop_rejects_oversized_output_before_encoding(): + class ExplodingEncodeString(str): + def encode(self, *args, **kwargs): + raise AssertionError("oversized output must be rejected before encoding") + + loop = BrokeredChatModelLoop( + _spec(tool_name="check-network-telemetry"), + [], + max_output_bytes=128, + ) + + with pytest.raises(AgentRunError) as exc_info: + asyncio.run( + loop.resume( + [{"role": "assistant", "content": None}], + call_id="call-1", + output=ExplodingEncodeString("x" * 129), + ) + ) + + assert exc_info.value.status == 413 + assert exc_info.value.code == "brokered_output_too_large" + + +def test_foundry_model_loop_rejects_lone_surrogate_tool_output_before_encoding(): + loop = BrokeredChatModelLoop( + _spec(tool_name="check-network-telemetry"), + [], + ) + + with pytest.raises(AgentRunError) as exc_info: + asyncio.run( + loop.resume( + [{"role": "assistant", "content": None}], + call_id="call-1", + output='{"value":"\ud800"}', + ) + ) + + assert exc_info.value.status == 400 + assert exc_info.value.code == "InvalidToolOutput" + assert "valid Unicode" in str(exc_info.value) + + +def test_foundry_brokered_model_loop_rejects_lone_surrogate_in_decoded_tool_arguments(): + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + request=request, + json=_chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "model-generated-call", + "type": "function", + "function": { + "name": "check-network-telemetry", + "arguments": '{"value":"\\ud800"}', + }, + } + ], + } + ), + ) + + app = _app( + _spec(tool_name="check-network-telemetry"), + brokered_model_loop_enabled=True, + brokered_model_http_client=httpx.AsyncClient(transport=httpx.MockTransport(respond)), + ) + + with TestClient(app, raise_server_exceptions=False) as client: + response = client.post("/responses", json={"input": "check-network-telemetry"}) + + assert response.status_code == 400 + assert response.json()["error"]["code"] == "InvalidToolArguments" + + + +def test_foundry_brokered_model_loop_rejects_decoded_lone_surrogate(): + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + request=request, + content=( + b'{"choices":[{"message":{"role":"assistant","content":"\\ud800"}}],' + b'"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}' + ), + headers={"content-type": "application/json"}, + ) + + model_client = httpx.AsyncClient(transport=httpx.MockTransport(respond)) + app = _app( + _spec(tool_name="check-network-telemetry"), + brokered_model_loop_enabled=True, + brokered_model_http_client=model_client, + ) + + with TestClient(app, raise_server_exceptions=False) as client: + response = client.post("/responses", json={"input": "check-network-telemetry"}) + + assert response.status_code == 502 + assert response.json()["error"] == { + "message": "model service returned an invalid JSON response", + "code": "InvalidModelResponse", + } + + +def test_foundry_brokered_model_loop_sanitizes_upstream_auth_failure_on_resume(): + spec_data = _spec(tool_name="check-network-telemetry").model_dump(by_alias=True) + internal_url = "https://private-model.internal.example/v1" + spec_data["model"]["baseURL"] = internal_url + spec = AgentSpec.model_validate(spec_data) + calls = 0 + + def respond(request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + if calls == 1: + return httpx.Response( + 200, + request=request, + json=_chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "model_generated_call_id", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": "{}"}, + } + ], + } + ), + ) + if calls == 2: + return httpx.Response(401, request=request, json={"error": "credential rejected"}) + return httpx.Response( + 200, + request=request, + json=_chat_response({"role": "assistant", "content": "Retry after auth recovery."}), + ) + + model_client = httpx.AsyncClient(transport=httpx.MockTransport(respond)) + app = _app( + spec, + brokered_model_loop_enabled=True, + brokered_model_http_client=model_client, + ) + + with TestClient(app) as client: + initial = client.post("/responses", json={"input": "check-network-telemetry"}) + call = _call(initial.json()) + payload = _continuation(initial.json()["id"], call["call_id"], {"approved": True, "output": {"ok": True}}) + failed = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + 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 internal_url not in failed.text + assert retried.status_code == 200, retried.text + assert _message_text(retried.json()) == "Retry after auth recovery." + + +def test_foundry_brokered_malformed_file_state_fails_startup_without_overwriting(tmp_path): state_file = tmp_path / "responses-state.json" + raw_state = "{not valid json" state_file.write_text(raw_state, encoding="utf-8") with pytest.raises(RuntimeError, match="invalid Foundry response state file"): @@ -2253,23 +2779,19 @@ def test_foundry_brokered_invalid_file_state_fails_startup_without_overwriting(t assert state_file.read_text(encoding="utf-8") == raw_state -@pytest.mark.parametrize("expires_at", [float("nan"), float("inf"), float("-inf")]) -def test_foundry_brokered_rejects_nonfinite_persisted_expiry(tmp_path, expires_at: float): + +def test_foundry_brokered_state_file_larger_than_limit_fails_startup_without_overwriting(tmp_path): state_file = tmp_path / "responses-state.json" - with TestClient(_app(response_state_file=state_file)) as client: - _start(client) - persisted = json.loads(state_file.read_text(encoding="utf-8")) - state = next(iter(persisted["states"].values())) - state["expiresAt"] = expires_at - raw_state = json.dumps(persisted, separators=(",", ":")) + raw_state = json.dumps({"states": {}, "padding": "x" * 8192}) state_file.write_text(raw_state, encoding="utf-8") - with pytest.raises(RuntimeError, match="stored expiresAt must be finite"): - _app(response_state_file=state_file) + with pytest.raises(RuntimeError, match="exceeds configured byte limit"): + _app(response_state_file=state_file, max_response_state_bytes=4 * 1024) assert state_file.read_text(encoding="utf-8") == raw_state + def test_foundry_brokered_file_state_is_written_with_private_permissions(tmp_path): state_file = tmp_path / "responses-state.json" app = _app(response_state_file=state_file) @@ -2281,10 +2803,1626 @@ def test_foundry_brokered_file_state_is_written_with_private_permissions(tmp_pat assert state_file.stat().st_mode & 0o777 == 0o600 -def test_foundry_brokered_file_state_recovers_unfinalized_accepted_continuation_after_restart(tmp_path): - state_file = tmp_path / "responses-state.json" +def test_foundry_brokered_failed_continuation_state_persist_is_retryable_without_stranding_progress(monkeypatch, tmp_path): + state_file = tmp_path / "responses-state.json" + spec = _spec(tool_name="check-network-telemetry") + fake = _FakeChatTransport( + [ + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "model_generated_call_id", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": "{}"}, + } + ], + } + ), + _chat_response({"role": "assistant", "content": "Recovered after storage retry."}), + ] + ) + app = _model_loop_app(spec, fake, response_state_file=state_file) + original_replace = Path.replace + failed_once = False + + def fail_first_resuming_replace(path: Path, target: Path) -> Path: + nonlocal failed_once + if not failed_once and path.name == f".{state_file.name}.tmp": + persisted = json.loads(path.read_text(encoding="utf-8")) + stored_state = next(iter(persisted["states"].values())) + if stored_state["status"] == "resuming": + failed_once = True + raise OSError("simulated continuation state storage failure") + return original_replace(path, target) + + with TestClient(app, raise_server_exceptions=False) as client: + initial = client.post("/responses", json={"input": "check-network-telemetry"}) + call = _call(initial.json()) + payload = _continuation(initial.json()["id"], call["call_id"], {"approved": True, "output": {"ok": True}}) + monkeypatch.setattr(Path, "replace", fail_first_resuming_replace) + failed = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + retried = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + + assert failed.status_code == 503 + assert failed.json()["error"] == { + "message": "brokered response state storage unavailable", + "code": "brokered_response_state_storage_error", + } + assert retried.status_code == 200, retried.text + assert _message_text(retried.json()) == "Recovered after storage retry." + assert len(fake.requests) == 2 + + +def test_foundry_brokered_failed_final_state_persist_retries_cached_completion_without_second_resume(monkeypatch, tmp_path): + state_file = tmp_path / "responses-state.json" + spec = _spec(tool_name="check-network-telemetry") + fake = _FakeChatTransport( + [ + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "model_generated_call_id", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": "{}"}, + } + ], + } + ), + _chat_response({"role": "assistant", "content": "Persist this exact completion."}), + _chat_response({"role": "assistant", "content": "A duplicate resume incorrectly ran."}), + ] + ) + app = _model_loop_app(spec, fake, response_state_file=state_file) + original_replace = Path.replace + failed_once = False + + def fail_first_completed_replace(path: Path, target: Path) -> Path: + nonlocal failed_once + if not failed_once and path.name == f".{state_file.name}.tmp": + persisted = json.loads(path.read_text(encoding="utf-8")) + stored_state = next(iter(persisted["states"].values())) + if stored_state["status"] == "completed": + failed_once = True + raise OSError("simulated completed state storage failure") + return original_replace(path, target) + + with TestClient(app, raise_server_exceptions=False) as client: + initial = client.post("/responses", json={"input": "check-network-telemetry"}) + call = _call(initial.json()) + payload = _continuation(initial.json()["id"], call["call_id"], {"approved": True, "output": {"ok": True}}) + monkeypatch.setattr(Path, "replace", fail_first_completed_replace) + failed = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + retried = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + + assert failed.status_code == 503 + assert failed.json()["error"]["code"] == "brokered_response_state_storage_error" + assert retried.status_code == 200, retried.text + assert _message_text(retried.json()) == "Persist this exact completion." + assert len(fake.requests) == 2 + persisted_state = json.loads(state_file.read_text(encoding="utf-8"))["states"][initial.json()["id"]] + assert persisted_state["status"] == "completed" + assert persisted_state["finalPayload"] == retried.json() + + +def test_foundry_brokered_unrelated_full_map_persist_marks_cached_completion_durable(monkeypatch, tmp_path): + state_file = tmp_path / "responses-state.json" + spec = _spec(tool_name="check-network-telemetry") + + def tool_request(call_id: str) -> dict[str, Any]: + return _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": call_id, + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": "{}"}, + } + ], + } + ) + + fake = _FakeChatTransport( + [ + tool_request("model_generated_call_id"), + _chat_response({"role": "assistant", "content": "Persisted by an unrelated transaction."}), + tool_request("unrelated_call_1"), + tool_request("unrelated_call_2"), + ] + ) + app = _model_loop_app(spec, fake, response_state_file=state_file, max_pending_responses=2) + original_replace = Path.replace + first_completed_write_failed = False + fail_next_write = False + second_failure_triggered = False + + def fail_selected_replaces(path: Path, target: Path) -> Path: + nonlocal first_completed_write_failed, second_failure_triggered + if path.name == f".{state_file.name}.tmp": + persisted = json.loads(path.read_text(encoding="utf-8")) + if not first_completed_write_failed and any(state["status"] == "completed" for state in persisted["states"].values()): + first_completed_write_failed = True + raise OSError("simulated completed state storage failure") + if fail_next_write: + second_failure_triggered = True + raise OSError("simulated later storage failure") + return original_replace(path, target) + + with TestClient(app, raise_server_exceptions=False) as client: + initial = client.post("/responses", json={"input": "check-network-telemetry"}) + call = _call(initial.json()) + payload = _continuation(initial.json()["id"], call["call_id"], {"approved": True, "output": {"ok": True}}) + monkeypatch.setattr(Path, "replace", fail_selected_replaces) + failed = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + unrelated = client.post("/responses", json={"input": "check-network-telemetry unrelated"}) + durable_state = json.loads(state_file.read_text(encoding="utf-8"))["states"][initial.json()["id"]] + fail_next_write = True + retried = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + fail_next_write = False + replacement = client.post("/responses", json={"input": "check-network-telemetry replacement"}) + + assert failed.status_code == 503 + assert unrelated.status_code == 200, unrelated.text + assert durable_state["status"] == "completed" + assert durable_state["finalPayload"] is not None + assert retried.status_code == 200, retried.text + assert _message_text(retried.json()) == "Persisted by an unrelated transaction." + assert second_failure_triggered is False + assert replacement.status_code == 200, replacement.text + assert _call(replacement.json()) + assert len(fake.requests) == 4 + + +def test_foundry_brokered_recovered_storage_persists_and_evicts_cached_completion_at_capacity(monkeypatch, tmp_path): + state_file = tmp_path / "responses-state.json" + spec = _spec(tool_name="check-network-telemetry") + fake = _FakeChatTransport( + [ + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "model_generated_call_id", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": "{}"}, + } + ], + } + ), + _chat_response({"role": "assistant", "content": "Completion cached while storage failed."}), + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "replacement_model_call", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": "{}"}, + } + ], + } + ), + ] + ) + app = _model_loop_app(spec, fake, response_state_file=state_file, max_pending_responses=1) + original_replace = Path.replace + failed_once = False + + def fail_first_completed_replace(path: Path, target: Path) -> Path: + nonlocal failed_once + if not failed_once and path.name == f".{state_file.name}.tmp": + persisted = json.loads(path.read_text(encoding="utf-8")) + if any(state["status"] == "completed" for state in persisted["states"].values()): + failed_once = True + raise OSError("simulated completed state storage failure") + return original_replace(path, target) + + with TestClient(app, raise_server_exceptions=False) as client: + initial = client.post("/responses", json={"input": "check-network-telemetry"}) + call = _call(initial.json()) + payload = _continuation(initial.json()["id"], call["call_id"], {"approved": True, "output": {"ok": True}}) + monkeypatch.setattr(Path, "replace", fail_first_completed_replace) + failed = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + replacement = client.post("/responses", json={"input": "check-network-telemetry replacement"}) + + assert failed.status_code == 503 + assert replacement.status_code == 200, replacement.text + replacement_body = replacement.json() + assert _call(replacement_body) + persisted_states = json.loads(state_file.read_text(encoding="utf-8"))["states"] + assert initial.json()["id"] not in persisted_states + assert replacement_body["id"] in persisted_states + assert len(fake.requests) == 3 + + +def test_foundry_brokered_file_state_recovers_unfinalized_accepted_continuation_after_restart(tmp_path): + state_file = tmp_path / "responses-state.json" + spec = _spec(tool_name="check-network-telemetry") + fake_first = _FakeChatTransport( + [ + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "model_generated_call_id", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": '{}'}, + } + ], + } + ) + ] + ) + first_app = _app(spec, brokered_model_loop_enabled=True, brokered_model_http_client=httpx.AsyncClient(transport=httpx.MockTransport(fake_first.handler)), response_state_file=state_file) + + with TestClient(first_app) as client: + initial = client.post("/responses", json={"input": "check-network-telemetry"}) + call = _call(initial.json()) + + payload = _continuation(initial.json()["id"], call["call_id"], {"approved": True, "output": {"ok": True}}) + state_data = json.loads(state_file.read_text(encoding="utf-8")) + state = state_data["states"][initial.json()["id"]] + state["acceptedOutputs"] = {call["call_id"]: payload["input"][0]["output"]} + state["status"] = "resuming" + state["finalPayload"] = None + state_file.write_text(json.dumps(state_data, separators=(",", ":"), sort_keys=True), encoding="utf-8") + + fake_second = _FakeChatTransport([_chat_response({"role": "assistant", "content": "Recovered."})]) + second_app = _app(spec, brokered_model_loop_enabled=True, brokered_model_http_client=httpx.AsyncClient(transport=httpx.MockTransport(fake_second.handler)), response_state_file=state_file) + with TestClient(second_app) as client: + retried = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + + assert retried.status_code == 200, retried.text + assert _message_text(retried.json()) == "Recovered." + + +def test_foundry_brokered_model_loop_accepts_integer_arguments_encoded_as_integral_float(): + spec = _spec(tool_name="retry-tool") + spec.brokered_tools[0].parameters = { + "type": "object", + "properties": {"retries": {"type": "integer"}}, + "required": ["retries"], + } + fake = _FakeChatTransport( + [ + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_model", + "type": "function", + "function": {"name": "retry-tool", "arguments": '{"retries":1.0}'}, + } + ], + } + ) + ] + ) + app = _model_loop_app(spec, fake) + + with TestClient(app) as client: + response = client.post("/responses", json={"input": "call retry-tool"}) + + assert response.status_code == 200, response.text + assert json.loads(_call(response.json())["arguments"]) == {"retries": 1.0} + + +def test_foundry_brokered_model_loop_unexpected_resume_failure_can_be_retried(): + spec = _spec(tool_name="check-network-telemetry") + + class FlakyResumeTransport: + def __init__(self) -> None: + self.requests: list[dict[str, Any]] = [] + self.resume_attempts = 0 + + def handler(self, request: httpx.Request) -> httpx.Response: + payload = json.loads(request.content.decode("utf-8")) + self.requests.append(payload) + if len(self.requests) == 1: + return httpx.Response( + 200, + json=_chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "model_generated_call_id", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": '{}'}, + } + ], + } + ), + ) + self.resume_attempts += 1 + if self.resume_attempts == 1: + raise RuntimeError("transient resume failure") + return httpx.Response(200, json=_chat_response({"role": "assistant", "content": "Recovered after retry."})) + + fake = FlakyResumeTransport() + app = _model_loop_app(spec, fake) + + with TestClient(app) as client: + initial = client.post("/responses", json={"input": "check-network-telemetry"}) + call = _call(initial.json()) + payload = _continuation(initial.json()["id"], call["call_id"], {"approved": True, "output": {"ok": True}}) + failed = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + 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 retried.status_code == 200, retried.text + assert _message_text(retried.json()) == "Recovered after retry." + + +def test_foundry_brokered_model_loop_failed_resume_can_be_retried(): + spec = _spec(tool_name="check-network-telemetry") + fake = _FakeChatTransport( + [ + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "model_generated_call_id", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": '{}'}, + } + ], + } + ), + {"choices": []}, + _chat_response({"role": "assistant", "content": "Retry worked."}), + ] + ) + app = _model_loop_app(spec, fake) + + with TestClient(app) as client: + initial = client.post("/responses", json={"input": "check-network-telemetry"}) + call = _call(initial.json()) + payload = _continuation(initial.json()["id"], call["call_id"], {"approved": True, "output": {"ok": True}}) + failed = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + retried = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + + assert failed.status_code == 502 + assert retried.status_code == 200, retried.text + assert _message_text(retried.json()) == "Retry worked." + + +def test_foundry_brokered_model_loop_bounds_resumed_final_text_without_installing_completion(tmp_path): + state_file = tmp_path / "responses-state.json" + oversized_text = "é" * 80 + assert len(oversized_text) < 128 < len(oversized_text.encode("utf-8")) + spec = _spec(tool_name="check-network-telemetry") + fake = _FakeChatTransport( + [ + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "model_generated_call_id", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": "{}"}, + } + ], + } + ), + _chat_response({"role": "assistant", "content": oversized_text}), + _chat_response({"role": "assistant", "content": "Retry stayed bounded."}), + ] + ) + app = _model_loop_app( + spec, + fake, + response_state_file=state_file, + max_brokered_output_bytes=128, + max_response_state_bytes=4 * 1024, + ) + + with TestClient(app) as client: + initial = client.post("/responses", json={"input": "check-network-telemetry"}) + call = _call(initial.json()) + payload = _continuation(initial.json()["id"], call["call_id"], {"approved": True, "output": {"ok": True}}) + failed = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + state_after_failure = json.loads(state_file.read_text(encoding="utf-8"))["states"][initial.json()["id"]] + retried = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + + assert failed.status_code == 502 + assert failed.json()["error"] == { + "message": "model response is too large to retain safely", + "code": "ModelResponseTooLarge", + } + assert state_after_failure["status"] == "pending" + assert state_after_failure["finalPayload"] is None + assert state_after_failure["acceptedOutputDigests"] == {} + assert retried.status_code == 200, retried.text + assert _message_text(retried.json()) == "Retry stayed bounded." + assert len(fake.requests) == 3 + + +def test_foundry_brokered_model_loop_oversized_completion_can_be_retried(): + spec = _spec(tool_name="check-network-telemetry") + fake = _FakeChatTransport( + [ + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "model_generated_call_id", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": "{}"}, + } + ], + } + ), + _chat_response({"role": "assistant", "content": "x" * 5_000}), + _chat_response({"role": "assistant", "content": "Retry stayed bounded."}), + ] + ) + app = _model_loop_app(spec, fake, max_response_state_bytes=4 * 1024) + + with TestClient(app) as client: + initial = client.post("/responses", json={"input": "check-network-telemetry"}) + call = _call(initial.json()) + payload = _continuation(initial.json()["id"], call["call_id"], {"approved": True, "output": {"ok": True}}) + failed = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + retried = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + + assert failed.status_code == 502 + assert failed.json()["error"] == { + "message": "model response is too large to retain safely", + "code": "ModelResponseTooLarge", + } + assert retried.status_code == 200, retried.text + assert _message_text(retried.json()) == "Retry stayed bounded." + assert len(fake.requests) == 3 + + +def test_foundry_brokered_model_loop_discards_resume_transcript_before_final_state_sizing(tmp_path): + state_file = tmp_path / "responses-state.json" + spec = _spec(tool_name="check-network-telemetry") + first_fake = _FakeChatTransport( + [ + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "model-generated-call", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": "{}"}, + } + ], + } + ) + ] + ) + with TestClient(_model_loop_app(spec, first_fake, response_state_file=state_file)) as client: + initial = client.post("/responses", json={"input": "check-network-telemetry " + "x" * 2_000}) + call = _call(initial.json()) + pending_state_bytes = state_file.stat().st_size + + second_fake = _FakeChatTransport([_chat_response({"role": "assistant", "content": "Small final answer."})]) + with TestClient( + _model_loop_app( + spec, + second_fake, + response_state_file=state_file, + max_response_state_bytes=pending_state_bytes + 256, + ) + ) as client: + completed = client.post( + "/responses", + headers=CONTINUATION_AUTH, + json=_continuation(initial.json()["id"], call["call_id"], {"approved": True, "output": {"ok": True}}), + ) + + assert completed.status_code == 200, completed.text + assert _message_text(completed.json()) == "Small final answer." + persisted = json.loads(state_file.read_text(encoding="utf-8"))["states"][initial.json()["id"]] + assert persisted["modelMessages"] is None + assert persisted["initialUsage"] == {} + + +def test_foundry_brokered_model_loop_aggregate_state_pressure_is_terminal_for_duplicate(tmp_path): + state_file = tmp_path / "responses-state.json" + spec = _spec(tool_name="check-network-telemetry") + + def tool_request(call_id: str) -> dict[str, Any]: + return _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": call_id, + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": "{}"}, + } + ], + } + ) + + fake = _FakeChatTransport( + [ + tool_request("first-model-call"), + tool_request("second-model-call"), + _chat_response({"role": "assistant", "content": "Computed once under capacity pressure."}), + _chat_response({"role": "assistant", "content": "must not be called for duplicate"}), + ] + ) + app = _model_loop_app( + spec, + fake, + response_state_file=state_file, + max_pending_responses=3, + max_response_state_bytes=2_500, + ) + + with TestClient(app) as client: + first = client.post("/responses", json={"input": "check-network-telemetry first"}) + first_call = _call(first.json()) + second = client.post("/responses", json={"input": "check-network-telemetry second"}) + _call(second.json()) + payload = _continuation(first.json()["id"], first_call["call_id"], {"approved": True, "output": {"ok": True}}) + failed = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + duplicate = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + + restart_fake = _FakeChatTransport([_chat_response({"role": "assistant", "content": "must not run after restart"})]) + with TestClient( + _model_loop_app( + spec, + restart_fake, + response_state_file=state_file, + max_pending_responses=3, + max_response_state_bytes=2_500, + ) + ) as client: + restarted_duplicate = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + + assert first.status_code == 200, first.text + assert second.status_code == 200, second.text + assert failed.status_code == 429 + assert failed.json()["error"]["code"] == "brokered_response_state_full" + assert duplicate.status_code == 429 + assert duplicate.json() == failed.json() + assert len(fake.requests) == 3 + assert restarted_duplicate.status_code == 429 + assert restarted_duplicate.json() == failed.json() + assert restart_fake.requests == [] + + +def test_foundry_brokered_bounds_initial_model_state_before_copying(monkeypatch): + spec = _spec(tool_name="check-network-telemetry") + fake = _FakeChatTransport( + [ + _chat_response( + { + "role": "assistant", + "content": "x" * 5_000, + "tool_calls": [ + { + "id": "model_generated_call_id", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": "{}"}, + } + ], + } + ) + ] + ) + app = _model_loop_app(spec, fake, max_response_state_bytes=4 * 1024) + original_deepcopy = foundry_module.deepcopy + + def reject_oversized_state_copy(value: Any, memo: dict[int, Any] | None = None) -> Any: + if type(value).__name__ == "_HostedResponseState" and value.model_messages: + if any(len(message.get("content") or "") > 4 * 1024 for message in value.model_messages): + raise AssertionError("oversized state must be rejected before deepcopy") + return original_deepcopy(value, memo) if memo is not None else original_deepcopy(value) + + monkeypatch.setattr(foundry_module, "deepcopy", reject_oversized_state_copy) + with TestClient(app, raise_server_exceptions=False) as client: + response = client.post("/responses", json={"input": "check-network-telemetry"}) + + assert response.status_code == 502 + assert response.json()["error"] == { + "message": "model response is too large to retain safely", + "code": "ModelResponseTooLarge", + } + + +def test_foundry_brokered_model_loop_bounds_immediate_final_text_in_utf8_bytes_without_state_install(): + oversized_text = "é" * 80 + assert len(oversized_text) < 128 < len(oversized_text.encode("utf-8")) + fake = _FakeChatTransport( + [ + _chat_response({"role": "assistant", "content": oversized_text}), + _chat_response({"role": "assistant", "content": "Retry stayed available."}), + ] + ) + app = _model_loop_app( + _spec(tool_name="check-network-telemetry"), + fake, + max_pending_responses=1, + max_brokered_output_bytes=128, + max_response_state_bytes=4 * 1024, + ) + + with TestClient(app) as client: + oversized = client.post("/responses", json={"input": "Say hello"}) + retry = client.post("/responses", json={"input": "Say hello again"}) + + assert oversized.status_code == 502 + assert oversized.json()["error"] == { + "message": "model response is too large to retain safely", + "code": "ModelResponseTooLarge", + } + assert retry.status_code == 200, retry.text + assert _message_text(retry.json()) == "Retry stayed available." + assert len(fake.requests) == 2 + + +def test_foundry_brokered_model_loop_can_return_final_message_without_tool_call(): + fake = _FakeChatTransport([_chat_response({"role": "assistant", "content": "No tool needed."})]) + app = _model_loop_app(_spec(tool_name="check-network-telemetry"), fake) + + with TestClient(app) as client: + response = client.post("/responses", json={"input": "Say hello"}) + + assert response.status_code == 200, response.text + assert _message_text(response.json()) == "No tool needed." + assert fake.requests[0]["tool_choice"] == "auto" + + +def test_foundry_brokered_model_loop_rejects_unsupported_pattern_deterministically(): + spec = _spec(tool_name="check-network-telemetry") + spec.brokered_tools[0].parameters = { + "type": "object", + "properties": {"site": {"type": "string", "pattern": "\\p{L}+"}}, + "required": ["site"], + } + fake = _FakeChatTransport( + [ + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_model", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": '{"site":"sfo"}'}, + } + ], + } + ) + ] + ) + app = _model_loop_app(spec, fake) + + with TestClient(app) as client: + response = client.post("/responses", json={"input": "call check-network-telemetry"}) + + assert response.status_code == 400 + assert response.json()["error"]["code"] == "UnsupportedBrokeredSchema" + + +def test_foundry_brokered_model_loop_validates_additional_properties_schema(): + spec = _spec(tool_name="check-network-telemetry") + spec.brokered_tools[0].parameters = { + "type": "object", + "properties": {"site": {"type": "string"}}, + "additionalProperties": {"type": "string"}, + } + fake = _FakeChatTransport( + [ + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_model", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": '{"site":"sfo","count":1}'}, + } + ], + } + ) + ] + ) + app = _model_loop_app(spec, fake) + + with TestClient(app) as client: + response = client.post("/responses", json={"input": "call check-network-telemetry"}) + + assert response.status_code == 400 + assert response.json()["error"]["code"] == "InvalidToolArguments" + + +def test_foundry_brokered_model_loop_rejects_unknown_model_tool_request(): + fake = _FakeChatTransport( + [ + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_model", + "type": "function", + "function": {"name": "unknown", "arguments": "{}"}, + } + ], + } + ) + ] + ) + app = _model_loop_app(_spec(tool_name="check-network-telemetry"), fake) + + with TestClient(app) as client: + response = client.post("/responses", json={"input": "call unknown"}) + + assert response.status_code == 400 + assert response.json()["error"]["code"] == "unknown_brokered_tool" + + +def test_foundry_brokered_model_loop_validates_schema_valued_additional_properties(): + spec = _spec(tool_name="flex-tool") + spec.brokered_tools[0].parameters = { + "type": "object", + "additionalProperties": {"type": "integer"}, + } + fake = _FakeChatTransport( + [ + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_model", + "type": "function", + "function": {"name": "flex-tool", "arguments": '{"safe":"not-int"}'}, + } + ], + } + ) + ] + ) + app = _model_loop_app(spec, fake) + + with TestClient(app) as client: + response = client.post("/responses", json={"input": "call flex-tool"}) + + assert response.status_code == 400 + assert response.json()["error"]["code"] == "InvalidToolArguments" + + +def test_foundry_brokered_model_loop_rejects_nonfinite_model_arguments(): + spec = _spec(tool_name="check-network-telemetry") + fake = _FakeChatTransport( + [ + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_model", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": '{"value":NaN}'}, + } + ], + } + ) + ] + ) + app = _model_loop_app(spec, fake) + + with TestClient(app) as client: + response = client.post("/responses", json={"input": "call check-network-telemetry"}) + + assert response.status_code == 400 + assert response.json()["error"]["code"] == "InvalidToolArguments" + + +def test_foundry_brokered_model_loop_uses_type_strict_const_and_enum_matching(): + for property_schema, arguments in [ + ({"const": 1}, '{"value":true}'), + ({"enum": [0]}, '{"value":false}'), + ]: + spec = _spec(tool_name="check-network-telemetry") + spec.brokered_tools[0].parameters = { + "type": "object", + "properties": {"value": property_schema}, + "required": ["value"], + } + fake = _FakeChatTransport( + [ + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_model", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": arguments}, + } + ], + } + ) + ] + ) + app = _model_loop_app(spec, fake) + + with TestClient(app) as client: + response = client.post("/responses", json={"input": "call check-network-telemetry"}) + + assert response.status_code == 400 + assert response.json()["error"]["code"] == "InvalidToolArguments" + + +def test_foundry_brokered_model_loop_rejects_arguments_that_violate_declared_schema(): + spec = _spec(tool_name="check-network-telemetry") + spec.brokered_tools[0].parameters = { + "type": "object", + "properties": {"site": {"type": "string"}}, + "required": ["site"], + "additionalProperties": False, + } + fake = _FakeChatTransport( + [ + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_model", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": '{"site":123,"extra":"nope"}'}, + } + ], + } + ) + ] + ) + app = _model_loop_app(spec, fake) + + with TestClient(app) as client: + response = client.post("/responses", json={"input": "call check-network-telemetry"}) + + assert response.status_code == 400 + assert response.json()["error"]["code"] == "InvalidToolArguments" + + +def test_foundry_brokered_model_loop_validates_large_integer_bounds_exactly(): + spec = _spec(tool_name="check-network-telemetry") + spec.brokered_tools[0].parameters = { + "type": "object", + "properties": {"count": {"type": "integer", "maximum": 9007199254740992}}, + "required": ["count"], + } + fake = _FakeChatTransport( + [ + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_model", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": '{"count":9007199254740993}'}, + } + ], + } + ) + ] + ) + app = _model_loop_app(spec, fake) + + with TestClient(app) as client: + response = client.post("/responses", json={"input": "call check-network-telemetry"}) + + assert response.status_code == 400 + assert response.json()["error"]["code"] == "InvalidToolArguments" + + +def test_foundry_brokered_model_loop_rejects_float_arguments_that_cannot_round_trip_exactly(): + spec = _spec(tool_name="check-network-telemetry") + spec.brokered_tools[0].parameters = { + "type": "object", + "properties": {"count": {"type": "integer", "maximum": 9007199254740992}}, + "required": ["count"], + } + fake = _FakeChatTransport( + [ + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_model", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": '{"count":9007199254740993.0}'}, + } + ], + } + ) + ] + ) + app = _model_loop_app(spec, fake) + + with TestClient(app) as client: + response = client.post("/responses", json={"input": "call check-network-telemetry"}) + + assert response.status_code == 400 + assert response.json()["error"]["code"] == "InvalidToolArguments" + + +def test_foundry_brokered_model_loop_rejects_unsafe_model_generated_arguments(): + fake = _FakeChatTransport( + [ + _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_model", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": '{"site":"sfo","tokenValue":"ghp_not_real"}'}, + } + ], + } + ) + ] + ) + app = _model_loop_app(_spec(tool_name="check-network-telemetry"), fake) + + with TestClient(app) as client: + response = client.post("/responses", json={"input": "call check-network-telemetry"}) + + assert response.status_code == 400 + assert response.json()["error"]["code"] == "UnsafeBrokeredArguments" + + +def test_foundry_brokered_rejects_tool_choice_none_instead_of_ignoring_it(): + app = _app() + + with TestClient(app) as client: + resp = client.post("/responses", json={"input": "hi", "tool_choice": "none"}) + + assert resp.status_code == 400 + assert resp.json()["error"]["code"] == "tool_choice_unsupported" + + +def test_foundry_brokered_rejects_request_supplied_tools_even_when_static_tools_exist(): + app = _app() + + with TestClient(app) as client: + resp = client.post("/responses", json={"input": "hi", "tools": [{"type": "function"}]}) + + assert resp.status_code == 400 + assert resp.json()["error"]["code"] == "tools_unsupported" + + +def _fixture(name: str) -> dict[str, Any]: + path = __import__("pathlib").Path(__file__).parent / "fixtures" / "foundry_brokered" / name + return json.loads(path.read_text(encoding="utf-8")) + + +def _normalize_initial_response(body: dict[str, Any]) -> dict[str, Any]: + normalized = deepcopy(body) + actual_response_id = normalized["id"] + normalized["id"] = "caresp_test" + normalized["created_at"] = 0 + item = normalized["output"][0] + item["id"] = "fc_test" + item["call_id"] = item["call_id"].replace(actual_response_id, "caresp_test") + item["response_id"] = "caresp_test" + return normalized + + +def _normalize_final_response(body: dict[str, Any], *, previous_response_id: str) -> dict[str, Any]: + normalized = deepcopy(body) + actual_response_id = normalized["id"] + normalized["id"] = "caresp_final" + normalized["created_at"] = 0 + normalized["previous_response_id"] = "caresp_test" + item = normalized["output"][0] + item["id"] = "msg_final" + item["response_id"] = "caresp_final" + assert previous_response_id + assert actual_response_id + return normalized + + +def _materialize_continuation(fixture: dict[str, Any], *, response_id: str, call_id: str) -> dict[str, Any]: + materialized = deepcopy(fixture) + materialized["previous_response_id"] = response_id + materialized["input"][0]["call_id"] = call_id + return materialized + + +def test_foundry_brokered_golden_fixtures_pin_function_call_loop_and_errors(): + app = _app() + + with TestClient(app) as client: + initial = client.post("/responses", json=_fixture("initial_request.json")) + assert initial.status_code == 200, initial.text + initial_body = initial.json() + call = _call(initial_body) + assert _normalize_initial_response(initial_body) == _fixture("function_call_response.json") + + continuation = _materialize_continuation( + _fixture("continuation_request.json"), + response_id=initial_body["id"], + call_id=call["call_id"], + ) + final = client.post("/responses", json=continuation, headers=CONTINUATION_AUTH) + assert final.status_code == 200, final.text + assert _normalize_final_response(final.json(), previous_response_id=initial_body["id"]) == _fixture("final_message_response.json") + + unknown_prev = client.post( + "/responses", + headers=CONTINUATION_AUTH, + json=_materialize_continuation( + _fixture("continuation_request.json"), + response_id="caresp_unknown", + call_id=call["call_id"], + ), + ) + assert unknown_prev.status_code == 404 + assert unknown_prev.json() == _fixture("unknown_previous_response_id_error.json") + + app_for_errors = _app() + with TestClient(app_for_errors) as client: + initial_body = _start(client) + call = _call(initial_body) + unknown_call = client.post( + "/responses", + headers=CONTINUATION_AUTH, + json=_materialize_continuation( + _fixture("continuation_request.json"), + response_id=initial_body["id"], + call_id="call_unknown", + ), + ) + assert unknown_call.status_code == 400 + assert unknown_call.json() == _fixture("unknown_call_id_error.json") + + multiple = _materialize_continuation( + _fixture("continuation_request.json"), + response_id=initial_body["id"], + call_id=call["call_id"], + ) + multiple["input"].append(dict(multiple["input"][0])) + multiple_resp = client.post("/responses", json=multiple, headers=CONTINUATION_AUTH) + assert multiple_resp.status_code == 400 + assert multiple_resp.json() == _fixture("multiple_function_calls_unsupported_error.json") + + +# Regression coverage retained from the merged brokered-continuation stack. + +@pytest.mark.parametrize("ttl_seconds", [float("nan"), float("inf"), float("-inf")]) +def test_foundry_brokered_nonfinite_explicit_state_ttl_uses_default(ttl_seconds: float): + app = _app(state_ttl_seconds=ttl_seconds) + + with TestClient(app) as client: + readiness = client.get("/readiness") + + assert readiness.status_code == 200 + assert readiness.json()["foundryResponses"]["stateTtlSeconds"] == 900.0 + +@pytest.mark.parametrize("raw_ttl", ["nan", "inf", "-inf"]) +def test_foundry_brokered_nonfinite_state_ttl_env_uses_default(monkeypatch, raw_ttl: str): + monkeypatch.setenv("AGENTKIT_FOUNDRY_RESPONSE_STATE_TTL_SECONDS", raw_ttl) + app = _app() + + with TestClient(app) as client: + readiness = client.get("/readiness") + + assert readiness.status_code == 200 + assert readiness.json()["foundryResponses"]["stateTtlSeconds"] == 900.0 + +def test_foundry_brokered_rejects_normal_followup_while_previous_response_is_resuming(monkeypatch): + stores: list[Any] = [] + original_store = foundry_module._FoundryResponseStateStore + + def capture_store(*args: Any, **kwargs: Any) -> Any: + store = original_store(*args, **kwargs) + stores.append(store) + return store + + monkeypatch.setattr(foundry_module, "_FoundryResponseStateStore", capture_store) + app = _app() + + with TestClient(app) as client: + initial = _start(client) + state = stores[0].get(initial["id"]) + state.status = "resuming" + stores[0].save(state) + resp = client.post("/responses", json={"previous_response_id": initial["id"], "input": "next question"}) + + assert resp.status_code == 409 + assert resp.json()["error"]["code"] == "response_pending_function_call_output" + +def test_foundry_brokered_integer_arguments_preserve_large_integer_bounds(): + huge = 10**100 + spec = _spec(tool_name="large-integer-bounds") + spec.brokered_tools[0].parameters = { + "type": "object", + "properties": { + "minimum": {"type": "integer", "minimum": huge}, + "exclusiveMinimum": {"type": "integer", "exclusiveMinimum": huge}, + "exclusiveMaximum": {"type": "integer", "exclusiveMaximum": -huge}, + }, + "required": ["minimum", "exclusiveMinimum", "exclusiveMaximum"], + } + app = _app(spec) + + with TestClient(app) as client: + response = client.post("/responses", json={"input": "large-integer-bounds"}) + + assert response.status_code == 200, response.text + assert json.loads(_call(response.json())["arguments"]) == { + "minimum": huge, + "exclusiveMinimum": huge + 1, + "exclusiveMaximum": -huge - 1, + } + +def test_foundry_brokered_integer_synthesis_rejects_unsupported_multiple_of(): + spec = _spec(tool_name="check-network-telemetry") + spec.brokered_tools[0].parameters = { + "type": "object", + "properties": {"n": {"type": "integer", "minimum": 1, "multipleOf": 2}}, + "required": ["n"], + } + app = _app(spec) + + with TestClient(app) as client: + response = client.post("/responses", json={"input": "check-network-telemetry"}) + + assert response.status_code == 400 + assert response.json()["error"]["code"] == "UnsupportedBrokeredSchema" + +def test_foundry_brokered_number_synthesis_handles_large_finite_bounds_without_overflow(): + spec = _spec(tool_name="large-number-bounds") + spec.brokered_tools[0].parameters = { + "type": "object", + "properties": { + "midpoint": {"type": "number", "minimum": 1e308, "maximum": 1.1e308}, + "above": {"type": "number", "exclusiveMinimum": 1e308}, + "below": {"type": "number", "exclusiveMaximum": -1e308}, + "exactInteger": {"type": "number", "minimum": 9007199254740993, "maximum": 9007199254740994}, + "roundedMidpoint": {"type": "number", "exclusiveMinimum": 0, "exclusiveMaximum": 0.9999999999999999}, + "narrowValue": {"type": "number", "exclusiveMinimum": 1.0, "maximum": 1.0000000000000002}, + }, + "required": ["midpoint", "above", "below", "exactInteger", "roundedMidpoint", "narrowValue"], + } + app = _app(spec) + + with TestClient(app) as client: + response = client.post("/responses", json={"input": "large-number-bounds"}) + + assert response.status_code == 200, response.text + arguments = json.loads(_call(response.json())["arguments"]) + assert arguments["midpoint"] == 1e308 + assert arguments["above"] == math.nextafter(1e308, math.inf) + assert arguments["below"] == math.nextafter(-1e308, -math.inf) + assert arguments["exactInteger"] == 9007199254740993 + assert 0 < arguments["roundedMidpoint"] < 0.9999999999999999 + assert arguments["narrowValue"] == 1.0000000000000002 + +def test_foundry_brokered_number_synthesis_rejects_unsupported_multiple_of(): + spec = _spec(tool_name="number-multiple") + spec.brokered_tools[0].parameters = { + "type": "object", + "properties": {"value": {"type": "number", "minimum": 0, "multipleOf": 0.5}}, + "required": ["value"], + } + app = _app(spec) + + with TestClient(app) as client: + response = client.post("/responses", json={"input": "number-multiple"}) + + assert response.status_code == 400 + assert response.json()["error"]["code"] == "UnsupportedBrokeredSchema" + +def test_foundry_brokered_number_synthesis_combines_all_declared_bounds(): + spec = _spec(tool_name="combined-number-bounds") + spec.brokered_tools[0].parameters = { + "type": "object", + "properties": { + "value": {"type": "number", "minimum": 0, "exclusiveMinimum": 0, "maximum": 1}, + }, + "required": ["value"], + } + app = _app(spec) + + with TestClient(app) as client: + response = client.post("/responses", json={"input": "combined-number-bounds"}) + + assert response.status_code == 200, response.text + assert json.loads(_call(response.json())["arguments"]) == {"value": 1} + +@pytest.mark.parametrize("bound", [float("nan"), float("inf"), float("-inf")]) +def test_foundry_brokered_number_synthesis_rejects_nonfinite_bounds(bound: float): + spec = _spec(tool_name="nonfinite-number-bound") + spec.brokered_tools[0].parameters = { + "type": "object", + "properties": {"value": {"type": "number", "minimum": bound}}, + "required": ["value"], + } + app = _app(spec) + + with TestClient(app) as client: + response = client.post("/responses", json={"input": "nonfinite-number-bound"}) + + assert response.status_code == 400 + assert response.json()["error"]["code"] == "UnsupportedBrokeredSchema" + +def test_foundry_brokered_number_synthesis_prefers_compact_zero_within_wide_bounds(): + spec = _spec(tool_name="compact-number-bounds") + properties = { + f"value{index}": {"type": "number", "minimum": -1e308, "maximum": 1e308} + for index in range(27) + } + spec.brokered_tools[0].parameters = { + "type": "object", + "properties": properties, + "required": list(properties), + } + app = _app(spec) + + with TestClient(app) as client: + response = client.post("/responses", json={"input": "compact-number-bounds"}) + + assert response.status_code == 200, response.text + assert json.loads(_call(response.json())["arguments"]) == {name: 0 for name in properties} + +def test_foundry_brokered_number_synthesis_prefers_compact_float_boundary_over_large_integer(): + spec = _spec(tool_name="compact-large-number-bounds") + properties = { + f"value{index}": {"type": "number", "minimum": 1e308, "maximum": 1.1e308} + for index in range(27) + } + spec.brokered_tools[0].parameters = { + "type": "object", + "properties": properties, + "required": list(properties), + } + app = _app(spec) + + with TestClient(app) as client: + response = client.post("/responses", json={"input": "compact-large-number-bounds"}) + + assert response.status_code == 200, response.text + assert json.loads(_call(response.json())["arguments"]) == {name: 1e308 for name in properties} + +def test_foundry_brokered_bounds_nested_array_synthesis_before_recursive_allocation(): + spec = _spec(tool_name="nested-array") + spec.brokered_tools[0].parameters = { + "type": "object", + "properties": { + "values": { + "type": "array", + "minItems": 32, + "items": { + "type": "array", + "minItems": 32, + "items": {"type": "string"}, + }, + } + }, + "required": ["values"], + } + app = _app(spec) + + with TestClient(app) as client: + response = client.post("/responses", json={"input": "nested-array"}) + + assert response.status_code == 413 + assert response.json()["error"]["code"] == "brokered_arguments_too_large" + +def test_foundry_brokered_counts_wide_object_members_against_synthesis_budget(): + item_properties = {f"field{index}": {"type": "string"} for index in range(32)} + spec = _spec(tool_name="wide-object-array") + spec.brokered_tools[0].parameters = { + "type": "object", + "properties": { + "values": { + "type": "array", + "minItems": 32, + "items": { + "type": "object", + "properties": item_properties, + "required": list(item_properties), + }, + } + }, + "required": ["values"], + } + app = _app(spec) + + with TestClient(app) as client: + response = client.post("/responses", json={"input": "wide-object-array"}) + + assert response.status_code == 413 + assert response.json()["error"]["code"] == "brokered_arguments_too_large" + +def test_foundry_brokered_rejects_overly_deep_deterministic_synthesis(): + child: dict[str, Any] = {"type": "string", "default": "leaf"} + for _ in range(129): + child = { + "type": "object", + "properties": {"child": child}, + "required": ["child"], + } + spec = _spec(tool_name="deep-synthesis") + spec.brokered_tools[0].parameters = { + "type": "object", + "properties": {"root": child}, + "required": ["root"], + } + app = _app(spec) + + with TestClient(app) as client: + response = client.post("/responses", json={"input": "deep-synthesis"}) + + assert response.status_code == 413 + assert response.json()["error"]["code"] == "brokered_arguments_too_large" + +def test_foundry_brokered_rejects_overly_deep_root_literal_synthesis(): + literal: dict[str, Any] = {} + for _ in range(129): + literal = {"nested": literal} + spec = _spec(tool_name="deep-root-literal") + spec.brokered_tools[0].parameters = { + "type": "object", + "const": literal, + } + app = _app(spec) + + with TestClient(app) as client: + response = client.post("/responses", json={"input": "deep-root-literal"}) + + assert response.status_code == 413 + assert response.json()["error"]["code"] == "brokered_arguments_too_large" + +@pytest.mark.parametrize( + ("literal", "expected"), + [ + ({"const": 2.0}, 2.0), + ({"default": 1.0}, 1.0), + ({"enum": [3.0]}, 3.0), + ], +) +def test_foundry_brokered_honors_integral_float_integer_literals( + literal: dict[str, Any], expected: float +): + spec = _spec(tool_name="integer-literal") + spec.brokered_tools[0].parameters = { + "type": "object", + "properties": {"count": {"type": "integer", **literal}}, + "required": ["count"], + } + app = _app(spec) + + with TestClient(app) as client: + response = client.post("/responses", json={"input": "integer-literal"}) + + assert response.status_code == 200, response.text + count = json.loads(_call(response.json())["arguments"])["count"] + assert count == expected + assert isinstance(count, float) + +def test_foundry_brokered_rejects_lossy_function_call_output_float_before_state_change(tmp_path): + state_file = tmp_path / "responses-state.json" + app = _app(response_state_file=state_file) + + with TestClient(app) as client: + initial = _start(client) + call = _call(initial) + persisted_before = state_file.read_bytes() + lossy = client.post( + "/responses", + headers=CONTINUATION_AUTH, + json={ + "previous_response_id": initial["id"], + "input": [ + { + "type": "function_call_output", + "call_id": call["call_id"], + "output": '{"approved":true,"output":{"id":9007199254740993.0}}', + } + ], + }, + ) + persisted_after_rejection = state_file.read_bytes() + exact = client.post( + "/responses", + headers=CONTINUATION_AUTH, + json={ + "previous_response_id": initial["id"], + "input": [ + { + "type": "function_call_output", + "call_id": call["call_id"], + "output": '{"approved":true,"output":{"id":9007199254740992.0}}', + } + ], + }, + ) + + assert lossy.status_code == 400 + assert lossy.json()["error"]["code"] == "invalid_function_call_output" + assert persisted_after_rejection == persisted_before + assert exact.status_code == 200, exact.text + assert _message_text(exact.json()).endswith('{"id":9007199254740992.0}') + +def test_foundry_brokered_rejects_duplicate_keys_in_function_call_output(tmp_path): + state_file = tmp_path / "responses-state.json" + app = _app(response_state_file=state_file) + + with TestClient(app) as client: + initial = _start(client) + call = _call(initial) + persisted_before = state_file.read_bytes() + response = client.post( + "/responses", + headers=CONTINUATION_AUTH, + json={ + "previous_response_id": initial["id"], + "input": [ + { + "type": "function_call_output", + "call_id": call["call_id"], + "output": '{"approved":false,"approved":true,"output":{}}', + } + ], + }, + ) + + assert response.status_code == 400 + assert response.json()["error"]["code"] == "invalid_function_call_output" + assert state_file.read_bytes() == persisted_before + +def test_foundry_brokered_rejects_object_valued_function_call_output_before_state_change(tmp_path): + state_file = tmp_path / "responses-state.json" + app = _app(response_state_file=state_file) + + with TestClient(app) as client: + initial = _start(client) + call = _call(initial) + persisted_before = state_file.read_bytes() + response = client.post( + "/responses", + headers=CONTINUATION_AUTH, + json={ + "previous_response_id": initial["id"], + "input": [ + { + "type": "function_call_output", + "call_id": call["call_id"], + "output": {"approved": True, "output": {"id": 9007199254740993.0}}, + } + ], + }, + ) + + assert response.status_code == 400 + assert response.json()["error"]["code"] == "invalid_function_call_output" + assert state_file.read_bytes() == persisted_before + +def test_foundry_brokered_rejects_oversized_raw_output_before_json_parsing(monkeypatch): + app = _app(max_brokered_output_bytes=64) + + def fail_if_parsed(_output: Any) -> dict[str, Any]: + raise AssertionError("oversized raw output must be rejected before JSON parsing") + + monkeypatch.setattr(foundry_module, "_json_object_from_output", fail_if_parsed) + with TestClient(app) as client: + initial = _start(client) + call = _call(initial) + response = client.post( + "/responses", + headers=CONTINUATION_AUTH, + json={ + "previous_response_id": initial["id"], + "input": [ + { + "type": "function_call_output", + "call_id": call["call_id"], + "output": '{"approved":true,"output":{"blob":"' + ("x" * 128) + '"}}', + } + ], + }, + ) + + assert response.status_code == 413 + assert response.json()["error"]["code"] == "brokered_output_too_large" + +def test_foundry_brokered_accepting_continuation_refreshes_state_expiry(monkeypatch): + stores: list[Any] = [] + original_store = foundry_module._FoundryResponseStateStore + + def capture_store(*args: Any, **kwargs: Any) -> Any: + store = original_store(*args, **kwargs) + stores.append(store) + return store + + monkeypatch.setattr(foundry_module, "_FoundryResponseStateStore", capture_store) + app = _app(state_ttl_seconds=60) + + with TestClient(app) as client: + initial = _start(client) + call = _call(initial) + state = stores[0].get(initial["id"]) + state.expires_at = time.time() + 1 + stores[0].save(state) + response = client.post( + "/responses", + headers=CONTINUATION_AUTH, + json=_continuation(initial["id"], call["call_id"], {"approved": True, "output": {"ok": True}}), + ) + + assert response.status_code == 200, response.text + refreshed = stores[0].get(initial["id"]) + assert refreshed.expires_at > time.time() + 50 + + +def test_foundry_brokered_replays_persisted_output_after_limit_is_lowered(tmp_path): + state_file = tmp_path / "responses-state.json" + + with TestClient(_app(response_state_file=state_file, max_brokered_output_bytes=1024)) as client: + initial = _start(client) + call = _call(initial) + payload = _continuation( + initial["id"], + call["call_id"], + {"approved": True, "output": {"blob": "x" * 256}}, + ) + completed = client.post("/responses", json=payload, headers=CONTINUATION_AUTH) + + with TestClient(_app(response_state_file=state_file, max_brokered_output_bytes=128)) as client: + replay = client.post("/responses", json=payload, headers=CONTINUATION_AUTH) + + assert completed.status_code == 200, completed.text + assert replay.status_code == 200, replay.text + assert replay.json() == completed.json() + + +def test_foundry_brokered_rejects_persisted_model_loop_continuation_when_loop_is_disabled(tmp_path): + state_file = tmp_path / "foundry-model-state.json" spec = _spec(tool_name="check-network-telemetry") - fake_first = _FakeChatTransport( + fake = _FakeChatTransport( [ _chat_response( { @@ -2292,217 +4430,257 @@ def test_foundry_brokered_file_state_recovers_unfinalized_accepted_continuation_ "content": None, "tool_calls": [ { - "id": "model_generated_call_id", + "id": "call_model", "type": "function", - "function": {"name": "check-network-telemetry", "arguments": '{}'}, + "function": {"name": "check-network-telemetry", "arguments": "{}"}, } ], } ) ] ) - first_app = _app(spec, brokered_model_loop_enabled=True, brokered_model_http_client=httpx.AsyncClient(transport=httpx.MockTransport(fake_first.handler)), response_state_file=state_file) + with TestClient(_model_loop_app(spec, fake, response_state_file=state_file)) as client: + initial = client.post("/responses", json={"input": "check-network-telemetry"}).json() + call = _call(initial) - with TestClient(first_app) as client: - initial = client.post("/responses", json={"input": "check-network-telemetry"}) - call = _call(initial.json()) + persisted_before = state_file.read_bytes() + with TestClient(_app(spec, response_state_file=state_file, brokered_model_loop_enabled=False)) as client: + response = client.post( + "/responses", + headers=CONTINUATION_AUTH, + json=_continuation(initial["id"], call["call_id"], {"approved": True, "output": {"ok": True}}), + ) - payload = _continuation(initial.json()["id"], call["call_id"], {"approved": True, "output": {"ok": True}}) - state_data = json.loads(state_file.read_text(encoding="utf-8")) - state = state_data["states"][initial.json()["id"]] - state["acceptedOutputs"] = {call["call_id"]: payload["input"][0]["output"]} - state["status"] = "resuming" - state["finalPayload"] = None - state_file.write_text(json.dumps(state_data, separators=(",", ":"), sort_keys=True), encoding="utf-8") + assert response.status_code == 503 + assert response.json()["error"]["code"] == "brokered_model_loop_unavailable" + assert state_file.read_bytes() == persisted_before - fake_second = _FakeChatTransport([_chat_response({"role": "assistant", "content": "Recovered."})]) - second_app = _app(spec, brokered_model_loop_enabled=True, brokered_model_http_client=httpx.AsyncClient(transport=httpx.MockTransport(fake_second.handler)), response_state_file=state_file) - with TestClient(second_app) as client: - retried = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) +def test_foundry_brokered_rejects_normal_followup_to_expired_pending_response(): + app = _app(state_ttl_seconds=0) - assert retried.status_code == 200, retried.text - assert _message_text(retried.json()) == "Recovered." + with TestClient(app) as client: + initial = _start(client) + time.sleep(0.01) + response = client.post( + "/responses", + json={"previous_response_id": initial["id"], "input": "next question"}, + ) + assert response.status_code == 410 + assert response.json()["error"]["code"] == "response_state_expired" -def test_foundry_brokered_model_loop_accepts_integer_arguments_encoded_as_integral_float(): - spec = _spec(tool_name="retry-tool") +def test_foundry_brokered_allows_normal_followup_after_completed_state_expires(monkeypatch): + stores: list[Any] = [] + original_store = foundry_module._FoundryResponseStateStore + + def capture_store(*args: Any, **kwargs: Any) -> Any: + store = original_store(*args, **kwargs) + stores.append(store) + return store + + monkeypatch.setattr(foundry_module, "_FoundryResponseStateStore", capture_store) + app = _app() + + with TestClient(app) as client: + initial = _start(client) + call = _call(initial) + completed = client.post( + "/responses", + headers=CONTINUATION_AUTH, + json=_continuation(initial["id"], call["call_id"], {"approved": True, "output": {"ok": True}}), + ) + state = stores[0].get(initial["id"]) + state.expires_at = time.time() - 1 + stores[0].save(state) + followup = client.post( + "/responses", + json={"previous_response_id": initial["id"], "input": "next question"}, + ) + + assert completed.status_code == 200, completed.text + assert followup.status_code == 200, followup.text + assert _call(followup.json()) + +@pytest.mark.parametrize("brokered_class", ["write", "coordination"]) +def test_foundry_brokered_refuses_multi_value_root_enum_side_effecting_arguments(brokered_class: str): + spec = _spec(tool_name="dispatch-work-order", brokered_class=brokered_class) spec.brokered_tools[0].parameters = { "type": "object", - "properties": {"retries": {"type": "integer"}}, - "required": ["retries"], + "enum": [ + {"operation": "delete"}, + {"operation": "create"}, + ], } - fake = _FakeChatTransport( - [ - _chat_response( - { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_model", - "type": "function", - "function": {"name": "retry-tool", "arguments": '{"retries":1.0}'}, - } - ], - } - ) - ] - ) - app = _model_loop_app(spec, fake) + app = _app(spec) with TestClient(app) as client: - response = client.post("/responses", json={"input": "call retry-tool"}) - - assert response.status_code == 200, response.text - assert json.loads(_call(response.json())["arguments"]) == {"retries": 1.0} + resp = client.post("/responses", json={"input": "dispatch-work-order"}) + assert resp.status_code == 400 + assert resp.json()["error"]["code"] == "UnsupportedBrokeredSchema" -def test_foundry_brokered_model_loop_unexpected_resume_failure_can_be_retried(): - spec = _spec(tool_name="check-network-telemetry") +def test_foundry_brokered_allows_single_value_root_enum_write_arguments(): + spec = _spec(tool_name="dispatch-work-order", brokered_class="write") + spec.brokered_tools[0].parameters = { + "type": "object", + "enum": [{"operation": "create"}], + } + app = _app(spec) - class FlakyResumeTransport: - def __init__(self) -> None: - self.requests: list[dict[str, Any]] = [] - self.resume_attempts = 0 + with TestClient(app) as client: + resp = client.post("/responses", json={"input": "dispatch-work-order"}) - def handler(self, request: httpx.Request) -> httpx.Response: - payload = json.loads(request.content.decode("utf-8")) - self.requests.append(payload) - if len(self.requests) == 1: - return httpx.Response( - 200, - json=_chat_response( - { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "model_generated_call_id", - "type": "function", - "function": {"name": "check-network-telemetry", "arguments": '{}'}, - } - ], - } - ), - ) - self.resume_attempts += 1 - if self.resume_attempts == 1: - raise RuntimeError("transient resume failure") - return httpx.Response(200, json=_chat_response({"role": "assistant", "content": "Recovered after retry."})) + assert resp.status_code == 200 + assert json.loads(_call(resp.json())["arguments"]) == {"operation": "create"} - fake = FlakyResumeTransport() - app = _model_loop_app(spec, fake) +@pytest.mark.parametrize("brokered_class", ["write", "coordination"]) +def test_foundry_brokered_refuses_nonliteral_optional_prompt_for_side_effecting_tools(brokered_class: str): + spec = _spec(tool_name="dispatch-work-order", brokered_class=brokered_class) + spec.brokered_tools[0].parameters = { + "type": "object", + "properties": {"prompt": {"type": "string"}}, + } + app = _app(spec) with TestClient(app) as client: - initial = client.post("/responses", json={"input": "check-network-telemetry"}) - call = _call(initial.json()) - payload = _continuation(initial.json()["id"], call["call_id"], {"approved": True, "output": {"ok": True}}) - failed = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) - retried = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + resp = client.post("/responses", json={"input": "dispatch-work-order with arbitrary payload"}) - assert failed.status_code == 502 - assert failed.json()["error"] == {"message": "model resume failed", "code": "ModelResumeError"} - assert retried.status_code == 200, retried.text - assert _message_text(retried.json()) == "Recovered after retry." + assert resp.status_code == 400 + assert resp.json()["error"]["code"] == "UnsupportedBrokeredSchema" +def test_foundry_brokered_allows_literal_optional_prompt_for_write_tools(): + spec = _spec(tool_name="dispatch-work-order", brokered_class="write") + spec.brokered_tools[0].parameters = { + "type": "object", + "properties": {"prompt": {"type": "string", "const": "fixed"}}, + } + app = _app(spec) -def test_foundry_brokered_model_loop_failed_resume_can_be_retried(): - spec = _spec(tool_name="check-network-telemetry") - fake = _FakeChatTransport( - [ - _chat_response( - { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "model_generated_call_id", - "type": "function", - "function": {"name": "check-network-telemetry", "arguments": '{}'}, - } - ], - } - ), - {"choices": []}, - _chat_response({"role": "assistant", "content": "Retry worked."}), - ] - ) - app = _model_loop_app(spec, fake) + with TestClient(app) as client: + resp = client.post("/responses", json={"input": "dispatch-work-order with arbitrary payload"}) + + assert resp.status_code == 200 + assert json.loads(_call(resp.json())["arguments"]) == {"prompt": "fixed"} + +def test_foundry_brokered_rejects_continuation_with_extra_non_output_item(): + app = _app() with TestClient(app) as client: - initial = client.post("/responses", json={"input": "check-network-telemetry"}) - call = _call(initial.json()) - payload = _continuation(initial.json()["id"], call["call_id"], {"approved": True, "output": {"ok": True}}) - failed = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) - retried = client.post("/responses", headers=CONTINUATION_AUTH, json=payload) + initial = _start(client) + call = _call(initial) + request = _continuation(initial["id"], call["call_id"], {"approved": True, "output": {"ok": True}}) + request["input"].append({"type": "message", "role": "user", "content": "ignored"}) + rejected = client.post("/responses", json=request, headers=CONTINUATION_AUTH) + accepted = client.post( + "/responses", + json=_continuation(initial["id"], call["call_id"], {"approved": True, "output": {"ok": True}}), + headers=CONTINUATION_AUTH, + ) - assert failed.status_code == 502 - assert retried.status_code == 200, retried.text - assert _message_text(retried.json()) == "Retry worked." + assert rejected.status_code == 400 + assert rejected.json()["error"]["code"] == "multiple_tool_outputs_unsupported" + assert accepted.status_code == 200, accepted.text +def test_foundry_brokered_model_loop_omits_authorization_when_auth_is_omitted(monkeypatch): + captured_headers: dict[str, str] = {} -def test_foundry_brokered_model_loop_can_return_final_message_without_tool_call(): - fake = _FakeChatTransport([_chat_response({"role": "assistant", "content": "No tool needed."})]) - app = _model_loop_app(_spec(tool_name="check-network-telemetry"), fake) + class FakeStream: + async def __aenter__(self) -> httpx.Response: + return httpx.Response( + 200, + request=httpx.Request("POST", "https://api.openai.com/v1/chat/completions"), + content=json.dumps(_chat_response({"role": "assistant", "content": "done"})).encode("utf-8"), + headers={"content-type": "application/json"}, + ) - with TestClient(app) as client: - response = client.post("/responses", json={"input": "Say hello"}) + async def __aexit__(self, *args: Any) -> None: + return None + + class FakeClient: + def __init__(self, *, headers: dict[str, str], timeout: int) -> None: + assert timeout == 60 + captured_headers.update(headers) - assert response.status_code == 200, response.text - assert _message_text(response.json()) == "No tool needed." - assert fake.requests[0]["tool_choice"] == "auto" + def stream(self, method: str, url: str, **kwargs: Any) -> FakeStream: + assert method == "POST" + assert url.endswith("/chat/completions") + return FakeStream() + async def aclose(self) -> None: + return None -def test_foundry_brokered_model_loop_returns_assistant_refusal_without_tool_call(): - fake = _FakeChatTransport( - [_chat_response({"role": "assistant", "content": None, "refusal": "I cannot help with that."})] - ) - app = _model_loop_app(_spec(tool_name="check-network-telemetry"), fake) + monkeypatch.setattr(foundry_model_loop_module.httpx, "AsyncClient", FakeClient) + app = _app(_spec(tool_name="check-network-telemetry"), brokered_model_loop_enabled=True) with TestClient(app) as client: - response = client.post("/responses", json={"input": "Do something unsafe"}) + response = client.post("/responses", json={"input": "check-network-telemetry"}) assert response.status_code == 200, response.text - assert _message_text(response.json()) == "I cannot help with that." + assert "Authorization" not in captured_headers @pytest.mark.parametrize( - "message", + "usage", [ - {"role": "assistant", "content": "\ud800"}, - {"role": "assistant", "content": None, "refusal": "\ud800"}, + {"prompt_tokens": {"unexpected": 1}}, + {"completion_tokens": "not-a-number"}, + {"prompt_tokens": "12"}, + {"total_tokens": 1.5}, ], ) -def test_foundry_brokered_model_loop_rejects_surrogate_final_text(message: dict[str, Any]): - fake = _FakeChatTransport([_chat_response(message)]) +def test_foundry_brokered_model_loop_normalizes_malformed_usage(usage: dict[str, Any]): + fake = _FakeChatTransport( + [ + { + "choices": [{"message": {"role": "assistant", "content": "done"}}], + "usage": usage, + } + ] + ) app = _model_loop_app(_spec(tool_name="check-network-telemetry"), fake) with TestClient(app) as client: - response = client.post("/responses", json={"input": "Say hello"}) + response = client.post("/responses", json={"input": "check-network-telemetry"}) assert response.status_code == 502 assert response.json()["error"]["code"] == "InvalidModelResponse" - -@pytest.mark.parametrize("content", [None, {"text": "not supported"}, ["not supported"]]) -def test_foundry_brokered_model_loop_rejects_non_string_final_content(content: Any): - fake = _FakeChatTransport([_chat_response({"role": "assistant", "content": content})]) - app = _model_loop_app(_spec(tool_name="check-network-telemetry"), fake) +def test_foundry_brokered_model_loop_preserves_total_only_usage_across_resume(): + spec = _spec(tool_name="check-network-telemetry") + initial_model_response = _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "model_generated_call_id", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": "{}"}, + } + ], + } + ) + initial_model_response["usage"] = {"total_tokens": 5} + final_model_response = _chat_response({"role": "assistant", "content": "done"}) + final_model_response["usage"] = {"total_tokens": 7} + app = _model_loop_app(spec, _FakeChatTransport([initial_model_response, final_model_response])) with TestClient(app) as client: - response = client.post("/responses", json={"input": "Say hello"}) - - assert response.status_code == 502 - assert response.json()["error"]["code"] == "InvalidModelResponse" + initial = client.post("/responses", json={"input": "check-network-telemetry"}) + call = _call(initial.json()) + final = client.post( + "/responses", + headers=CONTINUATION_AUTH, + json=_continuation(initial.json()["id"], call["call_id"], {"approved": True, "output": {"ok": True}}), + ) + assert initial.json()["usage"] == {"input_tokens": 0, "output_tokens": 0, "total_tokens": 5} + assert final.status_code == 200, final.text + assert final.json()["usage"] == {"input_tokens": 0, "output_tokens": 0, "total_tokens": 12} -def test_foundry_brokered_model_loop_rejects_unsupported_pattern_deterministically(): - spec = _spec(tool_name="check-network-telemetry") - spec.brokered_tools[0].parameters = { - "type": "object", - "properties": {"site": {"type": "string", "pattern": "\\p{L}+"}}, - "required": ["site"], - } +def test_foundry_brokered_model_loop_rejects_noncanonical_denied_output_before_resume(tmp_path): + state_file = tmp_path / "responses-state.json" + spec = _spec(tool_name="dispatch-work-order", brokered_class="write") fake = _FakeChatTransport( [ _chat_response( @@ -2511,85 +4689,173 @@ def test_foundry_brokered_model_loop_rejects_unsupported_pattern_deterministical "content": None, "tool_calls": [ { - "id": "call_model", + "id": "model_generated_call_id", "type": "function", - "function": {"name": "check-network-telemetry", "arguments": '{"site":"sfo"}'}, + "function": {"name": "dispatch-work-order", "arguments": "{}"}, } ], } ) ] ) - app = _model_loop_app(spec, fake) + app = _model_loop_app(spec, fake, response_state_file=state_file) with TestClient(app) as client: - response = client.post("/responses", json={"input": "call check-network-telemetry"}) + initial = client.post("/responses", json={"input": "dispatch-work-order"}) + call = _call(initial.json()) + persisted_before = state_file.read_bytes() + denied = client.post( + "/responses", + headers=CONTINUATION_AUTH, + json=_continuation( + initial.json()["id"], + call["call_id"], + { + "approved": False, + "error": {"code": "approval_declined", "message": "denied"}, + "output": {"sensitiveDiagnostic": "must-not-reach-model"}, + }, + ), + ) + + assert denied.status_code == 400 + assert denied.json()["error"]["code"] == "invalid_function_call_output" + assert state_file.read_bytes() == persisted_before + assert len(fake.requests) == 1 + +@pytest.mark.parametrize("payload", [{"approved": True}, {"approved": True, "output": None}]) +def test_foundry_brokered_rejects_approved_continuation_without_object_output(payload: dict[str, Any]): + app = _app() + + with TestClient(app) as client: + initial = _start(client) + call = _call(initial) + response = client.post( + "/responses", + headers=CONTINUATION_AUTH, + json=_continuation(initial["id"], call["call_id"], payload), + ) assert response.status_code == 400 - assert response.json()["error"]["code"] == "UnsupportedBrokeredSchema" + assert response.json()["error"]["code"] == "invalid_function_call_output" + +@pytest.mark.parametrize("payload", [{"approved": False}, {"approved": False, "error": None}]) +def test_foundry_brokered_rejects_denied_continuation_without_error_object(payload: dict[str, Any]): + app = _app() + with TestClient(app) as client: + initial = _start(client) + call = _call(initial) + response = client.post( + "/responses", + headers=CONTINUATION_AUTH, + json=_continuation(initial["id"], call["call_id"], payload), + ) -def test_foundry_brokered_model_loop_validates_additional_properties_schema(): - spec = _spec(tool_name="check-network-telemetry") - spec.brokered_tools[0].parameters = { - "type": "object", - "properties": {"site": {"type": "string"}}, - "additionalProperties": {"type": "string"}, - } - fake = _FakeChatTransport( - [ - _chat_response( - { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_model", - "type": "function", - "function": {"name": "check-network-telemetry", "arguments": '{"site":"sfo","count":1}'}, - } - ], - } - ) - ] + assert response.status_code == 400 + assert response.json()["error"]["code"] == "invalid_function_call_output" + +def test_foundry_brokered_rejects_deep_continuation_output_without_consuming_state(): + app = _app() + nested: dict[str, Any] = {} + for _ in range(200): + nested = {"nested": nested} + + with TestClient(app) as client: + initial = _start(client) + call = _call(initial) + rejected = client.post( + "/responses", + headers=CONTINUATION_AUTH, + json=_continuation(initial["id"], call["call_id"], {"approved": True, "output": nested}), + ) + accepted = client.post( + "/responses", + headers=CONTINUATION_AUTH, + json=_continuation(initial["id"], call["call_id"], {"approved": True, "output": {"ok": True}}), + ) + + assert rejected.status_code == 400 + assert rejected.json()["error"]["code"] == "invalid_function_call_output" + assert accepted.status_code == 200, accepted.text + +def test_foundry_brokered_bounds_request_body_before_model_call(): + fake = _FakeChatTransport([_chat_response({"role": "assistant", "content": "must not run"})]) + app = _model_loop_app( + _spec(tool_name="check-network-telemetry"), + fake, + max_request_body_bytes=64, ) - app = _model_loop_app(spec, fake) with TestClient(app) as client: - response = client.post("/responses", json={"input": "call check-network-telemetry"}) + response = client.post( + "/responses", + content=json.dumps({"input": "x" * 128}), + headers={"content-type": "application/json"}, + ) - assert response.status_code == 400 - assert response.json()["error"]["code"] == "InvalidToolArguments" + assert response.status_code == 413 + assert response.json()["error"]["code"] == "request_body_too_large" + assert fake.requests == [] +def test_foundry_brokered_model_message_size_matches_persistence_encoding(): + messages = [{"role": "user", "content": "😀" * 8}] -def test_foundry_brokered_model_loop_rejects_unknown_model_tool_request(): - fake = _FakeChatTransport( - [ - _chat_response( + measured = foundry_module._persistence_json_bytes(messages) + persisted = json.dumps( + messages, + allow_nan=False, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + utf8_compact = json.dumps( + messages, + allow_nan=False, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + + assert measured == persisted + assert len(measured) > len(utf8_compact) + +def test_foundry_brokered_bounds_persisted_model_messages_and_releases_reservation(): + tool_response = _chat_response( + { + "role": "assistant", + "content": None, + "tool_calls": [ { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_model", - "type": "function", - "function": {"name": "unknown", "arguments": "{}"}, - } - ], + "id": "call_model", + "type": "function", + "function": {"name": "check-network-telemetry", "arguments": "{}"}, } - ) - ] + ], + } + ) + fake = _FakeChatTransport([tool_response, deepcopy(tool_response)]) + app = _model_loop_app( + _spec(tool_name="check-network-telemetry"), + fake, + max_request_body_bytes=4096, + max_model_messages_bytes=512, + max_pending_responses=1, ) - app = _model_loop_app(_spec(tool_name="check-network-telemetry"), fake) with TestClient(app) as client: - response = client.post("/responses", json={"input": "call unknown"}) - - assert response.status_code == 400 - assert response.json()["error"]["code"] == "unknown_brokered_tool" + oversized = client.post( + "/responses", + json={"input": "check-network-telemetry " + ("x" * 1024)}, + ) + accepted = client.post("/responses", json={"input": "check-network-telemetry"}) + assert oversized.status_code == 413 + 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 -def test_foundry_brokered_model_loop_rejects_object_valued_tool_arguments(): +def test_foundry_brokered_model_loop_reserves_capacity_before_model_call(): fake = _FakeChatTransport( [ _chat_response( @@ -2600,26 +4866,31 @@ def test_foundry_brokered_model_loop_rejects_object_valued_tool_arguments(): { "id": "call_model", "type": "function", - "function": { - "name": "check-network-telemetry", - "arguments": {"id": 9007199254740993.0}, - }, + "function": {"name": "check-network-telemetry", "arguments": "{}"}, } ], } - ) + ), + _chat_response({"role": "assistant", "content": "must not be called"}), ] ) - app = _model_loop_app(_spec(tool_name="check-network-telemetry"), fake) + app = _model_loop_app( + _spec(tool_name="check-network-telemetry"), + fake, + max_pending_responses=1, + ) with TestClient(app) as client: - response = client.post("/responses", json={"input": "check-network-telemetry"}) - - assert response.status_code == 400 - assert response.json()["error"]["code"] == "InvalidToolArguments" + first = client.post("/responses", json={"input": "check-network-telemetry"}) + second = client.post("/responses", json={"input": "check-network-telemetry"}) + assert first.status_code == 200, first.text + assert _call(first.json()) + assert second.status_code == 429 + assert second.json()["error"]["code"] == "brokered_response_state_full" + assert len(fake.requests) == 1 -def test_foundry_brokered_model_loop_rejects_duplicate_tool_argument_keys(): +def test_foundry_brokered_model_loop_unused_reservation_preserves_completed_replay(): fake = _FakeChatTransport( [ _chat_response( @@ -2630,26 +4901,42 @@ def test_foundry_brokered_model_loop_rejects_duplicate_tool_argument_keys(): { "id": "call_model", "type": "function", - "function": { - "name": "check-network-telemetry", - "arguments": '{"site":"sfo","site":"sea"}', - }, + "function": {"name": "check-network-telemetry", "arguments": "{}"}, } ], } - ) + ), + _chat_response({"role": "assistant", "content": "First response completed."}), + _chat_response({"role": "assistant", "content": "No tool needed."}), ] ) - app = _model_loop_app(_spec(tool_name="check-network-telemetry"), fake) + app = _model_loop_app( + _spec(tool_name="check-network-telemetry"), + fake, + max_pending_responses=1, + ) with TestClient(app) as client: - response = client.post("/responses", json={"input": "check-network-telemetry"}) - - assert response.status_code == 400 - assert response.json()["error"]["code"] == "InvalidToolArguments" + initial = client.post("/responses", json={"input": "check-network-telemetry"}) + call = _call(initial.json()) + continuation = _continuation( + initial.json()["id"], + call["call_id"], + {"approved": True, "output": {"ok": True}}, + ) + completed = client.post("/responses", headers=CONTINUATION_AUTH, json=continuation) + direct = client.post("/responses", json={"input": "answer without a tool"}) + replayed = client.post("/responses", headers=CONTINUATION_AUTH, json=continuation) + assert completed.status_code == 200, completed.text + assert direct.status_code == 200, direct.text + assert replayed.status_code == 200, replayed.text + assert replayed.json() == completed.json() + assert len(fake.requests) == 3 -def test_foundry_brokered_model_loop_rejects_decoded_surrogate_arguments(): +def test_foundry_brokered_model_loop_rejects_oversized_output_before_resume_or_state_change(tmp_path): + state_file = tmp_path / "responses-state.json" + spec = _spec(tool_name="check-network-telemetry") fake = _FakeChatTransport( [ _chat_response( @@ -2658,131 +4945,108 @@ def test_foundry_brokered_model_loop_rejects_decoded_surrogate_arguments(): "content": None, "tool_calls": [ { - "id": "call_model", + "id": "model_generated_call_id", "type": "function", - "function": { - "name": "check-network-telemetry", - "arguments": '{"site":"\\ud800"}', - }, + "function": {"name": "check-network-telemetry", "arguments": "{}"}, } ], } ) ] ) - app = _model_loop_app(_spec(tool_name="check-network-telemetry"), fake) + app = _model_loop_app( + spec, + fake, + response_state_file=state_file, + max_brokered_output_bytes=128, + ) with TestClient(app) as client: - response = client.post("/responses", json={"input": "check-network-telemetry"}) + initial = client.post("/responses", json={"input": "check-network-telemetry"}) + call = _call(initial.json()) + persisted_before = state_file.read_bytes() + oversized = client.post( + "/responses", + headers=CONTINUATION_AUTH, + json=_continuation( + initial.json()["id"], + call["call_id"], + {"approved": True, "output": {"blob": "x" * 256}}, + ), + ) - assert response.status_code == 400 - assert response.json()["error"]["code"] == "InvalidToolArguments" + assert oversized.status_code == 413 + assert oversized.json()["error"]["code"] == "brokered_output_too_large" + assert state_file.read_bytes() == persisted_before + assert len(fake.requests) == 1 +@pytest.mark.parametrize("raw_state", ["{not valid json", "[]"]) +def test_foundry_brokered_invalid_file_state_fails_startup_without_overwriting(tmp_path, raw_state: str): + state_file = tmp_path / "responses-state.json" + state_file.write_text(raw_state, encoding="utf-8") -def test_foundry_brokered_model_loop_rejects_excessively_nested_arguments(): - nested = '"leaf"' - for _ in range(200): - nested = f"[{nested}]" - fake = _FakeChatTransport( - [ - _chat_response( - { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_model", - "type": "function", - "function": { - "name": "check-network-telemetry", - "arguments": '{"site":' + nested + "}", - }, - } - ], - } - ) - ] - ) - app = _model_loop_app(_spec(tool_name="check-network-telemetry"), fake) + with pytest.raises(RuntimeError, match="invalid Foundry response state file"): + _app(response_state_file=state_file) - with TestClient(app) as client: - response = client.post("/responses", json={"input": "check-network-telemetry"}) + assert state_file.read_text(encoding="utf-8") == raw_state - assert response.status_code == 400 - assert response.json()["error"]["code"] == "InvalidToolArguments" +@pytest.mark.parametrize("expires_at", [float("nan"), float("inf"), float("-inf")]) +def test_foundry_brokered_rejects_nonfinite_persisted_expiry(tmp_path, expires_at: float): + state_file = tmp_path / "responses-state.json" + with TestClient(_app(response_state_file=state_file)) as client: + _start(client) + persisted = json.loads(state_file.read_text(encoding="utf-8")) + state = next(iter(persisted["states"].values())) + state["expiresAt"] = expires_at + raw_state = json.dumps(persisted, separators=(",", ":")) + state_file.write_text(raw_state, encoding="utf-8") + + with pytest.raises(RuntimeError, match="stored expiresAt must be finite"): + _app(response_state_file=state_file) + assert state_file.read_text(encoding="utf-8") == raw_state -def test_foundry_brokered_model_loop_bounds_raw_arguments_before_parsing(monkeypatch): +def test_foundry_brokered_model_loop_returns_assistant_refusal_without_tool_call(): fake = _FakeChatTransport( - [ - _chat_response( - { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_model", - "type": "function", - "function": { - "name": "check-network-telemetry", - "arguments": '{"blob":"' + ("x" * 128) + '"}', - }, - } - ], - } - ) - ] + [_chat_response({"role": "assistant", "content": None, "refusal": "I cannot help with that."})] ) - - def fail_if_parsed(_raw: Any) -> dict[str, Any]: - raise AssertionError("oversized model arguments must be rejected before parsing") - - monkeypatch.setattr(foundry_model_loop_module, "_parse_arguments", fail_if_parsed) - app = _model_loop_app(_spec(tool_name="check-network-telemetry"), fake, max_brokered_argument_bytes=64) + app = _model_loop_app(_spec(tool_name="check-network-telemetry"), fake) with TestClient(app) as client: - response = client.post("/responses", json={"input": "check-network-telemetry"}) - - assert response.status_code == 413 - assert response.json()["error"]["code"] == "brokered_arguments_too_large" + response = client.post("/responses", json={"input": "Do something unsafe"}) + assert response.status_code == 200, response.text + assert _message_text(response.json()) == "I cannot help with that." -def test_foundry_brokered_model_loop_normalizes_json_parser_failures(): - fake = _FakeChatTransport( - [ - _chat_response( - { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_model", - "type": "function", - "function": { - "name": "check-network-telemetry", - "arguments": '{"value":' + ("9" * 5000) + "}", - }, - } - ], - } - ) - ] - ) +@pytest.mark.parametrize( + "message", + [ + {"role": "assistant", "content": "\ud800"}, + {"role": "assistant", "content": None, "refusal": "\ud800"}, + ], +) +def test_foundry_brokered_model_loop_rejects_surrogate_final_text(message: dict[str, Any]): + fake = _FakeChatTransport([_chat_response(message)]) app = _model_loop_app(_spec(tool_name="check-network-telemetry"), fake) with TestClient(app) as client: - response = client.post("/responses", json={"input": "check-network-telemetry"}) + response = client.post("/responses", json={"input": "Say hello"}) - assert response.status_code == 400 - assert response.json()["error"]["code"] == "InvalidToolArguments" + assert response.status_code == 502 + assert response.json()["error"]["code"] == "InvalidModelResponse" +@pytest.mark.parametrize("content", [None, {"text": "not supported"}, ["not supported"]]) +def test_foundry_brokered_model_loop_rejects_non_string_final_content(content: Any): + fake = _FakeChatTransport([_chat_response({"role": "assistant", "content": content})]) + app = _model_loop_app(_spec(tool_name="check-network-telemetry"), fake) -def test_foundry_brokered_model_loop_validates_schema_valued_additional_properties(): - spec = _spec(tool_name="flex-tool") - spec.brokered_tools[0].parameters = { - "type": "object", - "additionalProperties": {"type": "integer"}, - } + with TestClient(app) as client: + response = client.post("/responses", json={"input": "Say hello"}) + + assert response.status_code == 502 + assert response.json()["error"]["code"] == "InvalidModelResponse" + +def test_foundry_brokered_model_loop_rejects_object_valued_tool_arguments(): fake = _FakeChatTransport( [ _chat_response( @@ -2793,24 +5057,25 @@ def test_foundry_brokered_model_loop_validates_schema_valued_additional_properti { "id": "call_model", "type": "function", - "function": {"name": "flex-tool", "arguments": '{"safe":"not-int"}'}, + "function": { + "name": "check-network-telemetry", + "arguments": {"id": 9007199254740993.0}, + }, } ], } ) ] ) - app = _model_loop_app(spec, fake) + app = _model_loop_app(_spec(tool_name="check-network-telemetry"), fake) with TestClient(app) as client: - response = client.post("/responses", json={"input": "call flex-tool"}) + response = client.post("/responses", json={"input": "check-network-telemetry"}) assert response.status_code == 400 assert response.json()["error"]["code"] == "InvalidToolArguments" - -def test_foundry_brokered_model_loop_rejects_nonfinite_model_arguments(): - spec = _spec(tool_name="check-network-telemetry") +def test_foundry_brokered_model_loop_rejects_duplicate_tool_argument_keys(): fake = _FakeChatTransport( [ _chat_response( @@ -2821,67 +5086,25 @@ def test_foundry_brokered_model_loop_rejects_nonfinite_model_arguments(): { "id": "call_model", "type": "function", - "function": {"name": "check-network-telemetry", "arguments": '{"value":NaN}'}, + "function": { + "name": "check-network-telemetry", + "arguments": '{"site":"sfo","site":"sea"}', + }, } ], } ) ] ) - app = _model_loop_app(spec, fake) + app = _model_loop_app(_spec(tool_name="check-network-telemetry"), fake) with TestClient(app) as client: - response = client.post("/responses", json={"input": "call check-network-telemetry"}) + response = client.post("/responses", json={"input": "check-network-telemetry"}) assert response.status_code == 400 assert response.json()["error"]["code"] == "InvalidToolArguments" - -def test_foundry_brokered_model_loop_uses_type_strict_const_and_enum_matching(): - for property_schema, arguments in [ - ({"const": 1}, '{"value":true}'), - ({"enum": [0]}, '{"value":false}'), - ]: - spec = _spec(tool_name="check-network-telemetry") - spec.brokered_tools[0].parameters = { - "type": "object", - "properties": {"value": property_schema}, - "required": ["value"], - } - fake = _FakeChatTransport( - [ - _chat_response( - { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_model", - "type": "function", - "function": {"name": "check-network-telemetry", "arguments": arguments}, - } - ], - } - ) - ] - ) - app = _model_loop_app(spec, fake) - - with TestClient(app) as client: - response = client.post("/responses", json={"input": "call check-network-telemetry"}) - - assert response.status_code == 400 - assert response.json()["error"]["code"] == "InvalidToolArguments" - - -def test_foundry_brokered_model_loop_rejects_arguments_that_violate_declared_schema(): - spec = _spec(tool_name="check-network-telemetry") - spec.brokered_tools[0].parameters = { - "type": "object", - "properties": {"site": {"type": "string"}}, - "required": ["site"], - "additionalProperties": False, - } +def test_foundry_brokered_model_loop_rejects_decoded_surrogate_arguments(): fake = _FakeChatTransport( [ _chat_response( @@ -2892,29 +5115,28 @@ def test_foundry_brokered_model_loop_rejects_arguments_that_violate_declared_sch { "id": "call_model", "type": "function", - "function": {"name": "check-network-telemetry", "arguments": '{"site":123,"extra":"nope"}'}, + "function": { + "name": "check-network-telemetry", + "arguments": '{"site":"\\ud800"}', + }, } ], } ) ] ) - app = _model_loop_app(spec, fake) + app = _model_loop_app(_spec(tool_name="check-network-telemetry"), fake) with TestClient(app) as client: - response = client.post("/responses", json={"input": "call check-network-telemetry"}) + response = client.post("/responses", json={"input": "check-network-telemetry"}) assert response.status_code == 400 assert response.json()["error"]["code"] == "InvalidToolArguments" - -def test_foundry_brokered_model_loop_validates_large_integer_bounds_exactly(): - spec = _spec(tool_name="check-network-telemetry") - spec.brokered_tools[0].parameters = { - "type": "object", - "properties": {"count": {"type": "integer", "maximum": 9007199254740992}}, - "required": ["count"], - } +def test_foundry_brokered_model_loop_rejects_excessively_nested_arguments(): + nested = '"leaf"' + for _ in range(200): + nested = f"[{nested}]" fake = _FakeChatTransport( [ _chat_response( @@ -2925,29 +5147,25 @@ def test_foundry_brokered_model_loop_validates_large_integer_bounds_exactly(): { "id": "call_model", "type": "function", - "function": {"name": "check-network-telemetry", "arguments": '{"count":9007199254740993}'}, + "function": { + "name": "check-network-telemetry", + "arguments": '{"site":' + nested + "}", + }, } ], } ) ] ) - app = _model_loop_app(spec, fake) + app = _model_loop_app(_spec(tool_name="check-network-telemetry"), fake) with TestClient(app) as client: - response = client.post("/responses", json={"input": "call check-network-telemetry"}) + response = client.post("/responses", json={"input": "check-network-telemetry"}) assert response.status_code == 400 assert response.json()["error"]["code"] == "InvalidToolArguments" - -def test_foundry_brokered_model_loop_rejects_float_arguments_that_cannot_round_trip_exactly(): - spec = _spec(tool_name="check-network-telemetry") - spec.brokered_tools[0].parameters = { - "type": "object", - "properties": {"count": {"type": "integer", "maximum": 9007199254740992}}, - "required": ["count"], - } +def test_foundry_brokered_model_loop_bounds_raw_arguments_before_parsing(monkeypatch): fake = _FakeChatTransport( [ _chat_response( @@ -2958,23 +5176,30 @@ def test_foundry_brokered_model_loop_rejects_float_arguments_that_cannot_round_t { "id": "call_model", "type": "function", - "function": {"name": "check-network-telemetry", "arguments": '{"count":9007199254740993.0}'}, + "function": { + "name": "check-network-telemetry", + "arguments": '{"blob":"' + ("x" * 128) + '"}', + }, } ], } ) ] ) - app = _model_loop_app(spec, fake) - with TestClient(app) as client: - response = client.post("/responses", json={"input": "call check-network-telemetry"}) + def fail_if_parsed(_raw: Any) -> dict[str, Any]: + raise AssertionError("oversized model arguments must be rejected before parsing") - assert response.status_code == 400 - assert response.json()["error"]["code"] == "InvalidToolArguments" + monkeypatch.setattr(foundry_model_loop_module, "_parse_arguments", fail_if_parsed) + app = _model_loop_app(_spec(tool_name="check-network-telemetry"), fake, max_brokered_argument_bytes=64) + with TestClient(app) as client: + response = client.post("/responses", json={"input": "check-network-telemetry"}) -def test_foundry_brokered_model_loop_rejects_unsafe_model_generated_arguments(): + assert response.status_code == 413 + assert response.json()["error"]["code"] == "brokered_arguments_too_large" + +def test_foundry_brokered_model_loop_normalizes_json_parser_failures(): fake = _FakeChatTransport( [ _chat_response( @@ -2985,7 +5210,10 @@ def test_foundry_brokered_model_loop_rejects_unsafe_model_generated_arguments(): { "id": "call_model", "type": "function", - "function": {"name": "check-network-telemetry", "arguments": '{"site":"sfo","tokenValue":"ghp_not_real"}'}, + "function": { + "name": "check-network-telemetry", + "arguments": '{"value":' + ("9" * 5000) + "}", + }, } ], } @@ -2995,123 +5223,7 @@ def test_foundry_brokered_model_loop_rejects_unsafe_model_generated_arguments(): app = _model_loop_app(_spec(tool_name="check-network-telemetry"), fake) with TestClient(app) as client: - response = client.post("/responses", json={"input": "call check-network-telemetry"}) + response = client.post("/responses", json={"input": "check-network-telemetry"}) assert response.status_code == 400 - assert response.json()["error"]["code"] == "UnsafeBrokeredArguments" - - -def test_foundry_brokered_rejects_tool_choice_none_instead_of_ignoring_it(): - app = _app() - - with TestClient(app) as client: - resp = client.post("/responses", json={"input": "hi", "tool_choice": "none"}) - - assert resp.status_code == 400 - assert resp.json()["error"]["code"] == "tool_choice_unsupported" - - -def test_foundry_brokered_rejects_request_supplied_tools_even_when_static_tools_exist(): - app = _app() - - with TestClient(app) as client: - resp = client.post("/responses", json={"input": "hi", "tools": [{"type": "function"}]}) - - assert resp.status_code == 400 - assert resp.json()["error"]["code"] == "tools_unsupported" - - -def _fixture(name: str) -> dict[str, Any]: - path = __import__("pathlib").Path(__file__).parent / "fixtures" / "foundry_brokered" / name - return json.loads(path.read_text(encoding="utf-8")) - - -def _normalize_initial_response(body: dict[str, Any]) -> dict[str, Any]: - normalized = deepcopy(body) - actual_response_id = normalized["id"] - normalized["id"] = "caresp_test" - normalized["created_at"] = 0 - item = normalized["output"][0] - item["id"] = "fc_test" - item["call_id"] = item["call_id"].replace(actual_response_id, "caresp_test") - item["response_id"] = "caresp_test" - return normalized - - -def _normalize_final_response(body: dict[str, Any], *, previous_response_id: str) -> dict[str, Any]: - normalized = deepcopy(body) - actual_response_id = normalized["id"] - normalized["id"] = "caresp_final" - normalized["created_at"] = 0 - normalized["previous_response_id"] = "caresp_test" - item = normalized["output"][0] - item["id"] = "msg_final" - item["response_id"] = "caresp_final" - assert previous_response_id - assert actual_response_id - return normalized - - -def _materialize_continuation(fixture: dict[str, Any], *, response_id: str, call_id: str) -> dict[str, Any]: - materialized = deepcopy(fixture) - materialized["previous_response_id"] = response_id - materialized["input"][0]["call_id"] = call_id - return materialized - - -def test_foundry_brokered_golden_fixtures_pin_function_call_loop_and_errors(): - app = _app() - - with TestClient(app) as client: - initial = client.post("/responses", json=_fixture("initial_request.json")) - assert initial.status_code == 200, initial.text - initial_body = initial.json() - call = _call(initial_body) - assert _normalize_initial_response(initial_body) == _fixture("function_call_response.json") - - continuation = _materialize_continuation( - _fixture("continuation_request.json"), - response_id=initial_body["id"], - call_id=call["call_id"], - ) - final = client.post("/responses", json=continuation, headers=CONTINUATION_AUTH) - assert final.status_code == 200, final.text - assert _normalize_final_response(final.json(), previous_response_id=initial_body["id"]) == _fixture("final_message_response.json") - - unknown_prev = client.post( - "/responses", - headers=CONTINUATION_AUTH, - json=_materialize_continuation( - _fixture("continuation_request.json"), - response_id="caresp_unknown", - call_id=call["call_id"], - ), - ) - assert unknown_prev.status_code == 404 - assert unknown_prev.json() == _fixture("unknown_previous_response_id_error.json") - - app_for_errors = _app() - with TestClient(app_for_errors) as client: - initial_body = _start(client) - call = _call(initial_body) - unknown_call = client.post( - "/responses", - headers=CONTINUATION_AUTH, - json=_materialize_continuation( - _fixture("continuation_request.json"), - response_id=initial_body["id"], - call_id="call_unknown", - ), - ) - assert unknown_call.status_code == 400 - assert unknown_call.json() == _fixture("unknown_call_id_error.json") - - multiple = _materialize_continuation( - _fixture("continuation_request.json"), - response_id=initial_body["id"], - call_id=call["call_id"], - ) - multiple["input"].append(dict(multiple["input"][0])) - multiple_resp = client.post("/responses", json=multiple, headers=CONTINUATION_AUTH) - assert multiple_resp.status_code == 400 - assert multiple_resp.json() == _fixture("multiple_function_calls_unsupported_error.json") + assert response.json()["error"]["code"] == "InvalidToolArguments" diff --git a/runtimes/common/tests/test_foundry_protocol.py b/runtimes/common/tests/test_foundry_protocol.py index e7f18b5..146b4cc 100644 --- a/runtimes/common/tests/test_foundry_protocol.py +++ b/runtimes/common/tests/test_foundry_protocol.py @@ -1,5 +1,7 @@ from __future__ import annotations +import logging + from types import TracebackType from fastapi.testclient import TestClient @@ -7,7 +9,7 @@ from agentkit_serve_common.config import AgentSpec from agentkit_serve_common.conversation import RunRequest from agentkit_serve_common.foundry import create_foundry_app -from agentkit_serve_common.runtime import RunResult, RuntimeSession +from agentkit_serve_common.runtime import AgentRunError, RunResult, RuntimeSession def _spec() -> AgentSpec: @@ -98,22 +100,146 @@ def test_foundry_responses_tolerates_stream_flag_with_non_streaming_response(): assert resp.json()["output"][0]["content"][0]["text"] == "echo: hi" -def test_foundry_invocations_enforces_request_body_limit_before_runtime(): +def test_foundry_non_brokered_routes_reject_oversized_streamed_request_before_runtime(): factory = EchoFactory() app = create_foundry_app(_spec(), factory, max_request_body_bytes=64) + headers = {"content-type": "application/json"} with TestClient(app) as client: - response = client.post( + invocation = client.post( "/invocations", - content='{"message":"' + ("x" * 128) + '"}', + headers=headers, + content=iter([b'{"message":"', b"x" * 128, b'"}']), + ) + response = client.post( + "/responses", + headers=headers, + content=iter([b'{"input":"', b"x" * 128, b'"}']), + ) + + expected = { + "message": "Foundry request body is too large", + "code": "request_too_large", + } + assert invocation.status_code == 413 + assert invocation.json()["error"] == expected + assert response.status_code == 413 + assert response.json()["error"] == expected + assert factory.runtime.requests == [] + + +def test_foundry_non_brokered_request_limit_can_be_configured_from_environment(monkeypatch): + monkeypatch.setenv("AGENTKIT_FOUNDRY_REQUEST_BODY_MAX_BYTES", "64") + factory = EchoFactory() + app = create_foundry_app(_spec(), factory) + + with TestClient(app) as client: + response = client.post( + "/responses", headers={"content-type": "application/json"}, + content=b'{"input":"' + b"x" * 128 + b'"}', ) assert response.status_code == 413 - assert "too large" in response.text + assert response.json()["error"]["code"] == "request_too_large" assert factory.runtime.requests == [] +def test_foundry_non_brokered_runtime_5xx_is_logged_but_public_response_is_sanitized(caplog): + marker = "INTERNAL_RUNTIME_FAILURE_MARKER" + internal_url = "https://private-runtime.internal.example/v1" + + class FailingRuntime(EchoRuntime): + async def run(self, request: RunRequest) -> RunResult: + self.requests.append(request) + raise AgentRunError( + f"provider request to {internal_url} failed: {marker}", + status=503, + code="PrivateProviderError", + ) + + factory = EchoFactory() + factory.runtime = FailingRuntime() + app = create_foundry_app(_spec(), factory) + + with caplog.at_level(logging.WARNING, logger="agentkit_serve_common.foundry"): + with TestClient(app) as client: + invocation = client.post("/invocations", json={"message": "hello"}) + response = client.post("/responses", json={"input": "hello"}) + + expected = { + "message": "agent runtime request failed", + "code": "RuntimeFailure", + } + assert invocation.status_code == 503 + assert invocation.json()["error"] == expected + assert response.status_code == 503 + assert response.json()["error"] == expected + assert marker not in invocation.text + assert marker not in response.text + assert internal_url not in invocation.text + assert internal_url not in response.text + assert marker in caplog.text + assert internal_url in caplog.text + + +def test_foundry_non_brokered_unexpected_runtime_failure_is_sanitized_and_logged(caplog): + marker = "UNEXPECTED_RUNTIME_FAILURE_MARKER" + internal_url = "https://unexpected-runtime.internal.example/v1" + + class FailingRuntime(EchoRuntime): + async def run(self, request: RunRequest) -> RunResult: + self.requests.append(request) + raise RuntimeError(f"connection to {internal_url} failed: {marker}") + + factory = EchoFactory() + factory.runtime = FailingRuntime() + app = create_foundry_app(_spec(), factory) + + with caplog.at_level(logging.ERROR, logger="agentkit_serve_common.foundry"): + with TestClient(app) as client: + invocation = client.post("/invocations", json={"message": "hello"}) + response = client.post("/responses", json={"input": "hello"}) + + expected = { + "message": "agent runtime request failed", + "code": "RuntimeFailure", + } + assert invocation.status_code == 502 + assert invocation.json()["error"] == expected + assert response.status_code == 502 + assert response.json()["error"] == expected + assert marker not in invocation.text + assert marker not in response.text + assert internal_url not in invocation.text + assert internal_url not in response.text + assert marker in caplog.text + assert internal_url in caplog.text + + +def test_foundry_non_brokered_runtime_4xx_preserves_caller_actionable_detail(): + detail = "caller must provide a supported deployment name" + + class FailingRuntime(EchoRuntime): + async def run(self, request: RunRequest) -> RunResult: + self.requests.append(request) + raise AgentRunError(detail, status=422, code="InvalidDeployment") + + factory = EchoFactory() + factory.runtime = FailingRuntime() + app = create_foundry_app(_spec(), factory) + + with TestClient(app) as client: + invocation = client.post("/invocations", json={"message": "hello"}) + response = client.post("/responses", json={"input": "hello"}) + + expected = {"message": detail, "code": "InvalidDeployment"} + assert invocation.status_code == 422 + assert invocation.json()["error"] == expected + assert response.status_code == 422 + assert response.json()["error"] == expected + + def test_foundry_protocols_reject_non_object_json(): app = create_foundry_app(_spec(), EchoFactory()) with TestClient(app) as client: @@ -140,6 +266,63 @@ def test_foundry_protocol_forwards_session_header_and_query_param(): assert factory.runtime.requests[1].session_id == "session-2" +def test_foundry_invocations_prefers_hosted_session_identity_over_caller_fields(monkeypatch): + monkeypatch.setenv("FOUNDRY_AGENT_SESSION_ID", "hosted-session") + factory = EchoFactory() + app = create_foundry_app(_spec(), factory) + + with TestClient(app) as client: + response = client.post( + "/invocations?agent_session_id=query-session", + headers={ + "x-agent-session-id": "gateway-header-session", + "x-agentkit-session-id": "compatibility-header-session", + }, + json={"message": "hello"}, + ) + + assert response.status_code == 200 + assert factory.runtime.requests[0].session_id == "hosted-session" + + +def test_foundry_invocations_preserves_local_query_and_header_fallback(monkeypatch): + monkeypatch.delenv("FOUNDRY_AGENT_SESSION_ID", raising=False) + factory = EchoFactory() + app = create_foundry_app(_spec(), factory) + + with TestClient(app) as client: + query_response = client.post( + "/invocations?session_id=query-session", + headers={ + "x-agent-session-id": "gateway-header-session", + "x-agentkit-session-id": "compatibility-header-session", + }, + json={"message": "query"}, + ) + gateway_header_response = client.post( + "/invocations", + headers={ + "x-agent-session-id": "gateway-header-session", + "x-agentkit-session-id": "compatibility-header-session", + }, + json={"message": "gateway header"}, + ) + compatibility_header_response = client.post( + "/invocations", + headers={"x-agentkit-session-id": "compatibility-header-session"}, + json={"message": "compatibility header"}, + ) + + assert query_response.status_code == 200 + assert gateway_header_response.status_code == 200 + assert compatibility_header_response.status_code == 200 + assert [request.session_id for request in factory.runtime.requests] == [ + "query-session", + "gateway-header-session", + "compatibility-header-session", + ] + + def test_foundry_responses_prefers_body_session_id_over_local_query_and_header(monkeypatch): monkeypatch.delenv("FOUNDRY_AGENT_SESSION_ID", raising=False) factory = EchoFactory() @@ -407,3 +590,21 @@ def test_foundry_brokered_cli_rejects_agents_without_brokered_tools(tmp_path): assert "brokeredTools" in str(exc) else: # pragma: no cover - assertion path. raise AssertionError("expected missing brokeredTools to fail") + + +# Regression coverage retained from the merged brokered-continuation stack. + +def test_foundry_invocations_enforces_request_body_limit_before_runtime(): + factory = EchoFactory() + app = create_foundry_app(_spec(), factory, max_request_body_bytes=64) + + with TestClient(app) as client: + response = client.post( + "/invocations", + content='{"message":"' + ("x" * 128) + '"}', + headers={"content-type": "application/json"}, + ) + + assert response.status_code == 413 + assert "too large" in response.text + assert factory.runtime.requests == [] diff --git a/runtimes/common/tests/test_foundry_transcript_verifier.py b/runtimes/common/tests/test_foundry_transcript_verifier.py index 1f7932e..372bdfb 100644 --- a/runtimes/common/tests/test_foundry_transcript_verifier.py +++ b/runtimes/common/tests/test_foundry_transcript_verifier.py @@ -45,6 +45,8 @@ def _write_transcript(tmp_path: Path) -> Path: } ], } + if initial_response.get("agent_session_id") is not None: + continuation_request["agent_session_id"] = initial_response["agent_session_id"] continuation_response = client.post("/responses", json=continuation_request).json() files = { @@ -362,3 +364,36 @@ def build_runtime(self, spec): # noqa: ANN001 assert summary["call_id"].startswith("call_") assert summary["arguments"] == {"probe": True} + + +def test_verify_brokered_transcript_requires_returned_agent_session_id_on_continuation(tmp_path): + verifier = _load_verifier() + transcript = _write_transcript(tmp_path) + initial = json.loads((transcript / "02-initial-response.json").read_text(encoding="utf-8")) + initial["agent_session_id"] = "gateway-session" + (transcript / "02-initial-response.json").write_text(json.dumps(initial), encoding="utf-8") + continuation = json.loads((transcript / "03-continuation-request.json").read_text(encoding="utf-8")) + continuation.pop("agent_session_id", None) + (transcript / "03-continuation-request.json").write_text(json.dumps(continuation), encoding="utf-8") + + try: + verifier.verify_transcript(transcript) + except ValueError as exc: + assert "agent_session_id" in str(exc) + else: # pragma: no cover - assertion path. + raise AssertionError("expected a missing continuation agent_session_id to fail") + + +def test_verify_brokered_transcript_rejects_archived_continuation_proof(tmp_path): + verifier = _load_verifier() + transcript = _write_transcript(tmp_path) + continuation = json.loads((transcript / "03-continuation-request.json").read_text(encoding="utf-8")) + continuation["brokered_continuation_proof"] = "proof-must-not-be-archived" + (transcript / "03-continuation-request.json").write_text(json.dumps(continuation), encoding="utf-8") + + try: + verifier.verify_transcript(transcript) + except ValueError as exc: + assert "continuation proof" in str(exc) + else: # pragma: no cover - assertion path. + raise AssertionError("expected an archived continuation proof to fail") diff --git a/runtimes/common/tests/test_orka_protocol.py b/runtimes/common/tests/test_orka_protocol.py index f8c4c6d..0886e49 100644 --- a/runtimes/common/tests/test_orka_protocol.py +++ b/runtimes/common/tests/test_orka_protocol.py @@ -1,6 +1,8 @@ from __future__ import annotations +import asyncio import json +import threading import time from datetime import UTC, datetime, timedelta from types import TracebackType @@ -9,12 +11,15 @@ import pytest from fastapi.testclient import TestClient +import agentkit_serve_common.orka as orka_module from agentkit_serve_common.config import AgentSpec from agentkit_serve_common.conversation import RunRequest from agentkit_serve_common.orka import ORKA_HARNESS_VERSION, create_orka_app from agentkit_serve_common.runtime import ( + AgentRunError, BrokeredToolCall, BrokeredToolDefinition, + BrokeredToolResult, OfflineEchoRuntimeFactory, RunResult, RuntimeSession, @@ -23,6 +28,8 @@ AUTH = {"authorization": "Bearer test-token"} TERMINAL_TYPES = {"TurnCompleted", "TurnFailed", "TurnCancelled"} +EXPECTED_MAX_OUTPUT_BYTES = 512 * 1024 +ORKA_CLIENT_MAX_SSE_TOKEN_BYTES = 1 << 20 def _deadline() -> str: @@ -82,6 +89,43 @@ def build_runtime(self, spec: AgentSpec) -> RuntimeSession: return self.runtime +class StaticOutputRuntime: + def __init__(self, text: str) -> None: + self.text = text + + async def __aenter__(self) -> RuntimeSession: + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> bool | None: + return None + + async def run(self, request: RunRequest) -> RunResult: + return RunResult(text=self.text) + + +class StaticOutputFactory: + def __init__(self, text: str) -> None: + self.runtime = StaticOutputRuntime(text) + + def build_runtime(self, spec: AgentSpec) -> RuntimeSession: + return self.runtime + + +class RaisingRuntime(StaticOutputRuntime): + async def run(self, request: RunRequest) -> RunResult: + raise RuntimeError(self.text) + + +class RaisingFactory(StaticOutputFactory): + def __init__(self, message: str) -> None: + self.runtime = RaisingRuntime(message) + + def _start_payload(**overrides: Any) -> dict[str, Any]: payload: dict[str, Any] = { "version": ORKA_HARNESS_VERSION, @@ -198,6 +242,17 @@ def _wait_for_event_type(client: TestClient, turn_id: str, event_type: str) -> d raise AssertionError(f"timed out waiting for {event_type}") +def _wait_for_tool_call_id(client: TestClient, turn_id: str, tool_call_id: str) -> dict[str, Any]: + for _ in range(100): + events = client.app.state.turns[turn_id].events + for event in events: + frame = event.as_frame() + if frame["type"] == "ToolCallRequested" and frame["toolCallID"] == tool_call_id: + return frame + time.sleep(0.01) + raise AssertionError(f"timed out waiting for tool call {tool_call_id}") + + def _frames(resp_text: str) -> list[dict[str, Any]]: frames: list[dict[str, Any]] = [] for raw in resp_text.strip().split("\n\n"): @@ -209,6 +264,26 @@ def _frames(resp_text: str) -> list[dict[str, Any]]: return frames +def _assert_sse_lines_fit_orka_client(raw: bytes) -> None: + data_lines = [line for line in raw.splitlines() if line.startswith(b"data: ")] + assert data_lines + assert max(map(len, data_lines)) < ORKA_CLIENT_MAX_SSE_TOKEN_BYTES + + +def _compact_json_bytes(value: Any) -> bytes: + return json.dumps(value, ensure_ascii=False, separators=(",", ":"), sort_keys=True).encode() + + +def _json_object_with_size(size: int) -> dict[str, str]: + empty = _compact_json_bytes({"value": ""}) + remaining = size - len(empty) + assert remaining >= 0 + value = "é" * (remaining // len("é".encode())) + "x" * (remaining % len("é".encode())) + output = {"value": value} + assert len(_compact_json_bytes(output)) == size + return output + + def _create_turn(client: TestClient, **overrides: Any) -> str: payload = _start_payload(**overrides) resp = client.post("/v1/turns", json=payload, headers=AUTH) @@ -274,6 +349,7 @@ def test_orka_health_and_capabilities_are_open_and_match_contract(): "supportsSuspend": False, "supportsWorkspaceSnapshot": False, "maxConcurrentTurns": 1, + "maxOutputBytes": EXPECTED_MAX_OUTPUT_BYTES, "metadata": {"agentName": "orka-test", "model": "gpt-4o-mini", "agentkitProvider": "openai-compatible"}, } @@ -292,7 +368,6 @@ def test_orka_turn_lifecycle_streams_contract_frames_and_one_terminal(): _assert_frame_identity(frames[1], seq=2, typ="RuntimeOutput") assert frames[1]["contentText"] == "echo: hello" assert frames[1]["content"] == { - "message": "echo: hello", "usage": {"completion_tokens": 2, "prompt_tokens": 1, "total_tokens": 3}, } _assert_frame_identity(frames[2], seq=3, typ="TurnCompleted") @@ -316,6 +391,174 @@ def test_orka_events_support_after_seq_replay(): assert [frame["type"] for frame in _frames(replay.text)] == ["RuntimeOutput", "TurnCompleted"] +def test_orka_observed_output_uses_utf8_bytes_accepts_exact_limit_and_replays_safely(): + output = "é" * (EXPECTED_MAX_OUTPUT_BYTES // len("é".encode())) + assert len(output.encode()) == EXPECTED_MAX_OUTPUT_BYTES + app = create_orka_app(_spec(), StaticOutputFactory(output), AUTH["authorization"].removeprefix("Bearer ")) + + with TestClient(app) as client: + turn_id = _create_turn(client, turnID="turn-output-boundary") + response = client.get(f"/v1/turns/{turn_id}/events", headers=AUTH) + replay = client.get(f"/v1/turns/{turn_id}/events?afterSeq=1", headers=AUTH) + + frames = _frames(response.text) + assert [frame["type"] for frame in frames] == ["TurnStarted", "RuntimeOutput", "TurnCompleted"] + assert frames[1]["contentText"] == output + assert frames[1]["content"] == {"usage": {}} + assert frames[2]["completed"]["result"] == output + assert [frame["type"] for frame in _frames(replay.text)] == ["RuntimeOutput", "TurnCompleted"] + _assert_sse_lines_fit_orka_client(response.content) + _assert_sse_lines_fit_orka_client(replay.content) + + +def test_orka_observed_output_over_utf8_limit_fails_without_retaining_payload_and_replays_terminal(): + output = "é" * (EXPECTED_MAX_OUTPUT_BYTES // len("é".encode())) + "x" + output_bytes = len(output.encode()) + assert output_bytes == EXPECTED_MAX_OUTPUT_BYTES + 1 + app = create_orka_app(_spec(), StaticOutputFactory(output), AUTH["authorization"].removeprefix("Bearer ")) + + with TestClient(app) as client: + turn_id = _create_turn(client, turnID="turn-output-over-limit") + response = client.get(f"/v1/turns/{turn_id}/events", headers=AUTH) + replay = client.get(f"/v1/turns/{turn_id}/events?afterSeq=1", headers=AUTH) + retained = client.app.state.turns[turn_id].events + + frames = _frames(response.text) + assert [frame["type"] for frame in frames] == ["TurnStarted", "TurnFailed"] + assert frames[-1]["failed"] == { + "reason": "MaxOutputBytesExceeded", + "message": ( + f"runtime output is {output_bytes} UTF-8 bytes; " + f"maxOutputBytes is {EXPECTED_MAX_OUTPUT_BYTES}" + ), + "retryable": False, + } + assert [frame["type"] for frame in _frames(replay.text)] == ["TurnFailed"] + assert all(event.content_text != output for event in retained) + assert all(event.completed is None or event.completed.get("result") != output for event in retained) + _assert_sse_lines_fit_orka_client(response.content) + _assert_sse_lines_fit_orka_client(replay.content) + + +def test_orka_observed_output_that_expands_past_scanner_limit_fails_with_visible_terminal(): + output = "\x00" * 200_000 + assert len(output.encode()) < EXPECTED_MAX_OUTPUT_BYTES + app = create_orka_app(_spec(), StaticOutputFactory(output), AUTH["authorization"].removeprefix("Bearer ")) + + with TestClient(app) as client: + turn_id = _create_turn(client, turnID="turn-output-json-expansion") + response = client.get(f"/v1/turns/{turn_id}/events", headers=AUTH) + + frames = _frames(response.text) + assert [frame["type"] for frame in frames] == ["TurnStarted", "TurnFailed"] + assert frames[-1]["failed"]["reason"] == "HarnessFrameTooLarge" + assert "Orka client limit is 1048575" in frames[-1]["failed"]["message"] + _assert_sse_lines_fit_orka_client(response.content) + + +def test_orka_observed_output_with_unpaired_surrogate_fails_with_visible_terminal(): + app = create_orka_app( + _spec(), + StaticOutputFactory("\ud800"), + AUTH["authorization"].removeprefix("Bearer "), + ) + + with TestClient(app) as client: + turn_id = _create_turn(client, turnID="turn-output-invalid-utf8") + response = client.get(f"/v1/turns/{turn_id}/events", headers=AUTH) + + frames = _frames(response.text) + assert [frame["type"] for frame in frames] == ["TurnStarted", "TurnFailed"] + assert frames[-1]["failed"] == { + "reason": "InvalidOutputEncoding", + "message": "runtime output is not valid UTF-8", + "retryable": False, + } + _assert_sse_lines_fit_orka_client(response.content) + + +@pytest.mark.parametrize( + "message", + [pytest.param("\x00" * 200_000, id="oversized"), pytest.param("\ud800", id="invalid-utf8")], +) +def test_orka_runtime_failure_with_unstreamable_detail_uses_bounded_terminal_fallback(message: str): + app = create_orka_app( + _spec(), + RaisingFactory(message), + AUTH["authorization"].removeprefix("Bearer "), + ) + + with TestClient(app) as client: + turn_id = _create_turn(client, turnID="turn-unstreamable-failure-detail") + response = client.get(f"/v1/turns/{turn_id}/events", headers=AUTH) + + frames = _frames(response.text) + assert [frame["type"] for frame in frames] == ["TurnStarted", "TurnFailed"] + assert frames[-1]["failed"] == { + "reason": "TerminalFrameRejected", + "message": "terminal failure details could not be emitted safely", + "retryable": False, + } + _assert_sse_lines_fit_orka_client(response.content) + + +def test_orka_turn_state_rejects_nonterminal_events_after_terminal(): + async def exercise() -> list[str]: + state = orka_module.TurnState( + runtime_session_id="runtime-session-1", + turn_id="turn-terminal-finality", + correlation_id="corr-1", + ) + await state.append("TurnStarted", summary="turn started") + await state.append( + "TurnFailed", + summary="turn failed", + failed={"reason": "TestFailure", "message": "failed", "retryable": False}, + ) + with pytest.raises(orka_module.AgentRunError, match="already terminal"): + await state.append("RuntimeOutput", summary="late output", content_text="late") + return [event.type for event in state.events] + + assert asyncio.run(exercise()) == ["TurnStarted", "TurnFailed"] + + +def test_orka_rejects_an_unstreamable_start_frame_without_retaining_a_poisoned_turn(): + app = create_orka_app(_spec(), EchoFactory(), AUTH["authorization"].removeprefix("Bearer ")) + + with TestClient(app, raise_server_exceptions=False) as client: + rejected = client.post( + "/v1/turns", + json=_start_payload(turnID="turn-unstreamable-start", metadata={"large": "\x00" * 200_000}), + headers=AUTH, + ) + accepted = client.post( + "/v1/turns", + json=_start_payload(turnID="turn-after-unstreamable-start"), + headers=AUTH, + ) + + assert rejected.status_code == 413 + assert "TurnStarted SSE data line" in rejected.text + assert accepted.status_code == 202 + assert "turn-unstreamable-start" not in app.state.turns + + +def test_orka_rejects_start_frame_text_that_is_not_valid_utf8_without_retaining_turn(): + app = create_orka_app(_spec(), EchoFactory(), AUTH["authorization"].removeprefix("Bearer ")) + payload = _start_payload(turnID="turn-invalid-utf8-start", metadata={"invalid": "\ud800"}) + + with TestClient(app, raise_server_exceptions=False) as client: + rejected = client.post( + "/v1/turns", + content=json.dumps(payload, ensure_ascii=True).encode(), + headers={**AUTH, "content-type": "application/json"}, + ) + + assert rejected.status_code == 400 + assert rejected.json() == {"detail": "TurnStarted contains text that is not valid UTF-8"} + assert "turn-invalid-utf8-start" not in app.state.turns + + def test_orka_duplicate_turn_rejection_matches_orka_conformance_contract(): app = create_orka_app(_spec(), EchoFactory(delay=60), auth_token="test-token") @@ -437,16 +680,49 @@ def test_orka_protected_endpoints_require_bearer_token(): assert cancel.status_code == 401 -def test_orka_rejects_turn_ids_that_do_not_fit_route_path(): +@pytest.mark.parametrize("turn_id", ["", " ", ".", "..", " turn", "turn ", "turn/one", r"turn\one"]) +def test_orka_rejects_turn_ids_that_are_not_trimmed_single_path_segments(turn_id: str): app = create_orka_app(_spec(), EchoFactory(), auth_token="test-token") with TestClient(app) as client: - slash = client.post("/v1/turns", json=_start_payload(turnID="bad/id"), headers=AUTH) - query = client.post("/v1/turns", json=_start_payload(turnID="bad?id"), headers=AUTH) + response = client.post("/v1/turns", json=_start_payload(turnID=turn_id), headers=AUTH) + + assert response.status_code == 400 + assert "turnID" in response.text + + +@pytest.mark.parametrize( + ("turn_id", "escaped_turn_id"), + [ + pytest.param("turn:one", "turn:one", id="colon"), + pytest.param("turn one", "turn%20one", id="space"), + pytest.param("türn-雪", "t%C3%BCrn-%E9%9B%AA", id="unicode"), + pytest.param("turn$&+=@", "turn$&+=@", id="orka-path-safe-reserved"), + pytest.param("turn,one", "turn%2Cone", id="escaped-reserved"), + pytest.param("turn?one", "turn%3Fone", id="query-delimiter"), + pytest.param("\x1cturn\x1c", "%1Cturn%1C", id="go-non-space-control"), + pytest.param("t" * 1024, "t" * 1024, id="long-segment"), + ], +) +def test_orka_accepts_and_path_escapes_valid_turn_segments_exactly_like_orka(turn_id: str, escaped_turn_id: str): + app = create_orka_app(_spec(), EchoFactory(), auth_token=AUTH["authorization"].removeprefix("Bearer ")) - assert slash.status_code == 400 - assert query.status_code == 400 - assert "URL-safe" in slash.text + with TestClient(app) as client: + response = client.post("/v1/turns", json=_start_payload(turnID=turn_id), headers=AUTH) + if response.status_code == 202: + event_stream_path = response.json()["eventStreamPath"] + events = client.get(event_stream_path, headers=AUTH) + else: + event_stream_path = "" + events = None + + assert response.status_code == 202, response.text + assert event_stream_path == f"/v1/turns/{escaped_turn_id}/events" + assert events is not None + assert events.status_code == 200 + frames = _frames(events.text) + assert frames[0]["turnID"] == turn_id + assert frames[-1]["type"] == "TurnCompleted" def test_orka_cancel_accepts_contract_request_and_produces_cancelled_terminal_frame(): @@ -498,6 +774,39 @@ def test_orka_cancel_rejects_runtime_session_or_correlation_mismatch(): assert "correlationID" in wrong_correlation.text +@pytest.mark.parametrize( + ("field_name", "mismatched_value"), + [ + ("namespace", "other-namespace"), + ("taskName", "other-task"), + ("sessionName", "other-session-name"), + ], +) +def test_orka_cancel_rejects_turn_owner_mismatch_without_cancelling_turn(field_name: str, mismatched_value: str): + app = create_orka_app(_spec(), EchoFactory(delay=60), auth_token="test-token") + + with TestClient(app) as client: + turn_id = _create_turn(client, turnID=f"turn-cancel-{field_name}", input={"prompt": "slow", "contextRefs": [], "env": []}) + state = client.app.state.turns[turn_id] + mismatch = client.post( + f"/v1/turns/{turn_id}/cancel", + json=_cancel_payload(turnID=turn_id, **{field_name: mismatched_value}), + headers=AUTH, + ) + + assert mismatch.status_code == 400 + assert mismatch.json() == {"detail": "cancel namespace/taskName/sessionName must match turn"} + assert state.task is not None + assert state.task.cancelling() == 0 + assert state.terminal_event is None + + accepted = client.post(f"/v1/turns/{turn_id}/cancel", json=_cancel_payload(turnID=turn_id), headers=AUTH) + frames = _frames(client.get(f"/v1/turns/{turn_id}/events", headers=AUTH).text) + + assert accepted.status_code == 202 + assert frames[-1]["type"] == "TurnCancelled" + + def test_orka_cancel_rejects_body_turn_id_mismatch(): app = create_orka_app(_spec(), EchoFactory(delay=60), auth_token="test-token") @@ -555,6 +864,70 @@ def test_orka_terminal_turn_retention_is_bounded(): assert [frame["type"] for frame in _frames(kept.text)] == ["RuntimeOutput", "TurnCompleted"] +def test_orka_lifespan_awaits_cancelled_turn_and_terminal_callback_before_runtime_close(monkeypatch): + turn_started = threading.Event() + + class ShutdownRuntime: + def __init__(self) -> None: + self.state: Any = None + self.close_observations: list[tuple[bool, str | None]] = [] + + async def __aenter__(self) -> RuntimeSession: + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> bool | None: + terminal_type = self.state.terminal_event.type if self.state.terminal_event is not None else None + self.close_observations.append((self.state.task.done(), terminal_type)) + return None + + async def run(self, request: RunRequest) -> RunResult: + raise AssertionError("patched turn runner should own the active task") + + class ShutdownFactory: + def __init__(self) -> None: + self.runtime = ShutdownRuntime() + + def build_runtime(self, spec: AgentSpec) -> RuntimeSession: + return self.runtime + + async def uncaught_cancel_run_turn( + get_runtime, + turns, + terminal_order, + state, + run_request, + *, + max_terminal_turns, + brokered_tools=None, + ) -> None: + del turns, terminal_order, state, max_terminal_turns, brokered_tools + await get_runtime(run_request) + turn_started.set() + await asyncio.Event().wait() + + monkeypatch.setattr(orka_module, "_run_turn", uncaught_cancel_run_turn) + factory = ShutdownFactory() + app = create_orka_app(_spec(), factory, auth_token="test-token") + + with TestClient(app) as client: + turn_id = _create_turn(client, turnID="turn-shutdown", input={"prompt": "slow", "contextRefs": [], "env": []}) + assert turn_started.wait(timeout=2) + state = client.app.state.turns[turn_id] + factory.runtime.state = state + assert state.task is not None + assert not state.task.done() + + assert factory.runtime.close_observations == [(True, "TurnCancelled")] + assert state.task.done() + assert state.terminal_event is not None + assert state.terminal_event.type == "TurnCancelled" + + class EnvRuntime: def __init__(self, token: str) -> None: self.token = token @@ -595,6 +968,81 @@ def build_runtime(self, spec: AgentSpec) -> RuntimeSession: return runtime +class SlowCloseEnvRuntime(EnvRuntime): + def __init__( + self, + token: str, + *, + close_delay: float, + close_error: BaseException | None = None, + close_release: threading.Event | None = None, + ) -> None: + super().__init__(token) + self.close_delay = close_delay + self.close_error = close_error + self.close_release = close_release + self.close_calls = 0 + self.close_completed = 0 + self.close_cancelled = 0 + self.close_failed = 0 + self.close_started = threading.Event() + self.close_finished = threading.Event() + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> bool | None: + self.close_calls += 1 + self.close_started.set() + try: + if self.close_release is not None: + await asyncio.to_thread(self.close_release.wait) + else: + await asyncio.sleep(self.close_delay) + except asyncio.CancelledError: + self.close_cancelled += 1 + raise + if self.close_error is not None: + self.close_failed += 1 + self.close_finished.set() + raise self.close_error + self.close_completed += 1 + self.exited += 1 + self.close_finished.set() + return None + + +class SlowFirstCloseEnvFactory: + def __init__( + self, + *, + close_delay: float = 0.25, + close_error: BaseException | None = None, + close_release: threading.Event | None = None, + ) -> None: + self.close_delay = close_delay + self.close_error = close_error + self.close_release = close_release + self.runtimes: list[SlowCloseEnvRuntime] = [] + + def build_runtime(self, spec: AgentSpec) -> RuntimeSession: + import os + + delay = self.close_delay if not self.runtimes else 0 + close_error = self.close_error if not self.runtimes else None + close_release = self.close_release if not self.runtimes else None + runtime = SlowCloseEnvRuntime( + os.environ["MODEL_TOKEN"], + close_delay=delay, + close_error=close_error, + close_release=close_release, + ) + self.runtimes.append(runtime) + return runtime + + def test_orka_runtime_build_is_deferred_until_turn_env_is_available(monkeypatch): monkeypatch.delenv("MODEL_TOKEN", raising=False) factory = EnvFactory() @@ -649,6 +1097,283 @@ def test_orka_runtime_session_cache_is_bounded_and_closes_evicted_runtime(monkey assert factory.runtimes[1].exited == 1 +def test_orka_capacity_eviction_cleanup_survives_turn_deadline_and_blocks_new_runtime(monkeypatch): + monkeypatch.delenv("MODEL_TOKEN", raising=False) + close_release = threading.Event() + capacity_wait_started = threading.Event() + original_asyncio_wait = asyncio.wait + + async def observed_asyncio_wait(fs, *, timeout=None, return_when=asyncio.ALL_COMPLETED): + waitables = tuple(fs) + if return_when == asyncio.FIRST_COMPLETED and any( + isinstance(item, asyncio.Task) and item.get_name() == "agentkit-orka-runtime-close" for item in waitables + ): + capacity_wait_started.set() + return await original_asyncio_wait(waitables, timeout=timeout, return_when=return_when) + + monkeypatch.setattr(orka_module.asyncio, "wait", observed_asyncio_wait) + factory = SlowFirstCloseEnvFactory(close_release=close_release) + app = create_orka_app(_spec(), factory, auth_token=AUTH["authorization"].removeprefix("Bearer "), max_runtime_sessions=1) + + with TestClient(app) as client: + try: + first_id = _create_turn( + client, + turnID="turn-session-one", + runtimeSessionID="runtime-session-one", + correlationID="corr-one", + input={"prompt": "one", "contextRefs": [], "env": [{"name": "MODEL_TOKEN", "value": "one-token"}]}, + ) + first_frames = _frames(client.get(f"/v1/turns/{first_id}/events", headers=AUTH).text) + assert first_frames[-1]["type"] == "TurnCompleted" + + second_deadline = (datetime.now(UTC) + timedelta(milliseconds=500)).isoformat().replace("+00:00", "Z") + second_id = _create_turn( + client, + turnID="turn-session-two", + runtimeSessionID="runtime-session-two", + correlationID="corr-two", + deadline=second_deadline, + input={"prompt": "two", "contextRefs": [], "env": [{"name": "MODEL_TOKEN", "value": "two-token"}]}, + ) + second_frames = _frames(client.get(f"/v1/turns/{second_id}/events", headers=AUTH).text) + assert second_frames[-1]["type"] == "TurnFailed" + assert second_frames[-1]["failed"]["reason"] == "DeadlineExceeded" + + third_deadline = (datetime.now(UTC) + timedelta(milliseconds=500)).isoformat().replace("+00:00", "Z") + third_id = _create_turn( + client, + turnID="turn-session-three", + runtimeSessionID="runtime-session-three", + correlationID="corr-three", + deadline=third_deadline, + input={"prompt": "three", "contextRefs": [], "env": [{"name": "MODEL_TOKEN", "value": "three-token"}]}, + ) + third_frames = _frames(client.get(f"/v1/turns/{third_id}/events", headers=AUTH).text) + assert third_frames[-1]["type"] == "TurnFailed" + assert third_frames[-1]["failed"]["reason"] == "DeadlineExceeded" + assert len(factory.runtimes) == 1 + + capacity_wait_started.clear() + fourth_id = _create_turn( + client, + turnID="turn-session-four", + runtimeSessionID="runtime-session-four", + correlationID="corr-four", + input={"prompt": "four", "contextRefs": [], "env": [{"name": "MODEL_TOKEN", "value": "four-token"}]}, + ) + assert capacity_wait_started.wait(timeout=2) + close_release.set() + fourth_frames = _frames(client.get(f"/v1/turns/{fourth_id}/events", headers=AUTH).text) + assert fourth_frames[-1]["type"] == "TurnCompleted" + finally: + close_release.set() + assert factory.runtimes[0].close_finished.wait(timeout=2) + + assert len(factory.runtimes) == 2 + assert factory.runtimes[0].close_calls == 1 + assert factory.runtimes[0].close_cancelled == 0 + assert factory.runtimes[0].close_completed == 1 + assert factory.runtimes[1].close_calls == 1 + assert factory.runtimes[1].close_completed == 1 + + +def test_orka_env_replacement_cleanup_survives_turn_cancellation(monkeypatch): + monkeypatch.delenv("MODEL_TOKEN", raising=False) + close_release = threading.Event() + factory = SlowFirstCloseEnvFactory(close_release=close_release) + app = create_orka_app(_spec(), factory, auth_token=AUTH["authorization"].removeprefix("Bearer "), max_runtime_sessions=2) + + with TestClient(app) as client: + try: + first_id = _create_turn( + client, + turnID="turn-env-one", + runtimeSessionID="runtime-session-rotating", + correlationID="corr-one", + input={"prompt": "one", "contextRefs": [], "env": [{"name": "MODEL_TOKEN", "value": "one-token"}]}, + ) + first_frames = _frames(client.get(f"/v1/turns/{first_id}/events", headers=AUTH).text) + assert first_frames[-1]["type"] == "TurnCompleted" + + second_id = _create_turn( + client, + turnID="turn-env-two", + runtimeSessionID="runtime-session-rotating", + correlationID="corr-two", + input={"prompt": "two", "contextRefs": [], "env": [{"name": "MODEL_TOKEN", "value": "two-token"}]}, + ) + assert factory.runtimes[0].close_started.wait(timeout=2) + cancel = client.post( + f"/v1/turns/{second_id}/cancel", + json=_cancel_payload( + turnID=second_id, + runtimeSessionID="runtime-session-rotating", + correlationID="corr-two", + ), + headers=AUTH, + ) + second_frames = _frames(client.get(f"/v1/turns/{second_id}/events", headers=AUTH).text) + + assert cancel.status_code == 202 + assert second_frames[-1]["type"] == "TurnCancelled" + assert not factory.runtimes[0].close_finished.is_set() + + third_deadline = (datetime.now(UTC) + timedelta(milliseconds=500)).isoformat().replace("+00:00", "Z") + third_id = _create_turn( + client, + turnID="turn-env-three", + runtimeSessionID="runtime-session-rotating", + correlationID="corr-three", + deadline=third_deadline, + input={"prompt": "three", "contextRefs": [], "env": [{"name": "MODEL_TOKEN", "value": "two-token"}]}, + ) + third_frames = _frames(client.get(f"/v1/turns/{third_id}/events", headers=AUTH).text) + assert third_frames[-1]["type"] == "TurnFailed" + assert third_frames[-1]["failed"]["reason"] == "DeadlineExceeded" + assert len(factory.runtimes) == 1 + finally: + close_release.set() + assert factory.runtimes[0].close_finished.wait(timeout=2) + + fourth_id = _create_turn( + client, + turnID="turn-env-four", + runtimeSessionID="runtime-session-rotating", + correlationID="corr-four", + input={"prompt": "four", "contextRefs": [], "env": [{"name": "MODEL_TOKEN", "value": "two-token"}]}, + ) + fourth_frames = _frames(client.get(f"/v1/turns/{fourth_id}/events", headers=AUTH).text) + assert fourth_frames[-1]["type"] == "TurnCompleted" + + assert len(factory.runtimes) == 2 + assert factory.runtimes[0].close_calls == 1 + assert factory.runtimes[0].close_cancelled == 0 + assert factory.runtimes[0].close_completed == 1 + assert factory.runtimes[0].close_finished.is_set() + assert factory.runtimes[1].close_calls == 1 + assert factory.runtimes[1].close_completed == 1 + + +def test_orka_capacity_eviction_close_failure_fails_uncancelled_turn_without_double_close(monkeypatch): + monkeypatch.delenv("MODEL_TOKEN", raising=False) + factory = SlowFirstCloseEnvFactory(close_delay=0, close_error=RuntimeError("runtime close failed")) + app = create_orka_app(_spec(), factory, auth_token=AUTH["authorization"].removeprefix("Bearer "), max_runtime_sessions=1) + + with pytest.raises(AgentRunError, match="runtime cleanup failed"): + with TestClient(app) as client: + first_id = _create_turn( + client, + turnID="turn-close-failure-one", + runtimeSessionID="runtime-session-one", + correlationID="corr-one", + input={"prompt": "one", "contextRefs": [], "env": [{"name": "MODEL_TOKEN", "value": "one-token"}]}, + ) + first_frames = _frames(client.get(f"/v1/turns/{first_id}/events", headers=AUTH).text) + assert first_frames[-1]["type"] == "TurnCompleted" + + second_id = _create_turn( + client, + turnID="turn-close-failure-two", + runtimeSessionID="runtime-session-two", + correlationID="corr-two", + input={"prompt": "two", "contextRefs": [], "env": [{"name": "MODEL_TOKEN", "value": "two-token"}]}, + ) + second_frames = _frames(client.get(f"/v1/turns/{second_id}/events", headers=AUTH).text) + + assert second_frames[-1]["type"] == "TurnFailed" + assert second_frames[-1]["failed"] == { + "reason": "RuntimeCloseFailed", + "message": "runtime cleanup failed; restart required before opening another runtime session", + "retryable": False, + } + + third_id = _create_turn( + client, + turnID="turn-after-close-failure", + runtimeSessionID="runtime-session-three", + correlationID="corr-three", + input={"prompt": "three", "contextRefs": [], "env": [{"name": "MODEL_TOKEN", "value": "three-token"}]}, + ) + third_frames = _frames(client.get(f"/v1/turns/{third_id}/events", headers=AUTH).text) + assert third_frames[-1]["type"] == "TurnFailed" + assert third_frames[-1]["failed"] == { + "reason": "RuntimeCloseFailed", + "message": "runtime cleanup failed; restart required before opening another runtime session", + "retryable": False, + } + assert len(factory.runtimes) == 1 + + assert len(factory.runtimes) == 1 + assert factory.runtimes[0].close_calls == 1 + assert factory.runtimes[0].close_failed == 1 + assert factory.runtimes[0].close_completed == 0 + + +def test_orka_shutdown_closes_active_runtime_after_orphaned_close_failure(monkeypatch): + monkeypatch.delenv("MODEL_TOKEN", raising=False) + factory = SlowFirstCloseEnvFactory(close_delay=1.0, close_error=RuntimeError("orphaned close failed")) + app = create_orka_app(_spec(), factory, auth_token=AUTH["authorization"].removeprefix("Bearer "), max_runtime_sessions=2) + + with pytest.raises(AgentRunError, match="runtime cleanup failed"): + with TestClient(app) as client: + first_id = _create_turn( + client, + turnID="turn-orphan-failure-one", + runtimeSessionID="runtime-session-one", + correlationID="corr-one", + input={"prompt": "one", "contextRefs": [], "env": [{"name": "MODEL_TOKEN", "value": "one-token"}]}, + ) + first_frames = _frames(client.get(f"/v1/turns/{first_id}/events", headers=AUTH).text) + assert first_frames[-1]["type"] == "TurnCompleted" + + second_id = _create_turn( + client, + turnID="turn-orphan-failure-two", + runtimeSessionID="runtime-session-two", + correlationID="corr-two", + input={"prompt": "two", "contextRefs": [], "env": [{"name": "MODEL_TOKEN", "value": "two-token"}]}, + ) + second_frames = _frames(client.get(f"/v1/turns/{second_id}/events", headers=AUTH).text) + assert second_frames[-1]["type"] == "TurnCompleted" + + short_deadline = (datetime.now(UTC) + timedelta(milliseconds=500)).isoformat().replace("+00:00", "Z") + third_id = _create_turn( + client, + turnID="turn-orphan-failure-three", + runtimeSessionID="runtime-session-three", + correlationID="corr-three", + deadline=short_deadline, + input={"prompt": "three", "contextRefs": [], "env": [{"name": "MODEL_TOKEN", "value": "three-token"}]}, + ) + third_frames = _frames(client.get(f"/v1/turns/{third_id}/events", headers=AUTH).text) + assert third_frames[-1]["type"] == "TurnFailed" + assert third_frames[-1]["failed"]["reason"] == "DeadlineExceeded" + assert factory.runtimes[0].close_finished.wait(timeout=2) + + fourth_id = _create_turn( + client, + turnID="turn-after-orphan-failure", + runtimeSessionID="runtime-session-four", + correlationID="corr-four", + input={"prompt": "four", "contextRefs": [], "env": [{"name": "MODEL_TOKEN", "value": "four-token"}]}, + ) + fourth_frames = _frames(client.get(f"/v1/turns/{fourth_id}/events", headers=AUTH).text) + assert fourth_frames[-1]["type"] == "TurnFailed" + assert fourth_frames[-1]["failed"] == { + "reason": "RuntimeCloseFailed", + "message": "runtime cleanup failed; restart required before opening another runtime session", + "retryable": False, + } + assert len(factory.runtimes) == 2 + + assert len(factory.runtimes) == 2 + assert factory.runtimes[0].close_calls == 1 + assert factory.runtimes[0].close_failed == 1 + assert factory.runtimes[1].close_calls == 1 + assert factory.runtimes[1].close_completed == 1 + + def test_orka_required_env_must_be_supplied_per_turn_or_process(monkeypatch): spec = AgentSpec.model_validate( { @@ -954,6 +1679,7 @@ def test_orka_brokered_start_validates_safe_read_tool_schemas(): class CapturingBrokeredRuntime: def __init__(self) -> None: self.tools: list[BrokeredToolDefinition] = [] + self.results: list[BrokeredToolResult] = [] async def __aenter__(self) -> RuntimeSession: return self @@ -979,6 +1705,7 @@ async def run_brokered(self, request: RunRequest, tools: list[BrokeredToolDefini brokered_class="read", ) ) + self.results.append(result) return RunResult(text=f"captured {result.output}") @@ -993,6 +1720,524 @@ def build_runtime(self, spec: AgentSpec) -> RuntimeSession: return self.runtime +class NonAmplifyingBrokeredRuntime: + async def __aenter__(self) -> RuntimeSession: + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> bool | None: + return None + + async def run(self, request: RunRequest) -> RunResult: + raise AssertionError("brokered mode must not call direct run()") + + async def run_brokered(self, request: RunRequest, tools: list[BrokeredToolDefinition], broker: ToolBroker) -> RunResult: + await broker.request_tool( + BrokeredToolCall( + tool_call_id="tool-call-1", + name=tools[0].name, + arguments={"incident": "INC-1"}, + brokered_class="read", + ) + ) + return RunResult(text="brokered output accepted") + + +class NonAmplifyingBrokeredFactory: + def __init__(self) -> None: + self.runtime = NonAmplifyingBrokeredRuntime() + + def supports_brokered_read(self) -> bool: + return True + + def build_runtime(self, spec: AgentSpec) -> RuntimeSession: + return self.runtime + + +class TwoStepBrokeredRuntime(NonAmplifyingBrokeredRuntime): + async def run_brokered(self, request: RunRequest, tools: list[BrokeredToolDefinition], broker: ToolBroker) -> RunResult: + for tool_call_id in ("tool-call-1", "tool-call-2"): + await broker.request_tool( + BrokeredToolCall( + tool_call_id=tool_call_id, + name=tools[0].name, + arguments={"toolCallID": tool_call_id}, + brokered_class="read", + ) + ) + return RunResult(text="two-step brokered output accepted") + + +class TwoStepBrokeredFactory(NonAmplifyingBrokeredFactory): + def __init__(self) -> None: + self.runtime = TwoStepBrokeredRuntime() + + +class MutatingBrokeredRuntime(NonAmplifyingBrokeredRuntime): + async def run_brokered(self, request: RunRequest, tools: list[BrokeredToolDefinition], broker: ToolBroker) -> RunResult: + result = await broker.request_tool( + BrokeredToolCall( + tool_call_id="tool-call-1", + name=tools[0].name, + arguments={"probe": True}, + brokered_class="read", + ) + ) + result.output["nested"]["items"].append("mutated") + return RunResult(text="mutated adapter-local result") + + +class MutatingBrokeredFactory(NonAmplifyingBrokeredFactory): + def __init__(self) -> None: + self.runtime = MutatingBrokeredRuntime() + + +def test_orka_brokered_retained_event_is_an_immutable_json_snapshot(): + app = create_orka_app( + _spec(), + MutatingBrokeredFactory(), + AUTH["authorization"].removeprefix("Bearer "), + enable_brokered_read=True, + ) + output = {"nested": {"items": ["original"]}} + + with TestClient(app) as client: + turn_id = _create_turn(client, toolExecutionMode="brokered", input=_brokered_input()) + _wait_for_event_type(client, turn_id, "ToolCallRequested") + continuation = _continue_payload(turnID=turn_id) + continuation["toolResults"][0]["output"] = output + accepted = client.post(f"/v1/turns/{turn_id}/continue", json=continuation, headers=AUTH) + frames = _frames(client.get(f"/v1/turns/{turn_id}/events", headers=AUTH).text) + + assert accepted.status_code == 202 + assert frames[2]["type"] == "ToolResultReceived" + assert frames[2]["content"] == output + + +def test_orka_brokered_accepted_result_replays_while_later_tool_call_is_pending(): + app = create_orka_app( + _spec(), + TwoStepBrokeredFactory(), + AUTH["authorization"].removeprefix("Bearer "), + enable_brokered_read=True, + ) + + with TestClient(app) as client: + turn_id = _create_turn( + client, + turnID="turn-brokered-inflight-replay", + toolExecutionMode="brokered", + input=_brokered_input(), + ) + _wait_for_tool_call_id(client, turn_id, "tool-call-1") + first = _continue_payload(turnID=turn_id) + first["toolResults"][0]["turnID"] = turn_id + first["toolResults"][0]["idempotencyKey"] = f"runtime-session-1:{turn_id}:tool-call-1" + accepted = client.post(f"/v1/turns/{turn_id}/continue", json=first, headers=AUTH) + _wait_for_tool_call_id(client, turn_id, "tool-call-2") + replayed = client.post(f"/v1/turns/{turn_id}/continue", json=first, headers=AUTH) + second = _continue_payload(turnID=turn_id) + second_result = second["toolResults"][0] + second_result["turnID"] = turn_id + second_result["toolCallID"] = "tool-call-2" + second_result["idempotencyKey"] = f"runtime-session-1:{turn_id}:tool-call-2" + completed = client.post(f"/v1/turns/{turn_id}/continue", json=second, headers=AUTH) + frames = _frames(client.get(f"/v1/turns/{turn_id}/events", headers=AUTH).text) + + assert accepted.status_code == 202 + assert replayed.status_code == 202 + assert completed.status_code == 202 + assert frames[-1]["type"] == "TurnCompleted" + + +def test_orka_brokered_batch_validates_conflicts_before_oversized_members(): + app = create_orka_app( + _spec(), + TwoStepBrokeredFactory(), + AUTH["authorization"].removeprefix("Bearer "), + enable_brokered_read=True, + ) + + with TestClient(app) as client: + turn_id = _create_turn( + client, + turnID="turn-brokered-batch-validation", + toolExecutionMode="brokered", + input=_brokered_input(), + ) + _wait_for_tool_call_id(client, turn_id, "tool-call-1") + first = _continue_payload(turnID=turn_id) + first_result = first["toolResults"][0] + first_result["turnID"] = turn_id + first_result["idempotencyKey"] = f"runtime-session-1:{turn_id}:tool-call-1" + assert client.post(f"/v1/turns/{turn_id}/continue", json=first, headers=AUTH).status_code == 202 + _wait_for_tool_call_id(client, turn_id, "tool-call-2") + + oversized_second = dict(first_result) + oversized_second["toolCallID"] = "tool-call-2" + oversized_second["idempotencyKey"] = f"runtime-session-1:{turn_id}:tool-call-2" + oversized_second["output"] = _json_object_with_size(EXPECTED_MAX_OUTPUT_BYTES + 1) + conflicting_first = dict(first_result) + conflicting_first["output"] = {"different": True} + invalid_batch = _continue_payload(turnID=turn_id) + invalid_batch["toolResults"] = [oversized_second, conflicting_first] + invalid = client.post(f"/v1/turns/{turn_id}/continue", json=invalid_batch, headers=AUTH) + + valid_second = _continue_payload(turnID=turn_id) + valid_second_result = valid_second["toolResults"][0] + valid_second_result["turnID"] = turn_id + valid_second_result["toolCallID"] = "tool-call-2" + valid_second_result["idempotencyKey"] = f"runtime-session-1:{turn_id}:tool-call-2" + completed = client.post(f"/v1/turns/{turn_id}/continue", json=valid_second, headers=AUTH) + frames = _frames(client.get(f"/v1/turns/{turn_id}/events", headers=AUTH).text) + + assert invalid.status_code == 409 + assert "conflicting tool result" in invalid.text + assert completed.status_code == 202 + assert frames[-1]["type"] == "TurnCompleted" + + +def test_orka_brokered_json_output_accepts_exact_utf8_limit_replays_and_retains_one_copy(): + output = _json_object_with_size(EXPECTED_MAX_OUTPUT_BYTES) + app = create_orka_app( + _spec(), + NonAmplifyingBrokeredFactory(), + AUTH["authorization"].removeprefix("Bearer "), + enable_brokered_read=True, + ) + + with TestClient(app) as client: + turn_id = _create_turn( + client, + turnID="turn-brokered-output-boundary", + toolExecutionMode="brokered", + input=_brokered_input(), + ) + _wait_for_event_type(client, turn_id, "ToolCallRequested") + continuation = _continue_payload(turnID=turn_id) + continuation["toolResults"][0]["turnID"] = turn_id + continuation["toolResults"][0]["idempotencyKey"] = f"runtime-session-1:{turn_id}:tool-call-1" + continuation["toolResults"][0]["output"] = output + accepted = client.post(f"/v1/turns/{turn_id}/continue", json=continuation, headers=AUTH) + response = client.get(f"/v1/turns/{turn_id}/events", headers=AUTH) + replay = client.get(f"/v1/turns/{turn_id}/events?afterSeq=1", headers=AUTH) + duplicate = client.post(f"/v1/turns/{turn_id}/continue", json=continuation, headers=AUTH) + state = client.app.state.turns[turn_id] + + frames = _frames(response.text) + assert accepted.status_code == 202, accepted.text + assert duplicate.status_code == 202, duplicate.text + assert [frame["type"] for frame in frames] == [ + "TurnStarted", + "ToolCallRequested", + "ToolResultReceived", + "RuntimeOutput", + "TurnCompleted", + ] + assert frames[2]["content"] == output + assert [frame["type"] for frame in _frames(replay.text)] == [ + "ToolCallRequested", + "ToolResultReceived", + "RuntimeOutput", + "TurnCompleted", + ] + assert state.pending_tools == {} + assert sum(event.content == output for event in state.events) == 1 + _assert_sse_lines_fit_orka_client(response.content) + _assert_sse_lines_fit_orka_client(replay.content) + + +def test_orka_brokered_result_preflight_reserves_maximum_sequence_width(monkeypatch): + app = create_orka_app( + _spec(), + NonAmplifyingBrokeredFactory(), + AUTH["authorization"].removeprefix("Bearer "), + enable_brokered_read=True, + ) + original_ensure = orka_module._ensure_sse_frame_fits + tool_result_sequences: list[int] = [] + + def recording_ensure(event): + if event.type == "ToolResultReceived": + tool_result_sequences.append(event.seq) + return original_ensure(event) + + monkeypatch.setattr(orka_module, "_ensure_sse_frame_fits", recording_ensure) + + with TestClient(app) as client: + turn_id = _create_turn(client, toolExecutionMode="brokered", input=_brokered_input()) + _wait_for_event_type(client, turn_id, "ToolCallRequested") + accepted = client.post(f"/v1/turns/{turn_id}/continue", json=_continue_payload(turnID=turn_id), headers=AUTH) + client.get(f"/v1/turns/{turn_id}/events", headers=AUTH) + + assert accepted.status_code == 202 + assert tool_result_sequences[0] == 9_223_372_036_854_775_807 + assert tool_result_sequences[-1] < tool_result_sequences[0] + + +def test_orka_brokered_json_output_over_utf8_limit_returns_413_and_visible_terminal_failure(): + output = _json_object_with_size(EXPECTED_MAX_OUTPUT_BYTES + 1) + app = create_orka_app( + _spec(), + NonAmplifyingBrokeredFactory(), + AUTH["authorization"].removeprefix("Bearer "), + enable_brokered_read=True, + ) + + with TestClient(app) as client: + turn_id = _create_turn( + client, + turnID="turn-brokered-output-over-limit", + toolExecutionMode="brokered", + input=_brokered_input(), + ) + _wait_for_event_type(client, turn_id, "ToolCallRequested") + continuation = _continue_payload(turnID=turn_id) + continuation["toolResults"][0]["turnID"] = turn_id + continuation["toolResults"][0]["idempotencyKey"] = f"runtime-session-1:{turn_id}:tool-call-1" + continuation["toolResults"][0]["output"] = output + rejected = client.post(f"/v1/turns/{turn_id}/continue", json=continuation, headers=AUTH) + replayed_rejection = client.post(f"/v1/turns/{turn_id}/continue", json=continuation, headers=AUTH) + conflicting_continuation = _continue_payload(turnID=turn_id) + conflicting_continuation["toolResults"][0]["turnID"] = turn_id + conflicting_continuation["toolResults"][0]["idempotencyKey"] = f"runtime-session-1:{turn_id}:tool-call-1" + conflicting_continuation["toolResults"][0]["output"] = {"different": True} + conflicting = client.post( + f"/v1/turns/{turn_id}/continue", + json=conflicting_continuation, + headers=AUTH, + ) + response = client.get(f"/v1/turns/{turn_id}/events", headers=AUTH) + replay = client.get(f"/v1/turns/{turn_id}/events?afterSeq=1", headers=AUTH) + state = client.app.state.turns[turn_id] + + message = ( + f"brokered tool output is {EXPECTED_MAX_OUTPUT_BYTES + 1} UTF-8 bytes; " + f"maxOutputBytes is {EXPECTED_MAX_OUTPUT_BYTES}" + ) + frames = _frames(response.text) + assert rejected.status_code == 413 + assert rejected.json() == {"detail": message} + assert replayed_rejection.status_code == 413 + assert replayed_rejection.json() == rejected.json() + assert conflicting.status_code == 409 + assert [frame["type"] for frame in frames] == ["TurnStarted", "ToolCallRequested", "TurnFailed"] + assert frames[-1]["failed"] == {"reason": "MaxOutputBytesExceeded", "message": message, "retryable": False} + assert [frame["type"] for frame in _frames(replay.text)] == ["ToolCallRequested", "TurnFailed"] + assert state.pending_tools == {} + assert all(event.content != output for event in state.events) + _assert_sse_lines_fit_orka_client(response.content) + _assert_sse_lines_fit_orka_client(replay.content) + + +def test_orka_brokered_under_limit_output_with_unstreamable_error_uses_frame_failure_code(): + app = create_orka_app( + _spec(), + NonAmplifyingBrokeredFactory(), + AUTH["authorization"].removeprefix("Bearer "), + enable_brokered_read=True, + ) + + with TestClient(app) as client: + turn_id = _create_turn( + client, + turnID="turn-brokered-error-frame-over-limit", + toolExecutionMode="brokered", + input=_brokered_input(), + ) + _wait_for_event_type(client, turn_id, "ToolCallRequested") + continuation = _continue_payload(turnID=turn_id) + result = continuation["toolResults"][0] + result["turnID"] = turn_id + result["idempotencyKey"] = f"runtime-session-1:{turn_id}:tool-call-1" + result["output"] = {"small": True} + result["error"] = {"code": "HugeError", "message": "\x00" * 200_000, "retryable": False} + rejected = client.post(f"/v1/turns/{turn_id}/continue", json=continuation, headers=AUTH) + event_response = client.get(f"/v1/turns/{turn_id}/events", headers=AUTH) + frames = _frames(event_response.text) + + assert rejected.status_code == 413 + assert frames[-1]["type"] == "TurnFailed" + assert frames[-1]["failed"]["reason"] == "HarnessFrameTooLarge" + _assert_sse_lines_fit_orka_client(event_response.content) + + +def test_orka_brokered_output_rejection_is_atomic_with_a_racing_valid_continue(monkeypatch): + app = create_orka_app( + _spec(), + NonAmplifyingBrokeredFactory(), + AUTH["authorization"].removeprefix("Bearer "), + enable_brokered_read=True, + ) + rejection_locked = threading.Event() + release_rejection = threading.Event() + original_append_failure = orka_module._append_output_failure_locked + + def blocking_append_failure(state, message, code): + rejection_locked.set() + assert release_rejection.wait(timeout=5) + return original_append_failure(state, message, code) + + monkeypatch.setattr(orka_module, "_append_output_failure_locked", blocking_append_failure) + + with TestClient(app) as client: + turn_id = _create_turn( + client, + turnID="turn-brokered-output-race", + toolExecutionMode="brokered", + input=_brokered_input(), + ) + _wait_for_event_type(client, turn_id, "ToolCallRequested") + oversized = _continue_payload(turnID=turn_id) + oversized_result = oversized["toolResults"][0] + oversized_result["turnID"] = turn_id + oversized_result["idempotencyKey"] = f"runtime-session-1:{turn_id}:tool-call-1" + oversized_result["output"] = _json_object_with_size(EXPECTED_MAX_OUTPUT_BYTES + 1) + valid = _continue_payload(turnID=turn_id) + valid_result = valid["toolResults"][0] + valid_result["turnID"] = turn_id + valid_result["idempotencyKey"] = f"runtime-session-1:{turn_id}:tool-call-1" + + responses: dict[str, Any] = {} + oversized_thread = threading.Thread( + target=lambda: responses.setdefault( + "oversized", + client.post(f"/v1/turns/{turn_id}/continue", json=oversized, headers=AUTH), + ) + ) + valid_thread = threading.Thread( + target=lambda: responses.setdefault( + "valid", + client.post(f"/v1/turns/{turn_id}/continue", json=valid, headers=AUTH), + ) + ) + oversized_thread.start() + assert rejection_locked.wait(timeout=5) + valid_thread.start() + time.sleep(0.05) + assert valid_thread.is_alive() + release_rejection.set() + oversized_thread.join(timeout=5) + valid_thread.join(timeout=5) + assert not oversized_thread.is_alive() + assert not valid_thread.is_alive() + frames = _frames(client.get(f"/v1/turns/{turn_id}/events", headers=AUTH).text) + + assert responses["oversized"].status_code == 413 + assert responses["valid"].status_code == 409 + assert frames[-1]["type"] == "TurnFailed" + assert all(frame["type"] != "ToolResultReceived" for frame in frames) + + +@pytest.mark.parametrize( + "output", + [ + pytest.param({"answer": "ok"}, id="object"), + pytest.param([], id="array"), + pytest.param("", id="string"), + pytest.param(0, id="number"), + pytest.param(False, id="boolean"), + pytest.param(None, id="null"), + ], +) +def test_orka_brokered_tool_result_preserves_any_json_output_value(output: Any): + factory = CapturingBrokeredFactory() + app = create_orka_app( + _spec(), factory, AUTH["authorization"].removeprefix("Bearer "), enable_brokered_read=True + ) + + with TestClient(app) as client: + turn_id = _create_turn(client, toolExecutionMode="brokered", input=_brokered_input()) + _wait_for_event_type(client, turn_id, "ToolCallRequested") + continuation = _continue_payload(turnID=turn_id) + continuation["toolResults"][0]["output"] = output + response = client.post(f"/v1/turns/{turn_id}/continue", json=continuation, headers=AUTH) + if response.status_code == 202: + frames = _frames(client.get(f"/v1/turns/{turn_id}/events", headers=AUTH).text) + else: + frames = [] + + assert response.status_code == 202, response.text + assert len(factory.runtime.results) == 1 + assert type(factory.runtime.results[0].output) is type(output) + assert factory.runtime.results[0].output == output + assert frames[2]["type"] == "ToolResultReceived" + assert type(frames[2]["content"]) is type(output) + assert frames[2]["content"] == output + + +@pytest.mark.parametrize( + ("include_output", "expected_output_present"), + [ + pytest.param(False, False, id="absent"), + pytest.param(True, True, id="explicit-null"), + ], +) +def test_orka_brokered_tool_result_distinguishes_absent_output_from_explicit_null( + include_output: bool, expected_output_present: bool +): + factory = CapturingBrokeredFactory() + app = create_orka_app( + _spec(), factory, AUTH["authorization"].removeprefix("Bearer "), enable_brokered_read=True + ) + + with TestClient(app) as client: + turn_id = _create_turn(client, toolExecutionMode="brokered", input=_brokered_input()) + _wait_for_event_type(client, turn_id, "ToolCallRequested") + continuation = _continue_payload(turnID=turn_id) + if include_output: + continuation["toolResults"][0]["output"] = None + else: + continuation["toolResults"][0].pop("output") + continuation["toolResults"][0]["error"] = { + "code": "NoOutput", + "message": "tool completed without output", + "retryable": False, + } + response = client.post(f"/v1/turns/{turn_id}/continue", json=continuation, headers=AUTH) + if response.status_code == 202: + frames = _frames(client.get(f"/v1/turns/{turn_id}/events", headers=AUTH).text) + else: + frames = [] + + assert response.status_code == 202, response.text + assert len(factory.runtime.results) == 1 + result = factory.runtime.results[0] + assert result.output is None + assert result.output_present is expected_output_present + assert frames[2]["type"] == "ToolResultReceived" + if include_output: + assert "content" in frames[2] + assert frames[2]["content"] is None + else: + assert "content" not in frames[2] + assert frames[2]["error"] == { + "code": "NoOutput", + "message": "tool completed without output", + "retryable": False, + } + + +def test_orka_brokered_continue_rejects_absent_output_without_error(): + app = create_orka_app( + _spec(), OfflineEchoRuntimeFactory(), AUTH["authorization"].removeprefix("Bearer "), enable_brokered_read=True + ) + + with TestClient(app) as client: + turn_id = _create_turn(client, toolExecutionMode="brokered", input=_brokered_input()) + _wait_for_event_type(client, turn_id, "ToolCallRequested") + continuation = _continue_payload(turnID=turn_id) + continuation["toolResults"][0].pop("output") + response = client.post(f"/v1/turns/{turn_id}/continue", json=continuation, headers=AUTH) + + assert response.status_code == 400 + assert "output or error is required" in response.text + + def test_orka_brokered_runtime_receives_only_safe_tool_definition_fields(): factory = CapturingBrokeredFactory() app = create_orka_app(_spec(), factory, "test-token", enable_brokered_read=True) diff --git a/runtimes/langgraph/agentkit_serve/agent_factory.py b/runtimes/langgraph/agentkit_serve/agent_factory.py index 21cc289..52b5b65 100644 --- a/runtimes/langgraph/agentkit_serve/agent_factory.py +++ b/runtimes/langgraph/agentkit_serve/agent_factory.py @@ -39,6 +39,7 @@ from agentkit_serve_common.adapter_support import ( FORWARDED_ROLES, + AsyncExitStackLifecycle, AgentBuildError, declared_tool_env, normalize_agent_run_error, @@ -127,11 +128,12 @@ class LangGraphRuntime: def __init__(self, spec: AgentSpec) -> None: self.spec = spec self.stack = AsyncExitStack() + self.lifecycle = AsyncExitStackLifecycle(self.stack) self.graph: Any | None = None self.client: MultiServerMCPClient | None = None async def __aenter__(self) -> RuntimeSession: - try: + async def start() -> RuntimeSession: model = build_model(self.spec) tools = await self._load_tools() self.graph = create_agent( @@ -140,8 +142,12 @@ async def __aenter__(self) -> RuntimeSession: system_prompt=self.spec.instructions, ) return self - except Exception: - await self.stack.aclose() + + try: + return await self.lifecycle.enter(start) + except BaseException: + self.graph = None + self.client = None raise async def __aexit__( @@ -150,9 +156,11 @@ async def __aexit__( exc: BaseException | None, tb: TracebackType | None, ) -> bool | None: - self.graph = None - await self.stack.aclose() - return None + try: + return await self.lifecycle.exit(exc_type, exc, tb) + finally: + self.graph = None + self.client = None async def run(self, request: RunRequest) -> RunResult: return await run_agent(self, request) diff --git a/runtimes/langgraph/tests/test_lifecycle.py b/runtimes/langgraph/tests/test_lifecycle.py new file mode 100644 index 0000000..654ed92 --- /dev/null +++ b/runtimes/langgraph/tests/test_lifecycle.py @@ -0,0 +1,154 @@ +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from unittest import mock + +import pytest + +from agentkit_serve import agent_factory +from agentkit_serve_common.config import AgentSpec + + +def _runtime_with_tools(*names: str) -> agent_factory.LangGraphRuntime: + spec = AgentSpec.model_validate( + { + "abiVersion": "v0", + "metadata": {"name": "lifecycle-test"}, + "model": { + "provider": "openai-compatible", + "baseURL": "https://api.openai.com/v1", + "name": "gpt-4o-mini", + }, + "instructions": "Be helpful.", + "tools": [{"name": name, "command": ["fake-mcp"], "env": []} for name in names], + "expose": {"openai": True, "port": 8080}, + } + ) + return agent_factory.LangGraphRuntime(spec) + + +def test_second_mcp_session_failure_preserves_startup_error_while_unwinding_first(): + runtime = _runtime_with_tools("first", "second") + startup_error = RuntimeError("second session failed") + cleanup_error = RuntimeError("first session cleanup failed") + exited: list[str] = [] + owner_tasks: dict[str, asyncio.Task] = {} + + class _Session: + async def initialize(self) -> None: + return None + + class _SessionContext: + def __init__(self, name: str) -> None: + self.name = name + + async def __aenter__(self): + if self.name == "second": + raise startup_error + owner_tasks[self.name] = asyncio.current_task() + return _Session() + + async def __aexit__(self, exc_type, exc, tb): + assert asyncio.current_task() is owner_tasks[self.name] + exited.append(self.name) + raise cleanup_error + + class _FakeClient: + def __init__(self, connections, tool_name_prefix=False) -> None: + pass + + def session(self, server_name, auto_initialize=True): + return _SessionContext(server_name) + + async def _fake_load(session, *, server_name, tool_name_prefix): + return [SimpleNamespace(name=f"{server_name}_tool")] + + with ( + mock.patch("agentkit_serve.agent_factory.MultiServerMCPClient", _FakeClient), + mock.patch("agentkit_serve.agent_factory.load_mcp_tools", _fake_load), + ): + with pytest.raises(RuntimeError) as exc_info: + asyncio.run(runtime.__aenter__()) + + assert exc_info.value is startup_error + assert exited == ["first"] + + +def test_cancellation_while_initializing_second_mcp_session_finishes_cleanup(): + runtime = _runtime_with_tools("first", "second") + cleanup_error = RuntimeError("second session cleanup failed") + second_initialize_started = asyncio.Event() + second_exit_started = asyncio.Event() + allow_second_exit = asyncio.Event() + exited: list[str] = [] + owner_tasks: dict[str, asyncio.Task] = {} + + class _Session: + def __init__(self, name: str) -> None: + self.name = name + + async def initialize(self) -> None: + if self.name == "second": + second_initialize_started.set() + await asyncio.Event().wait() + + class _SessionContext: + def __init__(self, name: str) -> None: + self.name = name + + async def __aenter__(self): + owner_tasks[self.name] = asyncio.current_task() + return _Session(self.name) + + async def __aexit__(self, exc_type, exc, tb): + assert asyncio.current_task() is owner_tasks[self.name] + if self.name == "second": + second_exit_started.set() + await allow_second_exit.wait() + exited.append(self.name) + if self.name == "second": + raise cleanup_error + + class _FakeClient: + def __init__(self, connections, tool_name_prefix=False) -> None: + pass + + def session(self, server_name, auto_initialize=True): + return _SessionContext(server_name) + + async def _fake_load(session, *, server_name, tool_name_prefix): + return [] + + async def exercise() -> None: + task = asyncio.create_task(runtime.__aenter__()) + try: + await asyncio.wait_for(second_initialize_started.wait(), timeout=1) + task.cancel() + await asyncio.wait_for(second_exit_started.wait(), timeout=1) + + # A second cancellation must not interrupt the already-running cleanup. + task.cancel() + await asyncio.sleep(0) + assert not task.done() + + allow_second_exit.set() + with pytest.raises(asyncio.CancelledError) as exc_info: + await task + assert any( + "cleanup also failed with RuntimeError" in note + for note in getattr(exc_info.value, "__notes__", ()) + ) + finally: + allow_second_exit.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + with ( + mock.patch("agentkit_serve.agent_factory.MultiServerMCPClient", _FakeClient), + mock.patch("agentkit_serve.agent_factory.load_mcp_tools", _fake_load), + ): + asyncio.run(exercise()) + + assert exited == ["second", "first"] diff --git a/runtimes/microsoft-agent-framework/agentkit_serve/agent_factory.py b/runtimes/microsoft-agent-framework/agentkit_serve/agent_factory.py index b8801c3..4688bb0 100644 --- a/runtimes/microsoft-agent-framework/agentkit_serve/agent_factory.py +++ b/runtimes/microsoft-agent-framework/agentkit_serve/agent_factory.py @@ -9,9 +9,11 @@ import asyncio import importlib +import inspect import os import time -from contextlib import AsyncExitStack +from contextlib import AbstractAsyncContextManager, AsyncExitStack +from datetime import timedelta from types import TracebackType from urllib.parse import urlsplit @@ -26,8 +28,10 @@ SkillsProvider, ) from agent_framework.openai import OpenAIChatCompletionClient +from httpx import AsyncClient, URL from agentkit_serve_common.adapter_support import ( FORWARDED_ROLES, + AsyncExitStackLifecycle, AgentBuildError, declared_tool_env, normalize_agent_run_error, @@ -57,17 +61,18 @@ _CONTEXT_SOURCE_MCP = "mcp" _DEFAULT_SEARCH_AUDIENCE = "https://search.azure.com/.default" _DEFAULT_FOUNDRY_AUDIENCE = "https://ai.azure.com/.default" +_DEFAULT_MCP_REQUEST_TIMEOUT = 120 _DEFAULT_SESSION_CACHE_MAX = 256 -def _mcp_request_timeout() -> int | None: +def _mcp_request_timeout() -> int: """MCP request timeout (seconds), overridable via ``AGENTKIT_MCP_TIMEOUT``.""" - return positive_int_env(default=None) + return positive_int_env(default=_DEFAULT_MCP_REQUEST_TIMEOUT) or _DEFAULT_MCP_REQUEST_TIMEOUT def _remote_mcp_timeout() -> float: - """Bound network-backed MCP calls even when stdio timeout is left uncapped.""" - return float(_mcp_request_timeout() or 120) + """Use the shared bounded MCP timeout for network-backed calls.""" + return float(_mcp_request_timeout()) def _resolve_api_key(spec: AgentSpec) -> str: @@ -183,12 +188,35 @@ async def _provider() -> str: return _provider +def _uses_model_workload_identity_fallback(spec: AgentSpec) -> bool: + auth = spec.model.auth + return bool( + auth is not None + and auth.type == _AUTH_WORKLOAD_IDENTITY + and not ( + os.environ.get("AGENTKIT_MODEL_WORKLOAD_IDENTITY_TOKEN") + or os.environ.get("AGENTKIT_WORKLOAD_IDENTITY_TOKEN") + or os.environ.get("AGENTKIT_WORKLOAD_IDENTITY_TOKEN_COMMAND") + ) + ) + + def _disable_mcp_ping(mcp_tool: object) -> None: if hasattr(mcp_tool, "_ping_available"): setattr(mcp_tool, "_ping_available", False) -def build_client(spec: AgentSpec): +async def _close_resource(resource: object) -> None: + """Close a sync or async runtime-owned resource.""" + close = getattr(resource, "close", None) + if not callable(close): + return + result = close() + if inspect.isawaitable(result): + await result + + +def build_client(spec: AgentSpec, *, workload_identity_credential: object | None = None): """Construct the chat client for the configured model auth mode.""" auth = spec.model.auth if auth is not None and auth.type == _AUTH_WORKLOAD_IDENTITY: @@ -209,10 +237,15 @@ def build_client(spec: AgentSpec): raise AgentBuildError( "model workload identity auth requires agent-framework-foundry and azure-identity" ) from exc + credential = ( + workload_identity_credential + if workload_identity_credential is not None + else DefaultAzureCredential() + ) return FoundryChatClient( project_endpoint=_project_endpoint_from_openai_base_url(spec.model.base_url), model=spec.model.name, - credential=DefaultAzureCredential(), + credential=credential, ) return OpenAIChatCompletionClient( @@ -227,12 +260,10 @@ def _tool_env(tool: ToolSpec) -> dict[str, str]: return declared_tool_env(tool) -def build_tool(tool: ToolSpec): +def build_tool(tool: ToolSpec, *, stack: AsyncExitStack | None = None): """Create a stdio or Streamable HTTP MCP server for one tool spec.""" timeout = _mcp_request_timeout() if tool.url_env: - from httpx import AsyncClient, URL - url = resolve_tool_url(tool) target_url = URL(url) target_origin = (target_url.scheme, target_url.host, target_url.port) @@ -245,16 +276,22 @@ async def inject_headers(request): # noqa: ANN001 for key, value in (await asyncio.to_thread(resolve_tool_headers, tool)).items(): request.headers[key] = value + http_client = AsyncClient( + event_hooks={"request": [inject_headers]}, + follow_redirects=False, + timeout=remote_timeout, + ) + if stack is not None: + # MAF owns the MCP tool/session lifecycle, but the MCP SDK deliberately + # leaves caller-supplied HTTP clients open. Register the client before + # the Agent so LIFO cleanup disconnects the tool before closing HTTP. + stack.push_async_callback(http_client.aclose) kwargs: dict[str, object] = { "name": tool.name, "url": url, "tool_name_prefix": tool.name, "load_prompts": False, - "http_client": AsyncClient( - event_hooks={"request": [inject_headers]}, - follow_redirects=False, - timeout=remote_timeout, - ), + "http_client": http_client, } kwargs["request_timeout"] = int(remote_timeout) mcp_tool = MCPStreamableHTTPTool(**kwargs) @@ -272,18 +309,23 @@ async def inject_headers(request): # noqa: ANN001 "env": declared_tool_env(tool), "tool_name_prefix": tool.name, } - if timeout is not None: - kwargs["request_timeout"] = timeout + kwargs["request_timeout"] = timeout return MCPStdioTool(**kwargs) -def build_agent(spec: AgentSpec, *, context_providers=None) -> Agent: +def build_agent( + spec: AgentSpec, + *, + context_providers=None, + stack: AsyncExitStack | None = None, + client=None, +) -> Agent: """Assemble the MAF agent: client + system prompt + tools + context.""" return Agent( - client=build_client(spec), + client=client if client is not None else build_client(spec), instructions=spec.instructions, name=spec.metadata.name, - tools=[build_tool(t) for t in spec.tools], + tools=[build_tool(t, stack=stack) for t in spec.tools], context_providers=context_providers, ) @@ -294,19 +336,36 @@ class MAFRuntime: def __init__(self, spec: AgentSpec) -> None: self.spec = spec self.stack = AsyncExitStack() + self.lifecycle = AsyncExitStackLifecycle(self.stack) self.agent: Agent | None = None self.sessions: dict[str, AgentSession] = {} self.session_locks: dict[str, asyncio.Lock] = {} + # Claims cover both lock holders and queued waiters during eviction. + self.session_claims: dict[str, int] = {} + self.initialized_sessions: set[str] = set() + self.most_recent_session_id: str | None = None self.session_cache_max = _session_cache_max() async def __aenter__(self) -> RuntimeSession: - try: + async def start() -> RuntimeSession: context_providers = await self._build_context_providers() - self.agent = build_agent(self.spec, context_providers=context_providers) + client = await self._build_model_fallback_client() + self.agent = build_agent( + self.spec, + context_providers=context_providers, + stack=self.stack, + client=client, + ) + # Register before entering so a partially failed Agent.__aenter__ still + # unwinds the Agent's own internal AsyncExitStack. + self.stack.push_async_exit(self.agent) await self.agent.__aenter__() return self - except Exception: - await self.stack.aclose() + + try: + return await self.lifecycle.enter(start) + except BaseException: + self.agent = None raise async def __aexit__( @@ -315,47 +374,123 @@ async def __aexit__( exc: BaseException | None, tb: TracebackType | None, ) -> bool | None: - agent_result = None - if self.agent is not None: - agent_result = await self.agent.__aexit__(exc_type, exc, tb) + try: + return await self.lifecycle.exit(exc_type, exc, tb) + finally: self.agent = None - await self.stack.aclose() - return agent_result async def run(self, request: RunRequest) -> RunResult: if self.agent is None: raise AgentBuildError("MAF runtime session is not initialized") - session, existed, lock = self._session_for(request.session_id) - include_history = not (session is not None and existed) + session_id = request.session_id + session, lock = self._session_for(session_id) if lock is None: - return await run_agent(self.agent, request, session=session, include_history=include_history) - async with lock: - return await run_agent(self.agent, request, session=session, include_history=include_history) - - def _session_for(self, session_id: str | None) -> tuple[AgentSession | None, bool, asyncio.Lock | None]: + return await run_agent(self.agent, request, session=session, include_history=True) + assert session_id + try: + async with lock: + self._touch_session(session_id) + include_history = session_id not in self.initialized_sessions + result = await run_agent(self.agent, request, session=session, include_history=include_history) + self.initialized_sessions.add(session_id) + return result + finally: + self._release_session_claim(session_id) + + def _session_for(self, session_id: str | None) -> tuple[AgentSession | None, asyncio.Lock | None]: if not session_id: - return None, False, None - session = self.sessions.pop(session_id, None) - lock = self.session_locks.pop(session_id, None) - existed = session is not None + return None, None + session = self.sessions.get(session_id) + lock = self.session_locks.get(session_id) if session is None: session = AgentSession(session_id=session_id) + self.sessions[session_id] = session + self.initialized_sessions.discard(session_id) if lock is None: lock = asyncio.Lock() + self.session_locks[session_id] = lock + self.session_claims[session_id] = self.session_claims.get(session_id, 0) + 1 + self._evict_idle_sessions() + return session, lock + + def _touch_session(self, session_id: str) -> None: + session = self.sessions.pop(session_id) + lock = self.session_locks.pop(session_id) self.sessions[session_id] = session self.session_locks[session_id] = lock + self.most_recent_session_id = session_id self._evict_idle_sessions() - return session, existed, lock + + def _release_session_claim(self, session_id: str) -> None: + claims = self.session_claims.get(session_id, 0) + if claims <= 1: + self.session_claims.pop(session_id, None) + self._evict_idle_sessions() + else: + self.session_claims[session_id] = claims - 1 def _evict_idle_sessions(self) -> None: for session_id in list(self.sessions): if len(self.sessions) <= self.session_cache_max: return + if session_id == self.most_recent_session_id: + continue + if self.session_claims.get(session_id, 0) > 0: + continue lock = self.session_locks.get(session_id) if lock is not None and lock.locked(): continue self.sessions.pop(session_id, None) self.session_locks.pop(session_id, None) + self.session_claims.pop(session_id, None) + self.initialized_sessions.discard(session_id) + + async def _enter_owned_async_context(self, resource): + """Enter an adapter-owned async context with partial-enter cleanup.""" + enter = getattr(resource, "__aenter__", None) + exit_ = getattr(resource, "__aexit__", None) + if not callable(enter) or not callable(exit_): + return resource + self.stack.push_async_exit(resource) + return await enter() + + async def _build_model_fallback_client(self): + if not _uses_model_workload_identity_fallback(self.spec): + return None + try: + from azure.identity import DefaultAzureCredential + except ImportError as exc: # pragma: no cover - dependency guard. + raise AgentBuildError( + "model workload identity auth requires agent-framework-foundry and azure-identity" + ) from exc + + credential = DefaultAzureCredential() + if callable(getattr(credential, "close", None)): + self.stack.push_async_callback(_close_resource, credential) + client = build_client(self.spec, workload_identity_credential=credential) + # FoundryChatClient exposes its internally-created AIProjectClient, but + # MAF's Agent does not enter that project client. Own it here while leaving + # the framework chat client itself exclusively under Agent ownership. + project_client = getattr(client, "project_client", None) + if callable(getattr(project_client, "__aenter__", None)) and callable( + getattr(project_client, "__aexit__", None) + ): + await self._enter_owned_async_context(project_client) + + # Current FoundryChatClient is not an async context manager, while the + # AsyncOpenAI client it creates is. Close that model HTTP pool after the + # Agent but before the project client. If a future framework version owns + # the chat client lifecycle, leave its internals exclusively to Agent. + if not isinstance(client, AbstractAsyncContextManager): + model_http_client = getattr(client, "client", None) + if model_http_client is not None and model_http_client is not project_client: + enter = getattr(model_http_client, "__aenter__", None) + exit_ = getattr(model_http_client, "__aexit__", None) + if callable(enter) and callable(exit_): + await self._enter_owned_async_context(model_http_client) + elif callable(getattr(model_http_client, "close", None)): + self.stack.push_async_callback(_close_resource, model_http_client) + return client async def _build_context_providers(self): providers = [] @@ -368,7 +503,7 @@ async def _build_context_providers(self): elif provider.source == _CONTEXT_SOURCE_MCP: providers.append(await self._build_mcp_skills_provider(provider)) elif provider.type == _CONTEXT_TYPE_MEMORY: - providers.append(self._build_memory_provider(provider)) + providers.append(await self._build_memory_provider(provider)) return providers or None async def _build_search_provider(self, provider: ContextProviderSpec): @@ -387,16 +522,18 @@ async def _build_search_provider(self, provider: ContextProviderSpec): default_audience=_DEFAULT_SEARCH_AUDIENCE, async_credential=True, ) - close = getattr(credential, "close", None) - if callable(close): - self.stack.push_async_callback(close) - return search_provider( + if callable(getattr(credential, "close", None)): + self.stack.push_async_callback(_close_resource, credential) + # Agent invokes context providers but does not manage their async + # lifecycle, so the runtime must close their internally-created clients. + context_provider = search_provider( endpoint=endpoint, index_name=index, credential=credential, ) + return await self._enter_owned_async_context(context_provider) - def _build_memory_provider(self, provider: ContextProviderSpec): + async def _build_memory_provider(self, provider: ContextProviderSpec): try: foundry_mod = importlib.import_module("agent_framework.foundry") memory_provider = getattr(foundry_mod, "FoundryMemoryProvider") @@ -407,14 +544,20 @@ def _build_memory_provider(self, provider: ContextProviderSpec): endpoint = _env_required(provider.endpoint_env, field="context.providers[].endpointEnv") store_name = _env_required(provider.store_name_env, field="context.providers[].storeNameEnv") - return memory_provider( + credential = _credential_for_context(provider, default_audience=_DEFAULT_FOUNDRY_AUDIENCE) + if callable(getattr(credential, "close", None)): + self.stack.push_async_callback(_close_resource, credential) + # Entering the provider enters/closes its internally-created project + # client. Keep the credential below it on the stack so it closes last. + context_provider = memory_provider( source_id=provider.name or "memory", project_endpoint=endpoint, - credential=_credential_for_context(provider, default_audience=_DEFAULT_FOUNDRY_AUDIENCE), + credential=credential, memory_store_name=store_name, scope=_memory_scope(), update_delay=_memory_update_delay(), ) + return await self._enter_owned_async_context(context_provider) async def _build_mcp_skills_provider(self, provider: ContextProviderSpec): tool = next((t for t in self.spec.tools if t.name == provider.tool_ref), None) @@ -427,15 +570,25 @@ async def _build_mcp_skills_provider(self, provider: ContextProviderSpec): from mcp.client.streamable_http import streamablehttp_client url = resolve_tool_url(tool) + timeout = _remote_mcp_timeout() http_client_factory = same_origin_mcp_httpx_client_factory( tool, url, - timeout=_remote_mcp_timeout(), + timeout=timeout, ) + # This compatibility API creates the AsyncClient from the factory inside + # its own async context. Owning the transport context on our stack also + # owns that client; registering it separately would double-close it. read, write, _ = await self.stack.enter_async_context( streamablehttp_client(url=url, httpx_client_factory=http_client_factory) ) - session = await self.stack.enter_async_context(ClientSession(read, write)) + session = await self.stack.enter_async_context( + ClientSession( + read, + write, + read_timeout_seconds=timedelta(seconds=timeout), + ) + ) await session.initialize() return SkillsProvider(MCPSkillsSource(client=session)) diff --git a/runtimes/microsoft-agent-framework/tests/test_guardrails.py b/runtimes/microsoft-agent-framework/tests/test_guardrails.py index 9ff323e..6875f9f 100644 --- a/runtimes/microsoft-agent-framework/tests/test_guardrails.py +++ b/runtimes/microsoft-agent-framework/tests/test_guardrails.py @@ -492,6 +492,126 @@ async def fake_run_agent(agent, request, *, session=None, include_history=True): assert list(runtime.sessions) == ["s2", "s3"] +def test_session_cache_evicts_idle_entry_when_new_session_acquires(monkeypatch): + from agentkit_serve_common.config import AgentSpec + from agentkit_serve_common.conversation import RunRequest + from agentkit_serve_common.runtime import RunResult + + import asyncio + + b_started = asyncio.Event() + release_b = asyncio.Event() + seen_sessions = {} + + async def fake_run_agent(agent, request, *, session=None, include_history=True): + seen_sessions[session.session_id] = session + if session.session_id == "b": + b_started.set() + await release_b.wait() + return RunResult(text="ok") + + monkeypatch.setattr(agent_factory, "run_agent", fake_run_agent) + monkeypatch.setenv("AGENTKIT_SESSION_CACHE_MAX", "1") + spec = AgentSpec.model_validate({ + "abiVersion": "v0", + "metadata": {"name": "x"}, + "model": {"provider": "openai-compatible", "baseURL": "https://api.openai.com/v1", "name": "gpt-4o-mini"}, + "instructions": "hi", + "tools": [], + "expose": {"openai": True, "port": 8080}, + }) + runtime = agent_factory.MAFRuntime(spec) + runtime.agent = object() + + async def exercise(): + await runtime.run(RunRequest(prompt="one", session_id="a")) + b_task = asyncio.create_task(runtime.run(RunRequest(prompt="two", session_id="b"))) + await asyncio.wait_for(b_started.wait(), timeout=1) + state_while_b_runs = ( + list(runtime.sessions), + runtime.sessions.get("b") is seen_sessions["b"], + runtime.session_claims.get("b"), + runtime.session_locks["b"].locked(), + ) + release_b.set() + await b_task + return state_while_b_runs + + assert asyncio.run(exercise()) == (["b"], True, 1, True) + + +def test_session_cache_overflow_keeps_current_session_serialized(monkeypatch): + from agentkit_serve_common.config import AgentSpec + from agentkit_serve_common.conversation import RunRequest + from agentkit_serve_common.runtime import RunResult + + import asyncio + + busy_started = asyncio.Event() + release_busy = asyncio.Event() + first_current_started = asyncio.Event() + release_first_current = asyncio.Event() + second_current_started = asyncio.Event() + current_sessions = [] + + async def fake_run_agent(agent, request, *, session=None, include_history=True): + if session.session_id == "busy": + busy_started.set() + await release_busy.wait() + elif session.session_id == "current": + current_sessions.append(session) + if len(current_sessions) == 1: + first_current_started.set() + await release_first_current.wait() + else: + second_current_started.set() + return RunResult(text="ok") + + monkeypatch.setattr(agent_factory, "run_agent", fake_run_agent) + monkeypatch.setenv("AGENTKIT_SESSION_CACHE_MAX", "1") + spec = AgentSpec.model_validate({ + "abiVersion": "v0", + "metadata": {"name": "x"}, + "model": {"provider": "openai-compatible", "baseURL": "https://api.openai.com/v1", "name": "gpt-4o-mini"}, + "instructions": "hi", + "tools": [], + "expose": {"openai": True, "port": 8080}, + }) + runtime = agent_factory.MAFRuntime(spec) + runtime.agent = object() + + async def exercise(): + busy = asyncio.create_task(runtime.run(RunRequest(prompt="hold", session_id="busy"))) + await asyncio.wait_for(busy_started.wait(), timeout=1) + + first_current = asyncio.create_task(runtime.run(RunRequest(prompt="one", session_id="current"))) + await asyncio.wait_for(first_current_started.wait(), timeout=1) + cached_during_overflow = list(runtime.sessions) + + second_current = asyncio.create_task(runtime.run(RunRequest(prompt="two", session_id="current"))) + await asyncio.sleep(0) + await asyncio.sleep(0) + serialized_while_first_running = not second_current_started.is_set() + + release_first_current.set() + await asyncio.wait_for(second_current_started.wait(), timeout=1) + release_busy.set() + await asyncio.gather(busy, first_current, second_current) + cached_after_saturation = list(runtime.sessions) + await runtime.run(RunRequest(prompt="three", session_id="current")) + return cached_during_overflow, serialized_while_first_running, cached_after_saturation + + cached_during_overflow, serialized_while_first_running, cached_after_saturation = asyncio.run(exercise()) + same_session_entry = len(current_sessions) == 3 and all(session is current_sessions[0] for session in current_sessions) + + assert (cached_during_overflow, serialized_while_first_running, cached_after_saturation, same_session_entry) == ( + ["busy", "current"], + True, + ["current"], + True, + ) + + def test_context_credential_uses_async_default_for_search(monkeypatch): from agentkit_serve_common.config import ContextProviderSpec @@ -633,6 +753,53 @@ async def fake_run_agent(agent, request, *, session=None, include_history=True): assert include_history_values == [True, False] +def test_failed_first_turn_keeps_explicit_history_on_retry(monkeypatch): + from agentkit_serve_common.config import AgentSpec + from agentkit_serve_common.conversation import ConversationTurn, RunRequest + from agentkit_serve_common.runtime import RunResult + + include_history_values = [] + seen_sessions = [] + + async def fake_run_agent(agent, request, *, session=None, include_history=True): + include_history_values.append(include_history) + seen_sessions.append(session) + if len(include_history_values) == 1: + raise RuntimeError("provider unavailable") + return RunResult(text="ok") + + monkeypatch.setattr(agent_factory, "run_agent", fake_run_agent) + spec = AgentSpec.model_validate({ + "abiVersion": "v0", + "metadata": {"name": "x"}, + "model": {"provider": "openai-compatible", "baseURL": "https://api.openai.com/v1", "name": "gpt-4o-mini"}, + "instructions": "hi", + "tools": [], + "expose": {"openai": True, "port": 8080}, + }) + runtime = agent_factory.MAFRuntime(spec) + runtime.agent = object() + request = RunRequest( + prompt="current", + history=(ConversationTurn(role="user", text="old"),), + session_id="s1", + ) + + import asyncio + import pytest + + async def exercise(): + with pytest.raises(RuntimeError, match="provider unavailable"): + await runtime.run(request) + await runtime.run(request) + await runtime.run(request) + + asyncio.run(exercise()) + + assert include_history_values == [True, True, False] + assert seen_sessions[0] is seen_sessions[1] is seen_sessions[2] + + def test_remote_mcp_disables_ping(monkeypatch): tool = ToolSpec.model_validate({ "name": "toolbox", @@ -665,8 +832,9 @@ def fake_streamablehttp_client(**kwargs): return FakeTransport() class FakeClientSession: - def __init__(self, read, write): + def __init__(self, read, write, *, read_timeout_seconds): calls["session_args"] = (read, write) + calls["session_timeout"] = read_timeout_seconds async def __aenter__(self): return self @@ -683,6 +851,7 @@ async def initialize(self): monkeypatch.setattr(streamable_mod, "streamablehttp_client", fake_streamablehttp_client) monkeypatch.setattr(session_mod, "ClientSession", FakeClientSession) monkeypatch.setenv("TOOLBOX_ENDPOINT", "https://example.test/toolboxes/t/mcp") + monkeypatch.setenv("AGENTKIT_MCP_TIMEOUT", "7") spec = AgentSpec.model_validate({ "abiVersion": "v0", "metadata": {"name": "x"}, @@ -703,6 +872,7 @@ async def initialize(self): assert "httpx_client_factory" in calls assert "http_client" not in calls assert calls["initialized"] is True + assert calls["session_timeout"].total_seconds() == 7 diff --git a/runtimes/microsoft-agent-framework/tests/test_lifecycle.py b/runtimes/microsoft-agent-framework/tests/test_lifecycle.py new file mode 100644 index 0000000..8d77291 --- /dev/null +++ b/runtimes/microsoft-agent-framework/tests/test_lifecycle.py @@ -0,0 +1,1091 @@ +from __future__ import annotations + +import asyncio +import json +from unittest import mock + +import pytest +from mcp import McpError + +from agentkit_serve import agent_factory +from agentkit_serve_common.config import AgentSpec, ToolSpec + + +def _runtime() -> agent_factory.MAFRuntime: + spec = AgentSpec.model_validate( + { + "abiVersion": "v0", + "metadata": {"name": "lifecycle-test"}, + "model": { + "provider": "openai-compatible", + "baseURL": "https://api.openai.com/v1", + "name": "gpt-4o-mini", + }, + "instructions": "Be helpful.", + "tools": [], + "expose": {"openai": True, "port": 8080}, + } + ) + return agent_factory.MAFRuntime(spec) + + +def test_agent_partial_enter_failure_closes_agent_and_context_without_masking_startup_error(): + runtime = _runtime() + startup_error = RuntimeError("agent startup failed") + agent_cleanup_error = RuntimeError("agent cleanup failed") + events: list[str] = [] + owner_tasks: dict[str, asyncio.Task] = {} + + class _ContextResource: + async def __aenter__(self): + events.append("context-enter") + owner_tasks["context"] = asyncio.current_task() + return self + + async def __aexit__(self, exc_type, exc, tb): + assert asyncio.current_task() is owner_tasks["context"] + events.append("context-exit") + + class _PartialAgent: + async def __aenter__(self): + events.append("agent-enter") + owner_tasks["agent"] = asyncio.current_task() + raise startup_error + + async def __aexit__(self, exc_type, exc, tb): + assert asyncio.current_task() is owner_tasks["agent"] + events.append("agent-exit") + raise agent_cleanup_error + + async def _build_context_providers(): + await runtime.stack.enter_async_context(_ContextResource()) + return None + + runtime._build_context_providers = _build_context_providers # type: ignore[method-assign] + with mock.patch("agentkit_serve.agent_factory.build_agent", return_value=_PartialAgent()): + with pytest.raises(RuntimeError) as exc_info: + asyncio.run(runtime.__aenter__()) + + assert exc_info.value is startup_error + assert events == ["context-enter", "agent-enter", "agent-exit", "context-exit"] + assert runtime.agent is None + + +def test_runtime_exit_closes_context_stack_even_when_agent_exit_raises(): + runtime = _runtime() + agent_exit_error = RuntimeError("agent exit failed") + events: list[str] = [] + owner_tasks: dict[str, asyncio.Task] = {} + + class _ContextResource: + async def __aenter__(self): + events.append("context-enter") + owner_tasks["context"] = asyncio.current_task() + return self + + async def __aexit__(self, exc_type, exc, tb): + assert asyncio.current_task() is owner_tasks["context"] + events.append("context-exit") + + class _Agent: + async def __aenter__(self): + events.append("agent-enter") + owner_tasks["agent"] = asyncio.current_task() + return self + + async def __aexit__(self, exc_type, exc, tb): + assert asyncio.current_task() is owner_tasks["agent"] + events.append("agent-exit") + raise agent_exit_error + + async def _build_context_providers(): + await runtime.stack.enter_async_context(_ContextResource()) + return None + + async def exercise() -> None: + runtime._build_context_providers = _build_context_providers # type: ignore[method-assign] + with mock.patch("agentkit_serve.agent_factory.build_agent", return_value=_Agent()): + await runtime.__aenter__() + with pytest.raises(RuntimeError) as exc_info: + await runtime.__aexit__(None, None, None) + assert exc_info.value is agent_exit_error + + asyncio.run(exercise()) + + assert events == ["context-enter", "agent-enter", "agent-exit", "context-exit"] + assert runtime.agent is None + + +def test_cancellation_while_opening_second_context_resource_finishes_cleanup(): + runtime = _runtime() + second_opened = asyncio.Event() + second_exit_started = asyncio.Event() + allow_second_exit = asyncio.Event() + exited: list[str] = [] + owner_tasks: dict[str, asyncio.Task] = {} + + class _ContextResource: + def __init__(self, name: str) -> None: + self.name = name + + async def __aenter__(self): + owner_tasks[self.name] = asyncio.current_task() + return self + + async def __aexit__(self, exc_type, exc, tb): + assert asyncio.current_task() is owner_tasks[self.name] + if self.name == "second": + second_exit_started.set() + await allow_second_exit.wait() + exited.append(self.name) + + async def _build_context_providers(): + await runtime.stack.enter_async_context(_ContextResource("first")) + await runtime.stack.enter_async_context(_ContextResource("second")) + second_opened.set() + await asyncio.Event().wait() + + async def exercise() -> None: + runtime._build_context_providers = _build_context_providers # type: ignore[method-assign] + task = asyncio.create_task(runtime.__aenter__()) + try: + await asyncio.wait_for(second_opened.wait(), timeout=1) + task.cancel() + await asyncio.wait_for(second_exit_started.wait(), timeout=1) + + task.cancel() + await asyncio.sleep(0) + assert not task.done() + + allow_second_exit.set() + with pytest.raises(asyncio.CancelledError): + await task + finally: + allow_second_exit.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + asyncio.run(exercise()) + + assert exited == ["second", "first"] + + +def test_exit_cancellation_waits_for_owner_cleanup_and_keeps_cancellation_primary(): + runtime = _runtime() + cleanup_error = RuntimeError("context cleanup failed") + exit_started = asyncio.Event() + allow_exit = asyncio.Event() + events: list[str] = [] + owner_tasks: dict[str, asyncio.Task] = {} + + class _ContextResource: + async def __aenter__(self): + events.append("context-enter") + owner_tasks["context"] = asyncio.current_task() + return self + + async def __aexit__(self, exc_type, exc, tb): + assert asyncio.current_task() is owner_tasks["context"] + exit_started.set() + await allow_exit.wait() + events.append("context-exit") + raise cleanup_error + + class _Agent: + async def __aenter__(self): + events.append("agent-enter") + owner_tasks["agent"] = asyncio.current_task() + return self + + async def __aexit__(self, exc_type, exc, tb): + assert asyncio.current_task() is owner_tasks["agent"] + events.append("agent-exit") + + async def _build_context_providers(): + await runtime.stack.enter_async_context(_ContextResource()) + return None + + async def exercise() -> None: + runtime._build_context_providers = _build_context_providers # type: ignore[method-assign] + with mock.patch("agentkit_serve.agent_factory.build_agent", return_value=_Agent()): + await runtime.__aenter__() + + exit_task = asyncio.create_task(runtime.__aexit__(None, None, None)) + try: + await asyncio.wait_for(exit_started.wait(), timeout=1) + exit_task.cancel() + allow_exit.set() + with pytest.raises(asyncio.CancelledError) as exc_info: + await exit_task + assert exc_info.value.__cause__ is cleanup_error + finally: + allow_exit.set() + if not exit_task.done(): + exit_task.cancel() + await asyncio.gather(exit_task, return_exceptions=True) + + asyncio.run(exercise()) + + assert events == ["context-enter", "agent-enter", "agent-exit", "context-exit"] + + +@pytest.mark.parametrize( + ("raw", "expected"), + [(None, 120), ("garbage", 120), ("-1", 120), ("7.9", 7)], +) +def test_stdio_mcp_timeout_defaults_to_120_seconds_and_honors_override( + monkeypatch, + raw: str | None, + expected: int, +): + tool = ToolSpec(name="fetch", command=["fake-mcp"], env=[]) + + if raw is None: + monkeypatch.delenv("AGENTKIT_MCP_TIMEOUT", raising=False) + else: + monkeypatch.setenv("AGENTKIT_MCP_TIMEOUT", raw) + assert agent_factory.build_tool(tool).request_timeout == expected + + +def test_mcp_skills_initialize_timeout_closes_live_transport(monkeypatch): + request_bodies: list[bytes] = [] + active_writers: set[asyncio.StreamWriter] = set() + handler_tasks: set[asyncio.Task[None]] = set() + accepted = asyncio.Event() + disconnected = asyncio.Event() + + async def handle_client( + reader: asyncio.StreamReader, + writer: asyncio.StreamWriter, + ) -> None: + task = asyncio.current_task() + assert task is not None + handler_tasks.add(task) + active_writers.add(writer) + try: + header_bytes = await reader.readuntil(b"\r\n\r\n") + headers = header_bytes.decode("ascii").split("\r\n") + content_length = next( + int(line.split(":", 1)[1].strip()) + for line in headers + if line.lower().startswith("content-length:") + ) + request_bodies.append(await reader.readexactly(content_length)) + writer.write( + b"HTTP/1.1 202 Accepted\r\n" + b"Content-Length: 0\r\n" + b"Connection: keep-alive\r\n" + b"\r\n" + ) + await writer.drain() + accepted.set() + await reader.read() + finally: + writer.close() + await writer.wait_closed() + active_writers.discard(writer) + handler_tasks.discard(task) + disconnected.set() + + async def exercise() -> None: + server = await asyncio.start_server(handle_client, "127.0.0.1", 0) + assert server.sockets + port = server.sockets[0].getsockname()[1] + monkeypatch.setenv("AGENTKIT_MCP_TIMEOUT", "1") + monkeypatch.setenv("TOOLBOX_ENDPOINT", f"http://127.0.0.1:{port}/mcp") + spec = AgentSpec.model_validate( + { + "abiVersion": "v0", + "metadata": {"name": "mcp-skills-timeout"}, + "model": { + "provider": "openai-compatible", + "baseURL": "https://api.openai.com/v1", + "name": "gpt-4o-mini", + }, + "instructions": "Be helpful.", + "tools": [ + { + "name": "toolbox", + "type": "mcp", + "transport": "streamable-http", + "urlEnv": "TOOLBOX_ENDPOINT", + } + ], + "context": { + "providers": [ + { + "type": "skills", + "source": "mcp", + "toolRef": "toolbox", + } + ] + }, + "expose": {"openai": True, "port": 8080}, + } + ) + runtime = agent_factory.MAFRuntime(spec) + + try: + with pytest.raises(McpError, match="Timed out while waiting for response"): + await asyncio.wait_for(runtime.__aenter__(), timeout=3) + + await asyncio.wait_for(accepted.wait(), timeout=1) + await asyncio.wait_for(disconnected.wait(), timeout=1) + assert runtime.agent is None + assert not active_writers + assert json.loads(request_bodies[0])["method"] == "initialize" + finally: + server.close() + for writer in tuple(active_writers): + writer.close() + if active_writers: + await asyncio.gather( + *(writer.wait_closed() for writer in tuple(active_writers)), + return_exceptions=True, + ) + if handler_tasks: + await asyncio.gather(*tuple(handler_tasks), return_exceptions=True) + await server.wait_closed() + + asyncio.run(exercise()) + + +def test_runtime_owns_remote_mcp_http_client_and_closes_it_after_agent(monkeypatch): + monkeypatch.setenv("TOOLBOX_ENDPOINT", "http://127.0.0.1:8765/mcp") + spec = AgentSpec.model_validate( + { + "abiVersion": "v0", + "metadata": {"name": "remote-client-lifecycle"}, + "model": { + "provider": "openai-compatible", + "baseURL": "https://api.openai.com/v1", + "name": "gpt-4o-mini", + }, + "instructions": "Be helpful.", + "tools": [ + { + "name": "toolbox", + "type": "mcp", + "transport": "streamable-http", + "urlEnv": "TOOLBOX_ENDPOINT", + } + ], + "expose": {"openai": True, "port": 8080}, + } + ) + runtime = agent_factory.MAFRuntime(spec) + agents = [] + events: list[str] = [] + + class _Agent: + def __init__(self, **kwargs): + self.tools = kwargs["tools"] + agents.append(self) + + async def __aenter__(self): + events.append("agent-enter") + return self + + async def __aexit__(self, exc_type, exc, tb): + http_client = self.tools[0]._httpx_client + assert http_client is not None + assert http_client.is_closed is False + events.append("agent-exit") + + async def exercise() -> None: + with ( + mock.patch("agentkit_serve.agent_factory.Agent", _Agent), + mock.patch("agentkit_serve.agent_factory.build_client", return_value=object()), + ): + await runtime.__aenter__() + http_client = agents[0].tools[0]._httpx_client + assert http_client is not None + assert http_client.is_closed is False + await runtime.__aexit__(None, None, None) + assert http_client.is_closed is True + + asyncio.run(exercise()) + assert events == ["agent-enter", "agent-exit"] + + +def test_runtime_enters_async_context_provider_and_closes_it_before_credential(monkeypatch): + monkeypatch.setenv("SEARCH_ENDPOINT", "https://example.search.windows.net") + monkeypatch.setenv("SEARCH_INDEX", "knowledge") + spec = AgentSpec.model_validate( + { + "abiVersion": "v0", + "metadata": {"name": "context-lifecycle"}, + "model": { + "provider": "openai-compatible", + "baseURL": "https://api.openai.com/v1", + "name": "gpt-4o-mini", + }, + "instructions": "Be helpful.", + "tools": [], + "context": { + "providers": [ + { + "name": "knowledge", + "type": "search", + "endpointEnv": "SEARCH_ENDPOINT", + "indexEnv": "SEARCH_INDEX", + } + ] + }, + "expose": {"openai": True, "port": 8080}, + } + ) + runtime = agent_factory.MAFRuntime(spec) + events: list[str] = [] + + class _Credential: + async def close(self): + events.append("credential-close") + + class _Provider: + def __init__(self, **kwargs): + self.credential = kwargs["credential"] + events.append("provider-create") + + async def __aenter__(self): + events.append("provider-enter") + return self + + async def __aexit__(self, exc_type, exc, tb): + events.append("provider-exit") + + class _AzureModule: + AzureAISearchContextProvider = _Provider + + class _Agent: + async def __aenter__(self): + events.append("agent-enter") + return self + + async def __aexit__(self, exc_type, exc, tb): + events.append("agent-exit") + + def _build_agent(spec, *, context_providers=None, stack=None, client=None): + assert context_providers and isinstance(context_providers[0], _Provider) + return _Agent() + + async def exercise() -> None: + with ( + mock.patch("agentkit_serve.agent_factory._credential_for_context", return_value=_Credential()), + mock.patch("agentkit_serve.agent_factory.build_agent", side_effect=_build_agent), + mock.patch("agentkit_serve.agent_factory.importlib.import_module", return_value=_AzureModule), + ): + await runtime.__aenter__() + await runtime.__aexit__(None, None, None) + + asyncio.run(exercise()) + assert events == [ + "provider-create", + "provider-enter", + "agent-enter", + "agent-exit", + "provider-exit", + "credential-close", + ] + + +def test_runtime_owns_memory_provider_project_client_and_sync_credential(monkeypatch): + monkeypatch.setenv("MEMORY_ENDPOINT", "https://example.services.ai.azure.com/api/projects/proj") + monkeypatch.setenv("MEMORY_STORE_NAME", "agentkit-memory") + monkeypatch.setenv("AGENTKIT_MEMORY_SCOPE", "scope-1") + spec = AgentSpec.model_validate( + { + "abiVersion": "v0", + "metadata": {"name": "memory-lifecycle"}, + "model": { + "provider": "openai-compatible", + "baseURL": "https://api.openai.com/v1", + "name": "gpt-4o-mini", + }, + "instructions": "Be helpful.", + "tools": [], + "context": { + "providers": [ + { + "name": "memory", + "type": "memory", + "endpointEnv": "MEMORY_ENDPOINT", + "storeNameEnv": "MEMORY_STORE_NAME", + } + ] + }, + "expose": {"openai": True, "port": 8080}, + } + ) + runtime = agent_factory.MAFRuntime(spec) + events: list[str] = [] + + class _Credential: + def close(self): + events.append("credential-close") + + class _ProjectClient: + async def __aenter__(self): + events.append("project-enter") + return self + + async def __aexit__(self, exc_type, exc, tb): + events.append("project-exit") + + class _MemoryProvider: + def __init__(self, **kwargs): + self.credential = kwargs["credential"] + self.project_client = _ProjectClient() + events.append("provider-create") + + async def __aenter__(self): + events.append("provider-enter") + await self.project_client.__aenter__() + return self + + async def __aexit__(self, exc_type, exc, tb): + events.append("provider-exit-start") + await self.project_client.__aexit__(exc_type, exc, tb) + events.append("provider-exit") + + class _FoundryModule: + FoundryMemoryProvider = _MemoryProvider + + class _Agent: + async def __aenter__(self): + events.append("agent-enter") + return self + + async def __aexit__(self, exc_type, exc, tb): + events.append("agent-exit") + + def _build_agent(spec, *, context_providers=None, stack=None, client=None): + assert context_providers and isinstance(context_providers[0], _MemoryProvider) + return _Agent() + + async def exercise() -> None: + with ( + mock.patch("agentkit_serve.agent_factory._credential_for_context", return_value=_Credential()), + mock.patch("agentkit_serve.agent_factory.build_agent", side_effect=_build_agent), + mock.patch("agentkit_serve.agent_factory.importlib.import_module", return_value=_FoundryModule), + ): + await runtime.__aenter__() + await runtime.__aexit__(None, None, None) + + asyncio.run(exercise()) + assert events == [ + "provider-create", + "provider-enter", + "project-enter", + "agent-enter", + "agent-exit", + "provider-exit-start", + "project-exit", + "provider-exit", + "credential-close", + ] + + +def test_runtime_owns_model_fallback_credential_and_project_client_without_double_closing_framework_client( + monkeypatch, +): + for name in ( + "AGENTKIT_MODEL_WORKLOAD_IDENTITY_TOKEN", + "AGENTKIT_WORKLOAD_IDENTITY_TOKEN", + "AGENTKIT_WORKLOAD_IDENTITY_TOKEN_COMMAND", + ): + monkeypatch.delenv(name, raising=False) + spec = AgentSpec.model_validate( + { + "abiVersion": "v0", + "metadata": {"name": "model-lifecycle"}, + "model": { + "provider": "openai-compatible", + "baseURL": "https://example.services.ai.azure.com/api/projects/proj/openai/v1", + "name": "gpt-4.1-mini", + "auth": { + "type": "workload-identity-token", + "audience": "https://ai.azure.com/.default", + }, + }, + "instructions": "Be helpful.", + "tools": [], + "expose": {"openai": True, "port": 8080}, + } + ) + runtime = agent_factory.MAFRuntime(spec) + events: list[str] = [] + + class _Credential: + def __init__(self): + events.append("credential-create") + + def close(self): + events.append("credential-close") + + class _ProjectClient: + async def __aenter__(self): + events.append("project-enter") + return self + + async def __aexit__(self, exc_type, exc, tb): + events.append("project-exit") + + class _ModelHTTPClient: + async def close(self): + events.append("model-http-close") + + class _FoundryClient: + def __init__(self, **kwargs): + assert isinstance(kwargs["credential"], _Credential) + self.project_client = _ProjectClient() + self.client = _ModelHTTPClient() + events.append("client-create") + + async def __aenter__(self): + events.append("client-enter") + return self + + async def __aexit__(self, exc_type, exc, tb): + await self.client.close() + events.append("client-exit") + + class _Agent: + def __init__(self, **kwargs): + self.client = kwargs["client"] + + async def __aenter__(self): + events.append("agent-enter") + await self.client.__aenter__() + return self + + async def __aexit__(self, exc_type, exc, tb): + events.append("agent-exit-start") + await self.client.__aexit__(exc_type, exc, tb) + events.append("agent-exit") + + async def exercise() -> None: + with ( + mock.patch("azure.identity.DefaultAzureCredential", _Credential), + mock.patch("agent_framework.foundry.FoundryChatClient", _FoundryClient), + mock.patch("agentkit_serve.agent_factory.Agent", _Agent), + ): + await runtime.__aenter__() + await runtime.__aexit__(None, None, None) + + asyncio.run(exercise()) + assert events == [ + "credential-create", + "client-create", + "project-enter", + "agent-enter", + "client-enter", + "agent-exit-start", + "model-http-close", + "client-exit", + "agent-exit", + "project-exit", + "credential-close", + ] + + +def test_partial_agent_startup_closes_all_runtime_owned_resources_in_dependency_order(monkeypatch): + for name in ( + "AGENTKIT_MODEL_WORKLOAD_IDENTITY_TOKEN", + "AGENTKIT_WORKLOAD_IDENTITY_TOKEN", + "AGENTKIT_WORKLOAD_IDENTITY_TOKEN_COMMAND", + ): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("SEARCH_ENDPOINT", "https://example.search.windows.net") + monkeypatch.setenv("SEARCH_INDEX", "knowledge") + monkeypatch.setenv("TOOLBOX_ENDPOINT", "http://127.0.0.1:8765/mcp") + spec = AgentSpec.model_validate( + { + "abiVersion": "v0", + "metadata": {"name": "partial-startup-lifecycle"}, + "model": { + "provider": "openai-compatible", + "baseURL": "https://example.services.ai.azure.com/api/projects/proj/openai/v1", + "name": "gpt-4.1-mini", + "auth": { + "type": "workload-identity-token", + "audience": "https://ai.azure.com/.default", + }, + }, + "instructions": "Be helpful.", + "tools": [ + { + "name": "toolbox", + "type": "mcp", + "transport": "streamable-http", + "urlEnv": "TOOLBOX_ENDPOINT", + } + ], + "context": { + "providers": [ + { + "name": "knowledge", + "type": "search", + "endpointEnv": "SEARCH_ENDPOINT", + "indexEnv": "SEARCH_INDEX", + } + ] + }, + "expose": {"openai": True, "port": 8080}, + } + ) + runtime = agent_factory.MAFRuntime(spec) + startup_error = RuntimeError("agent startup failed") + events: list[str] = [] + + class _ContextCredential: + def __init__(self): + events.append("context-credential-create") + + async def close(self): + events.append("context-credential-close") + + class _ContextProvider: + def __init__(self, **kwargs): + events.append("context-provider-create") + + async def __aenter__(self): + events.append("context-provider-enter") + return self + + async def __aexit__(self, exc_type, exc, tb): + events.append("context-provider-exit") + + class _AzureModule: + AzureAISearchContextProvider = _ContextProvider + + class _ModelCredential: + def __init__(self): + events.append("model-credential-create") + + def close(self): + events.append("model-credential-close") + + class _ProjectClient: + async def __aenter__(self): + events.append("project-enter") + return self + + async def __aexit__(self, exc_type, exc, tb): + events.append("project-exit") + + class _ModelHTTPClient: + async def close(self): + events.append("model-http-close") + + class _FoundryClient: + def __init__(self, **kwargs): + self.project_client = _ProjectClient() + self.client = _ModelHTTPClient() + events.append("model-client-create") + + class _HTTPClient: + def __init__(self, **kwargs): + events.append("http-client-create") + + async def aclose(self): + events.append("http-client-close") + + class _PartialAgent: + def __init__(self, **kwargs): + pass + + async def __aenter__(self): + events.append("agent-enter") + raise startup_error + + async def __aexit__(self, exc_type, exc, tb): + events.append("agent-exit") + + async def exercise() -> None: + with ( + mock.patch( + "agentkit_serve.agent_factory._credential_for_context", + side_effect=lambda *args, **kwargs: _ContextCredential(), + ), + mock.patch("azure.identity.DefaultAzureCredential", _ModelCredential), + mock.patch("agent_framework.foundry.FoundryChatClient", _FoundryClient), + mock.patch("agentkit_serve.agent_factory.AsyncClient", _HTTPClient), + mock.patch("agentkit_serve.agent_factory.Agent", _PartialAgent), + mock.patch("agentkit_serve.agent_factory.importlib.import_module", return_value=_AzureModule), + ): + with pytest.raises(RuntimeError) as exc_info: + await runtime.__aenter__() + assert exc_info.value is startup_error + + asyncio.run(exercise()) + assert events == [ + "context-credential-create", + "context-provider-create", + "context-provider-enter", + "model-credential-create", + "model-client-create", + "project-enter", + "http-client-create", + "agent-enter", + "agent-exit", + "http-client-close", + "model-http-close", + "project-exit", + "model-credential-close", + "context-provider-exit", + "context-credential-close", + ] + assert runtime.agent is None + + +def test_startup_cancellation_waits_for_all_runtime_owned_resource_cleanup(monkeypatch): + for name in ( + "AGENTKIT_MODEL_WORKLOAD_IDENTITY_TOKEN", + "AGENTKIT_WORKLOAD_IDENTITY_TOKEN", + "AGENTKIT_WORKLOAD_IDENTITY_TOKEN_COMMAND", + ): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("SEARCH_ENDPOINT", "https://example.search.windows.net") + monkeypatch.setenv("SEARCH_INDEX", "knowledge") + monkeypatch.setenv("TOOLBOX_ENDPOINT", "http://127.0.0.1:8765/mcp") + spec = AgentSpec.model_validate( + { + "abiVersion": "v0", + "metadata": {"name": "cancelled-startup-lifecycle"}, + "model": { + "provider": "openai-compatible", + "baseURL": "https://example.services.ai.azure.com/api/projects/proj/openai/v1", + "name": "gpt-4.1-mini", + "auth": { + "type": "workload-identity-token", + "audience": "https://ai.azure.com/.default", + }, + }, + "instructions": "Be helpful.", + "tools": [ + { + "name": "toolbox", + "type": "mcp", + "transport": "streamable-http", + "urlEnv": "TOOLBOX_ENDPOINT", + } + ], + "context": { + "providers": [ + { + "name": "knowledge", + "type": "search", + "endpointEnv": "SEARCH_ENDPOINT", + "indexEnv": "SEARCH_INDEX", + } + ] + }, + "expose": {"openai": True, "port": 8080}, + } + ) + runtime = agent_factory.MAFRuntime(spec) + events: list[str] = [] + agent_started = asyncio.Event() + http_close_started = asyncio.Event() + allow_http_close = asyncio.Event() + + class _ContextCredential: + async def close(self): + events.append("context-credential-close") + + class _ContextProvider: + def __init__(self, **kwargs): + events.append("context-provider-create") + + async def __aenter__(self): + events.append("context-provider-enter") + return self + + async def __aexit__(self, exc_type, exc, tb): + events.append("context-provider-exit") + + class _AzureModule: + AzureAISearchContextProvider = _ContextProvider + + class _ModelCredential: + def close(self): + events.append("model-credential-close") + + class _ProjectClient: + async def __aenter__(self): + events.append("project-enter") + return self + + async def __aexit__(self, exc_type, exc, tb): + events.append("project-exit") + + class _ModelHTTPClient: + async def close(self): + events.append("model-http-close") + + class _FoundryClient: + def __init__(self, **kwargs): + self.project_client = _ProjectClient() + self.client = _ModelHTTPClient() + events.append("model-client-create") + + class _HTTPClient: + def __init__(self, **kwargs): + events.append("http-client-create") + + async def aclose(self): + events.append("http-client-close-start") + http_close_started.set() + await allow_http_close.wait() + events.append("http-client-close") + + class _BlockingAgent: + def __init__(self, **kwargs): + pass + + async def __aenter__(self): + events.append("agent-enter") + agent_started.set() + await asyncio.Event().wait() + + async def __aexit__(self, exc_type, exc, tb): + events.append("agent-exit") + + async def exercise() -> None: + with ( + mock.patch( + "agentkit_serve.agent_factory._credential_for_context", + side_effect=lambda *args, **kwargs: _ContextCredential(), + ), + mock.patch("azure.identity.DefaultAzureCredential", _ModelCredential), + mock.patch("agent_framework.foundry.FoundryChatClient", _FoundryClient), + mock.patch("agentkit_serve.agent_factory.AsyncClient", _HTTPClient), + mock.patch("agentkit_serve.agent_factory.Agent", _BlockingAgent), + mock.patch("agentkit_serve.agent_factory.importlib.import_module", return_value=_AzureModule), + ): + task = asyncio.create_task(runtime.__aenter__()) + try: + await asyncio.wait_for(agent_started.wait(), timeout=1) + task.cancel() + await asyncio.wait_for(http_close_started.wait(), timeout=1) + + task.cancel() + await asyncio.sleep(0) + assert task.done() is False + + allow_http_close.set() + with pytest.raises(asyncio.CancelledError): + await task + finally: + allow_http_close.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + asyncio.run(exercise()) + assert events == [ + "context-provider-create", + "context-provider-enter", + "model-client-create", + "project-enter", + "http-client-create", + "agent-enter", + "agent-exit", + "http-client-close-start", + "http-client-close", + "model-http-close", + "project-exit", + "model-credential-close", + "context-provider-exit", + "context-credential-close", + ] + assert runtime.agent is None + + +def test_runtime_closes_model_fallback_http_client_before_project_and_credential(monkeypatch): + for name in ( + "AGENTKIT_MODEL_WORKLOAD_IDENTITY_TOKEN", + "AGENTKIT_WORKLOAD_IDENTITY_TOKEN", + "AGENTKIT_WORKLOAD_IDENTITY_TOKEN_COMMAND", + ): + monkeypatch.delenv(name, raising=False) + spec = AgentSpec.model_validate( + { + "abiVersion": "v0", + "metadata": {"name": "model-http-lifecycle"}, + "model": { + "provider": "openai-compatible", + "baseURL": "https://example.services.ai.azure.com/api/projects/proj/openai/v1", + "name": "gpt-4.1-mini", + "auth": { + "type": "workload-identity-token", + "audience": "https://ai.azure.com/.default", + }, + }, + "instructions": "Be helpful.", + "tools": [], + "expose": {"openai": True, "port": 8080}, + } + ) + runtime = agent_factory.MAFRuntime(spec) + events: list[str] = [] + + class _Credential: + def __init__(self): + events.append("credential-create") + + def close(self): + events.append("credential-close") + + class _ProjectClient: + async def __aenter__(self): + events.append("project-enter") + return self + + async def __aexit__(self, exc_type, exc, tb): + events.append("project-exit") + + class _ModelHTTPClient: + async def close(self): + events.append("model-http-close") + + class _FoundryClient: + def __init__(self, **kwargs): + assert isinstance(kwargs["credential"], _Credential) + self.project_client = _ProjectClient() + self.client = _ModelHTTPClient() + events.append("client-create") + + class _Agent: + def __init__(self, **kwargs): + self.client = kwargs["client"] + + async def __aenter__(self): + events.append("agent-enter") + return self + + async def __aexit__(self, exc_type, exc, tb): + events.append("agent-exit") + + async def exercise() -> None: + with ( + mock.patch("azure.identity.DefaultAzureCredential", _Credential), + mock.patch("agent_framework.foundry.FoundryChatClient", _FoundryClient), + mock.patch("agentkit_serve.agent_factory.Agent", _Agent), + ): + await runtime.__aenter__() + await runtime.__aexit__(None, None, None) + + asyncio.run(exercise()) + assert events == [ + "credential-create", + "client-create", + "project-enter", + "agent-enter", + "agent-exit", + "model-http-close", + "project-exit", + "credential-close", + ] diff --git a/runtimes/pydantic-ai/agentkit_serve/agent_factory.py b/runtimes/pydantic-ai/agentkit_serve/agent_factory.py index a6509ea..8ce1e27 100644 --- a/runtimes/pydantic-ai/agentkit_serve/agent_factory.py +++ b/runtimes/pydantic-ai/agentkit_serve/agent_factory.py @@ -67,15 +67,15 @@ offline_orka_echo_enabled, ) -# Seconds to wait for a stdio MCP server's initialize handshake. pydantic-ai's -# default is 5s, which is too tight for a COLD `uvx`/`npx` tool: the first launch -# resolves, downloads, and installs the server package before it speaks MCP, which -# routinely exceeds 5s. Default generously and let operators tune via env. +# Seconds to wait for stdio MCP initialization and each subsequent request. +# pydantic-ai's 5s initialization default is too tight for a COLD `uvx`/`npx` +# tool: the first launch resolves, downloads, and installs the server package +# before it speaks MCP. Default generously and let operators tune both phases. _DEFAULT_MCP_INIT_TIMEOUT = 120.0 def _mcp_init_timeout() -> float: - """MCP stdio init timeout (seconds), overridable via AGENTKIT_MCP_TIMEOUT.""" + """MCP stdio init/read timeout, overridable via AGENTKIT_MCP_TIMEOUT.""" return positive_float_env(default=_DEFAULT_MCP_INIT_TIMEOUT) @@ -114,8 +114,8 @@ def build_tool_server(tool: ToolSpec) -> Any: command, args = split_tool_command(tool, example='["npx", "-y", "..."]') - # Generous init timeout: a cold uvx/npx tool installs its package before - # speaking MCP, which exceeds pydantic-ai's 5s default (see above). + # A single operator timeout bounds both cold initialization and later tool + # calls, preventing a blocked stdio server from hanging a request forever. env = declared_tool_env(tool) if MCPServerStdio is not None: @@ -124,6 +124,7 @@ def build_tool_server(tool: ToolSpec) -> Any: args=args, env=env, timeout=timeout, + read_timeout=timeout, # tool_prefix namespaces tool names so two servers can't collide. tool_prefix=tool.name, ) @@ -142,7 +143,7 @@ def build_tool_server(tool: ToolSpec) -> Any: # the stdio subprocess should be torn down instead of kept alive. keep_alive=False, ) - return MCPToolset(transport, init_timeout=timeout).prefixed(tool.name) + return MCPToolset(transport, init_timeout=timeout, read_timeout=timeout).prefixed(tool.name) def build_agent(spec: AgentSpec) -> Agent: diff --git a/runtimes/pydantic-ai/tests/test_lifecycle.py b/runtimes/pydantic-ai/tests/test_lifecycle.py new file mode 100644 index 0000000..3f1b402 --- /dev/null +++ b/runtimes/pydantic-ai/tests/test_lifecycle.py @@ -0,0 +1,138 @@ +from __future__ import annotations + +import asyncio +import os +import signal +import sys +import time +from pathlib import Path + +import pytest + +from agentkit_serve import agent_factory +from agentkit_serve_common.config import ToolSpec + + +def _process_exists(pid: int) -> bool: + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + return True + + +async def _wait_for_pid(pid_path: Path, *, timeout: float = 3.0) -> int: + deadline = asyncio.get_running_loop().time() + timeout + while asyncio.get_running_loop().time() < deadline: + if pid_path.exists(): + return int(pid_path.read_text(encoding="utf-8")) + await asyncio.sleep(0.01) + raise AssertionError("stdio MCP child did not publish its PID") + + +async def _wait_for_process_exit(pid: int, *, timeout: float = 3.0) -> None: + deadline = asyncio.get_running_loop().time() + timeout + while asyncio.get_running_loop().time() < deadline: + if not _process_exists(pid): + return + await asyncio.sleep(0.01) + raise AssertionError(f"stdio MCP child process {pid} was not cleaned up") + + +def test_stdio_tool_call_honors_mcp_timeout_and_cleans_up_child(monkeypatch, tmp_path): + server_script = tmp_path / "blocked_mcp_server.py" + pid_path = tmp_path / "blocked_mcp_server.pid" + server_script.write_text( + """ +import json +import os +import sys +import time +from pathlib import Path + +Path(sys.argv[1]).write_text(str(os.getpid()), encoding="utf-8") + +for line in sys.stdin: + message = json.loads(line) + method = message.get("method") + if method == "initialize": + response = { + "jsonrpc": "2.0", + "id": message["id"], + "result": { + "protocolVersion": message["params"]["protocolVersion"], + "capabilities": {"tools": {}}, + "serverInfo": {"name": "blocked-tool-test", "version": "1.0"}, + }, + } + print(json.dumps(response), flush=True) + elif method == "tools/list": + response = { + "jsonrpc": "2.0", + "id": message["id"], + "result": { + "tools": [ + { + "name": "block_forever", + "description": "Never returns", + "inputSchema": {"type": "object", "properties": {}}, + } + ] + }, + } + print(json.dumps(response), flush=True) + elif method == "tools/call": + while True: + time.sleep(3600) +""".lstrip(), + encoding="utf-8", + ) + monkeypatch.setenv("AGENTKIT_MCP_TIMEOUT", "0.3") + tool = ToolSpec( + name="blocked", + command=[sys.executable, str(server_script), str(pid_path)], + env=[], + ) + child_pid: int | None = None + + async def exercise() -> None: + nonlocal child_pid + toolset = agent_factory.build_tool_server(tool) + wrapped = toolset.wrapped + async with wrapped: + child_pid = await _wait_for_pid(pid_path) + started = time.monotonic() + with pytest.raises(Exception): + await asyncio.wait_for( + wrapped.direct_call_tool("block_forever", {}), + timeout=2.0, + ) + assert time.monotonic() - started < 1.2 + assert _process_exists(child_pid) + + await _wait_for_process_exit(child_pid) + + try: + asyncio.run(exercise()) + finally: + if child_pid is not None and _process_exists(child_pid): + os.kill(child_pid, signal.SIGKILL) + + +def test_legacy_stdio_server_receives_init_and_read_timeout(monkeypatch): + captured: dict[str, object] = {} + + class _LegacyServer: + def __init__(self, **kwargs): + captured.update(kwargs) + + monkeypatch.setattr(agent_factory, "MCPServerStdio", _LegacyServer) + monkeypatch.setenv("AGENTKIT_MCP_TIMEOUT", "1.25") + + server = agent_factory.build_tool_server( + ToolSpec(name="legacy", command=["legacy-mcp", "--stdio"], env=[]) + ) + + assert isinstance(server, _LegacyServer) + assert captured["timeout"] == 1.25 + assert captured["read_timeout"] == 1.25 diff --git a/runtimes/pydantic-ai/tests/test_orka_lifecycle.py b/runtimes/pydantic-ai/tests/test_orka_lifecycle.py new file mode 100644 index 0000000..c2723df --- /dev/null +++ b/runtimes/pydantic-ai/tests/test_orka_lifecycle.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +import asyncio +import json +from datetime import UTC, datetime, timedelta +from typing import Any + +from fastapi.testclient import TestClient +from pydantic_ai import Agent +from pydantic_ai.models.test import TestModel +from pydantic_ai.toolsets import AbstractToolset + +from agentkit_serve import agent_factory +from agentkit_serve_common.config import AgentSpec +from agentkit_serve_common.orka import ORKA_HARNESS_VERSION, create_orka_app +from agentkit_serve_common.runtime import RuntimeSession + +AUTH = {"authorization": "Bearer mock-token"} + + +def _spec() -> AgentSpec: + return AgentSpec.model_validate( + { + "abiVersion": "v0", + "metadata": {"name": "pydantic-orka-lifecycle"}, + "model": { + "provider": "openai-compatible", + "baseURL": "https://api.openai.com/v1", + "name": "gpt-4o-mini", + }, + "instructions": "Be helpful.", + "tools": [], + "env": [], + "expose": {"openai": True, "port": 8080}, + } + ) + + +def _start_payload(*, turn_id: str, runtime_session_id: str, correlation_id: str, deadline: str) -> dict[str, Any]: + return { + "version": ORKA_HARNESS_VERSION, + "namespace": "default", + "taskName": "task-1", + "sessionName": "session-1", + "runtimeSessionID": runtime_session_id, + "turnID": turn_id, + "correlationID": correlation_id, + "deadline": deadline, + "authIdentity": {"subject": "system:serviceaccount:default:orka"}, + "input": {"prompt": turn_id, "contextRefs": [], "env": []}, + "toolExecutionMode": "observed", + "metadata": {}, + } + + +def _frames(response_text: str) -> list[dict[str, Any]]: + return [json.loads(line.removeprefix("data: ")) for line in response_text.splitlines() if line.startswith("data: ")] + + +class SlowExitToolset(AbstractToolset[None]): + def __init__(self, toolset_id: str, *, close_delay: float, slow_close_call: int | None = None) -> None: + self.toolset_id = toolset_id + self.close_delay = close_delay + self.slow_close_call = slow_close_call + self.entered = 0 + self.close_calls = 0 + self.close_completed = 0 + self.close_cancelled = 0 + + @property + def id(self) -> str: + return self.toolset_id + + async def __aenter__(self) -> SlowExitToolset: + self.entered += 1 + return self + + async def __aexit__(self, *args: Any) -> bool | None: + self.close_calls += 1 + try: + if self.close_calls == self.slow_close_call: + await asyncio.sleep(self.close_delay) + except asyncio.CancelledError: + self.close_cancelled += 1 + raise + self.close_completed += 1 + return None + + async def get_tools(self, ctx: Any) -> dict[str, Any]: + return {} + + async def call_tool(self, name: str, tool_args: dict[str, Any], ctx: Any, tool: Any) -> Any: + raise AssertionError("the lifecycle reproduction exposes no callable tools") + + +class PydanticAgentFactory: + def __init__(self, *, first_close_delay: float = 1.0) -> None: + self.first_close_delay = first_close_delay + self.toolsets: list[SlowExitToolset] = [] + + def build_runtime(self, spec: AgentSpec) -> RuntimeSession: + del spec + index = len(self.toolsets) + toolset = SlowExitToolset( + f"toolset-{index}", + close_delay=self.first_close_delay if index == 0 else 0, + # Pydantic enters/exits the toolset for the first run, then exits it + # again when the long-lived Agent context is evicted from Orka. + slow_close_call=2 if index == 0 else None, + ) + self.toolsets.append(toolset) + agent = Agent(TestModel(custom_output_text=f"runtime-{index}"), instructions="Be helpful.", toolsets=[toolset]) + return agent_factory.PydanticRuntime(agent) + + +def test_orka_cache_limit_preserves_real_pydantic_agent_toolset_cleanup_after_deadline(): + factory = PydanticAgentFactory() + app = create_orka_app(_spec(), factory, auth_token=AUTH["authorization"].removeprefix("Bearer "), max_runtime_sessions=1) + long_deadline = (datetime.now(UTC) + timedelta(minutes=1)).isoformat().replace("+00:00", "Z") + + with TestClient(app) as client: + first = client.post( + "/v1/turns", + json=_start_payload( + turn_id="turn-pydantic-one", + runtime_session_id="runtime-session-one", + correlation_id="corr-one", + deadline=long_deadline, + ), + headers=AUTH, + ) + assert first.status_code == 202 + first_frames = _frames(client.get("/v1/turns/turn-pydantic-one/events", headers=AUTH).text) + assert first_frames[-1]["type"] == "TurnCompleted" + + short_deadline = (datetime.now(UTC) + timedelta(milliseconds=500)).isoformat().replace("+00:00", "Z") + second = client.post( + "/v1/turns", + json=_start_payload( + turn_id="turn-pydantic-two", + runtime_session_id="runtime-session-two", + correlation_id="corr-two", + deadline=short_deadline, + ), + headers=AUTH, + ) + assert second.status_code == 202 + second_frames = _frames(client.get("/v1/turns/turn-pydantic-two/events", headers=AUTH).text) + assert second_frames[-1]["type"] == "TurnFailed" + assert second_frames[-1]["failed"]["reason"] == "DeadlineExceeded" + + assert len(factory.toolsets) == 1 + assert factory.toolsets[0].entered >= 2 + assert factory.toolsets[0].close_calls == factory.toolsets[0].entered + assert factory.toolsets[0].close_cancelled == 0 + assert factory.toolsets[0].close_completed == factory.toolsets[0].close_calls diff --git a/scripts/live-copilot-agent-e2e.sh b/scripts/live-copilot-agent-e2e.sh index 6ab9054..6c41255 100755 --- a/scripts/live-copilot-agent-e2e.sh +++ b/scripts/live-copilot-agent-e2e.sh @@ -190,7 +190,7 @@ main() { docker buildx inspect --bootstrap fi make build-agentkit TAG="${tag}" - make build-serve-maf TAG="${tag}" + make build-serve-maf TAG="${tag}" PLATFORM="${platform}" log "Building live MAF agent image" docker buildx build ${buildx_args[@]+"${buildx_args[@]}"} . -f test/agentkitfile-maf-live.yaml \