From e904bf7cdafa3516156ce459c3ec665a7b0bd6f1 Mon Sep 17 00:00:00 2001 From: daniel Date: Tue, 8 Sep 2026 15:08:49 +0100 Subject: [PATCH 1/2] feat: support compose-managed multimodal subagents --- docs/plans/compose-managed-files.md | 10 +- docs/user/compose-and-local-tools.md | 56 ++ fixtures/mock-acp.py | 8 - src/acp_child.rs | 260 +++++- src/managed_files.rs | 330 ++++++- src/managed_files/tests.rs | 266 ++++++ src/managed_files/tests/faults.rs | 85 +- src/protocols/acp.rs | 406 +++++++- src/provider/adapter.rs | 575 +++++++++++- src/provider/adapter_native_tests.rs | 869 ++++++++++++++++++ src/runtime.rs | 109 ++- src/runtime/tests/managed_files.rs | 14 + .../tests/managed_files/native_subagent.rs | 174 ++++ src/session.rs | 95 +- src/tools/subagent.rs | 562 ++++++++++- src/tools/subagent/native-image-fixture.py | 71 ++ src/tools/subagent/native_images.rs | 751 +++++++++++++++ src/tools/subagent/tests.rs | 3 + tests/runtime.rs | 8 - 19 files changed, 4556 insertions(+), 96 deletions(-) create mode 100644 src/provider/adapter_native_tests.rs create mode 100644 src/runtime/tests/managed_files/native_subagent.rs create mode 100644 src/tools/subagent/native-image-fixture.py create mode 100644 src/tools/subagent/native_images.rs diff --git a/docs/plans/compose-managed-files.md b/docs/plans/compose-managed-files.md index 813fac4d..f7704f1b 100644 --- a/docs/plans/compose-managed-files.md +++ b/docs/plans/compose-managed-files.md @@ -71,7 +71,7 @@ stickered = subagent({ attachments: [cropped], output_schema: { type: "object", - properties: { result: { "$ref": "kit://schemas/file/v1" } }, + properties: { result: { "$ref": "kit://schemas/file/v1", "x-kit-image-index": 0 } }, required: ["result"], additionalProperties: false } @@ -91,6 +91,14 @@ The pipeline remains one compose invocation with final-return-only delivery, usi Extend the currently text-oriented ACP child prompt/output path to retain typed attachments and assistant media. Import native generated media into managed storage. Bind actual emitted media to a file-aware output contract; models must not invent file IDs. Specify single-image binding and reject ambiguous multiple output. Validate shape AND reference existence/access, with explicit contract failure rather than silent string fallback for the new file-aware contract. Parent outputs must survive child close; grants and promotion must preserve session isolation. +### Phase 3 implemented contract + +The subagent tool layer accepts optional typed `attachments` on `subagent`, `prompt`, and `fork`, resolving authority from `ToolRequest.session_id`. Native ACP image inputs and generated outputs use managed storage rather than model-created identities. Attachment grants cross working-directory stores explicitly; generated-image parent publication and final access validation complete before the existing success transition. Errors use the existing create cleanup, continuation retry-handle, and fork cleanup paths without adding shared-state instrumentation. + +File-aware `output_schema` uses the locally resolved `kit://schemas/file/v1` reference. The supported binding is exactly one root File or one required fixed nested object-property path. By default, Kit requires exactly one distinct native assistant image; an optional caller-fixed `x-kit-image-index` integer 0–7 beside the exact File `$ref` instead selects a distinct image in first-emission order. Kit rejects attempted model binding, validates surrounding JSON and the completed schema, and only constructs omitted surrounding objects when the complete result is valid. Arrays, unions, conditionals, indirect references and multiple bindings are unsupported. Every native image occurrence is independently validated and charged against occurrence, encoded/decoded byte, and aggregate pixel budgets before deduplication. Only equal MIME types and byte-identical image payloads from the current output collapse; visually identical images with different bytes remain distinct and require an explicit index to select one. Every occurrence, including unselected images, must validate and satisfy budgets before selection. Out-of-range indices fail without fallback, and only the selected image is imported/published. Kit does not strip signed metadata or deduplicate perceptually. No model identity, input image, or previous turn participates in that comparison. Ordinary text-only schema fallback is preserved. Native images without a file-aware schema have the explicit `output: { value, files }` surface, not base64 diagnostic updates. The [user guide](../user/compose-and-local-tools.md#attach-files-to-subagents-and-return-native-images) specifies the contract and final-return-only behavior. + +The complete read/rotate/crop → built-in ACP subagent → managed output → child close → `return output.result` pipeline has been verified against the canonical OpenRouter `google/gemini-3-pro-image` route. The caller explicitly selected distinct native output index 0; this does not claim that the backend emits only one image. The selected bytes were visually verified to contain the requested Hello Kitty sticker. Eligibility comes from exact-model and concrete-endpoint capability discovery, not a model allowlist. Other harness/provider routes require their own native-output support. Shared TUI presentation remains Phase 4 work. + ### Shared TUI presentation Generalize user-image rendering into reusable media presentation for user attachments, tool results, assistant-generated media, and Markdown image nodes. Share decode/cache/terminal protocols and budgets, and align live updates with history replay. Keep media out of text-only search/previews/logs. diff --git a/docs/user/compose-and-local-tools.md b/docs/user/compose-and-local-tools.md index 5dec8192..a0b3a5ff 100644 --- a/docs/user/compose-and-local-tools.md +++ b/docs/user/compose-and-local-tools.md @@ -167,6 +167,62 @@ References survive process restart and source modification or deletion. Authoriz There is no automatic managed-file garbage collection in this phase. Calls, cancellation, session close, and process exit do not delete these objects. Cancelled or failed imports can leave unreachable objects. Explicit removal of a session's managed-file directory invalidates its references; do not remove retained objects that you need after resume. Finalization can repeat against the same immutable references without importing again. A delivery error states that the compose program already completed and side effects may have occurred; it is not a rollback or an invitation to retry blindly. +## Attach files to subagents and return native images + +The hidden `subagent`, `prompt`, and `fork` tools accept optional `attachments: FileReference[]` (at most eight). Pass managed File values explicitly; a file ID, path, URL, or descriptor pasted into prompt text does not attach an image. Kit resolves attachments in the invoking session, checks the existing aggregate image budgets, grants durable copies to the child, and sends native ACP image blocks after the text prompt. The harness must advertise ACP image **input** support. That advertisement does not promise image generation. + +Use the exact local schema reference `{"$ref":"kit://schemas/file/v1"}` for a strict native-image output contract: + +```text +source = read_file({ path: "source.png" }) +edited = subagent({ + model: "openrouter:google/gemini-3-pro-image", + prompt: "Edit the attached image and emit exactly one distinct native assistant image.", + attachments: [source], + output_schema: { + type: "object", + properties: { result: { "$ref": "kit://schemas/file/v1" } }, + required: ["result"], + additionalProperties: false + } +}) +return edited.output.result +``` + +This example uses the built-in Kit ACP harness and a concrete OpenRouter image-output model. Select an eligible model explicitly; a vision model behind a text-only harness cannot generate native output through that harness. Standard ACP advertises image **input**, not image generation. Actual valid assistant image bytes, not the capability flag or a model name, establish output success. + +For the canonical OpenRouter endpoint, selecting a concrete model advertised with image output opts into generation and its additional provider cost. Kit checks the exact model catalogue and its concrete endpoint document, derives output modalities from advertised capabilities, and requires tools support so compose remains available. Automatic routing entries without concrete endpoints retain ordinary behavior. Native requests require providers to support all requested parameters; Kit does not silently remove compose or relax routing to make an incompatible model work. Editing also requires advertised image input. Custom endpoints do not inherit official OpenRouter capability assertions. Discovery failures do not manufacture support, and a required File contract fails if no native image arrives. + +This route uses bounded nonstreaming chat completions: at most 24 MiB of raw response, eight images, 8 MiB per image, 16 MiB of decoded image bytes and 32 megapixels in aggregate. Each image also passes the managed-file PNG/JPEG, animation, dimension, pixel and allocation checks. Only inline image bytes are accepted; Kit never fetches provider-generated HTTP or file URLs. Native media does not inject synthetic image-label text into the structured output. Generation has a 300-second attempt timeout and 310-second logical budget, with no automatic retries of ambiguous billable failures. Cancellation remains available. On continuation or replay, historical assistant images stay typed in canonical history and are projected into supported image-input blocks only in the outgoing provider request, after any complete parallel tool-result batch. The existing text-oriented behavior and tool-image user-message fallback remain unchanged for other routes. + +A file-aware schema supports **exactly one** File location: the root, or one fixed object-property path whose properties are required at every level. Arrays, unions, conditional binding, indirect references, multiple locations, and sibling keywords on the File `$ref` other than the optional `x-kit-image-index` annotation are rejected. Kit resolves the File schema locally. By default, the child must emit exactly one **distinct** native assistant image and return only the surrounding JSON fields, omitting the binding field. Kit independently validates every occurrence, including its declared MIME type, actual PNG/JPEG bytes, and pixels. Repeated occurrences collapse to one output only when both MIME type and actual bytes match exactly within this turn. Count, encoded/decoded byte, and aggregate pixel budgets count every occurrence before deduplication. Without an explicit selection index, different image bytes remain ambiguous even if they render identically; model IDs and File references do not determine identity. There is no cross-turn or input/output deduplication. Kit rejects model-written binding fields, including `null` placeholders and invented File IDs. A root binding requires empty text. Empty surrounding text is allowed only when Kit can construct the required object path and the resulting complete value validates; Kit does not invent other required fields or apply schema defaults. + +Missing images, multiple distinct images without an explicit index, out-of-range selection indices, capture errors, malformed image bytes, invalid surrounding JSON, failed schema validation, and inaccessible files fail the call explicitly. Tool-result images, thought images, resource links, and textual base64 are not native assistant output. Kit imports real PNG/JPEG bytes under the existing managed-file limits and publishes a durable parent-authorized copy before returning success. Outputs survive child close and different child working directories; unrelated sessions do not acquire access. Failed continuation calls retain the accepted handle generation for retry, as with existing text-only failures. + +To deliberately select one output from a backend that can emit multiple distinct images, fix `x-kit-image-index` beside the exact File `$ref` **before** starting the call. The annotation must be an integer from 0 through 7. It indexes distinct validated images in first-emission order, after byte-identical duplicates collapse. The default contract above remains strict: Kit never chooses among distinct outputs unless the caller supplies this annotation. + +```text +source = read_file({ path: "source.png" }) +edited = subagent({ + model: "openrouter:google/gemini-3-pro-image", + prompt: "Add the requested sticker to the attached image and emit native image output.", + attachments: [source], + output_schema: { + type: "object", + properties: { + result: { "$ref": "kit://schemas/file/v1", "x-kit-image-index": 0 } + }, + required: ["result"], + additionalProperties: false + } +}) +return edited.output.result +``` + +The root form is also supported: `output_schema: { "$ref": "kit://schemas/file/v1", "x-kit-image-index": 1 }` selects the second distinct image. Kit removes the annotation when expanding the local File schema. The model cannot supply or override the index or binding field. All occurrences—including unselected images—must pass validation and occurrence/byte/pixel budgets before selection. A requested index with no corresponding image fails explicitly, with no fallback. Only the selected image is imported and published; unselected images produce no File descriptors or diagnostic image updates. Kit does not strip signed metadata, compare images perceptually, or treat identical pixels with different PNG/JPEG bytes as duplicates. + +Without a file-aware schema, text-only results keep their existing behavior, including ordinary `output_schema` validation with string fallback. When native images accompany such a result, `output` is explicitly `{ value, files }`: `value` is the legacy text/JSON result and `files` contains imported File descriptors. Images are not hidden in diagnostic `updates`, and raw base64 is not included there. Returning `edited.output.files[0]` delivers that image; keeping it intermediate does not. All these tools remain callable only through compose, and final-return-only image delivery is unchanged. + ## Make exact file changes with `edit` `edit` operates on one file path with `op: "add"`, `"edit"`, or `"delete"`. Relative paths are resolved from Kit's working directory. Absolute paths, `..`, and paths through symlinks are accepted, so `edit` can change files outside the root when the Kit process has permission. Paths must be non-empty. diff --git a/fixtures/mock-acp.py b/fixtures/mock-acp.py index 6fdb3f10..b6394a1f 100644 --- a/fixtures/mock-acp.py +++ b/fixtures/mock-acp.py @@ -114,14 +114,6 @@ def prompt(request): "sessionUpdate": "agent_thought_chunk", "content": {"type": "text", "text": "internal"}, }, - { - "sessionUpdate": "agent_message_chunk", - "content": { - "type": "image", - "data": "aGVsbG8=", - "mimeType": "image/png", - }, - }, { "sessionUpdate": "tool_call", "toolCallId": "call-1", diff --git a/src/acp_child.rs b/src/acp_child.rs index 23781eba..35d7ed42 100644 --- a/src/acp_child.rs +++ b/src/acp_child.rs @@ -14,12 +14,14 @@ use std::{ use agent_client_protocol::{ByteStreams, schema::ProtocolVersion}; use agentkit_acp::{ - CancelNotification, CloseSessionRequest, ContentBlock, ForkSessionRequest, PermissionOption, - PermissionOptionKind, PromptResponse, RequestPermissionOutcome, RequestPermissionRequest, - RequestPermissionResponse, SelectedPermissionOutcome, SessionConfigKind, SessionId, - SessionNotification, SessionUpdate, SetSessionConfigOptionRequest, StopReason, + CancelNotification, CloseSessionRequest, ContentBlock, ForkSessionRequest, ImageContent, + PermissionOption, PermissionOptionKind, PromptResponse, RequestPermissionOutcome, + RequestPermissionRequest, RequestPermissionResponse, SelectedPermissionOutcome, + SessionConfigKind, SessionId, SessionNotification, SessionUpdate, + SetSessionConfigOptionRequest, StopReason, }; use agentkit_core::TurnCancellation; +use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64}; use futures_util::future::{Either, select}; use serde::Deserialize; use serde_json::Value; @@ -38,6 +40,10 @@ const PRE_HANDSHAKE_EXIT_SETTLE: Duration = Duration::from_millis(250); const CANCEL_SETTLE: Duration = Duration::from_secs(5); const MAX_CAPTURED_UPDATES: usize = 64; const MAX_CAPTURED_UPDATE_BYTES: usize = 64 * 1024; +pub(crate) const MAX_NATIVE_IMAGE_BYTES: usize = 8 * 1024 * 1024; +pub(crate) const MAX_NATIVE_IMAGES: usize = 8; +pub(crate) const MAX_NATIVE_IMAGE_TOTAL_BYTES: usize = 32 * 1024 * 1024; +pub(crate) const NATIVE_IMAGE_ERROR: &str = "ACP assistant image transport failed"; const FORK_PARENT_ID_META: &str = "kit.subagent.parent_id"; const FORK_PARENT_NAME_META: &str = "kit.subagent.parent_name"; pub const BUILTIN_HARNESS: &str = "acp.kit"; @@ -417,6 +423,7 @@ struct Prompt { serial: tokio::sync::OwnedMutexGuard<()>, session_id: SessionId, text: String, + attachments: Vec, cancellation: TurnCancellation, reply: oneshot::Sender>, } @@ -500,11 +507,43 @@ pub(crate) struct ChildOutput { pub text: String, pub updates: Vec, pub updates_truncated: bool, + pub(crate) images: Vec, + pub(crate) media_error: Option, + image_bytes: usize, update_bytes: usize, } impl ChildOutput { fn record(&mut self, update: SessionUpdate) { + if let SessionUpdate::Notice(notice) = &update + && notice.title == NATIVE_IMAGE_ERROR + { + self.media_error = Some(NATIVE_IMAGE_ERROR.into()); + return; + } + if let SessionUpdate::AgentMessageChunk(chunk) = &update + && let ContentBlock::Image(image) = &chunk.content + { + // Sticky rejection: later chunks cannot turn a partial image set into success. + if self.media_error.is_none() { + match native_image_size(image).and_then(|size| { + if self.images.len() >= MAX_NATIVE_IMAGES + || self.image_bytes + size > MAX_NATIVE_IMAGE_TOTAL_BYTES + { + Err("ACP assistant images exceed the output budget".into()) + } else { + Ok(size) + } + }) { + Ok(size) => { + self.images.push(image.clone()); + self.image_bytes += size; + } + Err(error) => self.media_error = Some(error), + } + } + return; + } if let SessionUpdate::AgentMessageChunk(chunk) = &update && let ContentBlock::Text(text) = &chunk.content { @@ -544,6 +583,23 @@ impl ChildOutput { } } +/// Validate encoded transport before allocating decoded bytes. File import validates pixels. +fn native_image_size(image: &ImageContent) -> Result { + if !matches!(image.mime_type.as_str(), "image/png" | "image/jpeg") { + return Err("ACP image MIME type must be PNG or JPEG".into()); + } + if image.data.len() > MAX_NATIVE_IMAGE_BYTES.div_ceil(3) * 4 { + return Err("ACP image exceeds the byte limit".into()); + } + let bytes = BASE64 + .decode(&image.data) + .map_err(|_| "ACP image contains invalid base64".to_string())?; + if bytes.is_empty() || bytes.len() > MAX_NATIVE_IMAGE_BYTES { + return Err("ACP image has empty or oversized bytes".into()); + } + Ok(bytes.len()) +} + fn deduplicate_tool_output(update: &mut Value) { let Some(object) = update.as_object_mut() else { return; @@ -680,6 +736,10 @@ impl ChildSession { } } + pub(crate) fn session_id(&self) -> &str { + self.session_id.0.as_ref() + } + pub fn is_closed(&self) -> bool { self.tx.is_closed() || *self.closed.borrow() } @@ -786,6 +846,35 @@ impl ChildSession { text: String, cancellation: TurnCancellation, ) -> Result { + self.prompt_with_attachments(text, Vec::new(), cancellation) + .await + } + + pub async fn prompt_with_attachments( + &self, + text: String, + attachments: Vec, + cancellation: TurnCancellation, + ) -> Result { + if !attachments.is_empty() && !self.capabilities.prompt_capabilities.image { + return Err(ChildError::Failed( + "ACP harness does not support image prompts".into(), + )); + } + if attachments.len() > MAX_NATIVE_IMAGES { + return Err(ChildError::Failed( + "ACP image attachments exceed the count limit".into(), + )); + } + let mut bytes = 0; + for image in &attachments { + bytes += native_image_size(image).map_err(ChildError::Failed)?; + if bytes > MAX_NATIVE_IMAGE_TOTAL_BYTES { + return Err(ChildError::Failed( + "ACP image attachments exceed the byte budget".into(), + )); + } + } // A one-shot admission race: an available gate may win concurrent // cancellation. The request retains cancellation after admission. let serial = match select( @@ -802,6 +891,7 @@ impl ChildSession { serial, session_id: self.session_id.clone(), text, + attachments, cancellation: cancellation.clone(), reply, }); @@ -1156,8 +1246,10 @@ async fn run( let session_id = prompt.session_id.clone(); let output = Arc::new(Mutex::new(ChildOutput::default())); if let Ok(mut routes) = routes.lock() { routes.insert(session_id.clone(), Arc::clone(&output)); } + let mut content = vec![ContentBlock::Text(agentkit_acp::TextContent::new(prompt.text))]; + content.extend(prompt.attachments.into_iter().map(ContentBlock::Image)); let request = connection.send_request(agentkit_acp::PromptRequest::new( - session_id.clone(), vec![ContentBlock::Text(agentkit_acp::TextContent::new(prompt.text))], + session_id.clone(), content, )).block_task(); tokio::pin!(request); // Response-first matches the original biased race. @@ -1636,6 +1728,7 @@ mod tests { serial, session_id: child.session_id.clone(), text: "queued".into(), + attachments: Vec::new(), cancellation: controller.handle().checkpoint(), reply, })) @@ -1915,13 +2008,146 @@ mod tests { }))); assert_eq!(output.text, "hello"); - assert_eq!(output.updates.len(), 3); - assert_eq!(output.updates[0]["content"]["type"], "image"); - assert_eq!(output.updates[1]["sessionUpdate"], "tool_call"); - assert_eq!(output.updates[2]["sessionUpdate"], "plan"); + assert_eq!(output.images.len(), 1); + assert!(output.media_error.is_none()); + assert_eq!(output.updates.len(), 2); + assert_eq!(output.updates[0]["sessionUpdate"], "tool_call"); + assert_eq!(output.updates[1]["sessionUpdate"], "plan"); assert!(!output.updates_truncated); } + #[test] + fn native_image_capture_has_sticky_explicit_limits() { + let image = |data: String| { + SessionUpdate::AgentMessageChunk(agentkit_acp::ContentChunk::new(ContentBlock::Image( + ImageContent::new(data, "image/png"), + ))) + }; + for bad in [ + String::new(), + "not base64!".into(), + "A".repeat(MAX_NATIVE_IMAGE_BYTES.div_ceil(3) * 4 + 1), + ] { + let mut output = ChildOutput::default(); + output.record(image(bad)); + output.record(image("AQID".into())); + assert!(output.media_error.is_some()); + assert!(output.images.is_empty()); + assert!(output.updates.is_empty()); + assert!(!output.updates_truncated); + } + let mut output = ChildOutput::default(); + for _ in 0..=MAX_NATIVE_IMAGES { + output.record(image("AQID".into())); + } + assert_eq!(output.images.len(), MAX_NATIVE_IMAGES); + assert!(output.media_error.is_some()); + assert!(output.updates.is_empty()); + } + + #[tokio::test] + async fn native_image_prompt_requires_advertised_capability_before_admission() { + let (session, mut requests) = admission_test_session(); + let error = session + .prompt_with_attachments( + "image".into(), + vec![ImageContent::new("AQID", "image/png")], + TurnCancellation::default(), + ) + .await + .unwrap_err(); + assert!(error.to_string().contains("does not support image")); + assert!(requests.try_recv().is_err()); + assert!(session.serial.try_lock().is_ok()); + } + + #[tokio::test] + async fn native_image_stdio_roundtrip_and_error_preserve_next_prompt() { + let root = tempfile::tempdir().unwrap(); + // A genuine external ACP peer: inspect typed input and echo native output + // before settling the JSON-RPC prompt response. + let script = r#"import json, sys +for line in sys.stdin: + r=json.loads(line); method=r.get('method'); p=r.get('params', {}) + if method=='initialize': out={'protocolVersion':1,'agentCapabilities':{'promptCapabilities':{'image':True}}} + elif method=='session/new': out={'sessionId':'image-session'} + elif method=='session/prompt': + blocks=p['prompt']; text=blocks[0]['text'] + if text=='image': + assert blocks[1]=={'type':'image','data':'AQID','mimeType':'image/png'}, blocks + for kind in ['agent_thought_chunk','agent_message_chunk']: + print(json.dumps({'jsonrpc':'2.0','method':'session/update','params':{'sessionId':p['sessionId'],'update':{'sessionUpdate':kind,'content':blocks[1]}}}), flush=True) + if text=='bad': + print(json.dumps({'jsonrpc':'2.0','method':'session/update','params':{'sessionId':p['sessionId'],'update':{'sessionUpdate':'agent_message_chunk','content':{'type':'image','data':'!','mimeType':'image/png'}}}}), flush=True) + print(json.dumps({'jsonrpc':'2.0','method':'session/update','params':{'sessionId':p['sessionId'],'update':{'sessionUpdate':'agent_message_chunk','content':{'type':'text','text':text}}}}), flush=True) + out={'stopReason':'end_turn'} + else: continue + print(json.dumps({'jsonrpc':'2.0','id':r['id'],'result':out}), flush=True) +"#; + let harnesses = AcpHarnesses::new(BTreeMap::from([( + "mock".into(), + AcpHarnessProfile { + command: "python3".into(), + args: vec!["-u".into(), "-c".into(), script.into()], + permissions: AcpPermissionPolicy::Deny, + }, + )])) + .unwrap(); + let config = ChildConfig { + root: root.path().to_path_buf(), + model: "unused".into(), + provider: Default::default(), + reasoning_effort: None, + openrouter_api_key: None, + configured_mcp_config: None, + configured_mcp_config_inherited: false, + legacy_mcp_config: false, + mcp_config: None, + credential_storage: Default::default(), + telemetry: Default::default(), + harnesses, + default_harness: "acp.mock".into(), + parent_id: None, + parent_name: None, + }; + let session = ChildSession::start( + config, + "acp.mock".into(), + None, + None, + 1, + TurnCancellation::default(), + ) + .await + .unwrap(); + let output = session + .prompt_with_attachments( + "image".into(), + vec![ImageContent::new("AQID", "image/png")], + TurnCancellation::default(), + ) + .await + .unwrap(); + assert_eq!(output.text, "image"); + assert_eq!(output.images.len(), 1); + assert_eq!(output.images[0].data, "AQID"); + assert!(output.updates.is_empty()); + assert!(output.media_error.is_none()); + let bad = session + .prompt("bad".into(), TurnCancellation::default()) + .await + .unwrap(); + assert!(bad.media_error.is_some()); + assert!(bad.updates.is_empty()); + let next = session + .prompt("next".into(), TurnCancellation::default()) + .await + .unwrap(); + assert_eq!(next.text, "next"); + assert!(next.images.is_empty()); + assert!(next.media_error.is_none()); + } + #[test] fn captured_tool_updates_drop_content_that_duplicates_raw_output() { let raw = json!({"exit_code": 0, "stdout": "done", "stderr": "", "success": true}); @@ -1981,28 +2207,20 @@ mod tests { #[test] fn rich_updates_are_bounded_by_count_and_bytes() { - let image = || { + let tool = |title: String| { update(json!({ - "sessionUpdate": "agent_message_chunk", - "content": {"type": "image", "data": "eA==", "mimeType": "image/png"} + "sessionUpdate": "tool_call", "toolCallId": "call", "title": title })) }; let mut counted = ChildOutput::default(); for _ in 0..=MAX_CAPTURED_UPDATES { - counted.record(image()); + counted.record(tool("Inspect".into())); } assert_eq!(counted.updates.len(), MAX_CAPTURED_UPDATES); assert!(counted.updates_truncated); let mut oversized = ChildOutput::default(); - oversized.record(update(json!({ - "sessionUpdate": "agent_message_chunk", - "content": { - "type": "image", - "data": "x".repeat(MAX_CAPTURED_UPDATE_BYTES), - "mimeType": "image/png" - } - }))); + oversized.record(tool("x".repeat(MAX_CAPTURED_UPDATE_BYTES))); assert!(oversized.updates.is_empty()); assert!(oversized.updates_truncated); } diff --git a/src/managed_files.rs b/src/managed_files.rs index e24ce5df..9faf1d57 100644 --- a/src/managed_files.rs +++ b/src/managed_files.rs @@ -64,6 +64,48 @@ pub(crate) struct FileStore { base: PathBuf, } +/// Unique cleanup owner for a fresh inherited authority directory. It follows +/// pending transcript creation into response publication; only commit retains it. +pub(crate) struct InheritedAuthority { + path: Option, + quarantine: PathBuf, + identity: fs::FileIdentity, + session: String, +} + +impl InheritedAuthority { + pub(crate) fn is_for_session(&self, session: &str) -> bool { + self.session == session + } + + pub(crate) fn commit(mut self) { + self.path = None; + } +} + +impl Drop for InheritedAuthority { + fn drop(&mut self) { + let Some(path) = &self.path else { return }; + let identity = |path: &Path| fs::Backend::identity(&fs::DiskBackend, path, false); + if identity(path).ok().flatten() != Some(self.identity) { + return; + } + // Detach before recursive removal. Verify again after the atomic move, + // so a replaced directory is never deleted, even across a pathname race. + if rename_no_replace(path, &self.quarantine).is_err() { + return; + } + if identity(&self.quarantine).ok().flatten() != Some(self.identity) { + let _ = rename_no_replace(&self.quarantine, path); + return; + } + let _ = fs::remove_dir_all(&self.quarantine); + if let Some(parent) = path.parent() { + let _ = fs::sync_directory(parent); + } + } +} + impl FileStore { pub(crate) fn new(root: &Path) -> Self { Self { @@ -121,6 +163,185 @@ impl FileStore { self.publish(session, bytes, name, mime_type, image, cancellation) } + /// Imports actual native image bytes, never a URI or descriptor surrogate. + pub(crate) fn import_bytes( + &self, + session: &str, + name: &str, + mime_type: &str, + bytes: &[u8], + cancellation: Option<&TurnCancellation>, + ) -> Result { + check_cancelled(cancellation)?; + if !valid_name(name) || bytes.is_empty() || bytes.len() as u64 > MAX_FILE_BYTES { + return Err("invalid image name or byte budget (1 byte to 8 MiB)".into()); + } + let (actual_mime, dimensions) = inspect_image(bytes)?; + if mime_type != actual_mime { + return Err("declared image MIME type does not match its bytes".into()); + } + self.publish( + session, + bytes.to_vec(), + name.into(), + actual_mime, + dimensions, + cancellation, + ) + } + + /// Resolves only the caller's authorized immutable snapshot for ACP input. + pub(crate) fn attachment_image( + &self, + session: &str, + reference: &FileReference, + cancellation: Option<&TurnCancellation>, + ) -> Result { + check_cancelled(cancellation)?; + let bytes = self.resolve(session, reference)?; + check_cancelled(cancellation)?; + Ok(Part::media( + Modality::Image, + reference.mime_type.clone(), + DataRef::InlineBytes(bytes), + )) + } + + /// Explicit durable replication, preserving identity without global lookup. + /// Existing destinations must match metadata and payload; never clobber. + pub(crate) fn grant_to( + &self, + session: &str, + reference: &FileReference, + destination_store: &Self, + destination_session: &str, + cancellation: Option<&TurnCancellation>, + ) -> Result { + check_cancelled(cancellation)?; + let bytes = self.resolve(session, reference)?; + check_cancelled(cancellation)?; + let directory = destination_store.session_directory(destination_session); + let destination = directory.join(&reference.id); + match fs::symlink_metadata(&destination) { + Ok(_) => { + // Existence only selects this branch; source and destination + // must both resolve with matching metadata, digest and bytes. + // Repair any outstanding directory durability after a prior + // publication attempt, without allocating a new disk object. + if destination_store.resolve(destination_session, reference)? != bytes { + return Err("existing managed grant has conflicting bytes".into()); + } + check_cancelled(cancellation)?; + fs::sync_directory(&directory).map_err(display)?; + fs::require_disk(&directory).map_err(display)?; + fs::sync_directory(&destination_store.base).map_err(display)?; + fs::require_disk(&destination).map_err(display)?; + check_cancelled(cancellation)?; + return Ok(reference.clone()); + } + Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} + Err(error) => return Err(display(error)), + } + fs::create_private_dir_all(&directory).map_err(display)?; + let mut random = [0_u8; 32]; + getrandom::fill(&mut random).map_err(display)?; + let staging_session = format!("grant-{}", blake3::Hash::from_bytes(random).to_hex()); + let staging_directory = destination_store.session_directory(&staging_session); + let staged = staging_directory.join(&reference.id); + let result = (|| { + destination_store.write_snapshot( + &staging_session, + reference.clone(), + &bytes, + cancellation, + )?; + check_cancelled(cancellation)?; + // Atomic exclusive rename preserves nlink == 1 throughout: readers + // and restart never observe a multiply linked published envelope. + match rename_no_replace(&staged, &destination) { + Ok(()) => {} + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => {} + Err(error) => return Err(display(error)), + } + if destination_store.resolve(destination_session, reference)? != bytes { + return Err("existing managed grant has conflicting bytes".into()); + } + fs::sync_directory(&staging_directory).map_err(display)?; + fs::require_disk(&staging_directory).map_err(display)?; + fs::sync_directory(&directory).map_err(display)?; + fs::require_disk(&directory).map_err(display)?; + fs::sync_directory(&destination_store.base).map_err(display)?; + fs::require_disk(&destination).map_err(display)?; + Ok(reference.clone()) + })(); + // A crash can leave unreachable staging bytes, never partial grants. + let _ = fs::remove_file(&staged); + let _ = std::fs::remove_dir(&staging_directory); + result + } + + /// Prepare a fresh inherited set. The caller must retain its cleanup owner + /// until session/response publication. Commit is infallible and does no I/O; + /// rollback must happen outside shared writer/registry locks. + pub(crate) fn prepare_inheritance( + &self, + source: &str, + destination: &str, + ) -> Result> { + let target = self.session_directory(destination); + match fs::symlink_metadata(&target) { + Ok(_) => return Err("fork destination already has managed file authority".into()), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} + Err(error) => return Err(display(error)), + } + let directory = self.session_directory(source); + let entries = match fs::read_dir(&directory) { + Ok(entries) => entries, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(error) => return Err(display(error)), + }; + let mut random = [0_u8; 32]; + getrandom::fill(&mut random).map_err(display)?; + let staging_session = format!("fork-{}", blake3::Hash::from_bytes(random).to_hex()); + let staging_directory = self.session_directory(&staging_session); + fs::create_private_dir_all(&staging_directory).map_err(display)?; + let identity = fs::Backend::identity(&fs::DiskBackend, &staging_directory, false) + .map_err(display)? + .ok_or("inherited authority requires native directory identity")?; + let mut authority = InheritedAuthority { + path: Some(staging_directory.clone()), + quarantine: staging_directory.with_extension("rollback"), + identity, + session: destination.into(), + }; + for entry in entries { + let entry = entry.map_err(display)?; + let name = entry.file_name(); + let name = name.to_str().ok_or("invalid managed object filename")?; + let mut file = fs::open_beneath(&directory, Path::new(name)).map_err(display)?; + let mut prefix = [0_u8; 12]; + file.read_exact(&mut prefix).map_err(display)?; + let length = u32::from_le_bytes(prefix[8..12].try_into().map_err(display)?) as usize; + if &prefix[..8] != MAGIC || length > MAX_HEADER_BYTES { + return Err("invalid inherited managed file envelope".into()); + } + let mut header = vec![0; length]; + file.read_exact(&mut header).map_err(display)?; + let header: Header = serde_json::from_slice(&header).map_err(display)?; + if header.file.id != name { + return Err("inherited file ID does not match its storage name".into()); + } + self.grant_to(source, &header.file, self, &staging_session, None)?; + } + // The preflight is not authority: even an empty destination created + // concurrently must survive unchanged at the atomic commit boundary. + rename_no_replace(&staging_directory, &target).map_err(display)?; + authority.path = Some(target.clone()); + fs::sync_directory(&self.base).map_err(display)?; + fs::require_disk(&target).map_err(display)?; + Ok(Some(authority)) + } + fn publish( &self, session: &str, @@ -143,9 +364,20 @@ impl FileStore { size_bytes: bytes.len() as u64, image, }; + self.write_snapshot(session, reference, &bytes, cancellation) + } + + fn write_snapshot( + &self, + session: &str, + reference: FileReference, + bytes: &[u8], + cancellation: Option<&TurnCancellation>, + ) -> Result { + check_cancelled(cancellation)?; let header = serde_json::to_vec(&Header { file: reference.clone(), - digest: blake3::hash(&bytes).to_hex().to_string(), + digest: blake3::hash(bytes).to_hex().to_string(), }) .map_err(display)?; if header.len() > MAX_HEADER_BYTES { @@ -175,7 +407,7 @@ impl FileStore { .write_all(&(header.len() as u32).to_le_bytes()) .map_err(display)?; output.write_all(&header).map_err(display)?; - output.write_all(&bytes).map_err(display)?; + output.write_all(bytes).map_err(display)?; output.sync_all().map_err(display)?; fs::require_disk(&destination).map_err(display)?; // Include newly created ancestor entries, not just the object contents. @@ -187,7 +419,7 @@ impl FileStore { Ok(reference) } - fn resolve(&self, session: &str, selected: &FileReference) -> Result> { + pub(crate) fn resolve(&self, session: &str, selected: &FileReference) -> Result> { selected.validate()?; let directory = self.session_directory(session); // A reference must never resolve an import still retained only in the @@ -285,13 +517,9 @@ impl FileStore { let mut parts = Vec::new(); for (label, reference) in selected { check_cancelled(cancellation)?; - let data = self.resolve(session, &reference)?; + let image = self.attachment_image(session, &reference, cancellation)?; parts.push(Part::text(label)); - parts.push(Part::media( - Modality::Image, - reference.mime_type, - DataRef::InlineBytes(data), - )); + parts.push(image); } check_cancelled(cancellation)?; Ok(parts) @@ -299,7 +527,7 @@ impl FileStore { } impl FileReference { - fn from_value(value: &Value) -> Result { + pub(crate) fn from_value(value: &Value) -> Result { // Bound descriptor strings before deserialization copies them. A marker // does not make an arbitrarily large object a bounded File value. let object = value @@ -410,6 +638,88 @@ impl Selection { } } +/// Native exclusive rename for both files and directories. Staging and target +/// are always in the same managed store. Never fall back to check-then-rename +/// or hard links: those break no-clobber or secure single-link resolution. +#[cfg(any(target_os = "linux", target_os = "macos"))] +fn rename_no_replace(source: &Path, destination: &Path) -> std::io::Result<()> { + use std::ffi::CString; + use std::os::unix::ffi::OsStrExt as _; + let source = CString::new(source.as_os_str().as_bytes())?; + let destination = CString::new(destination.as_os_str().as_bytes())?; + // SAFETY: both pointers reference live NUL-terminated path strings. These + // calls do not retain pointers. Flags request an atomic no-replace rename. + let result = unsafe { + #[cfg(target_os = "linux")] + { + libc::renameat2( + libc::AT_FDCWD, + source.as_ptr(), + libc::AT_FDCWD, + destination.as_ptr(), + libc::RENAME_NOREPLACE, + ) + } + #[cfg(target_os = "macos")] + { + libc::renamex_np(source.as_ptr(), destination.as_ptr(), libc::RENAME_EXCL) + } + }; + if result == 0 { + Ok(()) + } else { + Err(std::io::Error::last_os_error()) + } +} + +#[cfg(windows)] +fn rename_no_replace(source: &Path, destination: &Path) -> std::io::Result<()> { + use std::os::windows::ffi::OsStrExt as _; + #[link(name = "kernel32")] + unsafe extern "system" { + #[link_name = "MoveFileExW"] + fn move_file_ex(source: *const u16, destination: *const u16, flags: u32) -> i32; + } + let source: Vec = source.as_os_str().encode_wide().collect(); + let destination: Vec = destination.as_os_str().encode_wide().collect(); + if source.contains(&0) || destination.contains(&0) { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "path contains NUL", + )); + } + let source = [source, vec![0]].concat(); + let destination = [destination, vec![0]].concat(); + // SAFETY: live NUL-terminated UTF-16 paths; the call retains no pointers. + // Unlike std::fs::rename, flags 0 excludes MOVEFILE_REPLACE_EXISTING. + let result = unsafe { move_file_ex(source.as_ptr(), destination.as_ptr(), 0) }; + if result != 0 { + Ok(()) + } else { + Err(std::io::Error::last_os_error()) + } +} + +#[cfg(not(any(target_os = "linux", target_os = "macos", windows)))] +fn rename_no_replace(_source: &Path, _destination: &Path) -> std::io::Result<()> { + Err(std::io::Error::new( + std::io::ErrorKind::Unsupported, + "atomic exclusive managed-file rename is unavailable on this platform", + )) +} + +/// Validate native provider bytes without publishing them or granting authority. +pub(crate) fn validate_provider_image(bytes: &[u8]) -> Result<(String, u64)> { + if bytes.is_empty() || bytes.len() as u64 > MAX_FILE_BYTES { + return Err("image must contain 1 byte to 8 MiB".into()); + } + let (mime, dimensions) = inspect_image(bytes)?; + Ok(( + mime, + u64::from(dimensions.width) * u64::from(dimensions.height), + )) +} + fn inspect_image(bytes: &[u8]) -> Result<(String, ImageDimensions)> { let format = image::guess_format(bytes).map_err(display)?; let decoder = bounded_decoder(bytes, format)?; diff --git a/src/managed_files/tests.rs b/src/managed_files/tests.rs index d9545a14..619f400d 100644 --- a/src/managed_files/tests.rs +++ b/src/managed_files/tests.rs @@ -558,3 +558,269 @@ fn fresh_process_resolve_child() { mod faults; mod operations; + +#[test] +fn native_bytes_require_actual_image_and_matching_declared_mime() { + let f = Fixture::new(); + let source = f.source("native.png", ImageFormat::Png, 3, 2); + let bytes = disk::read(source).unwrap(); + for (mime, data) in [ + ("image/jpeg", bytes.as_slice()), + ("image/png", b"fake"), + ("image/png", b""), + ] { + assert!( + f.store + .import_bytes("session", "native.png", mime, data, None) + .is_err() + ); + } + let reference = f + .store + .import_bytes("session", "native.png", "image/png", &bytes, None) + .unwrap(); + assert_eq!(f.store.resolve("session", &reference).unwrap(), bytes); + assert!(matches!( + f.store + .attachment_image("session", &reference, None) + .unwrap(), + Part::Media(_) + )); + assert!(f.store.attachment_image("other", &reference, None).is_err()); + assert!( + f.store + .import_bytes( + "session", + "native.png", + "image/png", + &vec![0; MAX_FILE_BYTES as usize + 1], + None + ) + .is_err() + ); +} + +#[test] +fn grants_preserve_identity_survive_source_close_and_remain_isolated() { + let source = Fixture::new(); + let destination = Fixture::new(); + let reference = source.import("native.png"); + let bytes = source.store.resolve("session", &reference).unwrap(); + assert!( + source + .store + .grant_to("stranger", &reference, &destination.store, "parent", None) + .is_err() + ); + let mut forged = reference.clone(); + forged.name = "forged.png".into(); + assert!( + source + .store + .grant_to("session", &forged, &destination.store, "parent", None) + .is_err() + ); + for _ in 0..2 { + assert_eq!( + source + .store + .grant_to("session", &reference, &destination.store, "parent", None) + .unwrap(), + reference + ); + } + drop(source); + let reopened = FileStore { + base: destination.store.base.clone(), + }; + assert_eq!(reopened.resolve("parent", &reference).unwrap(), bytes); + assert!(reopened.resolve("sibling", &reference).is_err()); + reopened + .prepare_inheritance("parent", "fork") + .unwrap() + .unwrap() + .commit(); + disk::remove_file(reopened.session_directory("parent").join(&reference.id)).unwrap(); + assert_eq!(reopened.resolve("fork", &reference).unwrap(), bytes); + let later = reopened + .import_bytes("fork", "later.png", "image/png", &bytes, None) + .unwrap(); + assert!(reopened.resolve("parent", &later).is_err()); +} + +#[test] +fn grant_never_clobbers_corrupt_destination_or_accepts_missing_source() { + let f = Fixture::new(); + let reference = f.import("native.png"); + let destination = f.store.session_directory("parent"); + disk::create_dir_all(&destination).unwrap(); + let path = destination.join(&reference.id); + disk::write(&path, b"existing corrupt object").unwrap(); + assert!( + f.store + .grant_to("session", &reference, &f.store, "parent", None) + .is_err() + ); + assert_eq!(disk::read(&path).unwrap(), b"existing corrupt object"); + disk::remove_file(f.object(&reference)).unwrap(); + assert!( + f.store + .grant_to("session", &reference, &f.store, "new", None) + .is_err() + ); + assert!(!f.store.session_directory("new").exists()); +} + +#[test] +fn cancelled_grants_do_not_publish_authority() { + let f = Fixture::new(); + let reference = f.import("image.png"); + let controller = agentkit_core::CancellationController::new(); + let cancellation = controller.handle().checkpoint(); + controller.interrupt(); + assert!( + f.store + .grant_to( + "session", + &reference, + &f.store, + "cancelled", + Some(&cancellation) + ) + .is_err() + ); + assert!( + f.store + .attachment_image("session", &reference, Some(&cancellation)) + .is_err() + ); + assert!(!f.store.session_directory("cancelled").exists()); +} + +#[test] +fn failed_fork_does_not_publish_a_partial_authorized_set() { + let f = Fixture::new(); + let reference = f.import("good.png"); + let directory = f.store.session_directory("session"); + disk::write(directory.join(format!("file_{}", "0".repeat(64))), b"bad").unwrap(); + assert!( + f.store + .prepare_inheritance("session", "failed-fork") + .is_err() + ); + assert!(!f.store.session_directory("failed-fork").exists()); + assert!(f.store.resolve("failed-fork", &reference).is_err()); + assert!(f.store.resolve("session", &reference).is_ok()); +} + +#[test] +fn exclusive_rename_preserves_existing_files_and_empty_directories() { + let dir = tempfile::tempdir().unwrap(); + let source = dir.path().join("source"); + let destination = dir.path().join("destination"); + disk::write(&source, b"new").unwrap(); + disk::write(&destination, b"existing").unwrap(); + assert_eq!( + rename_no_replace(&source, &destination).unwrap_err().kind(), + std::io::ErrorKind::AlreadyExists + ); + assert_eq!(disk::read(&destination).unwrap(), b"existing"); + assert_eq!(disk::read(&source).unwrap(), b"new"); + disk::remove_file(&destination).unwrap(); + rename_no_replace(&source, &destination).unwrap(); + assert!(!source.exists()); + assert_eq!(disk::read(&destination).unwrap(), b"new"); + let source = dir.path().join("source-dir"); + let destination = dir.path().join("destination-dir"); + disk::create_dir(&source).unwrap(); + disk::create_dir(&destination).unwrap(); + disk::write(source.join("private"), b"private").unwrap(); + assert!(rename_no_replace(&source, &destination).is_err()); + assert!(source.join("private").exists()); + assert!(disk::read_dir(&destination).unwrap().next().is_none()); +} + +#[test] +fn concurrent_grants_publish_single_link_resolvable_envelopes() { + let f = Fixture::new(); + let reference = f.import("concurrent.png"); + let expected = f.store.resolve("session", &reference).unwrap(); + std::thread::scope(|scope| { + for _ in 0..8 { + scope.spawn(|| { + f.store + .grant_to("session", &reference, &f.store, "parent", None) + .unwrap(); + for _ in 0..8 { + assert_eq!(f.store.resolve("parent", &reference).unwrap(), expected); + } + }); + } + }); + #[cfg(unix)] + { + use std::os::unix::fs::MetadataExt as _; + let path = f.store.session_directory("parent").join(&reference.id); + assert_eq!(disk::metadata(&path).unwrap().nlink(), 1); + } + assert_eq!( + disk::read_dir(f.store.session_directory("parent")) + .unwrap() + .count(), + 1 + ); + let reopened = FileStore { + base: f.store.base.clone(), + }; + assert_eq!(reopened.resolve("parent", &reference).unwrap(), expected); +} + +#[test] +fn inherited_authority_owner_rolls_back_or_commits_without_touching_source() { + let f = Fixture::new(); + let reference = f.import("owner.png"); + let prepared = f + .store + .prepare_inheritance("session", "fork") + .unwrap() + .unwrap(); + assert!(f.store.resolve("fork", &reference).is_ok()); + drop(prepared); + assert!(!f.store.session_directory("fork").exists()); + assert!(f.store.resolve("session", &reference).is_ok()); + let prepared = f + .store + .prepare_inheritance("session", "fork") + .unwrap() + .unwrap(); + prepared.commit(); + let reopened = FileStore { + base: f.store.base.clone(), + }; + assert!(reopened.resolve("fork", &reference).is_ok()); + assert!(reopened.prepare_inheritance("session", "fork").is_err()); + assert!(reopened.resolve("fork", &reference).is_ok()); +} + +#[test] +fn inherited_authority_cleanup_does_not_delete_a_replacement_directory() { + let f = Fixture::new(); + let reference = f.import("owner.png"); + let prepared = f + .store + .prepare_inheritance("session", "fork") + .unwrap() + .unwrap(); + let target = f.store.session_directory("fork"); + let moved = target.with_extension("moved"); + disk::rename(&target, &moved).unwrap(); + disk::create_dir(&target).unwrap(); + disk::write(target.join("replacement"), b"not owned").unwrap(); + drop(prepared); + assert_eq!( + disk::read(target.join("replacement")).unwrap(), + b"not owned" + ); + assert!(moved.join(&reference.id).exists()); + assert!(f.store.resolve("session", &reference).is_ok()); +} diff --git a/src/managed_files/tests/faults.rs b/src/managed_files/tests/faults.rs index 370d3f51..df5deafa 100644 --- a/src/managed_files/tests/faults.rs +++ b/src/managed_files/tests/faults.rs @@ -96,7 +96,7 @@ impl FaultBackend { // Fs persists objects through sibling atomic-replacement files, not by // writing the final file_ basename. Scope faults to this session's file // contents so staged writes are covered without faulting directory work. - if path.parent() == Some(self.object_directory.as_path()) { + if path.starts_with(&self.object_directory) { Box::new(FaultFile { disk, mode: self.mode, @@ -271,3 +271,86 @@ fn fault_child() { ) .unwrap(); } + +#[test] +fn existing_grant_succeeds_when_new_object_writes_have_no_space() { + let f = Fixture::new(); + let reference = f.import("retry.png"); + f.store + .grant_to("session", &reference, &f.store, "parent", None) + .unwrap(); + let manifest = f.dir.path().join("grant-retry.json"); + disk::write( + &manifest, + serde_json::to_vec(&(f.store.base.clone(), reference)).unwrap(), + ) + .unwrap(); + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "managed_files::tests::faults::grant_retry_child", + "--ignored", + "--nocapture", + ]) + .env(MANIFEST_ENV, &manifest) + .output() + .unwrap(); + let stdout = String::from_utf8_lossy(&output.stdout); + assert!( + output.status.success(), + "{stdout}\n{}", + String::from_utf8_lossy(&output.stderr) + ); + assert!(stdout.contains("1 passed; 0 failed"), "{stdout}"); +} + +#[test] +#[ignore = "invoked by parent with immutable ENOSPC backend"] +fn grant_retry_child() { + let path = PathBuf::from(std::env::var_os(MANIFEST_ENV).unwrap()); + let (base, reference): (PathBuf, FileReference) = + serde_json::from_slice(&disk::read(path).unwrap()).unwrap(); + let store = FileStore { base }; + assert!( + fs::initialize_global(Fs::new(Arc::new(FaultBackend { + mode: Mode::NoSpace, + object_directory: store.base.clone(), + }))) + .is_ok() + ); + let granted = store + .grant_to("session", &reference, &store, "parent", None) + .unwrap(); + assert_eq!(granted, reference); + assert_eq!( + store.resolve("parent", &reference).unwrap(), + store.resolve("session", &reference).unwrap() + ); + let controller = agentkit_core::CancellationController::new(); + let cancellation = controller.handle().checkpoint(); + controller.interrupt(); + assert!( + store + .grant_to("session", &reference, &store, "parent", Some(&cancellation)) + .unwrap_err() + .contains("cancelled") + ); + // A genuinely new grant hits real filesystem write-back/durability failure. + assert!( + store + .grant_to("session", &reference, &store, "new", None) + .is_err() + ); + assert!(store.resolve("new", &reference).is_err()); + // Existing corrupt bytes must fail, even though retry is allocation-free. + disk::write( + store.session_directory("parent").join(&reference.id), + b"corrupt", + ) + .unwrap(); + assert!( + store + .grant_to("session", &reference, &store, "parent", None) + .is_err() + ); +} diff --git a/src/protocols/acp.rs b/src/protocols/acp.rs index c59c0bb4..aee15e35 100644 --- a/src/protocols/acp.rs +++ b/src/protocols/acp.rs @@ -176,10 +176,21 @@ fn transcript_replay( ) -> Vec { let mut replay = Vec::new(); for item in transcript { + // History stores items, not TurnFinished boundaries. Bound each durable + // assistant item independently: the live per-turn budget must not become + // a lifetime session quota, nor guess turn boundaries from user messages + // (automated turns need not have one). Individual image limits are shared + // with live output; live capture still enforces the aggregate turn budget. + let mut image_budget = NativeImageBudget::default(); for part in &item.parts { let update = match item.kind { ItemKind::User => user_replay_content(part).map(SessionUpdate::UserMessageChunk), - ItemKind::Assistant => assistant_replay_update(part), + ItemKind::Assistant => match part { + Part::Media(media) if media.modality == Modality::Image => { + Some(image_budget.update(media)) + } + _ => assistant_replay_update(part), + }, ItemKind::Tool => tool_replay_update(part), ItemKind::Developer if crate::compaction::is_compaction_summary(item) => { assistant_replay_update(part) @@ -231,6 +242,9 @@ fn assistant_replay_update(part: &Part) -> Option { .status(ToolCallStatus::Pending) .raw_input(call.input.clone()), )), + Part::Media(media) if media.modality == Modality::Image => { + Some(NativeImageBudget::default().update(media)) + } Part::Media(_) | Part::File(_) | Part::Structured(_) @@ -239,6 +253,38 @@ fn assistant_replay_update(part: &Part) -> Option { } } +/// A local budget, never shared across observers or guarded across notifications. +#[derive(Default)] +struct NativeImageBudget { + count: usize, + bytes: usize, +} + +impl NativeImageBudget { + fn update(&mut self, media: &MediaPart) -> SessionUpdate { + use crate::acp_child::{ + MAX_NATIVE_IMAGE_BYTES, MAX_NATIVE_IMAGE_TOTAL_BYTES, MAX_NATIVE_IMAGES, + NATIVE_IMAGE_ERROR, + }; + // URI and handle media are not byte output. In particular, never fetch a URL. + let DataRef::InlineBytes(bytes) = &media.data else { + return SessionUpdate::Notice(Notice::new(NoticeSeverity::Error, NATIVE_IMAGE_ERROR)); + }; + if !matches!(media.mime_type.as_str(), "image/png" | "image/jpeg") + || bytes.is_empty() + || bytes.len() > MAX_NATIVE_IMAGE_BYTES + || self.count >= MAX_NATIVE_IMAGES + || self.bytes + bytes.len() > MAX_NATIVE_IMAGE_TOTAL_BYTES + { + return SessionUpdate::Notice(Notice::new(NoticeSeverity::Error, NATIVE_IMAGE_ERROR)); + } + let image = ImageContent::new(BASE64.encode(bytes), media.mime_type.clone()); + self.count += 1; + self.bytes += bytes.len(); + SessionUpdate::AgentMessageChunk(ContentChunk::new(ContentBlock::Image(image))) + } +} + fn tool_replay_update(part: &Part) -> Option { let Part::ToolResult(result) = part else { return None; @@ -1050,6 +1096,29 @@ impl LoopObserver for ResponseInterruptionNoticeObserver { } return; } + // The pinned ACP adapter ignores media deltas. Forward only completed + // assistant items at this single boundary, before prompt settlement. + if let AgentEvent::TurnFinished(result) = &event.event + && result.finish_reason == FinishReason::Completed + { + let mut budget = NativeImageBudget::default(); + for item in &result.items { + if item.kind != ItemKind::Assistant { + continue; + } + for part in &item.parts { + if let Part::Media(media) = part + && media.modality == Modality::Image + && let Err(error) = self.client.notify_session(SessionNotification::new( + self.session_id.clone(), + budget.update(media), + )) + { + tracing::error!(%error, "failed to queue ACP assistant image output"); + } + } + } + } self.inner.handle_event(event); } } @@ -1895,6 +1964,7 @@ async fn session_actor(actor: SessionActor) { let mut transcript = driver.snapshot().transcript; crate::transcript::sanitize_forked_transcript(&mut transcript); Ok(AcpForkState { + source_session_id: session_id.to_string(), transcript, selection: adapter.selection().map_err(AcpRuntimeError::Loop)?, reasoning_effort: adapter @@ -4580,6 +4650,159 @@ pub(super) mod tests { )); } + #[tokio::test] + async fn native_images_forward_once_at_completed_assistant_boundary_and_replay() { + let integration = AcpIntegration::builder() + .name("native-image-test") + .approval_resolver(AutoDenyResolver) + .build() + .unwrap(); + let session_id = agentkit_acp::SessionId::new("native-image"); + let loop_session_id = AgentkitSessionId::new("native-image-loop"); + let (client, mut messages) = AcpClientHandle::channel(); + integration + .bind_session(AcpSessionBinding::new( + session_id.clone(), + loop_session_id.clone(), + client.clone(), + )) + .unwrap(); + let observer = ResponseInterruptionNoticeObserver::new( + integration, + client, + session_id.clone(), + test_activity(session_id.clone(), mpsc::unbounded_channel().0), + ); + let media = Part::media( + Modality::Image, + "image/png", + DataRef::inline_bytes([1, 2, 3]), + ); + let items = vec![ + Item::new(ItemKind::User, vec![media.clone()]), + Item::new(ItemKind::Tool, vec![media.clone()]), + Item::new( + ItemKind::Assistant, + vec![Part::text("answer"), media.clone()], + ), + ]; + let emit = |event| { + observer.handle_event(ObservedEvent { + session_id: Arc::new(loop_session_id.clone()), + event, + }) + }; + emit(AgentEvent::ContentDelta(Delta::CommitPart { + part: media.clone(), + })); + assert!( + messages.try_recv().is_err(), + "pinned adapter must not duplicate committed media" + ); + let result = agentkit_loop::TurnResult { + turn_id: agentkit_core::TurnId::new("image-turn"), + finish_reason: FinishReason::Cancelled, + items: items.clone(), + usage: None, + metadata: MetadataMap::default(), + }; + emit(AgentEvent::TurnFinished(result.clone())); + assert!( + messages.try_recv().is_err(), + "cancelled output must not be forwarded" + ); + emit(AgentEvent::TurnFinished(agentkit_loop::TurnResult { + finish_reason: FinishReason::Completed, + ..result + })); + let Some(AcpClientMessage::SessionNotification(notification)) = messages.recv().await + else { + panic!("expected native output") + }; + assert!( + matches!(¬ification.update, SessionUpdate::AgentMessageChunk(chunk) + if matches!(&chunk.content, ContentBlock::Image(image) if image.data == "AQID" && image.mime_type == "image/png")) + ); + assert!( + messages.try_recv().is_err(), + "only assistant media is forwarded, exactly once" + ); + let replay = transcript_replay(&session_id, &items); + let assistant_images = replay.iter().filter(|n| matches!(&n.update, + SessionUpdate::AgentMessageChunk(chunk) if matches!(chunk.content, ContentBlock::Image(_)))).collect::>(); + assert_eq!(assistant_images.len(), 1); + assert_eq!(assistant_images[0].update, notification.update); + } + + #[test] + fn native_image_replay_budget_is_not_a_session_lifetime_quota() { + use crate::acp_child::MAX_NATIVE_IMAGES; + let session = agentkit_acp::SessionId::new("many-image-turns"); + // No user items are required between durable assistant outputs, e.g. + // automated continuation turns. Both histories must replay all images. + for with_user_input in [false, true] { + let mut transcript = Vec::new(); + for _ in 0..=MAX_NATIVE_IMAGES { + if with_user_input { + transcript.push(Item::text(ItemKind::User, "next image")); + } + transcript.push(Item::new( + ItemKind::Assistant, + vec![Part::media( + Modality::Image, + "image/png", + DataRef::inline_bytes([1, 2, 3]), + )], + )); + } + let replay = transcript_replay(&session, &transcript); + assert_eq!(replay.iter().filter(|notification| matches!( + ¬ification.update, SessionUpdate::AgentMessageChunk(chunk) + if matches!(&chunk.content, ContentBlock::Image(image) if image.data == "AQID") + )).count(), MAX_NATIVE_IMAGES + 1); + assert!( + !replay + .iter() + .any(|notification| matches!(notification.update, SessionUpdate::Notice(_))) + ); + } + } + + #[test] + fn native_image_replay_reports_unsupported_sources_and_output_budgets() { + use crate::acp_child::{MAX_NATIVE_IMAGE_BYTES, MAX_NATIVE_IMAGES, NATIVE_IMAGE_ERROR}; + let session = agentkit_acp::SessionId::new("image-errors"); + for data in [ + DataRef::uri("https://example.invalid/image.png"), + DataRef::inline_bytes([]), + DataRef::inline_bytes(vec![0; MAX_NATIVE_IMAGE_BYTES + 1]), + ] { + let replay = transcript_replay( + &session, + &[Item::new( + ItemKind::Assistant, + vec![Part::media(Modality::Image, "image/png", data)], + )], + ); + assert!(matches!(&replay[0].update, SessionUpdate::Notice(notice) + if notice.severity == NoticeSeverity::Error && notice.title == NATIVE_IMAGE_ERROR)); + } + let replay = transcript_replay( + &session, + &[Item::new( + ItemKind::Assistant, + vec![ + Part::media(Modality::Image, "image/png", DataRef::inline_bytes([1])); + MAX_NATIVE_IMAGES + 1 + ], + )], + ); + assert!(matches!( + replay[MAX_NATIVE_IMAGES].update, + SessionUpdate::Notice(_) + )); + } + #[tokio::test] async fn response_interruption_marker_becomes_v1_warning_before_replacement() { let integration = AcpIntegration::builder() @@ -6268,6 +6491,19 @@ pub(super) mod tests { vec![Item::text(ItemKind::System, "system")], ) .unwrap(); + let store = crate::managed_files::FileStore::new(root.path()); + let source_id = crate::session::new_id(); + let image_path = root.path().join("inherited.png"); + image::DynamicImage::new_luma8(2, 2) + .save(&image_path) + .unwrap(); + let reference = store.import(&source_id, &image_path, None).unwrap(); + opened + .observer + .attach_inherited_authority( + store.prepare_inheritance(&source_id, &id).unwrap().unwrap(), + ) + .unwrap(); // A creation already owned by a prepared publication is rejected by // normal APIs, just as poison/write fencing is rejected in session tests. let held = if delivery == "rejected" { @@ -6321,6 +6557,29 @@ pub(super) mod tests { } drop(held); drop(opened); + assert!(store.resolve(&source_id, &reference).is_ok()); + let reopened_store = crate::managed_files::FileStore::new(root.path()); + assert_eq!( + reopened_store.resolve(&id, &reference).is_ok(), + delivery == "success" + ); + if delivery != "success" { + let retry = crate::session::open_uncommitted( + root.path(), + &id, + false, + vec![Item::text(ItemKind::System, "retry")], + ) + .unwrap(); + assert!(reopened_store.resolve(&id, &reference).is_err()); + drop(retry); + } + let base = crate::artifacts::base(root.path()).with_file_name("files"); + for session in [&source_id, &id] { + let _ = std::fs::remove_dir_all( + base.join(blake3::hash(session.as_bytes()).to_hex().as_str()), + ); + } assert_eq!( crate::session::load(root.path(), &id).is_ok(), delivery == "success" @@ -6328,6 +6587,151 @@ pub(super) mod tests { } } + #[tokio::test] + async fn native_kit_fork_inherits_file_authority_without_reattachment() { + let root = tempfile::tempdir().unwrap(); + let credentials = crate::credentials::CredentialStorage::Memory; + crate::provider::store_openrouter_test_credentials(&credentials); + let runtime = Runtime::new_with_provider_and_credentials( + root.path(), + "test-model", + crate::ProviderKind::OpenRouter, + credentials, + ) + .unwrap(); + let (client_transport, agent_transport) = Channel::duplex(); + let server = tokio::spawn(serve_transport(runtime, agent_transport)); + agent_client_protocol::Client + .builder() + .connect_with(client_transport, async move |connection| { + connection + .send_request(InitializeRequest::new(ProtocolVersion::V1)) + .block_task() + .await?; + let source = connection + .send_request(NewSessionRequest::new(root.path().to_path_buf())) + .block_task() + .await?; + let store = crate::managed_files::FileStore::new(root.path()); + let image_path = root.path().join("input.png"); + image::DynamicImage::new_rgb8(2, 2) + .save(&image_path) + .unwrap(); + let inherited = store + .import(&source.session_id.to_string(), &image_path, None) + .unwrap(); + let expected = store + .resolve(&source.session_id.to_string(), &inherited) + .unwrap(); + std::fs::remove_file(&image_path).unwrap(); + // The request carries only the source session ID, never attachments. + let fork = connection + .send_request(ForkSessionRequest::new( + source.session_id.clone(), + root.path().to_path_buf(), + )) + .block_task() + .await?; + assert_eq!( + store + .resolve(&fork.session_id.to_string(), &inherited) + .unwrap(), + expected + ); + // Imports after the authority snapshot are private to each branch. + image::DynamicImage::new_rgb8(3, 2) + .save(&image_path) + .unwrap(); + let later_source = store + .import(&source.session_id.to_string(), &image_path, None) + .unwrap(); + let later_branch = store + .import(&fork.session_id.to_string(), &image_path, None) + .unwrap(); + assert!( + store + .resolve(&fork.session_id.to_string(), &later_source) + .is_err() + ); + assert!( + store + .resolve(&source.session_id.to_string(), &later_branch) + .is_err() + ); + assert!( + store + .resolve(&source.session_id.to_string(), &later_source) + .is_ok() + ); + assert!( + store + .resolve(&fork.session_id.to_string(), &later_branch) + .is_ok() + ); + + // Corrupt a real storage envelope at the external filesystem boundary. + // Failed preparation must not publish a destination transcript/session. + let reference = serde_json::to_value(&inherited).unwrap(); + let object = crate::artifacts::base(root.path()) + .with_file_name("files") + .join( + blake3::hash(source.session_id.to_string().as_bytes()) + .to_hex() + .as_str(), + ) + .join(reference["id"].as_str().unwrap()); + let original = std::fs::read(&object).unwrap(); + std::fs::write(&object, b"corrupt").unwrap(); + let before = connection + .send_request(ListSessionsRequest::new().cwd(root.path().to_path_buf())) + .block_task() + .await?; + connection + .send_request(ForkSessionRequest::new( + source.session_id.clone(), + root.path().to_path_buf(), + )) + .block_task() + .await + .expect_err("corrupt authority must reject fork preparation"); + let after = connection + .send_request(ListSessionsRequest::new().cwd(root.path().to_path_buf())) + .block_task() + .await?; + assert_eq!(before.sessions.len(), after.sessions.len()); + assert_eq!( + store + .resolve(&fork.session_id.to_string(), &inherited) + .unwrap(), + expected + ); + std::fs::write(&object, original).unwrap(); + let retry = connection + .send_request(ForkSessionRequest::new( + source.session_id.clone(), + root.path().to_path_buf(), + )) + .block_task() + .await?; + assert_eq!( + store + .resolve(&retry.session_id.to_string(), &inherited) + .unwrap(), + expected + ); + for id in [retry.session_id, fork.session_id, source.session_id] { + connection + .send_request(CloseSessionRequest::new(id)) + .block_task() + .await?; + } + Ok(()) + }) + .await + .unwrap(); + server.await.unwrap().unwrap(); + } + #[tokio::test] async fn kit_server_advertises_supported_session_discovery_restoration_and_forking() { let root = tempfile::tempdir().unwrap(); diff --git a/src/provider/adapter.rs b/src/provider/adapter.rs index 04e60598..4946de7a 100644 --- a/src/provider/adapter.rs +++ b/src/provider/adapter.rs @@ -463,10 +463,11 @@ pub enum KitAdapter { #[derive(Clone)] pub struct OpenRouterKitAdapter { inner: OpenRouterAdapter, + config: Box, client: reqwest::Client, models_url: Option, model: String, - context_window: Arc>, + context_window: Arc>, } const SPEAKEASY_COMPLETIONS_URL: &str = "https://app.getgram.ai/chat/completions"; @@ -607,7 +608,7 @@ impl KitAdapter { )?; apply_openrouter_reasoning_effort(&mut config, reasoning_effort); let models_url = models_url(&config.base_url); - let inner = OpenRouterAdapter::new(config) + let inner = OpenRouterAdapter::new(config.clone()) .map_err(|error| error.to_string())? .with_resilience(agentkit_http::ResilienceConfig::default()); let client = reqwest::Client::builder() @@ -619,6 +620,7 @@ impl KitAdapter { .map_err(|_| "could not build OpenRouter model catalog client".to_owned())?; Ok(Self::OpenRouter(OpenRouterKitAdapter { inner, + config: Box::new(config), client, models_url, model, @@ -766,21 +768,53 @@ impl ModelAdapter for KitAdapter { .await .map(KitSession::OpenAiSubscription), Self::OpenRouter(adapter) => { - let session = adapter.inner.start_session(config).await?; - let context_window = match &adapter.models_url { + // The OnceCell publishes one complete immutable discovery result. Failed or + // cancelled discovery leaves it empty; no credentials or session state change. + let discovered = match &adapter.models_url { Some(url) => adapter .context_window .get_or_try_init(|| { - fetch_context_window(&adapter.client, url, &adapter.model) + fetch_openrouter_model(&adapter.client, url, &adapter.model) }) .await - .ok() - .copied(), + .inspect_err(|_| tracing::warn!("OpenRouter model discovery unavailable; retaining legacy routing without native capability assertions")) + .ok(), None => None, }; + let context_window = discovered.and_then(|info| info.context_window); + let native = discover_native_image( + &agentkit_http::Http::new(adapter.client.clone()), + &adapter.config, + discovered, + ) + .await?; + let session = if let Some(capability) = &native { + let native_config = + native_generation_config((*adapter.config).clone(), capability)?; + let client = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .connect_timeout(Duration::from_secs(10)) + .timeout(NATIVE_GENERATION_TIMEOUT) + .build() + .map_err(|_| { + LoopError::Provider("could not build native image client".into()) + })?; + CompletionsAdapter::with_client( + OpenRouterProvider::from(native_config), + agentkit_http::Http::new(BoundedImageClient { + inner: agentkit_http::Http::new(client), + }), + ) + .with_resilience(native_generation_resilience()) + .start_session(config) + .await? + } else { + adapter.inner.start_session(config).await? + }; Ok(KitSession::OpenRouter(OpenRouterKitSession { inner: session, context_window, + native, })) } Self::Speakeasy(adapter) => { @@ -816,6 +850,7 @@ pub enum KitSession { pub struct OpenRouterKitSession { inner: OpenRouterSession, context_window: Option, + native: Option, } pub struct SpeakeasyKitSession { @@ -824,8 +859,8 @@ pub struct SpeakeasyKitSession { } /// Project only the outbound request; the caller's canonical transcript stays typed. -/// Completions (including OpenRouter) stringify tool Parts, but encode ordinary -/// user Media as image_url content. Native transports retain typed tool images, +/// Completions (including OpenRouter) stringify tool Parts and reject assistant +/// Media, but encode ordinary user Media as image_url content. Native transports retain typed tool images, /// but detached notification images always need ordinary user attachments. pub(super) fn project_tool_output_images( mut request: TurnRequest, @@ -868,6 +903,9 @@ pub(super) fn project_tool_output_images( let mut transcript = Vec::with_capacity(request.transcript.len()); let mut outstanding = std::collections::HashSet::new(); let mut images = Vec::new(); + let mut assistant_image_bytes = 0; + let mut assistant_image_count = 0; + let mut assistant_image_pixels = 0; for mut item in request.transcript { // The loop has already answered detached calls with placeholders. Its // completion notification contains serialized ToolResultPart values, @@ -890,13 +928,55 @@ pub(super) fn project_tool_output_images( outstanding.insert(call.id.clone()); } } + // The canonical completed assistant item remains typed. Lift only the + // outbound copy's images, without synthetic text. Register all calls + // first so images wait for the complete parallel tool-result batch. + let mut lifted_assistant_images = false; + if !native && item.kind == agentkit_core::ItemKind::Assistant { + let mut supported = Vec::new(); + for part in std::mem::take(&mut item.parts) { + if let Part::Media(media) = &part + && media.modality == Modality::Image + { + let DataRef::InlineBytes(bytes) = &media.data else { + return Err(LoopError::InvalidState("selected-images-not-delivered: historical assistant images require inline PNG/JPEG bytes".into())); + }; + if bytes.len() > MAX_NATIVE_IMAGE_BYTES { + return Err(LoopError::InvalidState("selected-images-not-delivered: historical assistant image exceeds 8 MiB".into())); + } + let (mime, pixels) = crate::managed_files::validate_provider_image(bytes) + .map_err(|error| LoopError::InvalidState(format!("selected-images-not-delivered: invalid historical assistant image: {error}")))?; + if mime != media.mime_type { + return Err(LoopError::InvalidState("selected-images-not-delivered: historical image MIME disagrees with bytes".into())); + } + assistant_image_count += 1; + assistant_image_bytes += bytes.len(); + assistant_image_pixels += pixels; + if assistant_image_count > 8 + || assistant_image_bytes > MAX_NATIVE_DELIVERY_BYTES + || assistant_image_pixels > 32 * 1024 * 1024 + { + return Err(LoopError::InvalidState("selected-images-not-delivered: historical assistant images exceed 8 images, 16 MiB or 32 megapixels".into())); + } + images.push(part); + lifted_assistant_images = true; + } else { + supported.push(part); + } + } + item.parts = supported; + } for part in &mut item.parts { if let Part::ToolResult(result) = part { project_result_images(result, &mut images, native)?; outstanding.remove(&result.call_id); } } - transcript.push(item); + // An image-only assistant item has no supported content left. Do not + // emit an invalid empty assistant message before its user image block. + if !lifted_assistant_images || !item.parts.is_empty() { + transcript.push(item); + } if outstanding.is_empty() && !images.is_empty() { let mut attachment = agentkit_core::Item::new( agentkit_core::ItemKind::User, @@ -910,7 +990,7 @@ pub(super) fn project_tool_output_images( } if !images.is_empty() { return Err(LoopError::InvalidState( - "selected-images-not-delivered: cannot attach tool images before all outstanding tool calls are answered. The program may already have completed; do not retry or rerun the program.".into(), + "selected-images-not-delivered: cannot attach images before all outstanding tool calls are answered. The program may already have completed; do not retry or rerun the program.".into(), )); } request.transcript = transcript; @@ -1080,6 +1160,17 @@ impl ModelSession for KitSession { } else { project_tool_output_images(request, false)? }; + if let Self::OpenRouter(session) = self + && let Some(capability) = &session.native + && !capability.image_input + && request + .transcript + .iter() + .flat_map(|item| &item.parts) + .any(|part| matches!(part, Part::Media(media) if media.modality == Modality::Image)) + { + return Err(LoopError::Provider("selected-images-not-delivered: selected OpenRouter generation model does not support image input".into())); + } match self { Self::OpenAiSubscription(session) => session .begin_turn(request, cancellation) @@ -1095,6 +1186,7 @@ impl ModelSession for KitSession { context_window: session.context_window, media_part: None, next_media: 0, + native: session.native.is_some(), }) .map(KitTurn::OpenRouter), Self::Speakeasy(session) => session @@ -1106,6 +1198,7 @@ impl ModelSession for KitSession { context_window: session.context_window, media_part: None, next_media: 0, + native: false, }) .map(KitTurn::Speakeasy), } @@ -1139,6 +1232,7 @@ pub struct OpenRouterKitTurn { context_window: Option, media_part: Option, next_media: usize, + native: bool, } #[async_trait] @@ -1151,7 +1245,9 @@ impl ModelTurn for KitTurn { Self::OpenAiSubscription(turn) => turn.next_event(cancellation).await, Self::OpenRouter(turn) | Self::Speakeasy(turn) => { let mut event = turn.inner.next_event(cancellation).await?; - if let Some(ModelTurnEvent::Delta(delta)) = &mut event { + if turn.native { + normalize_native_event(&mut event)?; + } else if let Some(ModelTurnEvent::Delta(delta)) = &mut event { rewrite_openrouter_media(delta, &mut turn.media_part, &mut turn.next_media); } if let Some(context_window) = turn.context_window { @@ -1163,6 +1259,452 @@ impl ModelTurn for KitTurn { } } +// Native image generation is opt-in through exact catalogue model selection, not +// model-name heuristics. Official capabilities never apply to custom endpoints. +#[derive(Clone)] +struct NativeImageCapability { + image_input: bool, + modalities: Vec, +} + +struct OpenRouterModelInfo { + context_window: Option, + input: Vec, + output: Vec, + tools: bool, +} + +fn parse_openrouter_model(value: &Value, model: &str) -> Option { + let models = value.get("data")?.as_array()?; + if models.len() > MAX_MODELS { + return None; + } + let entry = models + .iter() + .find(|entry| entry["id"].as_str() == Some(model))?; + Some(OpenRouterModelInfo { + context_window: parse_context_window(value, model), + input: capability_strings(&entry["architecture"]["input_modalities"]).unwrap_or_default(), + output: capability_strings(&entry["architecture"]["output_modalities"]).ok()?, + tools: capability_strings(&entry["supported_parameters"]) + .unwrap_or_default() + .iter() + .any(|value| value == "tools"), + }) +} + +const MAX_NATIVE_ENDPOINTS: usize = 256; +const MAX_CAPABILITY_VALUES: usize = 128; + +fn capability_strings(value: &Value) -> Result, String> { + let values = value.as_array().ok_or("capability list must be an array")?; + if values.len() > MAX_CAPABILITY_VALUES { + return Err("capability list exceeds 128 entries".into()); + } + values + .iter() + .map(|value| { + value + .as_str() + .filter(|value| !value.is_empty() && value.len() <= 128) + .map(str::to_owned) + .ok_or_else(|| "invalid capability list entry".into()) + }) + .collect() +} + +fn native_endpoints_url(model: &str) -> Result { + let error = + || LoopError::Provider("native-image-discovery: invalid selected model path".into()); + if !valid_model_id(model) + || !model.contains('/') + || model + .split('/') + .any(|segment| matches!(segment, "" | "." | "..")) + { + return Err(error()); + } + let mut url = url::Url::parse(OPENROUTER_MODELS_URL).map_err(|_| error())?; + { + let mut path = url.path_segments_mut().map_err(|_| error())?; + for segment in model.split('/') { + path.push(segment); + } + path.push("endpoints"); + } + Ok(url) +} + +async fn discover_native_image( + client: &agentkit_http::Http, + config: &OpenRouterConfig, + catalog: Option<&OpenRouterModelInfo>, +) -> Result, LoopError> { + if !equivalent_openrouter_base_urls(&config.base_url, &OpenRouterConfig::new("", "").base_url) { + return Ok(None); + } + let Some(catalog) = catalog.filter(|info| info.output.iter().any(|value| value == "image")) + else { + return Ok(None); + }; + let url = native_endpoints_url(&config.model)?; + let response = client.get(url.as_str()).send().await.map_err(|_| { + LoopError::Provider( + "native-image-discovery: endpoint transport failed; generation eligibility is unknown" + .into(), + ) + })?; + if !response.status().is_success() { + return Err(LoopError::Provider(format!( + "native-image-discovery: endpoint catalog returned {}; generation eligibility is unknown", + response.status() + ))); + } + let mut body = Vec::new(); + let mut stream = response.bytes_stream(); + while let Some(chunk) = stream.next().await { + let chunk = chunk.map_err(|_| { + LoopError::Provider("native-image-discovery: endpoint body failed".into()) + })?; + if chunk.len() > MAX_MODELS_BYTES.saturating_sub(body.len()) { + return Err(LoopError::Provider( + "native-image-discovery: endpoint catalog exceeds 2 MiB".into(), + )); + } + body.extend_from_slice(&chunk); + } + let value = serde_json::from_slice(&body).map_err(|_| { + LoopError::Provider("native-image-discovery: endpoint catalog is not valid JSON".into()) + })?; + endpoint_image_capability(config, catalog, &value) +} + +fn endpoint_image_capability( + config: &OpenRouterConfig, + catalog: &OpenRouterModelInfo, + value: &Value, +) -> Result, LoopError> { + let invalid = |reason: &str| { + LoopError::Provider(format!( + "native-image-discovery: {reason}; generation eligibility is unknown" + )) + }; + let data = &value["data"]; + if data["id"].as_str() != Some(config.model.as_str()) { + return Err(invalid("endpoint model identity mismatch")); + } + let input = capability_strings(&data["architecture"]["input_modalities"]) + .map_err(|error| invalid(&error))?; + let output = capability_strings(&data["architecture"]["output_modalities"]) + .map_err(|error| invalid(&error))?; + let same_values = |a: &[String], b: &[String]| { + a.len() == b.len() + && a.iter().all(|value| b.contains(value)) + && b.iter().all(|value| a.contains(value)) + }; + if !same_values(&input, &catalog.input) || !same_values(&output, &catalog.output) { + return Err(invalid( + "endpoint architecture disagrees with model catalog", + )); + } + let endpoints = data["endpoints"] + .as_array() + .ok_or_else(|| invalid("missing endpoint list"))?; + if endpoints.len() > MAX_NATIVE_ENDPOINTS { + return Err(invalid("endpoint count exceeds 256")); + } + // Routing selectors can advertise broad aggregate modalities without having + // concrete generation providers. An empty list preserves legacy routing. + if endpoints.is_empty() { + return Ok(None); + } + let mut tools = false; + for endpoint in endpoints { + let parameters = capability_strings(&endpoint["supported_parameters"]) + .map_err(|error| invalid(&error))?; + tools |= parameters.iter().any(|parameter| parameter == "tools"); + } + if !tools { + return Err(LoopError::Provider("native-image-ineligible: no concrete OpenRouter endpoint supports tools; compose cannot be removed".into())); + } + let concrete = OpenRouterModelInfo { + context_window: catalog.context_window, + input, + output, + tools: catalog.tools && tools, + }; + native_image_capability(config, Some(&concrete)) +} + +fn native_image_capability( + config: &OpenRouterConfig, + info: Option<&OpenRouterModelInfo>, +) -> Result, LoopError> { + if !equivalent_openrouter_base_urls(&config.base_url, &OpenRouterConfig::new("", "").base_url) { + return Ok(None); + } + let Some(info) = info.filter(|info| info.output.iter().any(|value| value == "image")) else { + return Ok(None); + }; + if !info.tools { + return Err(LoopError::Provider("OpenRouter image model is not eligible for a Kit agent: catalogue does not advertise tools support; compose cannot be removed".into())); + } + if info + .output + .iter() + .any(|value| value != "image" && value != "text") + { + return Err(LoopError::Provider( + "OpenRouter image model advertises unsupported output modalities".into(), + )); + } + Ok(Some(NativeImageCapability { + image_input: info.input.iter().any(|value| value == "image"), + modalities: info.output.clone(), + })) +} + +fn native_generation_config( + mut config: OpenRouterConfig, + capability: &NativeImageCapability, +) -> Result { + let routing = config + .extra_body + .entry("provider".into()) + .or_insert_with(|| Value::Object(serde_json::Map::new())); + let routing = routing.as_object_mut().ok_or_else(|| { + LoopError::Provider( + "native-image-configuration: provider routing options must be an object".into(), + ) + })?; + routing.insert("require_parameters".into(), Value::Bool(true)); + // The proven generation contract is a complete JSON response. Buffer only + // after the transport enforces its raw-byte ceiling, before the decoder. + Ok(config.with_streaming(false).with_extra_body_value( + "modalities", + Value::Array( + capability + .modalities + .iter() + .cloned() + .map(Value::String) + .collect(), + ), + )) +} + +const NATIVE_GENERATION_TIMEOUT: Duration = Duration::from_secs(300); + +fn native_generation_resilience() -> agentkit_http::ResilienceConfig { + // A timed-out/nonstreaming generation may already have been billed. Keep + // authentication and cancellation under a finite logical deadline, but do + // not automatically replay ambiguous generation failures or HTTP statuses. + agentkit_http::ResilienceConfig { + max_retries: 0, + retry_budget: Duration::from_secs(310), + attempt_timeout: Some(NATIVE_GENERATION_TIMEOUT), + stream_idle_timeout: Some(NATIVE_GENERATION_TIMEOUT), + ..agentkit_http::ResilienceConfig::default() + } +} + +const MAX_NATIVE_RESPONSE_BYTES: usize = 24 * 1024 * 1024; +const MAX_NATIVE_IMAGE_BYTES: usize = 8 * 1024 * 1024; +const MAX_NATIVE_DELIVERY_BYTES: usize = 16 * 1024 * 1024; + +struct BoundedImageClient { + inner: agentkit_http::Http, +} + +#[async_trait] +impl agentkit_http::HttpClient for BoundedImageClient { + async fn execute( + &self, + request: agentkit_http::HttpRequest, + ) -> Result { + use agentkit_http::{HttpError, HttpResponse}; + let response = self.inner.execute(request).await?; + let status = response.status(); + let headers = response.headers().clone(); + let url = response.url().to_owned(); + let mut stream = response.bytes_stream(); + let mut body = Vec::new(); + while let Some(chunk) = stream.next().await { + let chunk = chunk?; + if chunk.len() > MAX_NATIVE_RESPONSE_BYTES.saturating_sub(body.len()) { + return Err(HttpError::Other( + "native-image-response-too-large: raw response exceeds 24 MiB".into(), + )); + } + body.extend_from_slice(&chunk); + } + if status.is_success() { + let value: Value = serde_json::from_slice(&body).map_err(|_| { + HttpError::Other("native-image-malformed: expected a complete JSON response".into()) + })?; + // The pinned decoder skips malformed image entries. Reject these + // before decoding so a valid sibling cannot mask invalid media. + let choices = value["choices"].as_array().ok_or_else(|| { + HttpError::Other("native-image-malformed: missing choices".into()) + })?; + if choices.len() != 1 { + return Err(HttpError::Other( + "native-image-malformed: expected exactly one completion choice".into(), + )); + } + let mut image_count = 0_usize; + let mut encoded_bytes = 0_usize; + for choice in choices { + for key in ["message", "delta"] { + let message = &choice[key]; + let images = match message.get("images") { + Some(images) => images + .as_array() + .ok_or_else(|| { + HttpError::Other( + "native-image-malformed: images must be an array".into(), + ) + })? + .as_slice(), + None => &[], + }; + // Both native message.images and standard content image_url + // parts pass the same strict checks, including skipped entries. + let content_images = message["content"] + .as_array() + .into_iter() + .flatten() + .filter(|part| part["type"] == "image_url"); + for image in images.iter().chain(content_images) { + let uri = image["image_url"]["url"].as_str().ok_or_else(|| { + HttpError::Other("native-image-malformed: missing image URL".into()) + })?; + let (_, payload) = native_image_payload(uri) + .map_err(|error| HttpError::Other(error.to_string()))?; + image_count += 1; + encoded_bytes += payload.len(); + if image_count > 8 + || encoded_bytes > MAX_NATIVE_DELIVERY_BYTES.div_ceil(3) * 4 + 8 * 4 + { + return Err(HttpError::Other( + "native-image-too-large: aggregate images exceed delivery budget" + .into(), + )); + } + } + } + } + } + Ok(HttpResponse::new( + status, + headers, + url, + Box::pin(futures_util::stream::once(async move { + Ok(bytes::Bytes::from(body)) + })), + )) + } +} + +fn native_image_payload(uri: &str) -> Result<(&str, &str), LoopError> { + let malformed = || { + LoopError::Provider("native-image-malformed: expected inline base64 PNG or JPEG; remote and file URIs are not fetched".into()) + }; + let (header, payload) = uri.split_once(',').ok_or_else(malformed)?; + let mime = header + .strip_prefix("data:") + .and_then(|value| value.strip_suffix(";base64")) + .ok_or_else(malformed)?; + if !matches!(mime, "image/png" | "image/jpeg" | "image/*") || payload.is_empty() { + return Err(malformed()); + } + if payload.len() > MAX_NATIVE_IMAGE_BYTES.div_ceil(3) * 4 { + return Err(LoopError::Provider( + "native-image-too-large: encoded image exceeds 8 MiB decoded budget".into(), + )); + } + Ok((mime, payload)) +} + +fn normalize_native_part(part: &mut Part) -> Result<(usize, u64), LoopError> { + use base64::Engine as _; + let Part::Media(media) = part else { + return Ok((0, 0)); + }; + if media.modality != Modality::Image { + return Err(LoopError::Provider( + "native-image-malformed: unsupported media modality".into(), + )); + } + let DataRef::Uri(uri) = &media.data else { + return Err(LoopError::Provider( + "native-image-malformed: expected inline image URI".into(), + )); + }; + let (declared, payload) = native_image_payload(uri)?; + let bytes = base64::engine::general_purpose::STANDARD + .decode(payload) + .map_err(|_| LoopError::Provider("native-image-malformed: invalid base64".into()))?; + if bytes.len() > MAX_NATIVE_IMAGE_BYTES { + return Err(LoopError::Provider( + "native-image-too-large: decoded image exceeds 8 MiB".into(), + )); + } + let (mime, pixels) = + crate::managed_files::validate_provider_image(&bytes).map_err(|error| { + // The shared validator returns strings; preserve the bounded decoder's + // size failures separately from invalid encoding/format failures. + let code = match error.as_str() { + "image exceeds decoded pixel or allocation budget" + | "Memory limit exceeded" + | "Image size exceeds limit" => "native-image-too-large", + _ => "native-image-malformed", + }; + LoopError::Provider(format!("{code}: {error}")) + })?; + if declared != "image/*" && declared != mime { + return Err(LoopError::Provider( + "native-image-malformed: declared MIME disagrees with image bytes".into(), + )); + } + let size = bytes.len(); + media.mime_type = mime; + media.data = DataRef::InlineBytes(bytes); + Ok((size, pixels)) +} + +fn normalize_native_event(event: &mut Option) -> Result<(), LoopError> { + match event { + Some(ModelTurnEvent::Delta(Delta::CommitPart { part })) => { + normalize_native_part(part)?; + } + Some(ModelTurnEvent::Finished(result)) => { + let mut count = 0; + let mut bytes = 0; + let mut pixels = 0; + for part in result + .output_items + .iter_mut() + .flat_map(|item| &mut item.parts) + { + let (size, image_pixels) = normalize_native_part(part)?; + count += usize::from(size > 0); + bytes += size; + pixels += image_pixels; + if count > 8 || bytes > MAX_NATIVE_DELIVERY_BYTES || pixels > 32 * 1024 * 1024 { + return Err(LoopError::Provider("native-image-too-large: delivery exceeds 8 images, 16 MiB or 32 megapixels".into())); + } + } + } + _ => {} + } + Ok(()) +} + +#[cfg(test)] +#[path = "adapter_native_tests.rs"] +mod native_tests; + fn rewrite_openrouter_media( delta: &mut Delta, media_part: &mut Option, @@ -1244,11 +1786,11 @@ fn models_url(completions_url: &str) -> Option { .map(|prefix| format!("{prefix}/models")) } -async fn fetch_context_window( +async fn fetch_openrouter_model( client: &reqwest::Client, url: &str, model: &str, -) -> Result { +) -> Result { let response = client .get(url) .send() @@ -1271,8 +1813,8 @@ async fn fetch_context_window( } let value: Value = serde_json::from_slice(&body) .map_err(|_| "OpenRouter model catalog is not valid JSON".to_owned())?; - parse_context_window(&value, model) - .ok_or_else(|| format!("OpenRouter model catalog omitted context length for {model:?}")) + parse_openrouter_model(&value, model) + .ok_or_else(|| "OpenRouter model catalog omitted selected model".to_owned()) } fn parse_context_window(value: &Value, model: &str) -> Option { @@ -1972,6 +2514,7 @@ mod tests { KitSession::OpenRouter(OpenRouterKitSession { inner, context_window: None, + native: None, }) } diff --git a/src/provider/adapter_native_tests.rs b/src/provider/adapter_native_tests.rs new file mode 100644 index 00000000..80eff6a5 --- /dev/null +++ b/src/provider/adapter_native_tests.rs @@ -0,0 +1,869 @@ +#![allow(clippy::disallowed_methods, clippy::disallowed_macros)] +use super::*; +use agentkit_core::{Item, ItemKind, MediaPart, MetadataMap, SessionId, TurnId}; +use agentkit_http::{ + HeaderMap, Http, HttpClient, HttpError, HttpRequest, HttpResponse, StatusCode, +}; +use base64::Engine as _; +use serde_json::json; + +fn info(input: &[&str], output: &[&str], tools: bool) -> OpenRouterModelInfo { + parse_openrouter_model( + &json!({"data": [{"id":"test/image", "architecture": { + "input_modalities": input, "output_modalities": output + }, "supported_parameters": if tools { vec!["tools"] } else { vec![] }}]}), + "test/image", + ) + .unwrap() +} + +#[test] +fn exact_catalog_capabilities_not_names_or_custom_endpoints() { + let config = OpenRouterConfig::new("test", "test/image"); + let image = info(&["image", "text"], &["image", "text"], true); + assert!( + native_image_capability(&config, Some(&image)) + .unwrap() + .unwrap() + .image_input + ); + assert!(native_image_capability(&config, None).unwrap().is_none()); + assert!( + native_image_capability(&config, Some(&info(&["text"], &["text"], true))) + .unwrap() + .is_none() + ); + assert!(native_image_capability(&config, Some(&info(&["text"], &["image"], false))).is_err()); + let custom = config.with_base_url("https://custom.invalid/chat/completions"); + assert!( + native_image_capability(&custom, Some(&image)) + .unwrap() + .is_none() + ); + assert!( + parse_openrouter_model(&json!({"data":[{"id":"different/image"}]}), "test/image").is_none() + ); +} + +fn png() -> Vec { + let mut bytes = std::io::Cursor::new(Vec::new()); + image::DynamicImage::new_rgba8(2, 2) + .write_to(&mut bytes, image::ImageFormat::Png) + .unwrap(); + bytes.into_inner() +} + +fn image_part(uri: String) -> Part { + Part::Media(MediaPart::new( + Modality::Image, + "image/*", + DataRef::Uri(uri), + )) +} + +#[test] +fn normalization_validates_bytes_and_never_fetches_remote_media() { + let bytes = png(); + let mut part = image_part(format!( + "data:image/*;base64,{}", + base64::engine::general_purpose::STANDARD.encode(&bytes) + )); + assert_eq!(normalize_native_part(&mut part).unwrap(), (bytes.len(), 4)); + let Part::Media(media) = part else { panic!() }; + assert_eq!(media.mime_type, "image/png"); + assert_eq!(media.data, DataRef::InlineBytes(bytes)); + for uri in [ + "https://example.invalid/image.png", + "file:///tmp/image.png", + "data:image/png;base64,???", + "data:image/png;base64,YWJj", + "data:image/gif;base64,YWJj", + ] { + let error = normalize_native_part(&mut image_part(uri.into())) + .unwrap_err() + .to_string(); + assert!(error.contains("native-image-malformed"), "{error}"); + } + let mut wrong_mime = image_part(format!( + "data:image/jpeg;base64,{}", + base64::engine::general_purpose::STANDARD.encode(png()) + )); + assert!(normalize_native_part(&mut wrong_mime).is_err()); + let oversized = format!( + "data:image/png;base64,{}", + "A".repeat(MAX_NATIVE_IMAGE_BYTES.div_ceil(3) * 4 + 4) + ); + assert!( + normalize_native_part(&mut image_part(oversized)) + .unwrap_err() + .to_string() + .contains("native-image-too-large") + ); +} + +struct FakeHttp { + body: Vec, + sent: tokio::sync::mpsc::UnboundedSender, +} + +#[async_trait] +impl HttpClient for FakeHttp { + async fn execute(&self, request: HttpRequest) -> Result { + self.sent + .send(serde_json::from_slice(request.body.as_ref().unwrap()).unwrap()) + .unwrap(); + Ok(HttpResponse::new( + StatusCode::OK, + HeaderMap::new(), + request.url, + Box::pin(futures_util::stream::iter( + self.body.clone().into_iter().map(Ok), + )), + )) + } +} + +fn request(with_image: bool) -> TurnRequest { + let mut parts = vec![Part::text("make a sticker")]; + if with_image { + parts.push(Part::Media(MediaPart::new( + Modality::Image, + "image/png", + DataRef::InlineBytes(png()), + ))); + } + TurnRequest { + session_id: SessionId::new("native"), + turn_id: TurnId::new("turn"), + transcript: vec![Item::new(ItemKind::User, parts)], + available_tools: vec![agentkit_tools_core::ToolSpec { + name: "compose".into(), + description: "Execute a program".into(), + input_schema: json!({"type":"object"}), + output_schema: None, + annotations: Default::default(), + metadata: MetadataMap::new(), + }], + cache: None, + metadata: MetadataMap::new(), + } +} + +async fn session( + body: Vec, + image_input: bool, +) -> (KitSession, tokio::sync::mpsc::UnboundedReceiver) { + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + let capability = NativeImageCapability { + image_input, + modalities: vec!["image".into(), "text".into()], + }; + let config = + native_generation_config(OpenRouterConfig::new("test", "test/image"), &capability).unwrap(); + let adapter = CompletionsAdapter::with_client( + OpenRouterProvider::from(config), + Http::new(BoundedImageClient { + inner: Http::new(FakeHttp { body, sent: tx }), + }), + ) + .with_resilience(native_generation_resilience()); + let inner = adapter + .start_session(SessionConfig::new("native")) + .await + .unwrap(); + ( + KitSession::OpenRouter(OpenRouterKitSession { + inner, + context_window: None, + native: Some(capability), + }), + rx, + ) +} + +fn response(images: Value) -> bytes::Bytes { + serde_json::to_vec( + &json!({"id":"gen-test", "model":"test/image", "choices":[{"index":0,"message":{ + "role":"assistant", "content":"{\"sticker\":true}", "images":images + }, "finish_reason":"stop"}]}), + ) + .unwrap() + .into() +} + +#[tokio::test] +async fn native_http_request_and_completed_bytes_preserve_real_text_without_labels() { + let bytes = png(); + let uri = format!( + "data:image/png;base64,{}", + base64::engine::general_purpose::STANDARD.encode(&bytes) + ); + let (mut session, mut sent) = + session(vec![response(json!([{"image_url":{"url":uri}}]))], true).await; + let mut turn = session.begin_turn(request(true), None).await.unwrap(); + let wire = sent.try_recv().unwrap(); + assert_eq!(wire["modalities"], json!(["image", "text"])); + assert_eq!(wire["stream"], false); + assert_eq!(wire["provider"]["require_parameters"], true); + assert_eq!(wire["tools"][0]["function"]["name"], "compose"); + assert!( + wire["messages"][0]["content"] + .as_array() + .unwrap() + .iter() + .any(|value| value["type"] == "image_url") + ); + let mut finished = false; + let mut committed_image = false; + while let Some(event) = turn.next_event(None).await.unwrap() { + match event { + ModelTurnEvent::Delta(Delta::CommitPart { + part: Part::Media(media), + }) => { + assert_eq!(media.mime_type, "image/png"); + assert_eq!(media.data, DataRef::InlineBytes(bytes.clone())); + committed_image = true; + } + ModelTurnEvent::Delta(Delta::AppendText { chunk, .. }) => { + assert!(!chunk.contains("[Image")) + } + ModelTurnEvent::Finished(result) => { + let parts: Vec<_> = result + .output_items + .iter() + .flat_map(|item| &item.parts) + .collect(); + assert!(parts.iter().any( + |part| matches!(part, Part::Text(text) if text.text == "{\"sticker\":true}") + )); + assert!(parts.iter().any(|part| matches!(part, Part::Media(media) if media.data == DataRef::InlineBytes(bytes.clone()) && media.mime_type == "image/png"))); + finished = true; + } + _ => {} + } + } + assert!(finished); + assert!(committed_image); + assert!( + sent.try_recv().is_err(), + "successful generation must not be replayed" + ); +} + +#[tokio::test] +async fn unsupported_image_input_fails_before_http() { + let (mut session, mut sent) = session(vec![], false).await; + assert!( + session + .begin_turn(request(true), None) + .await + .err() + .unwrap() + .to_string() + .contains("does not support image input") + ); + assert!(sent.try_recv().is_err()); +} + +#[tokio::test] +async fn transport_bounds_raw_json_and_rejects_malformed_siblings() { + let chunk = bytes::Bytes::from(vec![b' '; 1024 * 1024]); + let (mut session, _sent) = session(vec![chunk; 25], true).await; + // Keep capture channel alive: the fake is an actual HTTP boundary, not instrumentation. + let error = session + .begin_turn(request(false), None) + .await + .err() + .unwrap() + .to_string(); + assert!(error.contains("native-image-response-too-large"), "{error}"); + for images in [ + json!([{}]), + json!([{"image_url":{"url":"https://example.invalid/x"}}]), + json!({}), + ] { + let (mut session, _sent) = self::session(vec![response(images)], true).await; + let error = session + .begin_turn(request(false), None) + .await + .err() + .unwrap() + .to_string(); + assert!(error.contains("native-image-malformed"), "{error}"); + } +} + +#[tokio::test] +async fn custom_and_text_models_do_not_send_generation_modalities() { + for config in [ + OpenRouterConfig::new("test", "test/image") + .with_base_url("https://custom.invalid/chat/completions"), + OpenRouterConfig::new("test", "test/text"), + ] { + let reported = if config.model == "test/text" { + info(&["text"], &["text"], true) + } else { + info(&["image", "text"], &["image", "text"], true) + }; + assert!( + native_image_capability(&config, Some(&reported)) + .unwrap() + .is_none() + ); + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); + let adapter = CompletionsAdapter::with_client( + OpenRouterProvider::from(config.with_streaming(false)), + Http::new(FakeHttp { + body: vec![response(json!([]))], + sent: tx, + }), + ); + let mut session = adapter + .start_session(SessionConfig::new("native")) + .await + .unwrap(); + let mut turn = session.begin_turn(request(false), None).await.unwrap(); + let wire = rx.try_recv().unwrap(); + assert!(wire.get("modalities").is_none()); + assert_eq!(wire["tools"][0]["function"]["name"], "compose"); + while let Some(event) = turn.next_event(None).await.unwrap() { + if let ModelTurnEvent::Finished(result) = event { + assert!(result.output_items.iter().flat_map(|item| &item.parts).any( + |part| matches!(part, Part::Text(text) if text.text == "{\"sticker\":true}") + )); + } + } + } +} + +#[tokio::test] +async fn malformed_content_images_and_sse_are_strict_transport_errors() { + let malformed_content = + json!({"choices":[{"message":{"role":"assistant","content":[{"type":"image_url"}]}}]}); + for body in [ + serde_json::to_vec(&malformed_content).unwrap(), + b"data: {\"choices\":[]}\n\n".to_vec(), + br#"{"choices":[]}"#.to_vec(), + ] { + let (mut session, _sent) = session(vec![body.into()], true).await; + let error = session + .begin_turn(request(false), None) + .await + .err() + .unwrap() + .to_string(); + assert!(error.contains("native-image-malformed"), "{error}"); + } +} + +#[test] +fn excessive_dimensions_are_size_errors_not_malformed_media() { + let mut bytes = std::io::Cursor::new(Vec::new()); + image::DynamicImage::new_rgba8(8193, 1) + .write_to(&mut bytes, image::ImageFormat::Png) + .unwrap(); + let uri = format!( + "data:image/png;base64,{}", + base64::engine::general_purpose::STANDARD.encode(bytes.into_inner()) + ); + let error = normalize_native_part(&mut image_part(uri)) + .unwrap_err() + .to_string(); + assert!(error.contains("native-image-too-large"), "{error}"); +} + +fn endpoint_document(model: &str, endpoints: Value) -> Value { + json!({"data":{"id":model,"architecture":{"input_modalities":["image","text"],"output_modalities":["image","text"]},"endpoints":endpoints}}) +} + +struct EndpointHttp { + body: Vec, + status: StatusCode, + sent: tokio::sync::mpsc::UnboundedSender, +} + +#[async_trait] +impl HttpClient for EndpointHttp { + async fn execute(&self, request: HttpRequest) -> Result { + assert!(request.body.is_none()); + self.sent.send(request.url.clone()).unwrap(); + Ok(HttpResponse::new( + self.status, + HeaderMap::new(), + request.url, + Box::pin(futures_util::stream::iter( + self.body.clone().into_iter().map(Ok), + )), + )) + } +} + +fn endpoint_client( + body: Vec, + status: StatusCode, +) -> (Http, tokio::sync::mpsc::UnboundedReceiver) { + let (sent, received) = tokio::sync::mpsc::unbounded_channel(); + (Http::new(EndpointHttp { body, status, sent }), received) +} + +#[tokio::test] +async fn routing_selectors_with_empty_endpoints_remain_legacy_without_name_heuristics() { + for model in [ + "openrouter/auto", + "openrouter/auto-beta", + "other/arbitrary-router", + ] { + let document = endpoint_document(model, json!([])); + let (client, mut requests) = endpoint_client( + vec![serde_json::to_vec(&document).unwrap().into()], + StatusCode::OK, + ); + let config = OpenRouterConfig::new("test", model); + // Empty endpoints remain legacy even when the aggregate entry lacks tools. + let catalog = info(&["image", "text"], &["image", "text"], false); + assert!( + discover_native_image(&client, &config, Some(&catalog)) + .await + .unwrap() + .is_none() + ); + assert_eq!( + requests.try_recv().unwrap(), + format!("https://openrouter.ai/api/v1/models/{model}/endpoints") + ); + } +} + +#[tokio::test] +async fn concrete_mixed_endpoints_enable_generation_and_preserve_routing_privacy() { + let catalog = info(&["image", "text"], &["image", "text"], true); + let document = endpoint_document( + "test/image", + json!([ + {"supported_parameters":["temperature"]}, + {"supported_parameters":["tools","tool_choice"]} + ]), + ); + let (client, _requests) = endpoint_client( + vec![serde_json::to_vec(&document).unwrap().into()], + StatusCode::OK, + ); + let config = OpenRouterConfig::new("test", "test/image").with_extra_body_value("provider", json!({ + "data_collection":"deny", "order":["Google AI Studio"], "allow_fallbacks":false, "require_parameters":false + })); + let capability = discover_native_image(&client, &config, Some(&catalog)) + .await + .unwrap() + .unwrap(); + let native = native_generation_config(config, &capability).unwrap(); + assert_eq!(native.extra_body["modalities"], json!(["image", "text"])); + assert_eq!( + native.extra_body["provider"], + json!({ + "data_collection":"deny", "order":["Google AI Studio"], "allow_fallbacks":false, "require_parameters":true + }) + ); + let invalid = + OpenRouterConfig::new("test", "test/image").with_extra_body_value("provider", "invalid"); + assert!(native_generation_config(invalid, &capability).is_err()); +} + +#[tokio::test] +async fn endpoint_ineligibility_malformed_and_failed_discovery_do_not_fabricate_capability() { + let config = OpenRouterConfig::new("test", "test/image"); + let catalog = info(&["image", "text"], &["image", "text"], true); + let no_tools = endpoint_document( + "test/image", + json!([{"supported_parameters":["temperature"]}]), + ); + let mut wrong_architecture = endpoint_document("test/image", json!([])); + wrong_architecture["data"]["architecture"]["output_modalities"] = json!(["text"]); + let mut too_many = endpoint_document("test/image", json!([])); + too_many["data"]["endpoints"] = json!(vec![ + json!({"supported_parameters":["tools"]}); + MAX_NATIVE_ENDPOINTS + 1 + ]); + let cases = [ + (no_tools, "native-image-ineligible"), + ( + endpoint_document("wrong/model", json!([])), + "native-image-discovery", + ), + (wrong_architecture, "native-image-discovery"), + ( + endpoint_document("test/image", json!([{}])), + "native-image-discovery", + ), + ( + endpoint_document( + "test/image", + json!([{"supported_parameters":vec!["tools"; MAX_CAPABILITY_VALUES + 1]}]), + ), + "native-image-discovery", + ), + (too_many, "native-image-discovery"), + ( + json!({"data":{"id":"test/image"}}), + "native-image-discovery", + ), + ]; + for (document, code) in cases { + let (client, _requests) = endpoint_client( + vec![serde_json::to_vec(&document).unwrap().into()], + StatusCode::OK, + ); + let error = discover_native_image(&client, &config, Some(&catalog)) + .await + .err() + .unwrap() + .to_string(); + assert!(error.contains(code), "{error}"); + } + for (body, status) in [ + (vec![bytes::Bytes::from_static(b"not JSON")], StatusCode::OK), + ( + vec![ + bytes::Bytes::from(vec![b' '; MAX_MODELS_BYTES]), + bytes::Bytes::from_static(b"x"), + ], + StatusCode::OK, + ), + (vec![], StatusCode::SERVICE_UNAVAILABLE), + ] { + let (client, _requests) = endpoint_client(body, status); + assert!( + discover_native_image(&client, &config, Some(&catalog)) + .await + .err() + .unwrap() + .to_string() + .contains("native-image-discovery") + ); + } +} + +#[tokio::test] +async fn custom_and_nonimage_routes_never_fetch_official_endpoints() { + let (client, mut requests) = endpoint_client(vec![], StatusCode::SERVICE_UNAVAILABLE); + let config = OpenRouterConfig::new("test", "test/image"); + let catalog = info(&["image", "text"], &["image", "text"], true); + let custom = config + .clone() + .with_base_url("https://custom.invalid/chat/completions"); + assert!( + discover_native_image(&client, &custom, Some(&catalog)) + .await + .unwrap() + .is_none() + ); + let text = info(&["text"], &["text"], true); + assert!( + discover_native_image(&client, &config, Some(&text)) + .await + .unwrap() + .is_none() + ); + assert!( + discover_native_image(&client, &config, None) + .await + .unwrap() + .is_none() + ); + assert!(requests.try_recv().is_err()); +} + +#[test] +fn endpoint_model_paths_cannot_inject_query_fragment_or_traversal() { + for model in [ + "author/model?x=y", + "author/model#fragment", + "author/../model", + "author//model", + "author/%2e%2e", + "model", + "author/.", + ] { + assert!(native_endpoints_url(model).is_err(), "{model}"); + } + assert_eq!( + native_endpoints_url("author/model:free").unwrap().as_str(), + "https://openrouter.ai/api/v1/models/author/model:free/endpoints" + ); +} + +struct FailedEndpointHttp; +#[async_trait] +impl HttpClient for FailedEndpointHttp { + async fn execute(&self, _request: HttpRequest) -> Result { + Err(HttpError::Other("fixture transport unavailable".into())) + } +} + +#[tokio::test] +async fn failed_endpoint_transport_is_unknown_not_legacy_or_ineligible() { + let client = Http::new(FailedEndpointHttp); + let config = OpenRouterConfig::new("test", "test/image"); + let catalog = info(&["image", "text"], &["image", "text"], true); + let error = discover_native_image(&client, &config, Some(&catalog)) + .await + .err() + .unwrap() + .to_string(); + assert!(error.contains("native-image-discovery")); + assert!(error.contains("eligibility is unknown")); +} + +async fn completed_items(turn: &mut KitTurn) -> Vec { + let mut output = None; + while let Some(event) = turn.next_event(None).await.unwrap() { + if let ModelTurnEvent::Finished(result) = event { + output = Some(result.output_items); + } + } + output.expect("completed native turn") +} + +fn assert_historical_image_wire(wire: &Value, parallel_tools: bool) { + let messages = wire["messages"].as_array().unwrap(); + let image_position = messages + .iter() + .position(|message| { + message["role"] == "user" + && message["content"] + .as_array() + .is_some_and(|parts| parts.iter().any(|part| part["type"] == "image_url")) + }) + .expect("historical generated image must be encoded as actual image input"); + let image = messages[image_position]["content"] + .as_array() + .unwrap() + .iter() + .find(|part| part["type"] == "image_url") + .unwrap(); + assert_eq!( + image["image_url"]["url"], + format!( + "data:image/png;base64,{}", + base64::engine::general_purpose::STANDARD.encode(png()) + ) + ); + assert!(messages.iter().any( + |message| message["role"] == "assistant" && message["content"] == "{\"sticker\":true}" + )); + if parallel_tools { + let positions: Vec<_> = messages + .iter() + .enumerate() + .filter(|(_, message)| message["role"] == "tool") + .map(|(index, _)| index) + .collect(); + assert_eq!(positions.len(), 2); + assert!(positions.iter().all(|index| *index < image_position)); + assert_eq!(messages[positions[0]]["tool_call_id"], "a"); + assert_eq!(messages[positions[1]]["tool_call_id"], "b"); + } +} + +#[tokio::test] +async fn native_generated_history_supports_second_prompt_tool_roundtrip_and_restored_sessions() { + use agentkit_core::{ToolOutput, ToolResultPart}; + for parallel_tools in [false, true] { + let uri = format!( + "data:image/png;base64,{}", + base64::engine::general_purpose::STANDARD.encode(png()) + ); + let mut body: Value = + serde_json::from_slice(&response(json!([{"image_url":{"url":uri}}]))).unwrap(); + if parallel_tools { + body["choices"][0]["finish_reason"] = json!("tool_calls"); + body["choices"][0]["message"]["tool_calls"] = json!([ + {"id":"a","type":"function","function":{"name":"compose","arguments":"{}"}}, + {"id":"b","type":"function","function":{"name":"compose","arguments":"{}"}} + ]); + } + let response_bytes: bytes::Bytes = serde_json::to_vec(&body).unwrap().into(); + let (mut live, mut sent) = session(vec![response_bytes.clone()], true).await; + let first_request = request(false); + let mut first = live.begin_turn(first_request.clone(), None).await.unwrap(); + sent.try_recv().unwrap(); + let first_output = completed_items(&mut first).await; + let original_output = serde_json::to_value(&first_output).unwrap(); + let mut continuation = first_request; + continuation.transcript.extend(first_output.clone()); + if parallel_tools { + for id in ["a", "b"] { + continuation.transcript.push(Item::new( + ItemKind::Tool, + vec![Part::ToolResult(ToolResultPart::success( + id, + ToolOutput::Text(format!("result {id}")), + ))], + )); + } + } + continuation.transcript.push(Item::new( + ItemKind::User, + vec![Part::text("Refine that sticker")], + )); + let history = serde_json::to_value(&continuation).unwrap(); + let mut second = live.begin_turn(continuation.clone(), None).await.unwrap(); + assert_historical_image_wire(&sent.try_recv().unwrap(), parallel_tools); + completed_items(&mut second).await; + assert_eq!( + serde_json::to_value(&first_output).unwrap(), + original_output + ); + assert_eq!(serde_json::to_value(&continuation).unwrap(), history); + + // Fork/resume reconstruct new provider sessions from canonical history, + // not a provider-private cache or the already projected outbound copy. + for _ in ["fork", "resume"] { + let restored: TurnRequest = serde_json::from_value(history.clone()).unwrap(); + let (mut reconstructed, mut sent) = session(vec![response_bytes.clone()], true).await; + let mut turn = reconstructed.begin_turn(restored, None).await.unwrap(); + assert_historical_image_wire(&sent.try_recv().unwrap(), parallel_tools); + completed_items(&mut turn).await; + } + } +} + +#[tokio::test] +async fn image_only_assistant_history_has_no_empty_outbound_assistant_message() { + let mut history = request(false); + history.transcript.push(Item::new( + ItemKind::Assistant, + vec![Part::Media(MediaPart::new( + Modality::Image, + "image/png", + DataRef::InlineBytes(png()), + ))], + )); + let original = history.clone(); + let (mut session, mut sent) = session(vec![response(json!([]))], true).await; + session.begin_turn(history.clone(), None).await.unwrap(); + let wire = sent.try_recv().unwrap(); + assert!( + wire["messages"] + .as_array() + .unwrap() + .iter() + .all(|message| message["role"] != "assistant") + ); + assert_eq!(history, original); +} + +#[test] +fn historical_assistant_projection_is_bounded_and_waits_for_all_tool_results() { + use agentkit_core::ToolCallPart; + let image = Part::Media(MediaPart::new( + Modality::Image, + "image/png", + DataRef::InlineBytes(png()), + )); + let mut history = request(false); + history.transcript.push(Item::new( + ItemKind::Assistant, + vec![ + image.clone(), + Part::ToolCall(ToolCallPart::new("pending", "compose", json!({}))), + ], + )); + assert!( + project_tool_output_images(history.clone(), false) + .unwrap_err() + .to_string() + .contains("outstanding tool calls") + ); + history.transcript.last_mut().unwrap().parts = vec![image; 9]; + assert!( + project_tool_output_images(history, false) + .unwrap_err() + .to_string() + .contains("historical assistant images exceed") + ); +} + +#[test] +fn native_generation_policy_is_finite_and_never_replays_ambiguous_billable_work() { + let policy = native_generation_resilience(); + assert_eq!(policy.max_retries, 0); + assert_eq!(policy.attempt_timeout, Some(Duration::from_secs(300))); + assert_eq!(policy.stream_idle_timeout, Some(Duration::from_secs(300))); + assert_eq!(policy.retry_budget, Duration::from_secs(310)); + assert_eq!(NATIVE_GENERATION_TIMEOUT, Duration::from_secs(300)); +} + +struct PendingGenerationHttp(tokio::sync::mpsc::UnboundedSender<()>); +#[async_trait] +impl HttpClient for PendingGenerationHttp { + async fn execute(&self, _request: HttpRequest) -> Result { + self.0.send(()).unwrap(); + futures_util::future::pending().await + } +} + +struct AmbiguousGenerationHttp(tokio::sync::mpsc::UnboundedSender<()>); +#[async_trait] +impl HttpClient for AmbiguousGenerationHttp { + async fn execute(&self, _request: HttpRequest) -> Result { + self.0.send(()).unwrap(); + Err(HttpError::Timeout { + operation: "fixture accepted generation", + timeout: NATIVE_GENERATION_TIMEOUT, + }) + } +} + +async fn session_with_http(http: Http) -> KitSession { + let capability = NativeImageCapability { + image_input: true, + modalities: vec!["image".into(), "text".into()], + }; + let config = + native_generation_config(OpenRouterConfig::new("test", "test/image"), &capability).unwrap(); + let adapter = CompletionsAdapter::with_client( + OpenRouterProvider::from(config), + Http::new(BoundedImageClient { inner: http }), + ) + .with_resilience(native_generation_resilience()); + let inner = adapter + .start_session(SessionConfig::new("native")) + .await + .unwrap(); + KitSession::OpenRouter(OpenRouterKitSession { + inner, + context_window: None, + native: Some(capability), + }) +} + +#[tokio::test] +async fn native_pending_http_is_cancellable_and_ambiguous_timeout_is_not_replayed() { + // This is only a deadlock guard; no wall-clock performance assertion or sleep. + tokio::time::timeout(Duration::from_secs(5), async { + let (started, mut received) = tokio::sync::mpsc::unbounded_channel(); + let mut session = session_with_http(Http::new(PendingGenerationHttp(started))).await; + let controller = agentkit_core::CancellationController::new(); + let checkpoint = controller.handle().checkpoint(); + let (result, _) = tokio::join!( + session.begin_turn(request(false), Some(checkpoint)), + async { + received.recv().await.unwrap(); + controller.interrupt(); + } + ); + assert!(matches!(result, Err(LoopError::Cancelled))); + assert!(received.try_recv().is_err()); + + let (started, mut received) = tokio::sync::mpsc::unbounded_channel(); + let mut session = session_with_http(Http::new(AmbiguousGenerationHttp(started))).await; + assert!(session.begin_turn(request(false), None).await.is_err()); + received.try_recv().unwrap(); + assert!( + received.try_recv().is_err(), + "ambiguous billable generation must not be retried" + ); + }) + .await + .unwrap(); +} diff --git a/src/runtime.rs b/src/runtime.rs index e0088b9b..eae39813 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -100,6 +100,7 @@ mod test_support { fn rejected_fork_deferral_cleans_transcripts_before_releasing_identity_for_retry() { for kind in ["configured", "generated", "load"] { let root = tempfile::tempdir().unwrap(); + let selected_id = crate::session::new_id(); let runtime = if kind == "generated" { Runtime::new(root.path(), "gpt-5.4").unwrap() } else { @@ -107,7 +108,7 @@ mod test_support { root.path(), "gpt-5.4", SessionRequest { - id: "selected".into(), + id: selected_id.clone(), resume: false, force: false, }, @@ -115,7 +116,7 @@ mod test_support { .unwrap() }; let mut claim = if kind == "load" { - runtime.claim_session_load("selected").unwrap() + runtime.claim_session_load(&selected_id).unwrap() } else { runtime.claim_session().unwrap() }; @@ -127,11 +128,26 @@ mod test_support { vec![Item::text(agentkit_core::ItemKind::System, "system")], ) .unwrap(); + let store = crate::managed_files::FileStore::new(root.path()); + let source_id = crate::session::new_id(); + let image_path = root.path().join("inherited.png"); + image::DynamicImage::new_luma8(2, 2) + .save(&image_path) + .unwrap(); + let reference = store.import(&source_id, &image_path, None).unwrap(); + opened + .observer + .attach_inherited_authority( + store.prepare_inheritance(&source_id, &id).unwrap().unwrap(), + ) + .unwrap(); claim.guard_uncommitted_transcript(&opened.observer); drop(opened); assert!(crate::session::load(root.path(), &id).is_ok()); assert!(claim.defer_fork_commit().is_err()); assert!(crate::session::load(root.path(), &id).is_err()); + assert!(store.resolve(&id, &reference).is_err()); + assert!(store.resolve(&source_id, &reference).is_ok()); let mut retry = runtime.claim_session().unwrap(); assert_eq!(retry.id(), id); assert!(!retry.request.resume); @@ -143,12 +159,80 @@ mod test_support { ) .unwrap(); retry.guard_uncommitted_transcript(&recreated.observer); + assert!(store.resolve(&id, &reference).is_err()); retry.commit().unwrap(); drop(recreated); + let base = crate::artifacts::base(root.path()).with_file_name("files"); + let _ = std::fs::remove_dir_all( + base.join(blake3::hash(source_id.as_bytes()).to_hex().as_str()), + ); assert!(crate::session::load(root.path(), &id).is_ok()); } } + #[derive(Clone)] + struct DiscardLoopEvents; + impl LoopObserver for DiscardLoopEvents { + fn handle_event(&self, _event: agentkit_loop::ObservedEvent) {} + } + + #[tokio::test] + async fn native_fork_setup_failure_rolls_back_inherited_authority_with_claim() { + let root = tempfile::tempdir().unwrap(); + let runtime = Runtime::new(root.path(), "gpt-5.4").unwrap(); + let store = crate::managed_files::FileStore::new(root.path()); + let source_id = crate::session::new_id(); + let image_path = root.path().join("setup.png"); + image::DynamicImage::new_luma8(2, 2) + .save(&image_path) + .unwrap(); + let reference = store.import(&source_id, &image_path, None).unwrap(); + let mut claim = runtime.claim_session_fork().unwrap(); + let destination = claim.id().to_owned(); + let controller = CancellationController::new(); + let result = runtime + .start_acp_driver_with_initial( + AcpDriverContext { + cwd: root.path().to_path_buf(), + additional_directories: Vec::new(), + integration: Arc::new(DiscardLoopEvents), + cancellation: controller.handle(), + response_attempt_replacement: false, + }, + &mut claim, + Some(AcpForkState { + source_session_id: source_id.clone(), + transcript: vec![Item::text(ItemKind::System, "fork")], + // The real adapter constructor rejects this after inheritance. + selection: ModelSelection::new(ProviderKind::OpenRouter, ""), + reasoning_effort: None, + parent_context: None, + }), + ) + .await; + assert!( + matches!(result, Err(AcpRuntimeError::Loop(ref error)) if error.contains("model name")) + ); + assert!(store.resolve(&destination, &reference).is_ok()); + drop(claim); + assert!(crate::session::load(root.path(), &destination).is_err()); + assert!(store.resolve(&destination, &reference).is_err()); + assert!(store.resolve(&source_id, &reference).is_ok()); + let retry = crate::session::open_uncommitted( + root.path(), + &destination, + false, + vec![Item::text(ItemKind::System, "retry")], + ) + .unwrap(); + assert!(store.resolve(&destination, &reference).is_err()); + drop(retry); + let base = crate::artifacts::base(root.path()).with_file_name("files"); + let _ = std::fs::remove_dir_all( + base.join(blake3::hash(source_id.as_bytes()).to_hex().as_str()), + ); + } + impl Runtime { pub(crate) fn mcp_for_test(&self) -> &crate::tools::mcp::McpRuntime { &self.mcp @@ -311,6 +395,8 @@ impl SessionSelection { } pub(crate) struct AcpForkState { + /// Trusted actor storage identity, never supplied by transcript metadata. + pub source_session_id: String, pub transcript: Vec, pub selection: ModelSelection, pub reasoning_effort: Option, @@ -1477,13 +1563,14 @@ impl Runtime { } let request = claim.request.clone(); let session_id = request.id.clone(); - let (forked_transcript, selected, parent_context) = match forked { + let (forked_transcript, selected, parent_context, source_session_id) = match forked { Some(forked) => ( Some(forked.transcript), Some((forked.selection, forked.reasoning_effort)), forked.parent_context, + Some(forked.source_session_id), ), - None => (None, None, None), + None => (None, None, None, None), }; let is_fork = forked_transcript.is_some(); let initial = if let Some(transcript) = forked_transcript { @@ -1514,6 +1601,20 @@ impl Runtime { if is_fork || !request.resume { claim.guard_uncommitted_transcript(&opened.observer); } + if let Some(source_session_id) = source_session_id { + // The claim already guards the destination transcript. Authority must + // be prepared before any driver/actor or successful fork publication; + // failure leaves the existing uncommitted cleanup owner intact. + if let Some(authority) = crate::managed_files::FileStore::new(&self.root) + .prepare_inheritance(&source_session_id, &request.id) + .map_err(AcpRuntimeError::Loop)? + { + opened + .observer + .attach_inherited_authority(authority) + .map_err(AcpRuntimeError::Loop)?; + } + } // Every ACP route owns its model selection. Changing one session // cannot redirect another session served by the same runtime. let (selection, reasoning_effort) = selected.unwrap_or_else(|| { diff --git a/src/runtime/tests/managed_files.rs b/src/runtime/tests/managed_files.rs index dffeaa0c..b475d3ed 100644 --- a/src/runtime/tests/managed_files.rs +++ b/src/runtime/tests/managed_files.rs @@ -1,3 +1,5 @@ +mod native_subagent; + use super::*; use agentkit_core::{DataRef, Item, Modality, ToolResultPart}; use agentkit_http::Authentication; @@ -41,6 +43,18 @@ impl Fixture { // Reconstruct the runtime on every call to exercise restart-independent // storage rather than an in-memory resolver registry. let runtime = Runtime::new(self.root.path(), "gpt-5.4").unwrap(); + self.execute_with_runtime(runtime, script, input, background, outcome) + .await + } + + async fn execute_with_runtime( + &self, + runtime: Arc, + script: &str, + input: Value, + background: bool, + outcome: bool, + ) -> Result { let compose = runtime.compose(0); assert_eq!(compose.specs().len(), 1); assert_eq!(compose.specs()[0].name.0, "compose"); diff --git a/src/runtime/tests/managed_files/native_subagent.rs b/src/runtime/tests/managed_files/native_subagent.rs new file mode 100644 index 00000000..f5d391d6 --- /dev/null +++ b/src/runtime/tests/managed_files/native_subagent.rs @@ -0,0 +1,174 @@ +use super::*; +use crate::acp_child::{AcpHarnessProfile, AcpHarnesses}; +use std::collections::BTreeMap; + +const PIPELINE: &str = r#" +source = read_file({ path: "image.png" }) +rotated = image_rotate({ image: source, degrees: 90 }) +cropped = image_crop({ image: rotated, aspect_ratio: { width: 1, height: 1 }, anchor: "center" }) +child = subagent({ + prompt: input.prompt, + model: input.model, + attachments: [cropped], + output_schema: input.schema +}) +_ = close(child) +return child.output.result +"#; + +fn output_schema(index: Option) -> Value { + let mut binding = json!({"$ref":"kit://schemas/file/v1"}); + if let Some(index) = index { + binding["x-kit-image-index"] = json!(index); + } + json!({ + "type":"object", + "properties":{"result":binding}, + "required":["result"], + "additionalProperties":false + }) +} + +fn configured_runtime( + fixture: &Fixture, + name: &str, + profile: AcpHarnessProfile, + provider: crate::ProviderKind, + model: &str, +) -> Arc { + let runtime = Runtime::new_with_provider(fixture.root.path(), model, provider).unwrap(); + Runtime::with_acp_harnesses( + runtime, + AcpHarnesses::new(BTreeMap::from([(name.into(), profile)])).unwrap(), + format!("acp.{name}"), + ) + .unwrap() +} + +fn delivered_image(output: &ToolOutput) -> &[u8] { + let ToolOutput::Parts(parts) = output else { + panic!("pipeline did not deliver native parts: {output:?}"); + }; + let images = parts + .iter() + .filter_map(|part| match part { + Part::Media(media) if media.modality == Modality::Image => Some(media), + _ => None, + }) + .collect::>(); + assert_eq!( + images.len(), + 1, + "only the final selected image is delivered" + ); + let DataRef::InlineBytes(bytes) = &images[0].data else { + panic!("native bytes required"); + }; + bytes +} + +#[tokio::test] +async fn compose_transform_native_subagent_output_survives_close() { + let fixture = Fixture::new(); + let log = fixture.root.path().join("wire.jsonl"); + let runtime = configured_runtime( + &fixture, + "image-fixture", + AcpHarnessProfile { + command: "python3".into(), + args: vec![ + format!( + "{}/src/tools/subagent/native-image-fixture.py", + env!("CARGO_MANIFEST_DIR") + ), + base64::engine::general_purpose::STANDARD.encode(&fixture.bytes), + log.display().to_string(), + ], + permissions: Default::default(), + }, + crate::ProviderKind::OpenAiSubscription, + "gpt-5.4", + ); + let output = fixture + .execute_with_runtime( + runtime, + PIPELINE, + json!({"prompt":"root", "model":null, "schema":output_schema(None)}), + false, + true, + ) + .await + .unwrap(); + assert_eq!(delivered_image(&output), fixture.bytes); + let request: Value = serde_json::from_str( + std::fs::read_to_string(log) + .unwrap() + .lines() + .next() + .unwrap(), + ) + .unwrap(); + assert_eq!(request["prompt"][1]["type"], "image"); + let attached = base64::engine::general_purpose::STANDARD + .decode(request["prompt"][1]["data"].as_str().unwrap()) + .unwrap(); + let cropped = image::load_from_memory(&attached).unwrap(); + assert_eq!((cropped.width(), cropped.height()), (2, 2)); + assert_ne!( + attached, fixture.bytes, + "child receives the transformed snapshot" + ); +} + +/// Explicit opt-in only: this performs a billable request to the selected provider. +/// Build this worktree's Kit binary, then set KIT_LIVE_KIT_BINARY and +/// KIT_LIVE_IMAGE_MODEL to an explicitly selected image-output model. The caller +/// explicitly selects distinct output index 0; the backend may emit alternatives. +#[tokio::test] +#[ignore = "requires an explicitly selected live image model and provider credentials"] +async fn live_compose_sticker_pipeline() { + let binary = + std::env::var("KIT_LIVE_KIT_BINARY").expect("set this worktree's built Kit binary"); + let model = + std::env::var("KIT_LIVE_IMAGE_MODEL").expect("select a verified image-output model"); + let mut fixture = Fixture::new(); + let source = image::DynamicImage::ImageRgb8(image::RgbImage::from_pixel( + 64, + 96, + image::Rgb([153, 204, 255]), + )); + let mut encoded = std::io::Cursor::new(Vec::new()); + source + .write_to(&mut encoded, image::ImageFormat::Png) + .unwrap(); + fixture.bytes = encoded.into_inner(); + std::fs::write(fixture.root.path().join("image.png"), &fixture.bytes).unwrap(); + let runtime = configured_runtime( + &fixture, + "kit", + AcpHarnessProfile { + command: binary, + args: vec!["acp".into()], + permissions: Default::default(), + }, + crate::ProviderKind::OpenRouter, + model + .strip_prefix("openrouter:") + .expect("select an OpenRouter image model"), + ); + let output = fixture.execute_with_runtime(runtime, PIPELINE, json!({ + "model": model, + "schema": output_schema(Some(0)), + "prompt": "Add a small Hello Kitty sticker in the center of this blue image. Return exactly one edited image, no prose." + }), false, true).await.unwrap(); + let bytes = delivered_image(&output); + let image = image::load_from_memory(bytes).unwrap(); + assert!(image.width() > 0 && image.height() > 0); + assert_ne!( + bytes, fixture.bytes, + "native generation must not echo the source bytes" + ); + if let Ok(path) = std::env::var("KIT_LIVE_IMAGE_EVIDENCE") { + std::fs::write(path, bytes).unwrap(); + } +} diff --git a/src/session.rs b/src/session.rs index 66443203..5b8d81fd 100644 --- a/src/session.rs +++ b/src/session.rs @@ -116,6 +116,7 @@ struct CreatedTranscript { filesystem: Fs, path: Option, cleanup: Option, + inherited_authority: Option, keep: bool, } @@ -131,7 +132,7 @@ impl PreparedCreation { /// Called only after response submission. Both steps are infallible and /// require no shared lock or callback; readers resume after cleanup is kept. pub(crate) fn commit(mut self) { - self.created.keep = true; + self.created.mark_kept(); self.published.store(true, Ordering::Release); } } @@ -148,13 +149,21 @@ impl CreatedTranscript { path, filesystem, cleanup, + inherited_authority: None, keep: false, } } - fn keep(mut self) { + fn mark_kept(&mut self) { + if let Some(authority) = self.inherited_authority.take() { + authority.commit(); + } self.keep = true; } + + fn keep(mut self) { + self.mark_kept(); + } } impl Drop for CreatedTranscript { @@ -234,9 +243,15 @@ pub(crate) fn clone_completed_in( transcript, InitialTranscriptOptions { stamp_items: false, - commit_creation: true, + commit_creation: false, }, )?; + if let Some(authority) = + crate::managed_files::FileStore::new(root).prepare_inheritance(source, destination)? + { + opened.observer.attach_inherited_authority(authority)?; + } + opened.observer.prepare_creation()?.commit(); drop(opened); Ok(()) } @@ -598,6 +613,31 @@ fn finish_open( } impl SessionObserver { + /// Transfer private authority cleanup into the pending transcript owner. + /// Rejection drops the input only after the writer guard has been released. + pub(crate) fn attach_inherited_authority( + &self, + authority: crate::managed_files::InheritedAuthority, + ) -> Result<(), String> { + let mut writer = self + .0 + .lock() + .map_err(|_| "session transcript writer poisoned".to_string())?; + writer.check_ownership()?; + if !authority.is_for_session(&writer.session_id) { + return Err("inherited authority belongs to another session".into()); + } + let created = writer + .created + .as_mut() + .ok_or("session creation is not available for inherited authority")?; + if created.inherited_authority.is_some() { + return Err("session creation already owns inherited authority".into()); + } + created.inherited_authority = Some(authority); + Ok(()) + } + pub(crate) fn prepare_creation(&self) -> Result { let published = Arc::new(AtomicBool::new(false)); let mut writer = self @@ -2504,6 +2544,55 @@ mod tests { assert_eq!(cloned[1].created_at, Some(Timestamp(77))); } + #[test] + fn cloning_inherits_managed_files_without_later_branch_access() { + let root = tempfile::tempdir().unwrap(); + let source_id = format!( + "files-source-{}", + blake3::hash(root.path().as_os_str().as_encoded_bytes()).to_hex() + ); + let destination_id = format!( + "files-fork-{}", + blake3::hash(root.path().as_os_str().as_encoded_bytes()).to_hex() + ); + let store = crate::managed_files::FileStore::new(&project_root(root.path())); + let mut bytes = std::io::Cursor::new(Vec::new()); + image::DynamicImage::new_luma8(2, 2) + .write_to(&mut bytes, image::ImageFormat::Png) + .unwrap(); + let reference = store + .import_bytes(&source_id, "fork.png", "image/png", bytes.get_ref(), None) + .unwrap(); + let source = open( + root.path(), + &source_id, + false, + false, + vec![Item::text( + ItemKind::Assistant, + serde_json::to_string(&reference).unwrap(), + )], + ) + .unwrap(); + drop(source); + clone_completed(root.path(), &source_id, &destination_id).unwrap(); + let reopened = crate::managed_files::FileStore::new(&project_root(root.path())); + assert_eq!( + reopened.resolve(&destination_id, &reference).unwrap(), + *bytes.get_ref() + ); + let later = reopened + .import_bytes(&source_id, "later.png", "image/png", bytes.get_ref(), None) + .unwrap(); + assert!(reopened.resolve(&destination_id, &later).is_err()); + assert!(item_text(&load(root.path(), &destination_id).unwrap()[0]).contains("fork.png")); + let file_base = crate::artifacts::base(&project_root(root.path())).with_file_name("files"); + for session in [&source_id, &destination_id] { + fs::remove_dir_all(file_base.join(blake3::hash(session.as_bytes()).to_hex().as_str())) + .unwrap(); + } + } + #[test] fn cloning_sanitizes_session_bound_continuation_metadata() { let root = tempfile::tempdir().unwrap(); diff --git a/src/tools/subagent.rs b/src/tools/subagent.rs index f6353e50..5d3e2ebd 100644 --- a/src/tools/subagent.rs +++ b/src/tools/subagent.rs @@ -4,11 +4,14 @@ use std::{ sync::{Arc, Mutex}, }; +use crate::managed_files::{FileReference, FileStore}; +use agentkit_acp::ImageContent; use agentkit_core::{ToolOutput, ToolResultPart, TurnCancellation}; use agentkit_tools_core::{ Tool, ToolAnnotations, ToolContext, ToolError, ToolName, ToolRequest, ToolResult, ToolSpec, }; use async_trait::async_trait; +use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64}; use futures_util::future::{Either, select}; use serde::{Deserialize, Serialize}; #[cfg(test)] @@ -180,22 +183,495 @@ struct SubagentListing { task: String, } +const FILE_SCHEMA_REF: &str = "kit://schemas/file/v1"; +const IMAGE_INDEX_ANNOTATION: &str = "x-kit-image-index"; +const MEDIA_GUIDANCE: &str = " Attach images explicitly with attachments. A file-aware output_schema uses $ref kit://schemas/file/v1 at one root or required fixed object path; emit exactly one distinct native assistant image and omit the Kit-bound field. An optional x-kit-image-index integer 0..7 beside the File $ref selects that distinct image in emission order after all images validate; otherwise exactly one distinct image is required. Native images without this schema return output {value, files}. Only Files in the final compose return deliver pixels."; + +struct MediaContext { + store: FileStore, + session: String, + attachments: Vec, +} + +fn attachments_schema() -> Value { + Value::Object(Map::from_iter([ + ("type".into(), Value::from("array")), + ("maxItems".into(), Value::from(8)), + ("items".into(), super::read_file::file_schema()), + ])) +} + +// File bindings deliberately support only a fixed required object-property path. +// Scan first so refs hidden in unsupported applicators cannot silently fall back. +fn file_binding(schema: &mut Value) -> Result>, String> { + // Walk schema positions only, iteratively. Ordinary schemas retain the + // validator's existing limits: a deep/wide non-File schema must not acquire + // the narrower File contract's limits. Once a File ref is found, enforce + // bounds including all positions visited before that ref. + let mut pending = vec![(&*schema, 0usize)]; + let mut count = 0; + let mut nodes = 0; + let mut exceeded = false; + let mut unsupported = false; + while let Some((value, depth)) = pending.pop() { + nodes += 1; + exceeded |= nodes > 10_000 || depth > 64; + if let Value::Object(object) = value { + if object.contains_key(IMAGE_INDEX_ANNOTATION) + && object.get("$ref").and_then(Value::as_str) != Some(FILE_SCHEMA_REF) + { + return Err("x-kit-image-index is only allowed beside the exact File $ref".into()); + } + if let Some(reference) = object.get("$ref") { + if reference.as_str() == Some(FILE_SCHEMA_REF) { + count += 1; + } else { + unsupported = true; + } + } + for (key, value) in object { + unsupported |= matches!( + key.as_str(), + "anyOf" + | "oneOf" + | "allOf" + | "if" + | "then" + | "else" + | "not" + | "items" + | "prefixItems" + | "contains" + | "$dynamicRef" + | "$recursiveRef" + ); + match key.as_str() { + "properties" | "patternProperties" | "$defs" | "definitions" + | "dependentSchemas" | "dependencies" => { + if let Value::Object(properties) = value { + pending.extend(properties.values().map(|v| (v, depth + 1))); + } + } + "allOf" | "anyOf" | "oneOf" | "prefixItems" => { + if let Value::Array(schemas) = value { + pending.extend(schemas.iter().map(|v| (v, depth + 1))); + } + } + "items" => match value { + Value::Array(schemas) => { + pending.extend(schemas.iter().map(|v| (v, depth + 1))); + } + _ => pending.push((value, depth + 1)), + }, + "additionalProperties" + | "additionalItems" + | "unevaluatedProperties" + | "unevaluatedItems" + | "propertyNames" + | "contains" + | "not" + | "if" + | "then" + | "else" + | "contentSchema" => { + pending.push((value, depth + 1)); + } + // Annotations and literal data (const, enum, examples, + // defaults) are not schemas, even when they contain $ref. + _ => {} + } + } + } + if count > 0 && exceeded { + return Err("file-aware output_schema exceeds traversal limits".into()); + } + } + if count == 0 { + return Ok(None); + } + if count != 1 || unsupported { + return Err("file-aware output_schema requires exactly one fixed File binding; arrays, unions and other references are unsupported".into()); + } + fn locate(schema: &mut Value, path: &mut Vec) -> Result { + if schema.get("$ref").and_then(Value::as_str) == Some(FILE_SCHEMA_REF) { + if schema.as_object().is_none_or(|object| { + object + .keys() + .any(|key| key != "$ref" && key != IMAGE_INDEX_ANNOTATION) + }) { + return Err("File $ref only permits the optional x-kit-image-index sibling".into()); + } + if let Some(index) = schema.get(IMAGE_INDEX_ANNOTATION) + && index.as_u64().is_none_or(|index| index > 7) + { + return Err("x-kit-image-index must be an integer from 0 through 7".into()); + } + // Expand locally without passing Kit's selection annotation to the validator. + *schema = super::read_file::file_schema(); + return Ok(true); + } + let required = schema + .get("required") + .and_then(Value::as_array) + .cloned() + .unwrap_or_default(); + let object_type = schema.get("type").and_then(Value::as_str) == Some("object"); + if let Some(properties) = schema.get_mut("properties").and_then(Value::as_object_mut) { + for (key, child) in properties { + path.push(key.clone()); + if locate(child, path)? { + if !object_type || !required.contains(&Value::String(key.clone())) { + return Err( + "every File path property must be required in a fixed object schema" + .into(), + ); + } + return Ok(true); + } + path.pop(); + } + } + Ok(false) + } + let mut path = Vec::new(); + if !locate(schema, &mut path)? { + return Err("File binding must be at the root or a required object-property path".into()); + } + Ok(Some(path)) +} + +impl OutputContract { + fn for_request( + schema: Option, + attachments: Vec, + root: &std::path::Path, + session: &str, + ) -> Result { + if attachments.len() > 8 { + return Err(ToolError::InvalidInput( + "at most eight attachments are supported".into(), + )); + } + let explicit_schema = schema.is_some(); + let mut contract = Self::new(schema.unwrap_or(Value::Bool(true)))?; + contract.explicit_schema = explicit_schema; + contract.media = Some(MediaContext { + store: FileStore::new(root), + session: session.to_owned(), + attachments, + }); + Ok(contract) + } + + fn bind(&self, text: &str, file: Value) -> Result { + let path = self + .binding + .as_ref() + .ok_or_else(|| ChildError::Failed("missing File binding".into()))?; + let mut value = if text.trim().is_empty() { + Value::Object(Map::new()) + } else { + serde_json::from_str(text.trim()).map_err(|_| { + ChildError::Failed("file-aware output contains invalid surrounding JSON".into()) + })? + }; + if path.is_empty() { + if !text.trim().is_empty() { + return Err(ChildError::Failed( + "model must not supply the root File binding".into(), + )); + } + value = file; + } else { + let mut target = &mut value; + for (index, key) in path.iter().enumerate() { + let object = target.as_object_mut().ok_or_else(|| { + ChildError::Failed("File binding requires surrounding JSON objects".into()) + })?; + if index + 1 == path.len() { + if object.contains_key(key) { + return Err(ChildError::Failed( + "model must not supply the File binding field".into(), + )); + } + object.insert(key.clone(), file); + break; + } + target = object + .entry(key.clone()) + .or_insert_with(|| Value::Object(Map::new())); + } + } + self.validator.validate(&value).map_err(|_| { + ChildError::Failed("bound File output does not match output_schema".into()) + })?; + Ok(value) + } +} + +// Charge every occurrence before deduplication. These are the managed delivery +// budgets; the ACP transport independently enforces its capture budgets. +const MAX_OUTPUT_IMAGE_BYTES: usize = 16 * 1024 * 1024; +const MAX_OUTPUT_IMAGE_PIXELS: u64 = 32 * 1024 * 1024; + +fn distinct_native_images<'a>( + images: &'a [ImageContent], + cancellation: &TurnCancellation, +) -> Result)>, ChildError> { + use crate::acp_child::{MAX_NATIVE_IMAGE_BYTES, MAX_NATIVE_IMAGES}; + if images.len() > MAX_NATIVE_IMAGES { + return Err(ChildError::Failed( + "native assistant images exceed the occurrence budget".into(), + )); + } + let mut encoded_bytes = 0; + for image in images { + if image.data.len() > MAX_NATIVE_IMAGE_BYTES.div_ceil(3) * 4 { + return Err(ChildError::Failed( + "native assistant image exceeds the encoded byte budget".into(), + )); + } + encoded_bytes += image.data.len(); + } + // Each independently encoded occurrence can have up to two padding bytes. + if encoded_bytes > (MAX_OUTPUT_IMAGE_BYTES + 2 * images.len()).div_ceil(3) * 4 { + return Err(ChildError::Failed( + "native assistant images exceed the aggregate encoded byte budget".into(), + )); + } + let mut decoded_bytes = 0; + let mut pixels = 0; + let mut distinct: Vec<(&str, Vec)> = Vec::new(); + for image in images { + if cancellation.is_cancelled() { + return Err(ChildError::Cancelled); + } + let bytes = BASE64.decode(&image.data).map_err(|_| { + ChildError::Failed("native assistant image contains invalid base64".into()) + })?; + decoded_bytes += bytes.len(); + if decoded_bytes > MAX_OUTPUT_IMAGE_BYTES { + return Err(ChildError::Failed( + "native assistant images exceed the aggregate decoded byte budget".into(), + )); + } + // Validate every occurrence, including a repeated payload with a forged + // MIME declaration, before it can qualify as an exact duplicate. + let (mime, occurrence_pixels) = + crate::managed_files::validate_provider_image(&bytes).map_err(ChildError::Failed)?; + if mime != image.mime_type { + return Err(ChildError::Failed( + "native assistant image MIME type does not match its bytes".into(), + )); + } + pixels += occurrence_pixels; + if pixels > MAX_OUTPUT_IMAGE_PIXELS { + return Err(ChildError::Failed( + "native assistant images exceed the aggregate pixel budget".into(), + )); + } + if !distinct.iter().any(|(previous_mime, previous_bytes)| { + *previous_mime == image.mime_type && previous_bytes == &bytes + }) { + distinct.push((&image.mime_type, bytes)); + } + } + if cancellation.is_cancelled() { + return Err(ChildError::Cancelled); + } + Ok(distinct) +} + +async fn run_turn( + child: &ChildSession, + root: &std::path::Path, + child_id: &str, + kit: bool, + prompt: String, + contract: Option<&OutputContract>, + cancellation: TurnCancellation, +) -> Result<(Value, Option), ChildError> { + let media = contract.and_then(|c| c.media.as_ref()); + let child_store = FileStore::new(root); + // Kit's ACP session ID is its durable storage identity, including native + // forks. External harness IDs are not Kit identities (and may collide), so + // those use the unique parent-owned handle namespace instead. + let child_id = if kit { child.session_id() } else { child_id }; + let mut attachments = Vec::new(); + if let Some(media) = media { + let selection = serde_json::to_value(&media.attachments) + .map_err(|e| ChildError::Failed(e.to_string()))?; + // Preflight all descriptor and aggregate image budgets before encoding. + let parts = media + .store + .selected_parts(&media.session, &selection, Some(&cancellation)) + .map_err(ChildError::Failed)?; + let mut granted = Vec::new(); + for reference in &media.attachments { + // selected_parts already checked every occurrence's metadata. + // Deduplicate complete references, not unvalidated IDs. + if granted.contains(&reference) { + continue; + } + media + .store + .grant_to( + &media.session, + reference, + &child_store, + child_id, + Some(&cancellation), + ) + .map_err(ChildError::Failed)?; + granted.push(reference); + } + for part in parts { + if let agentkit_core::Part::Media(part) = part + && let agentkit_core::DataRef::InlineBytes(bytes) = part.data + { + attachments.push(ImageContent::new(BASE64.encode(bytes), part.mime_type)); + } + } + } + let prompt = structured_prompt(prompt, contract); + let output = if attachments.is_empty() { + child.prompt(prompt, cancellation.clone()).await? + } else { + child + .prompt_with_attachments(prompt, attachments, cancellation.clone()) + .await? + }; + if let Some(error) = &output.media_error { + return Err(ChildError::Failed(format!( + "native assistant image capture failed: {error}" + ))); + } + let mut images = distinct_native_images(&output.images, &cancellation)?; + let strict = contract.is_some_and(|c| c.binding.is_some()); + if let Some(index) = contract.and_then(|contract| contract.image_index) { + if index >= images.len() { + return Err(ChildError::Failed(format!( + "x-kit-image-index {index} is out of range for {} distinct native assistant images", + images.len() + ))); + } + // Selection is caller-fixed, and happens only after all occurrences have + // passed validation and budgets. No unselected image is imported/granted. + images = vec![images.swap_remove(index)]; + } else if strict && images.len() != 1 { + return Err(ChildError::Failed(format!( + "file-aware output requires exactly one distinct native assistant image (received {})", + images.len() + ))); + } + if output.images.is_empty() { + return Ok(turn_output(output, contract)); + } + let media = media.ok_or_else(|| { + ChildError::Failed("native image output requires an invoking session".into()) + })?; + let mut files = Vec::new(); + for (index, (mime_type, bytes)) in images.into_iter().enumerate() { + let extension = if mime_type == "image/png" { + "png" + } else { + "jpg" + }; + let file = child_store + .import_bytes( + child_id, + &format!("assistant-{}.{}", index + 1, extension), + mime_type, + &bytes, + Some(&cancellation), + ) + .map_err(ChildError::Failed)?; + child_store + .grant_to( + child_id, + &file, + &media.store, + &media.session, + Some(&cancellation), + ) + .map_err(ChildError::Failed)?; + files.push(serde_json::to_value(file).map_err(|e| ChildError::Failed(e.to_string()))?); + } + let (value, updates) = if let Some(contract) = contract.filter(|c| c.binding.is_some()) { + let value = contract.bind(&output.text, files.remove(0))?; + let (_, updates) = turn_output(output, None); + (value, updates) + } else { + let (value, updates) = turn_output(output, contract); + ( + Value::Object(Map::from_iter([ + ("value".into(), value), + ("files".into(), Value::Array(files)), + ])), + updates, + ) + }; + // Publication and complete access validation precede the success transition. + media + .store + .selected_parts(&media.session, &value, Some(&cancellation)) + .map_err(ChildError::Failed)?; + Ok((value, updates)) +} + struct OutputContract { schema: String, validator: jsonschema::Validator, + binding: Option>, + image_index: Option, + media: Option, + explicit_schema: bool, } impl OutputContract { fn new(schema: Value) -> Result { - let validator = jsonschema::validator_for(&schema) + let mut resolved = schema.clone(); + let binding = file_binding(&mut resolved).map_err(ToolError::InvalidInput)?; + let image_index = binding.as_ref().and_then(|path| { + path.iter() + .fold(&schema, |node, key| &node["properties"][key]) + .get(IMAGE_INDEX_ANNOTATION) + .and_then(Value::as_u64) + .map(|index| index as usize) + }); + let validator = jsonschema::validator_for(&resolved) .map_err(|error| ToolError::InvalidInput(format!("invalid output_schema: {error}")))?; let schema = serde_json::to_string(&schema).map_err(|error| { ToolError::ExecutionFailed(format!("failed to serialize output_schema: {error}")) })?; - Ok(Self { schema, validator }) + Ok(Self { + schema, + validator, + binding, + image_index, + media: None, + explicit_schema: true, + }) } fn prompt(&self, prompt: String) -> String { + if !self.explicit_schema { + return prompt; + } + if let Some(path) = &self.binding { + if let Some(index) = self.image_index { + return format!( + "{prompt}\n\nEmit native assistant images. The caller fixed distinct-image index {index} in first-emission order; Kit validates every occurrence and binds only that selected image at {}. Do not supply or override the index or binding field, a placeholder, or a File ID. Return only surrounding JSON fields; omit text if none are needed. Output schema: {}", + serde_json::to_string(path).unwrap_or_default(), + self.schema + ); + } + return format!( + "{prompt}\n\nEmit exactly one distinct native assistant image. Kit imports and binds it at {}. Do not write that field, a placeholder, or a File ID. Return only surrounding JSON fields; omit text if none are needed. Output schema: {}", + serde_json::to_string(path).unwrap_or_default(), + self.schema + ); + } format!( "{prompt}\n\nReturn only a JSON value matching this JSON Schema. Do not wrap it in Markdown or add commentary:\n{}", self.schema @@ -203,6 +679,9 @@ impl OutputContract { } fn parse(&self, output: &str) -> Option { + if !self.explicit_schema || self.binding.is_some() { + return None; + } let value: Value = serde_json::from_str(output.trim()).ok()?; self.validator.validate(&value).ok()?; Some(value) @@ -395,7 +874,7 @@ impl Subagents { let child_config = self .config .clone() - .with_root(root) + .with_root(root.clone()) .with_parent_context(id.clone(), state.lock().await.name.clone()); { let locked = state.lock().await; @@ -436,10 +915,7 @@ impl Subagents { self.emit_event(locked.runtime_event(id.clone())); } self.monitor_child_exit(id.clone(), &state, &child); - let output = match child - .prompt(structured_prompt(prompt, contract), cancellation) - .await - { + let output = match run_turn(&child, &root, &id, kit, prompt, contract, cancellation).await { Ok(output) => output, Err(error) => { self.fail_removed_and_remove(&id, &state).await; @@ -447,7 +923,7 @@ impl Subagents { return Err(error); } }; - let (output, updates) = turn_output(output, contract); + let (output, updates) = output; let mut locked = state.lock().await; self.check_active(&locked)?; locked.status = SubagentStatus::Idle; @@ -508,15 +984,24 @@ impl Subagents { .clone() .ok_or_else(|| ChildError::Failed("subagent session is still starting".into()))?; let name = locked.name.clone(); + let root = locked.root.clone(); + let kit = locked.kit; let event = locked.runtime_event(prior.id.clone()); drop(locked); self.emit_event(event); - match child - .prompt(structured_prompt(prompt, contract), cancellation) - .await + match run_turn( + &child, + &root, + &prior.id, + kit, + prompt, + contract, + cancellation, + ) + .await { Ok(output) => { - let (output, updates) = turn_output(output, contract); + let (output, updates) = output; let mut locked = state.lock().await; self.check_active(&locked)?; locked.status = SubagentStatus::Idle; @@ -742,7 +1227,7 @@ impl Subagents { let child_config = self .config .clone() - .with_root(root) + .with_root(root.clone()) .with_parent_context(id.clone(), branch_name.clone()); let child_result = if native_fork { let parent = kit.then(|| (id.clone(), branch_name)); @@ -804,9 +1289,16 @@ impl Subagents { .cleanup_installed_child(&id, &state, &child, ChildError::Cancelled) .await); } - let output = match child - .prompt(structured_prompt(prompt, contract.as_deref()), cancellation) - .await + let output = match run_turn( + &child, + &root, + &id, + kit, + prompt, + contract.as_deref(), + cancellation, + ) + .await { Ok(output) => output, Err(error) => { @@ -815,7 +1307,7 @@ impl Subagents { .await); } }; - let (output, updates) = turn_output(output, contract.as_deref()); + let (output, updates) = output; let mut locked = state.lock().await; if reply.is_closed() { drop(locked); @@ -1490,6 +1982,7 @@ fn continuation_schema() -> serde_json::Value { "prompt".into(), Value::Object(Map::from_iter([("type".into(), Value::from("string"))])), ), + ("attachments".into(), attachments_schema()), ( "output_schema".into(), Value::Object(Map::from_iter([( @@ -1570,7 +2063,7 @@ impl SubagentTool { ) }; let description = format!( - "Start a parent-owned configured ACP harness, preferably assign a concise role-oriented display name, prompt it, and return its reusable session value. {usage}Omit `harness` and `model` unless the user or active workflow explicitly supplies the exact override or a configured alias. Never choose an override based on your own model, provider, publisher, familiarity, cost, or perceived quality; advertised choices indicate availability, not preference." + "Start a parent-owned configured ACP harness, preferably assign a concise role-oriented display name, prompt it, and return its reusable session value. {usage}Omit `harness` and `model` unless the user or active workflow explicitly supplies the exact override or a configured alias. Never choose an override based on your own model, provider, publisher, familiarity, cost, or perceived quality; advertised choices indicate availability, not preference.{MEDIA_GUIDANCE}" ); let input_schema = Value::Object(Map::from_iter([ ("type".into(), Value::from("object")), @@ -1624,6 +2117,7 @@ impl SubagentTool { ), ])), ), + ("attachments".into(), attachments_schema()), ( "output_schema".into(), Value::Object(Map::from_iter([( @@ -1660,7 +2154,7 @@ impl PromptTool { manager, spec: ToolSpec::new( ToolName::new("prompt"), - "Re-prompt the same completed ACP subagent session using a prior subagent value.", + format!("Re-prompt the same completed ACP subagent session using a prior subagent value.{MEDIA_GUIDANCE}"), continuation_schema(), ) .with_output_schema(value_schema()) @@ -1676,7 +2170,7 @@ impl ForkTool { "" }; let description = format!( - "Fork a completed ACP subagent session using native capability support or the isolated Kit fallback, preferably assign the fork a concise role-oriented display name, prompt it, and return the new session value.{usage}" + "Fork a completed ACP subagent session using native capability support or the isolated Kit fallback, preferably assign the fork a concise role-oriented display name, prompt it, and return the new session value.{usage}{MEDIA_GUIDANCE}" ); Self { manager, @@ -1698,6 +2192,7 @@ impl ForkTool { )])), ), ("name".into(), display_name_schema()), + ("attachments".into(), attachments_schema()), ( "output_schema".into(), Value::Object(Map::from_iter([( @@ -1764,6 +2259,8 @@ impl CloseTool { #[derive(Deserialize)] #[serde(deny_unknown_fields)] struct Input { + #[serde(default)] + attachments: Vec, prompt: String, name: Option, harness: Option, @@ -1775,6 +2272,8 @@ struct Input { #[derive(Deserialize)] #[serde(deny_unknown_fields)] struct Continuation { + #[serde(default)] + attachments: Vec, subagent: SubagentValue, prompt: String, #[serde(default, deserialize_with = "deserialize_output_schema")] @@ -1783,6 +2282,8 @@ struct Continuation { #[derive(Deserialize)] #[serde(deny_unknown_fields)] struct ForkInput { + #[serde(default)] + attachments: Vec, subagent: SubagentValue, prompt: String, name: Option, @@ -1902,7 +2403,12 @@ impl Tool for SubagentTool { ) -> Result { let input: Input = serde_json::from_value(request.input.clone()) .map_err(|e| ToolError::InvalidInput(e.to_string()))?; - let contract = input.output_schema.map(OutputContract::new).transpose()?; + let contract = Some(OutputContract::for_request( + input.output_schema, + input.attachments, + &self.manager.config.root, + &request.session_id.0, + )?); result( request, self.manager @@ -1935,7 +2441,12 @@ impl Tool for PromptTool { ) -> Result { let input: Continuation = serde_json::from_value(request.input.clone()) .map_err(|e| ToolError::InvalidInput(e.to_string()))?; - let contract = input.output_schema.map(OutputContract::new).transpose()?; + let contract = Some(OutputContract::for_request( + input.output_schema, + input.attachments, + &self.manager.config.root, + &request.session_id.0, + )?); result( request, self.manager @@ -1962,7 +2473,12 @@ impl Tool for ForkTool { ) -> Result { let input: ForkInput = serde_json::from_value(request.input.clone()) .map_err(|e| ToolError::InvalidInput(e.to_string()))?; - let contract = input.output_schema.map(OutputContract::new).transpose()?; + let contract = Some(OutputContract::for_request( + input.output_schema, + input.attachments, + &self.manager.config.root, + &request.session_id.0, + )?); result( request, self.manager diff --git a/src/tools/subagent/native-image-fixture.py b/src/tools/subagent/native-image-fixture.py new file mode 100644 index 00000000..33d561e5 --- /dev/null +++ b/src/tools/subagent/native-image-fixture.py @@ -0,0 +1,71 @@ +"""Test-only ACP boundary for real image attachment/output contracts.""" +import base64 +import json +import struct +import sys +import zlib + +image, log = sys.argv[1:] +sequence = 0 + +def chunk(kind, payload): + return struct.pack(">I", len(payload)) + kind + payload + struct.pack(">I", zlib.crc32(kind + payload)) + +def alternate_png(width=2, height=3): + header = struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0) + return base64.b64encode(b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", header) + chunk(b"IDAT", zlib.compress((b"\x00" + b"\xff" * (width * 3)) * height)) + chunk(b"IEND", b"")).decode() + + +def send(value): + print(json.dumps(dict(jsonrpc="2.0", **value)), flush=True) + +for line in sys.stdin: + request = json.loads(line) + method = request.get("method") + params = request.get("params", {}) + result = {} + if method == "initialize": + result = {"protocolVersion": 1, "agentCapabilities": { + "promptCapabilities": {"image": True}, + "sessionCapabilities": {"fork": {}, "close": {}}}} + elif method in ("session/new", "session/fork"): + sequence += 1 + result = {"sessionId": "native-" + str(sequence)} + elif method == "session/prompt": + with open(log, "a") as output: + output.write(json.dumps(params) + "\n") + text = params["prompt"][0]["text"].split("\n")[0] + count = 0 if text == "zero" else 9 if text == "overcount" else 4 if text == "overpixels" else 3 if text == "duplicate-distinct" else 2 if text in ("multiple", "multiple-root", "duplicate", "duplicate-root", "invalid-mime", "invalid-png", "same-pixels") else 1 + for index in range(count): + payload = "invalid!" if text == "corrupt" else image + mime = "image/png" + if index == 1: + if text in ("multiple", "multiple-root"): + payload = alternate_png() + elif text == "invalid-mime": + mime = "image/jpeg" + elif text == "invalid-png": + payload = base64.b64encode(b"not a PNG").decode() + elif text == "same-pixels": + original = base64.b64decode(image) + payload = base64.b64encode(original[:-12] + chunk(b"tEXt", b"Comment\x00alternate encoding") + original[-12:]).decode() + if text == "duplicate-distinct" and index == 2: + payload = alternate_png() + if text == "overpixels" and index > 0: + payload = alternate_png(4096, 4096) + with open(log + ".images", "a") as emitted: + emitted.write(json.dumps({"data": payload, "mime": mime}) + "\n") + send({"method": "session/update", "params": { + "sessionId": params["sessionId"], "update": { + "sessionUpdate": "agent_message_chunk", "content": { + "type": "image", "data": payload, + "mimeType": mime}}}}) + surrounding = '{"result": null}' if text == "attempt" else "" if text in ("root", "zero", "duplicate-root", "multiple-root") else '{"caption":"native"}' + if surrounding: + send({"method": "session/update", "params": { + "sessionId": params["sessionId"], "update": { + "sessionUpdate": "agent_message_chunk", + "content": {"type": "text", "text": surrounding}}}}) + result = {"stopReason": "end_turn"} + if "id" in request: + send({"id": request["id"], "result": result}) diff --git a/src/tools/subagent/native_images.rs b/src/tools/subagent/native_images.rs new file mode 100644 index 00000000..772d23f7 --- /dev/null +++ b/src/tools/subagent/native_images.rs @@ -0,0 +1,751 @@ +use super::*; + +fn png() -> Vec { + let mut bytes = std::io::Cursor::new(Vec::new()); + image::DynamicImage::new_rgb8(2, 3) + .write_to(&mut bytes, image::ImageFormat::Png) + .unwrap(); + bytes.into_inner() +} + +fn file_schema() -> Value { + json!({"$ref": FILE_SCHEMA_REF}) +} +fn nested_schema() -> Value { + json!({"type":"object", "properties":{"result":file_schema(), "caption":{"type":"string"}}, "required":["result", "caption"], "additionalProperties":false}) +} +fn real_file(root: &Path) -> Value { + serde_json::to_value( + FileStore::new(root) + .import_bytes("parent", "test.png", "image/png", &png(), None) + .unwrap(), + ) + .unwrap() +} + +#[test] +fn strict_file_schema_supports_root_and_fixed_required_paths() { + let root = tempfile::tempdir().unwrap(); + let file = real_file(root.path()); + let contract = OutputContract::new(file_schema()).unwrap(); + assert_eq!(contract.bind("", file.clone()).unwrap(), file); + assert!(contract.bind("{}", file.clone()).is_err()); + let contract = OutputContract::new(nested_schema()).unwrap(); + assert_eq!( + contract.bind(r#"{"caption":"ok"}"#, file.clone()).unwrap()["result"], + file + ); + for text in ["", "not json", "[]", r#"{"result":null,"caption":"ok"}"#] { + assert!(contract.bind(text, file.clone()).is_err(), "{text}"); + } + let contract = OutputContract::new(json!({"type":"object", "properties":{"outer":{"type":"object", "properties":{"image":file_schema()}, "required":["image"]}}, "required":["outer"]})).unwrap(); + assert_eq!( + contract.bind("", file.clone()).unwrap(), + json!({"outer":{"image":file}}) + ); +} + +#[test] +fn strict_file_schema_rejects_ambiguous_optional_and_indirect_bindings() { + for schema in [ + json!({"type":"array", "items":file_schema()}), + json!({"anyOf":[file_schema(), {"type":"string"}]}), + json!({"type":"object", "properties":{"result":file_schema()}}), + json!({"type":"object", "properties":{"a":file_schema(), "b":file_schema()}, "required":["a","b"]}), + json!({"$defs":{"image":file_schema()}, "$ref":"#/$defs/image"}), + json!({"$ref":FILE_SCHEMA_REF, "description":"siblings unsupported"}), + ] { + assert!(OutputContract::new(schema.clone()).is_err(), "{schema}"); + } + // Property names are not schema keywords. + assert!( + OutputContract::new( + json!({"type":"object", "properties":{"items":file_schema()}, "required":["items"]}) + ) + .is_ok() + ); +} + +fn native_manager(root: &Path, log: &Path) -> Subagents { + let mut manager = manager_with_generic_harness(root, Vec::new()); + manager.config.harnesses = + crate::acp_child::AcpHarnesses::new(std::collections::BTreeMap::from([( + "generic".into(), + crate::acp_child::AcpHarnessProfile { + command: "python3".into(), + args: vec![ + format!( + "{}/src/tools/subagent/native-image-fixture.py", + env!("CARGO_MANIFEST_DIR") + ), + BASE64.encode(png()), + log.display().to_string(), + ], + permissions: Default::default(), + }, + )])) + .unwrap(); + manager +} + +#[tokio::test] +async fn native_files_are_attached_bound_and_durably_published_across_roots() { + let parent = tempfile::tempdir().unwrap(); + let child_root = tempfile::tempdir().unwrap(); + let log = parent.path().join("prompts.jsonl"); + let manager = native_manager(parent.path(), &log); + let file = real_file(parent.path()); + let contract = OutputContract::for_request( + Some(nested_schema()), + vec![serde_json::from_value(file.clone()).unwrap(); 2], + parent.path(), + "parent", + ) + .unwrap(); + let handle = manager + .create( + "native".into(), + CreateOptions { + cwd: Some(child_root.path().to_owned()), + ..Default::default() + }, + 0, + TurnCancellation::default(), + Some(&contract), + ) + .await + .unwrap(); + let request: Value = serde_json::from_str( + std::fs::read_to_string(&log) + .unwrap() + .lines() + .next() + .unwrap(), + ) + .unwrap(); + assert_eq!(request["prompt"].as_array().unwrap().len(), 2); + assert_eq!(request["prompt"][1]["type"], "image"); + assert_eq!( + BASE64 + .decode(request["prompt"][1]["data"].as_str().unwrap()) + .unwrap(), + png() + ); + assert_eq!(handle.output["caption"], "native"); + assert!(handle.updates.as_ref().is_none_or(|u| { + !serde_json::to_string(u) + .unwrap() + .contains(&BASE64.encode(png())) + })); + let child_store = FileStore::new(child_root.path()); + assert!( + !child_store + .selected_parts(&handle.id, &file, None) + .unwrap() + .is_empty() + ); + manager + .close(&handle.id, &TurnCancellation::default()) + .await + .unwrap(); + let store = FileStore::new(parent.path()); + assert!( + !store + .selected_parts("parent", &handle.output, None) + .unwrap() + .is_empty() + ); + assert!( + store + .selected_parts("sibling", &handle.output, None) + .is_err() + ); +} + +#[tokio::test] +async fn strict_native_errors_preserve_continuation_handle_ownership() { + let root = tempfile::tempdir().unwrap(); + let manager = native_manager(root.path(), &root.path().join("prompts.jsonl")); + let contract = + OutputContract::for_request(Some(file_schema()), Vec::new(), root.path(), "parent") + .unwrap(); + let handle = manager + .create( + "root".into(), + CreateOptions::default(), + 0, + TurnCancellation::default(), + Some(&contract), + ) + .await + .unwrap(); + for prompt in ["zero", "multiple", "corrupt", "attempt"] { + assert!( + manager + .prompt( + handle.clone(), + prompt.into(), + TurnCancellation::default(), + Some(&contract) + ) + .await + .is_err() + ); + let state = manager.lookup(&handle).unwrap(); + let locked = state.lock().await; + assert_eq!(locked.status, SubagentStatus::Idle); + assert_eq!(locked.handle_generation, handle.generation); + assert_eq!(locked.outcome, Some(GenerationOutcome::Failed)); + } + let next = manager + .prompt( + handle.clone(), + "root".into(), + TurnCancellation::default(), + Some(&contract), + ) + .await + .unwrap(); + assert!(next.generation > handle.generation); + assert!( + manager + .prompt( + handle, + "root".into(), + TurnCancellation::default(), + Some(&contract) + ) + .await + .is_err() + ); + manager + .close(&next.id, &TurnCancellation::default()) + .await + .unwrap(); +} + +#[tokio::test] +async fn native_images_without_schema_have_an_explicit_descriptor_surface() { + let root = tempfile::tempdir().unwrap(); + let manager = native_manager(root.path(), &root.path().join("prompts.jsonl")); + let contract = OutputContract::for_request(None, Vec::new(), root.path(), "parent").unwrap(); + let handle = manager + .create( + "native".into(), + CreateOptions::default(), + 0, + TurnCancellation::default(), + Some(&contract), + ) + .await + .unwrap(); + assert_eq!(handle.output["value"], r#"{"caption":"native"}"#); + assert_eq!(handle.output["files"].as_array().unwrap().len(), 1); + manager + .close(&handle.id, &TurnCancellation::default()) + .await + .unwrap(); +} + +#[tokio::test] +async fn attachment_authority_is_caller_scoped_and_failed_create_releases_capacity() { + let root = tempfile::tempdir().unwrap(); + let log = root.path().join("prompts.jsonl"); + let manager = native_manager(root.path(), &log); + let file = real_file(root.path()); + let contract = OutputContract::for_request( + None, + vec![serde_json::from_value(file).unwrap()], + root.path(), + "stranger", + ) + .unwrap(); + let error = manager + .create( + "root".into(), + CreateOptions::default(), + 0, + TurnCancellation::default(), + Some(&contract), + ) + .await + .unwrap_err(); + assert!(error.to_string().contains("inaccessible")); + assert!(!log.exists()); + wait_for_available_permits(&manager, MAX_LIVE_SUBAGENTS).await; + assert!( + manager + .list(&TurnCancellation::default()) + .await + .unwrap() + .is_empty() + ); +} + +#[tokio::test] +async fn strict_fork_failure_does_not_consume_source_and_success_publishes_branch() { + let root = tempfile::tempdir().unwrap(); + let manager = native_manager(root.path(), &root.path().join("prompts.jsonl")); + let contract = || { + Arc::new( + OutputContract::for_request(Some(file_schema()), Vec::new(), root.path(), "parent") + .unwrap(), + ) + }; + let handle = manager + .create( + "root".into(), + CreateOptions::default(), + 0, + TurnCancellation::default(), + Some(&contract()), + ) + .await + .unwrap(); + assert!( + manager + .fork( + handle.clone(), + "zero".into(), + None, + 0, + TurnCancellation::default(), + Some(contract()) + ) + .await + .is_err() + ); + let branch = manager + .fork( + handle.clone(), + "root".into(), + None, + 0, + TurnCancellation::default(), + Some(contract()), + ) + .await + .unwrap(); + assert_ne!(branch.id, handle.id); + manager + .close(&branch.id, &TurnCancellation::default()) + .await + .unwrap(); + assert!( + !FileStore::new(root.path()) + .selected_parts("parent", &branch.output, None) + .unwrap() + .is_empty() + ); + let next = manager + .prompt( + handle, + "root".into(), + TurnCancellation::default(), + Some(&contract()), + ) + .await + .unwrap(); + manager + .close(&next.id, &TurnCancellation::default()) + .await + .unwrap(); +} + +#[tokio::test] +async fn cancelled_native_continuation_keeps_retry_handle() { + let root = tempfile::tempdir().unwrap(); + let manager = native_manager(root.path(), &root.path().join("prompts.jsonl")); + let contract = + OutputContract::for_request(Some(file_schema()), Vec::new(), root.path(), "parent") + .unwrap(); + let handle = manager + .create( + "root".into(), + CreateOptions::default(), + 0, + TurnCancellation::default(), + Some(&contract), + ) + .await + .unwrap(); + let controller = agentkit_core::CancellationController::new(); + let cancel = controller.handle().checkpoint(); + controller.interrupt(); + assert!( + manager + .prompt(handle.clone(), "root".into(), cancel, Some(&contract)) + .await + .is_err() + ); + let next = manager + .prompt( + handle, + "root".into(), + TurnCancellation::default(), + Some(&contract), + ) + .await + .unwrap(); + manager + .close(&next.id, &TurnCancellation::default()) + .await + .unwrap(); +} + +#[test] +fn all_subagent_input_schemas_accept_optional_typed_attachments() { + let root = tempfile::tempdir().unwrap(); + let file = real_file(root.path()); + let manager = manager_with_generic_harness(root.path(), Vec::new()); + let create = SubagentTool::new(manager.clone(), 0); + let continuation = PromptTool::new(manager.clone()); + let fork = ForkTool::new(manager, 0); + for (schema, mut input) in [ + (&create.spec.input_schema, json!({"prompt":"work"})), + ( + &continuation.spec.input_schema, + json!({"prompt":"work","subagent":{"id":"s","output":null,"generation":1}}), + ), + ( + &fork.spec.input_schema, + json!({"prompt":"work","subagent":{"id":"s","output":null,"generation":1}}), + ), + ] { + let validator = jsonschema::validator_for(schema).unwrap(); + assert!(validator.is_valid(&input)); + input["attachments"] = json!([file]); + assert!(validator.is_valid(&input)); + input["attachments"] = Value::Null; + assert!(!validator.is_valid(&input)); + input["attachments"] = json!(["file_fabricated"]); + assert!(!validator.is_valid(&input)); + input["attachments"] = Value::Array(vec![file.clone(); 9]); + assert!(!validator.is_valid(&input)); + } +} + +#[test] +fn legacy_schema_depth_and_width_are_not_subject_to_file_binding_limits() { + let mut schema = json!({"type":"string"}); + let mut value = json!("leaf"); + for _ in 0..40 { + schema = json!({"type":"object", "properties":{"child":schema}, "required":["child"]}); + value = json!({"child":value}); + } + let contract = OutputContract::new(schema).unwrap(); + assert!(contract.binding.is_none()); + assert_eq!( + contract.parse(&serde_json::to_string(&value).unwrap()), + Some(value) + ); + + let properties = (0..10_001) + .map(|index| (format!("field{index}"), json!({"type":"string"}))) + .collect::>(); + let contract = OutputContract::new(json!({"type":"object", "properties":properties})).unwrap(); + assert!(contract.binding.is_none()); + assert_eq!(contract.parse("{}"), Some(json!({}))); +} + +#[test] +fn literal_and_annotation_file_refs_do_not_activate_strict_binding() { + for keyword in ["const", "default", "examples", "enum", "x-extension"] { + let value = file_schema(); + let literal = if matches!(keyword, "examples" | "enum") { + json!([value]) + } else { + value.clone() + }; + let contract = OutputContract::new(json!({keyword:literal})).unwrap(); + assert!(contract.binding.is_none(), "{keyword}"); + assert_eq!( + contract.parse(&serde_json::to_string(&value).unwrap()), + Some(value) + ); + } +} + +#[test] +fn deep_file_binding_still_has_a_strict_limit_even_after_deep_legacy_prefixes() { + let mut schema = file_schema(); + for _ in 0..65 { + schema = json!({"type":"object", "properties":{"child":schema}, "required":["child"]}); + } + assert!( + file_binding(&mut schema) + .unwrap_err() + .contains("traversal limits") + ); + // A literal annotation is ignored even alongside the real binding. + let mut schema = json!({"type":"object", "properties":{"image":file_schema()}, "required":["image"], "examples":[{"$ref":FILE_SCHEMA_REF}]}); + assert_eq!( + file_binding(&mut schema).unwrap(), + Some(vec!["image".into()]) + ); +} + +#[tokio::test] +async fn repeated_attachment_metadata_is_checked_before_deduplicating_grants() { + let root = tempfile::tempdir().unwrap(); + let log = root.path().join("prompts.jsonl"); + let manager = native_manager(root.path(), &log); + let original = real_file(root.path()); + let mut forged = original.clone(); + forged["name"] = json!("different.png"); + let contract = OutputContract::for_request( + None, + vec![ + serde_json::from_value(original).unwrap(), + serde_json::from_value(forged).unwrap(), + ], + root.path(), + "parent", + ) + .unwrap(); + assert!( + manager + .create( + "root".into(), + CreateOptions::default(), + 0, + TurnCancellation::default(), + Some(&contract) + ) + .await + .is_err() + ); + assert!(!log.exists()); + wait_for_available_permits(&manager, MAX_LIVE_SUBAGENTS).await; +} + +#[tokio::test] +async fn byte_identical_native_duplicates_bind_once_and_preserve_surrounding_text() { + let root = tempfile::tempdir().unwrap(); + let manager = native_manager(root.path(), &root.path().join("prompts.jsonl")); + for schema in [Some(nested_schema()), None] { + let contract = + OutputContract::for_request(schema.clone(), Vec::new(), root.path(), "parent").unwrap(); + let handle = manager + .create( + "duplicate".into(), + CreateOptions::default(), + 0, + TurnCancellation::default(), + Some(&contract), + ) + .await + .unwrap(); + if schema.is_some() { + assert_eq!(handle.output["caption"], "native"); + assert_eq!(handle.output["result"]["$kit"], "file"); + } else { + assert_eq!(handle.output["value"], r#"{"caption":"native"}"#); + assert_eq!(handle.output["files"].as_array().unwrap().len(), 1); + } + manager + .close(&handle.id, &TurnCancellation::default()) + .await + .unwrap(); + let parts = FileStore::new(root.path()) + .selected_parts("parent", &handle.output, None) + .unwrap(); + assert_eq!( + parts + .iter() + .filter(|p| matches!(p, agentkit_core::Part::Media(_))) + .count(), + 1 + ); + } +} + +#[tokio::test] +async fn every_native_occurrence_is_validated_and_byte_distinct_outputs_are_ambiguous() { + let root = tempfile::tempdir().unwrap(); + let manager = native_manager(root.path(), &root.path().join("prompts.jsonl")); + let contract = + OutputContract::for_request(Some(nested_schema()), Vec::new(), root.path(), "parent") + .unwrap(); + for (prompt, expected) in [ + ("invalid-mime", "MIME"), + ("invalid-png", "image format"), + ("overcount", "budget"), + ("same-pixels", "exactly one distinct"), + ("multiple", "exactly one distinct"), + ] { + let error = manager + .create( + prompt.into(), + CreateOptions::default(), + 0, + TurnCancellation::default(), + Some(&contract), + ) + .await + .unwrap_err(); + assert!(error.to_string().contains(expected), "{prompt}: {error}"); + wait_for_available_permits(&manager, MAX_LIVE_SUBAGENTS).await; + } +} + +#[test] +fn repeated_native_occurrences_do_not_evade_count_or_pixel_budgets() { + let image = ImageContent::new(BASE64.encode(png()), "image/png"); + let error = distinct_native_images(&vec![image; 9], &TurnCancellation::default()).unwrap_err(); + assert!(error.to_string().contains("occurrence budget")); + let mut bytes = std::io::Cursor::new(Vec::new()); + image::DynamicImage::new_rgb8(4096, 4096) + .write_to(&mut bytes, image::ImageFormat::Png) + .unwrap(); + let image = ImageContent::new(BASE64.encode(bytes.into_inner()), "image/png"); + let error = distinct_native_images(&vec![image; 3], &TurnCancellation::default()).unwrap_err(); + assert!(error.to_string().contains("aggregate pixel budget")); +} + +fn indexed_schema(index: Value, nested: bool) -> Value { + let file = json!({"$ref":FILE_SCHEMA_REF, "x-kit-image-index":index}); + if nested { + json!({"type":"object", "properties":{"result":file,"caption":{"type":"string"}}, "required":["result","caption"], "additionalProperties":false}) + } else { + file + } +} + +#[test] +fn image_index_is_only_an_integer_annotation_on_the_exact_binding() { + for index in [ + json!(true), + json!("0"), + json!(-1), + json!(0.5), + json!(8), + Value::Null, + ] { + assert!( + OutputContract::new(indexed_schema(index.clone(), false)).is_err(), + "{index}" + ); + } + for index in [0, 7] { + for nested in [false, true] { + let contract = OutputContract::new(indexed_schema(json!(index), nested)).unwrap(); + assert_eq!(contract.image_index, Some(index)); + } + } + for schema in [ + json!({"type":"object", "x-kit-image-index":0}), + json!({"$ref":FILE_SCHEMA_REF,"x-kit-image-index":0,"description":"not allowed"}), + json!({"type":"object","properties":{"a":indexed_schema(json!(0),false),"b":file_schema()},"required":["a","b"]}), + json!({"type":"array","items":indexed_schema(json!(0),false)}), + ] { + assert!(OutputContract::new(schema.clone()).is_err(), "{schema}"); + } +} + +#[tokio::test] +async fn explicit_indices_promote_only_selected_exact_bytes_in_distinct_emission_order() { + for (prompt, index, nested, occurrence) in [ + ("multiple", 0, true, 0), + ("multiple", 1, true, 1), + ("duplicate-distinct", 1, true, 2), + ("multiple-root", 1, false, 1), + ] { + let root = tempfile::tempdir().unwrap(); + let log = root.path().join("prompts.jsonl"); + let parent_session = session::new_id(); + let manager = native_manager(root.path(), &log); + let contract = OutputContract::for_request( + Some(indexed_schema(json!(index), nested)), + Vec::new(), + root.path(), + &parent_session, + ) + .unwrap(); + let handle = manager + .create( + prompt.into(), + CreateOptions::default(), + 0, + TurnCancellation::default(), + Some(&contract), + ) + .await + .unwrap(); + let emitted: Vec = std::fs::read_to_string(log.with_extension("jsonl.images")) + .unwrap() + .lines() + .map(|line| serde_json::from_str(line).unwrap()) + .collect(); + let expected = BASE64 + .decode(emitted[occurrence]["data"].as_str().unwrap()) + .unwrap(); + let file: FileReference = serde_json::from_value(if nested { + handle.output["result"].clone() + } else { + handle.output.clone() + }) + .unwrap(); + let store = FileStore::new(root.path()); + assert_eq!(store.resolve(&parent_session, &file).unwrap(), expected); + if nested { + assert_eq!(handle.output["caption"], "native"); + } + for session in [parent_session.as_str(), handle.id.as_str()] { + let directory = crate::artifacts::base(root.path()) + .with_file_name("files") + .join(blake3::hash(session.as_bytes()).to_hex().as_str()); + // Durable output objects, not internal execution instrumentation: + // neither child nor parent receives any unselected snapshot. + assert_eq!(std::fs::read_dir(directory).unwrap().count(), 1); + } + assert!( + handle + .updates + .as_ref() + .is_none_or(|updates| !serde_json::to_string(updates).unwrap().contains("file_")) + ); + manager + .close(&handle.id, &TurnCancellation::default()) + .await + .unwrap(); + assert_eq!(store.resolve(&parent_session, &file).unwrap(), expected); + } +} + +#[tokio::test] +async fn indexed_selection_never_bypasses_unselected_validation_or_missing_index_errors() { + for (prompt, index, nested, expected) in [ + ("root", 1, false, "out of range"), + ("duplicate-root", 1, false, "out of range"), + ("invalid-mime", 0, true, "MIME"), + ("invalid-png", 0, true, "image format"), + ("overcount", 0, true, "budget"), + ("overpixels", 0, true, "pixel budget"), + ("attempt", 0, true, "binding field"), + ("multiple", 0, false, "root File binding"), + ("multiple-root", 0, true, "does not match output_schema"), + ] { + let root = tempfile::tempdir().unwrap(); + let manager = native_manager(root.path(), &root.path().join("prompts.jsonl")); + let contract = OutputContract::for_request( + Some(indexed_schema(json!(index), nested)), + Vec::new(), + root.path(), + "parent", + ) + .unwrap(); + let error = manager + .create( + prompt.into(), + CreateOptions::default(), + 0, + TurnCancellation::default(), + Some(&contract), + ) + .await + .unwrap_err(); + assert!(error.to_string().contains(expected), "{prompt}: {error}"); + wait_for_available_permits(&manager, MAX_LIVE_SUBAGENTS).await; + } +} diff --git a/src/tools/subagent/tests.rs b/src/tools/subagent/tests.rs index 31d30583..db2e405c 100644 --- a/src/tools/subagent/tests.rs +++ b/src/tools/subagent/tests.rs @@ -1896,3 +1896,6 @@ async fn generic_harness_without_native_fork_returns_unsupported() { "ACP harness \"acp.generic\" does not advertise session/fork; transcript fallback is only available for Kit" ); } + +#[path = "native_images.rs"] +mod native_images; diff --git a/tests/runtime.rs b/tests/runtime.rs index 7a978e70..10b52bec 100644 --- a/tests/runtime.rs +++ b/tests/runtime.rs @@ -378,14 +378,6 @@ return { output: child.output, updates: child.updates }"#, "output": "rich done", "updates": { "items": [ - { - "sessionUpdate": "agent_message_chunk", - "content": { - "type": "image", - "data": "aGVsbG8=", - "mimeType": "image/png" - } - }, { "sessionUpdate": "tool_call", "toolCallId": "call-1", From d84baa24c80aee482f49ad445a05b6da45009bd5 Mon Sep 17 00:00:00 2001 From: daniel Date: Tue, 8 Sep 2026 15:56:45 +0100 Subject: [PATCH 2/2] fix: preserve safe imports and multimodal continuations --- docs/user/compose-and-local-tools.md | 6 +- src/managed_files.rs | 95 +++++++++------ src/managed_files/tests.rs | 71 +++++++++-- src/managed_files/tests/faults.rs | 83 +++++++------ src/provider/adapter.rs | 29 ++++- src/provider/adapter_native_tests.rs | 170 +++++++++++++++++++++++++++ 6 files changed, 375 insertions(+), 79 deletions(-) diff --git a/docs/user/compose-and-local-tools.md b/docs/user/compose-and-local-tools.md index a0b3a5ff..ff8b5503 100644 --- a/docs/user/compose-and-local-tools.md +++ b/docs/user/compose-and-local-tools.md @@ -161,7 +161,7 @@ The canonical tool result retains typed images. Supported terminal graphics rend File descriptors reserve `"$kit": "file"` and use schema version 1. They contain an opaque ID, a bounded display name, MIME type, encoded byte count, and image dimensions. Unknown fields, unknown versions, altered metadata, missing objects, and inaccessible IDs fail rather than appearing as successful text-only image delivery. Do not edit descriptors or invent IDs. -Kit stores immutable snapshots under `~/.kit/files//` (or `/.kit/files` when HOME is unavailable). A descriptor is returned only after the binary object and directory entries cross the disk durability barrier. Storage failure returns no usable descriptor. Each versioned binary envelope has a bounded metadata header and a digest-verified payload; truncated or corrupted objects are rejected. +Kit stores immutable snapshots under `~/.kit/files//` (or `/.kit/files` when HOME is unavailable). A descriptor is returned only after the binary object and directory entries cross the disk durability barrier. Storage failure returns no usable descriptor. Each versioned binary envelope has a bounded metadata header and a digest-verified payload; truncated or corrupted objects are rejected. Imports first complete and sync a private staging object outside session authority, then publish it with an atomic no-replace rename and sync both directories. A failure before publication cannot leave a partial object in the inherited set. A failure after rename may leave a complete object even though no descriptor was returned; Kit does not classify it as garbage. References survive process restart and source modification or deletion. Authorization comes from the calling session, not from possession of a marker or an arbitrary filesystem path. Copying a descriptor to a fork or another session does **not** grant access. Cross-session grants are not part of this reader. @@ -171,6 +171,8 @@ There is no automatic managed-file garbage collection in this phase. Calls, canc The hidden `subagent`, `prompt`, and `fork` tools accept optional `attachments: FileReference[]` (at most eight). Pass managed File values explicitly; a file ID, path, URL, or descriptor pasted into prompt text does not attach an image. Kit resolves attachments in the invoking session, checks the existing aggregate image budgets, grants durable copies to the child, and sends native ACP image blocks after the text prompt. The harness must advertise ACP image **input** support. That advertisement does not promise image generation. +Native Kit forks inherit the source session's complete managed-file authority, including valid version-1 objects written by older Kit versions. An empty, truncated, or otherwise malformed historical object blocks inheritance: Kit cannot distinguish an interrupted old import from damage to a previously published object, and transcript absence is not evidence that removal is safe. The source remains unchanged and destination authority is not committed. Restore the object, or investigate and explicitly remove it only if losing its references is acceptable. Kit never automatically skips or classifies these objects as garbage. + Use the exact local schema reference `{"$ref":"kit://schemas/file/v1"}` for a strict native-image output contract: ```text @@ -193,7 +195,7 @@ This example uses the built-in Kit ACP harness and a concrete OpenRouter image-o For the canonical OpenRouter endpoint, selecting a concrete model advertised with image output opts into generation and its additional provider cost. Kit checks the exact model catalogue and its concrete endpoint document, derives output modalities from advertised capabilities, and requires tools support so compose remains available. Automatic routing entries without concrete endpoints retain ordinary behavior. Native requests require providers to support all requested parameters; Kit does not silently remove compose or relax routing to make an incompatible model work. Editing also requires advertised image input. Custom endpoints do not inherit official OpenRouter capability assertions. Discovery failures do not manufacture support, and a required File contract fails if no native image arrives. -This route uses bounded nonstreaming chat completions: at most 24 MiB of raw response, eight images, 8 MiB per image, 16 MiB of decoded image bytes and 32 megapixels in aggregate. Each image also passes the managed-file PNG/JPEG, animation, dimension, pixel and allocation checks. Only inline image bytes are accepted; Kit never fetches provider-generated HTTP or file URLs. Native media does not inject synthetic image-label text into the structured output. Generation has a 300-second attempt timeout and 310-second logical budget, with no automatic retries of ambiguous billable failures. Cancellation remains available. On continuation or replay, historical assistant images stay typed in canonical history and are projected into supported image-input blocks only in the outgoing provider request, after any complete parallel tool-result batch. The existing text-oriented behavior and tool-image user-message fallback remain unchanged for other routes. +This route uses bounded nonstreaming chat completions: at most 24 MiB of raw response, eight images, 8 MiB per image, 16 MiB of decoded image bytes and 32 megapixels in aggregate. Each image also passes the managed-file PNG/JPEG, animation, dimension, pixel and allocation checks. Only inline image bytes are accepted; Kit never fetches provider-generated HTTP or file URLs. Native media does not inject synthetic image-label text into the structured output. Generation has a 300-second attempt timeout and 310-second logical budget, with no automatic retries of ambiguous billable failures. Cancellation remains available. On continuation or replay, historical assistant images stay typed in canonical history and are projected into supported image-input blocks only in the outgoing provider request, after any complete parallel tool-result batch. Delivery quotas apply separately to each assistant item, not cumulatively across history. Before base64 encoding on the Completions fallback route, a separate 64 MiB budget bounds retained assistant-image payloads. This is not a whole-request, user-attachment, or model-context limit. An oversized history fails explicitly; compact history or start a fresh session with selected attachments. The existing text-oriented behavior and tool-image user-message fallback remain unchanged for other routes. A file-aware schema supports **exactly one** File location: the root, or one fixed object-property path whose properties are required at every level. Arrays, unions, conditional binding, indirect references, multiple locations, and sibling keywords on the File `$ref` other than the optional `x-kit-image-index` annotation are rejected. Kit resolves the File schema locally. By default, the child must emit exactly one **distinct** native assistant image and return only the surrounding JSON fields, omitting the binding field. Kit independently validates every occurrence, including its declared MIME type, actual PNG/JPEG bytes, and pixels. Repeated occurrences collapse to one output only when both MIME type and actual bytes match exactly within this turn. Count, encoded/decoded byte, and aggregate pixel budgets count every occurrence before deduplication. Without an explicit selection index, different image bytes remain ambiguous even if they render identically; model IDs and File references do not determine identity. There is no cross-turn or input/output deduplication. Kit rejects model-written binding fields, including `null` placeholders and invented File IDs. A root binding requires empty text. Empty surrounding text is allowed only when Kit can construct the required object path and the resulting complete value validates; Kit does not invent other required fields or apply schema defaults. diff --git a/src/managed_files.rs b/src/managed_files.rs index 9faf1d57..788826b9 100644 --- a/src/managed_files.rs +++ b/src/managed_files.rs @@ -318,20 +318,30 @@ impl FileStore { let entry = entry.map_err(display)?; let name = entry.file_name(); let name = name.to_str().ok_or("invalid managed object filename")?; - let mut file = fs::open_beneath(&directory, Path::new(name)).map_err(display)?; - let mut prefix = [0_u8; 12]; - file.read_exact(&mut prefix).map_err(display)?; - let length = u32::from_le_bytes(prefix[8..12].try_into().map_err(display)?) as usize; - if &prefix[..8] != MAGIC || length > MAX_HEADER_BYTES { - return Err("invalid inherited managed file envelope".into()); - } - let mut header = vec![0; length]; - file.read_exact(&mut header).map_err(display)?; - let header: Header = serde_json::from_slice(&header).map_err(display)?; - if header.file.id != name { - return Err("inherited file ID does not match its storage name".into()); - } - self.grant_to(source, &header.file, self, &staging_session, None)?; + let inherit = || -> Result<()> { + let mut file = fs::open_beneath(&directory, Path::new(name)).map_err(display)?; + let mut prefix = [0_u8; 12]; + file.read_exact(&mut prefix).map_err(display)?; + let length = + u32::from_le_bytes(prefix[8..12].try_into().map_err(display)?) as usize; + if &prefix[..8] != MAGIC || length > MAX_HEADER_BYTES { + return Err("invalid inherited managed file envelope".into()); + } + let mut header = vec![0; length]; + file.read_exact(&mut header).map_err(display)?; + let header: Header = serde_json::from_slice(&header).map_err(display)?; + if header.file.id != name { + return Err("inherited file ID does not match its storage name".into()); + } + self.grant_to(source, &header.file, self, &staging_session, None)?; + Ok(()) + }; + inherit().map_err(|error| { + format!( + "cannot inherit managed object {}: {error}. Source unchanged; destination authority uncommitted. Restore the object, or investigate and explicitly remove it only if loss is acceptable", + directory.join(name).display() + ) + })?; } // The preflight is not authority: even an empty destination created // concurrently must survive unchanged at the atomic commit boundary. @@ -393,30 +403,49 @@ impl FileStore { } fs::create_private_dir_all(&directory).map_err(display)?; let destination = directory.join(&reference.id); - // No descriptor is published until all durability barriers succeed. - // Partial and cancelled imports may leave unreachable objects, never a - // reference that claims an in-memory-only snapshot survived a restart. + // Stage outside every session authority directory: inheritance enumerates + // all entries, including objects whose descriptor was never returned. + let mut random = [0_u8; 32]; + getrandom::fill(&mut random).map_err(display)?; + let staged = self.base.join(format!( + ".import-{}", + blake3::Hash::from_bytes(random).to_hex() + )); let mut output = fs::OpenOptions::new() .write(true) .create_new(true) .private(true) - .open(&destination) - .map_err(display)?; - output.write_all(MAGIC).map_err(display)?; - output - .write_all(&(header.len() as u32).to_le_bytes()) + .open(&staged) .map_err(display)?; - output.write_all(&header).map_err(display)?; - output.write_all(bytes).map_err(display)?; - output.sync_all().map_err(display)?; - fs::require_disk(&destination).map_err(display)?; - // Include newly created ancestor entries, not just the object contents. - for ancestor in durability_directories { - fs::sync_directory(ancestor).map_err(display)?; - fs::require_disk(ancestor).map_err(display)?; - } - check_cancelled(cancellation)?; - Ok(reference) + let result = (|| { + output.write_all(MAGIC).map_err(display)?; + output + .write_all(&(header.len() as u32).to_le_bytes()) + .map_err(display)?; + output.write_all(&header).map_err(display)?; + output.write_all(bytes).map_err(display)?; + output.sync_all().map_err(display)?; + fs::require_disk(&staged).map_err(display)?; + check_cancelled(cancellation)?; + // Reuse the grant publication primitive; never overwrite a prior + // object. A failure after this boundary may leave a complete object. + rename_no_replace(&staged, &destination).map_err(display)?; + fs::sync_directory(&self.base).map_err(display)?; + fs::require_disk(&self.base).map_err(display)?; + // Include newly created ancestor entries, not just object contents. + for ancestor in durability_directories { + fs::sync_directory(ancestor).map_err(display)?; + fs::require_disk(ancestor).map_err(display)?; + } + fs::require_disk(&destination).map_err(display)?; + check_cancelled(cancellation)?; + Ok(reference) + })(); + drop(output); + // Only our create-new staging file is eligible for cleanup. Never remove + // a destination, including when exclusive publication finds a collision. + let _ = fs::remove_file(&staged); + result } pub(crate) fn resolve(&self, session: &str, selected: &FileReference) -> Result> { diff --git a/src/managed_files/tests.rs b/src/managed_files/tests.rs index 619f400d..50e37079 100644 --- a/src/managed_files/tests.rs +++ b/src/managed_files/tests.rs @@ -698,19 +698,76 @@ fn cancelled_grants_do_not_publish_authority() { } #[test] -fn failed_fork_does_not_publish_a_partial_authorized_set() { +fn legacy_v1_envelopes_resolve_and_fork_without_rewriting_source() { + let f = Fixture::new(); + let source = f.source("legacy.png", ImageFormat::Png, 3, 2); + let payload = disk::read(source).unwrap(); + // Construct the historical envelope directly, independently of the writer. + let reference: FileReference = serde_json::from_value(json!({ + "$kit": "file", "version": 1, "id": format!("file_{}", "a".repeat(64)), + "name": "legacy.png", "mime_type": "image/png", "size_bytes": payload.len(), + "image": {"width": 3, "height": 2} + })) + .unwrap(); + let header = serde_json::to_vec(&json!({ + "file": reference, "digest": blake3::hash(&payload).to_hex().to_string() + })) + .unwrap(); + let mut envelope = b"KITFILE1".to_vec(); + envelope.extend_from_slice(&(header.len() as u32).to_le_bytes()); + envelope.extend_from_slice(&header); + envelope.extend_from_slice(&payload); + disk::create_dir_all(f.store.session_directory("session")).unwrap(); + disk::write(f.object(&reference), &envelope).unwrap(); + assert_eq!(f.store.resolve("session", &reference).unwrap(), payload); + f.store + .prepare_inheritance("session", "fork") + .unwrap() + .unwrap() + .commit(); + assert_eq!(f.store.resolve("fork", &reference).unwrap(), payload); + assert_eq!(disk::read(f.object(&reference)).unwrap(), envelope); +} + +#[test] +fn malformed_legacy_fork_fails_closed_without_source_mutation() { + for malformed in [b"".as_slice(), b"KITFILE1", b"KITFILE1\x10\x00\x00\x00{}"] { + let f = Fixture::new(); + let reference = f.import("good.png"); + let valid = disk::read(f.object(&reference)).unwrap(); + let directory = f.store.session_directory("session"); + let path = directory.join(format!("file_{}", "0".repeat(64))); + disk::write(&path, malformed).unwrap(); + let error = f + .store + .prepare_inheritance("session", "failed-fork") + .err() + .unwrap(); + assert!(error.contains("Source unchanged; destination authority uncommitted")); + assert!(error.contains("Restore the object")); + assert!(error.contains("only if loss is acceptable")); + assert!(error.contains(path.to_str().unwrap())); + assert!(!f.store.session_directory("failed-fork").exists()); + assert!(f.store.resolve("failed-fork", &reference).is_err()); + assert_eq!(disk::read(&path).unwrap(), malformed); + assert_eq!(disk::read(f.object(&reference)).unwrap(), valid); + assert_eq!(disk::read_dir(&directory).unwrap().count(), 2); + } +} + +#[test] +fn snapshot_publication_collision_preserves_existing_object() { let f = Fixture::new(); let reference = f.import("good.png"); - let directory = f.store.session_directory("session"); - disk::write(directory.join(format!("file_{}", "0".repeat(64))), b"bad").unwrap(); + let original = disk::read(f.object(&reference)).unwrap(); + let bytes = f.store.resolve("session", &reference).unwrap(); assert!( f.store - .prepare_inheritance("session", "failed-fork") + .write_snapshot("session", reference.clone(), &bytes, None) .is_err() ); - assert!(!f.store.session_directory("failed-fork").exists()); - assert!(f.store.resolve("failed-fork", &reference).is_err()); - assert!(f.store.resolve("session", &reference).is_ok()); + assert_eq!(disk::read(f.object(&reference)).unwrap(), original); + assert_eq!(disk::read_dir(&f.store.base).unwrap().count(), 1); } #[test] diff --git a/src/managed_files/tests/faults.rs b/src/managed_files/tests/faults.rs index df5deafa..6691e386 100644 --- a/src/managed_files/tests/faults.rs +++ b/src/managed_files/tests/faults.rs @@ -94,8 +94,8 @@ impl BackendFile for FaultFile { impl FaultBackend { fn wrap(&self, path: &Path, disk: Box) -> Box { // Fs persists objects through sibling atomic-replacement files, not by - // writing the final file_ basename. Scope faults to this session's file - // contents so staged writes are covered without faulting directory work. + // writing the final file_ basename. Scope faults to the store so private + // import staging is covered without faulting directory work. if path.starts_with(&self.object_directory) { Box::new(FaultFile { disk, @@ -201,6 +201,18 @@ fn write_sync_and_pending_recovery_fail_without_publishing_a_descriptor() { serde_json::from_slice(&disk::read(path.with_extension("result.json")).unwrap()) .unwrap(); assert_eq!(error, mode.expected_error()); + // With the faulting process gone, future real imports and inheritance + // work without investigating or deleting anything from source authority. + let reference = f.import("recovered.png"); + f.store + .prepare_inheritance("session", "recovered-fork") + .unwrap() + .unwrap() + .commit(); + assert_eq!( + f.store.resolve("recovered-fork", &reference).unwrap(), + f.store.resolve("session", &reference).unwrap() + ); } } @@ -216,48 +228,53 @@ fn fault_child() { assert!( fs::initialize_global(Fs::new(Arc::new(FaultBackend { mode: manifest.mode, - object_directory: store.session_directory("session"), + object_directory: store.base.clone(), }))) .is_ok(), "filesystem must be initialized only in this child" ); - // Err is the publication contract: unreachable partial objects may remain, - // but the caller must never receive a FileReference after any barrier fails. + // Every injected failure remains an error, including ENOSPC memory fallback. + // Failed envelopes must never enter the enumerable inherited authority set. let error = store.import("session", &manifest.source, None).unwrap_err(); assert_eq!(error, manifest.mode.expected_error()); let directory = store.session_directory("session"); - let object = disk::read_dir(&directory) - .unwrap() - .map(|entry| entry.unwrap().path()) - .find(|path| { - path.file_name() + assert_eq!(disk::read_dir(&directory).unwrap().count(), 0); + assert_eq!(fs::read_dir(&directory).unwrap().count(), 0); + match manifest.mode { + Mode::Write | Mode::Sync => { + assert!(disk::read_dir(&store.base).unwrap().all(|entry| { + !entry + .unwrap() + .file_name() + .to_string_lossy() + .starts_with(".import-") + })); + store + .prepare_inheritance("session", "fork-after-failure") .unwrap() - .to_str() .unwrap() - .starts_with("file_") - }) - .expect("native managed object was created before the fault"); - match manifest.mode { - Mode::Write => assert!(disk::read(&object).unwrap().is_empty()), - Mode::Sync => { - // The native payload was written before the injected sync error. - // Fs then abandons the unpublished replacement, retaining only the - // previously durable empty object, not unsynced payload contents. - assert!(disk::read(&object).unwrap().is_empty()); + .commit(); } Mode::NoSpace => { - // Show that writes were accepted as a complete memory-backed envelope - // while native storage cannot retain it. This is real Fs behavior, - // not a callback asserting which implementation method was called. - let accepted = fs::read(&object).unwrap(); - assert!(accepted.starts_with(MAGIC)); - let header_len = u32::from_le_bytes(accepted[8..12].try_into().unwrap()) as usize; - let header: Header = serde_json::from_slice(&accepted[12..12 + header_len]).unwrap(); - let payload = &accepted[12 + header_len..]; - assert_eq!(payload, disk::read(&manifest.source).unwrap()); - assert_eq!(payload.len() as u64, header.file.size_bytes); - assert_eq!(blake3::hash(payload).to_hex().as_str(), header.digest); - assert!(disk::metadata(&object).unwrap().len() < accepted.len() as u64); + // Recovery can also block best-effort cleanup, but only private + // staging can retain the memory-backed envelope, never authority. + let object = disk::read_dir(&store.base) + .unwrap() + .map(|entry| entry.unwrap().path()) + .find(|path| { + path.file_name() + .unwrap() + .to_string_lossy() + .starts_with(".import-") + }) + .unwrap(); + // Cleanup hides the volatile object immediately, but its queued + // unlink cannot reach disk while the earlier write still has ENOSPC. + assert_eq!( + fs::read(&object).unwrap_err().kind(), + io::ErrorKind::NotFound + ); + assert!(disk::read(&object).unwrap().is_empty()); assert!(fs::global().status().pending_operations > 0); assert_eq!( fs::require_disk(&object).unwrap_err().raw_os_error(), diff --git a/src/provider/adapter.rs b/src/provider/adapter.rs index 4946de7a..2d867bd9 100644 --- a/src/provider/adapter.rs +++ b/src/provider/adapter.rs @@ -858,6 +858,11 @@ pub struct SpeakeasyKitSession { context_window: Option, } +// Bound retained assistant-image payloads before the Completions encoder expands +// them to base64. This is not a limit on the whole request, user attachments, +// or model context; per-delivery validation remains separate. +const MAX_OUTGOING_ASSISTANT_IMAGE_BYTES: usize = 64 * 1024 * 1024; + /// Project only the outbound request; the caller's canonical transcript stays typed. /// Completions (including OpenRouter) stringify tool Parts and reject assistant /// Media, but encode ordinary user Media as image_url content. Native transports retain typed tool images, @@ -871,6 +876,7 @@ pub(super) fn project_tool_output_images( const MAX_NODES: usize = 100_000; const MAX_DEPTH: usize = 64; let mut visited = 0; + let mut retained_assistant_image_bytes = 0_usize; let mut pending = Vec::new(); for item in &request.transcript { visited += 1; @@ -889,6 +895,20 @@ pub(super) fn project_tool_output_images( if visited > MAX_NODES { return Err(tool_image_traversal_error()); } + if !native + && item.kind == agentkit_core::ItemKind::Assistant + && let Part::Media(media) = part + && media.modality == Modality::Image + && let DataRef::InlineBytes(bytes) = &media.data + { + retained_assistant_image_bytes = + retained_assistant_image_bytes.saturating_add(bytes.len()); + if retained_assistant_image_bytes > MAX_OUTGOING_ASSISTANT_IMAGE_BYTES { + return Err(LoopError::InvalidState( + "selected-images-not-delivered: outgoing request/history budget of 64 MiB for retained assistant-image payloads exceeded; compact history or start a fresh session with selected attachments".into(), + )); + } + } if let Part::ToolResult(result) = part && let ToolOutput::Parts(parts) = &result.output { @@ -903,9 +923,6 @@ pub(super) fn project_tool_output_images( let mut transcript = Vec::with_capacity(request.transcript.len()); let mut outstanding = std::collections::HashSet::new(); let mut images = Vec::new(); - let mut assistant_image_bytes = 0; - let mut assistant_image_count = 0; - let mut assistant_image_pixels = 0; for mut item in request.transcript { // The loop has already answered detached calls with placeholders. Its // completion notification contains serialized ToolResultPart values, @@ -933,6 +950,10 @@ pub(super) fn project_tool_output_images( // first so images wait for the complete parallel tool-result batch. let mut lifted_assistant_images = false; if !native && item.kind == agentkit_core::ItemKind::Assistant { + // These are delivery quotas, not cumulative transcript quotas. + let mut assistant_image_bytes = 0; + let mut assistant_image_count = 0; + let mut assistant_image_pixels = 0; let mut supported = Vec::new(); for part in std::mem::take(&mut item.parts) { if let Part::Media(media) = &part @@ -1285,7 +1306,7 @@ fn parse_openrouter_model(value: &Value, model: &str) -> Option = wire["messages"] + .as_array() + .unwrap() + .iter() + .filter_map(|message| message["content"].as_array()) + .flatten() + .filter(|part| part["type"] == "image_url") + .collect(); + assert_eq!(images.len(), 9); + assert!(images.iter().all(|part| part["image_url"]["url"] == uri)); + for _ in ["resume", "fork"] { + let restored: TurnRequest = + serde_json::from_value(serde_json::to_value(&history).unwrap()).unwrap(); + let (mut reconstructed, mut sent) = session(vec![body.clone()], true).await; + let mut turn = reconstructed.begin_turn(restored, None).await.unwrap(); + assert_eq!(sent.try_recv().unwrap(), wire); + completed_items(&mut turn).await; + } + assert_eq!(history, canonical); +} + +#[tokio::test] +async fn historical_image_limits_fail_before_http_and_leave_canonical_history_unchanged() { + let valid = Part::Media(MediaPart::new( + Modality::Image, + "image/png", + DataRef::InlineBytes(png()), + )); + let malformed = Part::Media(MediaPart::new( + Modality::Image, + "image/png", + DataRef::InlineBytes(vec![0; 8]), + )); + let oversized = Part::Media(MediaPart::new( + Modality::Image, + "image/png", + DataRef::InlineBytes(vec![0; MAX_NATIVE_IMAGE_BYTES + 1]), + )); + for (parts, expected) in [ + (vec![valid; 9], "historical assistant images exceed"), + (vec![malformed], "invalid historical assistant image"), + (vec![oversized], "historical assistant image exceeds 8 MiB"), + ] { + let mut history = request(false); + history + .transcript + .push(Item::new(ItemKind::Assistant, parts)); + let canonical = history.clone(); + let (mut live, mut sent) = session(vec![response(json!([]))], true).await; + let Err(error) = live.begin_turn(history.clone(), None).await else { + panic!("invalid history accepted") + }; + assert!(error.to_string().contains(expected), "{error}"); + assert!(matches!( + sent.try_recv(), + Err(tokio::sync::mpsc::error::TryRecvError::Empty) + )); + assert_eq!(history, canonical); + } + // Each item independently fits delivery limits; only retained history is too large. + // PNG permits trailing bytes, keeping the fixture's decoded allocation tiny. + let mut bytes = png(); + bytes.resize(MAX_NATIVE_IMAGE_BYTES, 0); + crate::managed_files::validate_provider_image(&bytes).unwrap(); + let mut history = request(false); + for _ in 0..9 { + history.transcript.push(Item::new( + ItemKind::Assistant, + vec![Part::Media(MediaPart::new( + Modality::Image, + "image/png", + DataRef::InlineBytes(bytes.clone()), + ))], + )); + } + let (mut live, mut sent) = session(vec![response(json!([]))], true).await; + let Err(error) = live.begin_turn(history.clone(), None).await else { + panic!("overbudget history accepted") + }; + let error = error.to_string(); + assert!( + error.contains("outgoing request/history budget of 64 MiB"), + "{error}" + ); + assert!(error.contains("compact history or start a fresh session with selected attachments")); + assert!(matches!( + sent.try_recv(), + Err(tokio::sync::mpsc::error::TryRecvError::Empty) + )); + assert_eq!(history.transcript.len(), 10); + for item in &history.transcript[1..] { + assert_eq!( + item.parts, + vec![Part::Media(MediaPart::new( + Modality::Image, + "image/png", + DataRef::InlineBytes(bytes.clone()) + ))] + ); + } +}