diff --git a/docs/plans/compose-managed-files.md b/docs/plans/compose-managed-files.md new file mode 100644 index 0000000..a9fe800 --- /dev/null +++ b/docs/plans/compose-managed-files.md @@ -0,0 +1,132 @@ +# Compose-managed files and multimodal tools + +## Goal and non-negotiable constraints + +Keep **compose as the only model-exposed tool**. Build composable file references for image reads, image transformations, multimodal subagents, and consistent TUI rendering. Do not add a directly exposed read-image tool. + +Implement as four stacked PRs, each owned by a distinct subagent in a distinct worktree created using `wt`. Phase 1 carries this plan. Start each later phase only after its predecessor passes the review gate. Never merge, enable auto-merge, enqueue, deploy, or change release versions. + +Base: current main (`origin/main` at setup: `0a9291d7d3ccdd6837903d0521ab2b040ab54044`). Phase-1 worktree was created using `wt switch --create feat/compose-files-phase-1 --base main --no-hooks --format json`, then fast-forwarded to fetched origin/main. The user's original dirty worktree must remain untouched. + +## Architecture + +### Managed File value + +Tools exchange bounded ordinary JSON descriptors, not base64 strings or raw filesystem paths. Illustrative descriptor (final schema requires design review): + +```json +{ + "$kit": "file", + "version": 1, + "id": "file_opaque_id", + "name": "my_image.png", + "mime_type": "image/png", + "size_bytes": 184230, + "image": { "width": 1024, "height": 768 } +} +``` + +Bytes are immutable snapshots in Kit-managed storage. Descriptors survive session resume independently of the original path. Resolution requires session access; metadata is validated against the stored object. A marker is not authorization. Forged IDs, cross-session references, and traversal must fail safely. Future fork/subagent grants must be explicit. Use the existing artifact/resilient filesystem infrastructure where appropriate, but do not stretch the UTF-8 artifact reader into a binary contract. Define version handling, retention, cleanup, restart behavior, and storage-failure semantics. + +### Compose delivery boundary + +```text +read_file / image transform / subagent -> managed bytes + JSON File descriptor + -> Runlet ordinary JSON -> final returned JSON + -> Kit reference validation and resolution + -> ToolOutput::Parts(text/structured result + typed image parts) + -> canonical tool result -> provider adapter -> model + -> ACP -> TUI +``` + +Only references reachable from the final returned value deliver content to the parent model. Intermediate images remain private to the operations consuming them. Resolve nested references with bounded depth/count, deterministic order, identity deduplication, and labels identifying return-value positions. Invalid/inaccessible/unsupported selected references must never appear as successful text-only image delivery. + +Integrate resolution with compose result finalization in `BackgroundableCompose`, sharing behavior across invoke, invoke_outcome, foreground, and background completion. JSON markers should permit the first implementation without changing Runlet or agentkit-tool-compose. Re-check current main before relying on earlier traced line numbers. + +Extract references before text spilling. Apply the 8 KiB spill policy only to the text/JSON portion; preserve selected typed images and enforce separate image budgets. Define interruption, retry/replay, cancellation, and ownership lifetime. + +### Provider contract + +The canonical transcript retains typed tool-result images. **Phase 1 MUST include the non-native fallback**, not defer it to a later phase. Use native multimodal tool output on verified routes; otherwise project text/metadata tool results with a pointer to an immediately following user-role image message through the existing user attachment encoder. Place that message after the full parallel tool-result batch, never between unanswered results. This is provider-request-only: never persist fake user turns, rerun compose, duplicate delivery, or stringify pixels. Preserve original result text, labels, diagnostics, call/result pairing, normalization budgets, background completion, and replay/continuations. Lack of native image tool-output support is not a fatal gate when ordinary image input transport is available; model vision capability remains a provider/model constraint, not an arbitrary Kit allowlist. Verify actual private Responses and Completions/OpenRouter request encoders. + +The acceptance criterion is an actual image block in the next outgoing provider request, not a base64 string or marker. Replay must use snapshotted content after the original file changes or disappears. + +### Hidden tool suite + +Initial read_file imports local content and returns a managed File. Image operations consume references and create new immutable references. Use flat callable names compatible with current Runlet. Export is explicit, not an implicit overwrite. Define format sniffing, bounded read/decode/allocation, cancellation, decoded-pixel/output budgets, animation policy, orientation, and metadata handling. No HTTP/SVG/PDF/OCR support is needed for the initial reader. + +Proposed end-state example (new APIs, not existing syntax contracts): + +```text +image = read_file({ path: "my_image.png" }) +rotated = image_rotate({ image, degrees: 90 }) +cropped = image_crop({ + image: rotated, + aspect_ratio: { width: 1, height: 1 }, + anchor: "center" +}) +stickered = subagent({ + model: "", + prompt: "Add a Hello Kitty sticker to the attached image.", + attachments: [cropped], + output_schema: { + type: "object", + properties: { result: { "$ref": "kit://schemas/file/v1" } }, + required: ["result"], + additionalProperties: false + } +}) +return stickered.output.result +``` + +Explicit attachment arguments, not concatenating a reference into prompt text, determine image input. A model capable of images behind a text-only harness is still unsupported. + +### Multimodal subagents + +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. + +### 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. + +Markdown image nodes such as `![Edited image](kit-file://file_opaque_id)` resolve through the managed file resolver. Local paths follow filesystem permission policy. Remote images require explicit network policy, asynchronous bounded fetch/decode, redirect/address validation, and no implicit credentials. Do not fetch arbitrary model-supplied URLs without the applicable authorization. Parse real image nodes, not regexes scanning code fences. Retain alt text, source links, placeholders, and terminal fallback. Ordinary links remain links unless explicitly previewed. TUI display never automatically attaches pixels to model context. + +## Milestones and PR stack + +| Phase | Branch / PR base | Owner | Scope and completion milestone | +| --- | --- | --- | --- | +| 1 | `feat/compose-files-phase-1` -> `main` | Dedicated phase-1 subagent | Versioned managed storage/reference contract; hidden read_file; compose finalizer; media-aware spill; native provider delivery plus mandatory user-image transport fallback; basic tool-result TUI rendering. `return read_file({ path: "screenshot.png" })` delivers real pixels to model and TUI. | +| 2 | `feat/compose-files-phase-2` -> phase-1 branch | New dedicated phase-2 subagent | Rotate/crop/resize/export, immutable references, bounded transforms, docs/tests. A pipeline works in one compose invocation; only returned files reach parent model. | +| 3 | `feat/compose-files-phase-3` -> phase-2 branch | New dedicated phase-3 subagent | Typed subagent attachments/output, file-aware schema/binding, capability checks, parent/child grants and output promotion. Sticker-editing pipeline is supported end to end. | +| 4 | `feat/compose-files-phase-4` -> phase-3 branch | New dedicated phase-4 subagent | Unified user/tool/assistant/Markdown rendering, resolver policy, live/replay consistency, caching and terminal fallback. All image origins render safely. | + +Each successor worktree is created with `wt` from its predecessor's CURRENT review-clean tip, invoked from the phase-1 worktree; do not fork all phases independently from main. Verify `git merge-base --is-ancestor HEAD`, and set PR base to the parent branch. The plan is inherited through the stack. Changes to a published parent require pausing descendants and asking for restack/history-rewrite approval. + +## Per-phase owner instructions and review gate + +1. Read this plan, current repository instructions, and relevant skills. Reinspect code on current main/parent; the original investigation was on a different dirty worktree. +2. Get independent design consensus before implementation; retain an explicit reviewer session ID in the orchestration ledger, not committed docs. Resolve material disagreements with the same reviewer. +3. Load shared-state skill for lock/shared-state changes, updating-artifact-schema for persistent schema changes, secure-rust-dependency-changes for dependency surface changes, and pr for PR creation. Apply ship-issue/shepherd with stack/no-merge overrides; this is not a Linear task. +4. Implement only your phase, preserving later extensibility. Keep production free of test-only instrumentation. Run smallest useful checks during iteration, then repository-required checks. No release version changes. +5. Self-review for reuse, quality, efficiency; obtain independent current-head review. Open a Conventional Commit PR with correct base, plan link, testing and limitations. +6. KEEP GOING through CI/reviewer findings. Fix valid findings, rerun tests, push follow-up commits, respond/resolve threads, and obtain fresh review evidence for the latest head. Never self-approve or misrepresent absent approval. +7. Stop successfully ONLY when local checks pass, CI passes with no pending/unknown required check, reviewers have approved the current head with no unresolved threads, and GitHub reports conflict-free/mergeable against its current base. A skipped stacked-base workflow is not a pass without repository-approved equivalent current-head evidence. +8. Do not merge, enable auto-merge, enqueue, deploy, or launch a successor yourself. Return the gate evidence to the parent orchestrator. If external approval/auth/infra genuinely blocks progress, report the blocker honestly; parent does not launch next phase. + +The parent launches each successor only after verifying the previous gate. Keep an uncommitted ledger outside tracked plan content with phase, worktree, branch/base/tip, agent/reviewer session IDs, PR URL, CI/check evidence, approval SHA, unresolved threads, mergeability, and blockers. + +## Acceptance test matrix + +- Real provider request contains image input after compose return, natively or in the mandatory request-only user-image fallback after the full tool-result batch. No descriptor/base64 text substituted for pixels. +- Nested returns, multiple references, deduplication, deterministic labels, and intermediate non-delivery. +- Text spill does not hide or corrupt image delivery; bounded traversal and independent byte/pixel budgets. +- Malformed/forged/stale/cross-session references, unsupported versions, missing files, corrupt formats, directories, permissions, oversized/decompression-bomb input. +- Foreground/background completion, interruption/resume, cancellation, replay, storage errors. +- Session resume/fork and child close preserve authorized output independently of original path. +- Transform geometry, format/animation/metadata policies, immutable source, explicit export. +- Subagent capability/contract failures, native output import, ambiguous binding, grant isolation. +- Live and replayed TUI tool/assistant/user images, Markdown policy, async placeholders/fallback, no credentialed arbitrary fetches. + +## Initial status + +Plan persisted before implementation. Phase 1 is next; phases 2–4 are blocked on predecessor review gates. Operational status and agent IDs live in the uncommitted orchestration ledger. diff --git a/docs/user/compose-and-local-tools.md b/docs/user/compose-and-local-tools.md index c68b9a1..8939540 100644 --- a/docs/user/compose-and-local-tools.md +++ b/docs/user/compose-and-local-tools.md @@ -97,6 +97,39 @@ The result contains `content`, `next_offset`, `total_bytes`, and `eof`. Continue An artifact-storage error does not turn an already-completed tool into a failed tool call. Kit returns a bounded preview with `artifact_error` when output cannot be retained. Do not repeat a side-effecting tool merely to obtain its output again. +## Return images with `read_file` + +`read_file` is a hidden compose callable, not another model-exposed tool. It imports a local image into Kit-managed storage and returns a small JSON **File reference**, not a path or a base64 string: + +```text +image = read_file({ path: "screenshot.png" }) +return { screenshot: image } +``` + +Only File references reachable from the final return deliver pixels to the parent model. A reference used only in an intermediate binding does not attach its image. Arrays and nested objects work; repeated references deliver one image, labeled with its first position as an escaped JSON Pointer. Every occurrence must have valid metadata, including duplicates. The whole selection is validated before any image is delivered. + +The initial reader supports **nonanimated PNG and JPEG**. It sniffs content rather than trusting the extension, rejects corrupt images and nonregular files, and reads at most 8 MiB. An image can have at most 8,192 pixels on either axis, 16 megapixels, and a 64 MiB decoder allocation. GIF, WebP, animated PNG, SVG, PDF, URLs, and text files are unsupported. Use `shell` for ordinary text reads. Relative paths resolve from Kit's working directory; absolute paths follow the Kit process's filesystem access, not a new project sandbox. + +Imports preserve the original bytes, metadata, and orientation. Width and height describe the encoded raster. There is no automatic transformation or export, and importing never overwrites the source. + +### Delivery limits and provider support + +A final compose return can select at most 8 distinct images, 16 MiB of encoded image bytes, and 32 megapixels in total. Selection traversal is bounded to 100,000 JSON nodes, depth 64, and 64 reference occurrences, with position labels bounded to 2 KiB each and 4 KiB in total. These limits are separate from the **8 KiB text-output budget**. Large returned JSON spills to a text artifact without hiding the selected image parts or their labels; media bytes do not enter the text artifact. + +Kit retains typed image tool results in the canonical transcript. The verified `openai-subscription:gpt-5.4` route sends native image tool output. Other subscription models (including `gpt-6-astra`) and the OpenRouter/Speakeasy adapters use their ordinary user-image input encoding: the tool result keeps its text and points to a following user-role image message, placed after the complete tool-result batch. This message exists only in the outgoing provider request, not as a synthetic user turn in session history. Replay and provider switching project the retained images again without rerunning compose. This transport fallback does not imply that every model supports vision; select a model with image input support. Provider image limits and normalization still apply. Do not rerun a side-effecting compose program merely because output delivery failed. + +The canonical tool result retains typed images. Supported terminal graphics render them in expanded tool cards using the existing bounded image cache; disabled graphics or decoding failures leave a text fallback. Displaying a tool image does not create a synthetic user message. + +### File identity, durability, and lifetime + +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. + +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. + +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. + ## 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/src/compose_output.rs b/src/compose_output.rs index b57ce77..84390f7 100644 --- a/src/compose_output.rs +++ b/src/compose_output.rs @@ -1,14 +1,106 @@ use std::{path::Path, sync::Arc}; -use agentkit_core::ToolOutput; +use agentkit_core::{Part, ToolOutput, TurnCancellation}; use agentkit_tools_core::ToolError; use serde_json::{Map, Value}; const MAX_MODEL_OUTPUT_BYTES: usize = 8 * 1024; +/// Resolve only the final Runlet JSON value. Both foreground and detached +/// completions use this boundary; resolving a reference never reimports its path. +pub(crate) async fn finalize( + root: &Path, + session: &str, + artifact_directory: &Path, + output: ToolOutput, + cancellation: Option, +) -> Result { + let output = if let ToolOutput::Structured(value) = output { + let store = crate::managed_files::FileStore::new(root); + let session = session.to_owned(); + let (value, selected) = tokio::task::spawn_blocking(move || { + let selected = store.selected_parts(&session, &value, cancellation.as_ref()); + (value, selected) + }) + .await + .map_err(|error| delivery_failed(error.to_string()))?; + let selected = selected.map_err(delivery_failed)?; + if selected.is_empty() { + ToolOutput::Structured(value) + } else { + // Keep each position label adjacent to its image, even when the + // ordinary JSON spills. Reserve label bytes from the text budget. + let label_bytes = selected + .iter() + .filter_map(|part| match part { + Part::Text(text) => Some(text.text.len()), + _ => None, + }) + .sum::(); + if label_bytes > MAX_MODEL_OUTPUT_BYTES / 2 { + return Err(delivery_failed( + "selected image labels exceed the 4 KiB label budget".into(), + )); + } + let text = guard_text( + artifact_directory, + ToolOutput::Structured(value), + MAX_MODEL_OUTPUT_BYTES - label_bytes, + ) + .await?; + let mut parts = match text { + ToolOutput::Structured(value) => vec![Part::structured(value)], + ToolOutput::Text(text) => vec![Part::text(text)], + ToolOutput::Parts(parts) => parts, + ToolOutput::Files(files) => files.into_iter().map(Part::File).collect(), + }; + parts.extend(selected); + return Ok(ToolOutput::Parts(parts)); + } + } else { + output + }; + guard(artifact_directory, output).await +} + +fn delivery_failed(detail: String) -> ToolError { + ToolError::ExecutionFailed(format!( + "compose program completed; selected File delivery failed: {detail}. Side effects may have occurred; do not rerun blindly." + )) +} + pub(crate) async fn guard( artifact_directory: &Path, output: ToolOutput, +) -> Result { + // Never serialize typed media into a spill artifact or count it as text. + let (output, media) = match output { + ToolOutput::Parts(parts) => { + let (media, text): (Vec<_>, Vec<_>) = parts + .into_iter() + .partition(|part| matches!(part, Part::Media(_))); + (ToolOutput::Parts(text), media) + } + output => (output, Vec::new()), + }; + let text = guard_text(artifact_directory, output, MAX_MODEL_OUTPUT_BYTES).await?; + if media.is_empty() { + return Ok(text); + } + let mut parts = match text { + ToolOutput::Text(text) => vec![Part::text(text)], + ToolOutput::Structured(value) => vec![Part::structured(value)], + ToolOutput::Parts(parts) => parts, + ToolOutput::Files(files) => files.into_iter().map(Part::File).collect(), + }; + parts.extend(media); + Ok(ToolOutput::Parts(parts)) +} + +async fn guard_text( + artifact_directory: &Path, + output: ToolOutput, + budget: usize, ) -> Result { let body = match &output { @@ -19,7 +111,7 @@ pub(crate) async fn guard( .map_err(|error| ToolError::Internal(error.to_string()))?, }; let original_bytes = body.len(); - if original_bytes <= MAX_MODEL_OUTPUT_BYTES { + if original_bytes <= budget { return Ok(output); } @@ -46,7 +138,7 @@ pub(crate) async fn guard( "\n...[tool completed; output truncated: {original_bytes} bytes; artifact storage failed]...\n" ) }; - let mut preview_budget = MAX_MODEL_OUTPUT_BYTES; + let mut preview_budget = budget; loop { let preview = preview(&body, &marker, preview_budget); let replacement = Value::Object(Map::from_iter([ @@ -61,11 +153,11 @@ pub(crate) async fn guard( let replacement_bytes = serde_json::to_vec(&replacement) .map_err(|error| ToolError::Internal(error.to_string()))? .len(); - if replacement_bytes <= MAX_MODEL_OUTPUT_BYTES { + if replacement_bytes <= budget { return Ok(ToolOutput::structured(replacement)); } let next_budget = preview_budget - .saturating_mul(MAX_MODEL_OUTPUT_BYTES) + .saturating_mul(budget) .checked_div(replacement_bytes) .unwrap_or(0) .min(preview_budget.saturating_sub(1)) diff --git a/src/lib.rs b/src/lib.rs index 80d5ab5..fe17f4c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -27,6 +27,7 @@ mod fatal; mod file_search; #[path = "resilient_fs/mod.rs"] mod filesystem; +mod managed_files; pub mod plugins; pub(crate) mod process_tree; pub mod protocols; diff --git a/src/managed_files.rs b/src/managed_files.rs new file mode 100644 index 0000000..92cb7a9 --- /dev/null +++ b/src/managed_files.rs @@ -0,0 +1,452 @@ +//! Durable, session-authorized immutable image snapshots. File values are JSON +//! capabilities only within the session that imported them, never filesystem paths. +use std::{ + collections::HashMap, + io::{Cursor, Read as _, Write as _}, + path::{Path, PathBuf}, +}; + +use agentkit_core::{DataRef, Modality, Part, TurnCancellation}; +use image::{DynamicImage, ImageDecoder, ImageFormat, ImageReader, Limits}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use crate::resilient_fs as fs; + +const MAX_FILE_BYTES: u64 = 8 * 1024 * 1024; +const MAX_HEADER_BYTES: usize = 4096; +const MAX_DIMENSION: u32 = 8192; +const MAX_PIXELS: u64 = 16 * 1024 * 1024; +const MAX_DECODE_BYTES: u64 = 64 * 1024 * 1024; +const MAX_DELIVERY_BYTES: u64 = 16 * 1024 * 1024; +const MAX_DELIVERY_PIXELS: u64 = 32 * 1024 * 1024; +const MAX_IMAGES: usize = 8; +pub(crate) const MAX_LABEL_BYTES: usize = 4 * 1024; +const MAX_OCCURRENCES: usize = 64; +const MAX_NODES: usize = 100_000; +const MAX_DEPTH: usize = 64; +const MAGIC: &[u8; 8] = b"KITFILE1"; + +type Result = std::result::Result; + +#[derive(Clone, Debug, Deserialize, PartialEq, Eq, Serialize)] +#[serde(deny_unknown_fields)] +pub(crate) struct FileReference { + #[serde(rename = "$kit")] + marker: String, + version: u32, + id: String, + name: String, + mime_type: String, + size_bytes: u64, + image: ImageDimensions, +} + +#[derive(Clone, Debug, Deserialize, PartialEq, Eq, Serialize)] +#[serde(deny_unknown_fields)] +struct ImageDimensions { + width: u32, + height: u32, +} + +#[derive(Deserialize, Serialize)] +#[serde(deny_unknown_fields)] +struct Header { + file: FileReference, + digest: String, +} + +#[derive(Clone)] +pub(crate) struct FileStore { + base: PathBuf, +} + +impl FileStore { + pub(crate) fn new(root: &Path) -> Self { + Self { + base: crate::artifacts::base(root).with_file_name("files"), + } + } + + fn session_directory(&self, session: &str) -> PathBuf { + self.base + .join(blake3::hash(session.as_bytes()).to_hex().as_str()) + } + + pub(crate) fn import( + &self, + session: &str, + path: &Path, + cancellation: Option<&TurnCancellation>, + ) -> Result { + check_cancelled(cancellation)?; + // Check before opening as well as on the opened descriptor. Nonblocking + // open on Unix prevents a replacement FIFO from blocking between checks. + if !std::fs::metadata(path).map_err(display)?.is_file() { + return Err("read_file requires a regular local file".into()); + } + let mut options = std::fs::OpenOptions::new(); + options.read(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt as _; + options.custom_flags(libc::O_NONBLOCK); + } + let file = options.open(path).map_err(display)?; + let metadata = file.metadata().map_err(display)?; + if !metadata.is_file() || metadata.len() > MAX_FILE_BYTES { + return Err("read_file requires a regular file no larger than 8 MiB".into()); + } + let mut bytes = Vec::new(); + file.take(MAX_FILE_BYTES + 1) + .read_to_end(&mut bytes) + .map_err(display)?; + if bytes.is_empty() || bytes.len() as u64 > MAX_FILE_BYTES { + return Err("read_file image must contain 1 byte to 8 MiB".into()); + } + check_cancelled(cancellation)?; + let (mime_type, image) = inspect_image(&bytes)?; + check_cancelled(cancellation)?; + let name = path + .file_name() + .and_then(|name| name.to_str()) + .filter(|name| valid_name(name)) + .ok_or( + "file name must be nonempty UTF-8, at most 255 bytes, without control characters", + )? + .to_owned(); + let mut random = [0_u8; 32]; + getrandom::fill(&mut random).map_err(display)?; + let id = format!("file_{}", blake3::Hash::from_bytes(random).to_hex()); + let reference = FileReference { + marker: "file".into(), + version: 1, + id, + name, + mime_type, + size_bytes: bytes.len() as u64, + image, + }; + let header = serde_json::to_vec(&Header { + file: reference.clone(), + digest: blake3::hash(&bytes).to_hex().to_string(), + }) + .map_err(display)?; + if header.len() > MAX_HEADER_BYTES { + return Err("managed file header is too large".into()); + } + let directory = self.session_directory(session); + let mut durability_directories = Vec::new(); + for ancestor in directory.ancestors() { + durability_directories.push(ancestor); + if fs::try_exists(ancestor).map_err(display)? { + break; + } + } + 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. + 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()) + .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) + } + + 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 + // resilient filesystem's volatile write-back layer. + fs::require_disk(directory.join(&selected.id)).map_err(display)?; + let mut file = fs::open_beneath(&directory, Path::new(&selected.id)).map_err(|error| { + format!("managed file is missing or inaccessible in this session: {error}") + })?; + let length = file.metadata().map_err(display)?.len(); + if length > MAX_FILE_BYTES + MAX_HEADER_BYTES as u64 + 12 { + return Err("managed file exceeds its storage budget".into()); + } + let mut prefix = [0_u8; 12]; + file.read_exact(&mut prefix).map_err(display)?; + if &prefix[..8] != MAGIC { + return Err("unsupported managed file envelope".into()); + } + let header_length = + u32::from_le_bytes([prefix[8], prefix[9], prefix[10], prefix[11]]) as usize; + if header_length == 0 || header_length > MAX_HEADER_BYTES { + return Err("invalid managed file header length".into()); + } + let mut header = vec![0; header_length]; + file.read_exact(&mut header).map_err(display)?; + let header: Header = serde_json::from_slice(&header).map_err(display)?; + header.file.validate()?; + if &header.file != selected { + return Err("selected file metadata does not match its stored object".into()); + } + if length != 12 + header_length as u64 + selected.size_bytes { + return Err("managed file envelope length does not match its payload".into()); + } + let mut bytes = Vec::new(); + file.take(MAX_FILE_BYTES + 1) + .read_to_end(&mut bytes) + .map_err(display)?; + if bytes.len() as u64 != selected.size_bytes + || blake3::hash(&bytes).to_hex().as_str() != header.digest + { + return Err("managed file payload is truncated or corrupt".into()); + } + // Digest and metadata bind the already validated immutable import. No + // repeated pixel decode is necessary for each selection or replay. + Ok(bytes) + } + + pub(crate) fn selected_parts( + &self, + session: &str, + value: &Value, + cancellation: Option<&TurnCancellation>, + ) -> Result> { + let mut selection = Selection::default(); + selection.walk(value, String::new(), 0)?; + let mut validated: HashMap = HashMap::new(); + let mut selected = Vec::new(); + let mut label_bytes = 0_usize; + let mut bytes = 0_u64; + let mut pixels = 0_u64; + for (pointer, reference) in selection.files { + check_cancelled(cancellation)?; + reference.validate()?; + if let Some(previous) = validated.get(&reference.id) { + if previous != &reference { + return Err("duplicate file reference has contradictory metadata".into()); + } + continue; + } + bytes += reference.size_bytes; + pixels += reference.pixels(); + if validated.len() >= MAX_IMAGES + || bytes > MAX_DELIVERY_BYTES + || pixels > MAX_DELIVERY_PIXELS + { + return Err( + "selected images exceed delivery budget (8 images, 16 MiB, 32 megapixels)" + .into(), + ); + } + let label = format!( + "Image #{} at JSON Pointer {}: {}", + validated.len() + 1, + serde_json::to_string(&pointer).map_err(display)?, + reference.name + ); + label_bytes += label.len(); + if label_bytes > MAX_LABEL_BYTES { + return Err("selected image labels exceed the 4 KiB label budget".into()); + } + validated.insert(reference.id.clone(), reference.clone()); + selected.push((label, reference)); + } + // Preflight all occurrence metadata and descriptor/label budgets before + // reading any payload. Nothing is emitted until every object resolves. + let mut parts = Vec::new(); + for (label, reference) in selected { + check_cancelled(cancellation)?; + let data = self.resolve(session, &reference)?; + parts.push(Part::text(label)); + parts.push(Part::media( + Modality::Image, + reference.mime_type, + DataRef::InlineBytes(data), + )); + } + check_cancelled(cancellation)?; + Ok(parts) + } +} + +impl FileReference { + 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 + .as_object() + .ok_or("File reference must be an object")?; + if object.len() != 7 + || object + .get("id") + .and_then(Value::as_str) + .is_none_or(|id| id.len() != 69) + || !object + .get("name") + .and_then(Value::as_str) + .is_some_and(valid_name) + || object + .get("mime_type") + .and_then(Value::as_str) + .is_none_or(|mime| mime.len() > 10) + { + return Err("invalid managed File reference shape or string budget".into()); + } + let reference = Self::deserialize(value).map_err(display)?; + reference.validate()?; + Ok(reference) + } + + fn pixels(&self) -> u64 { + u64::from(self.image.width) * u64::from(self.image.height) + } + + fn validate(&self) -> Result<()> { + let suffix = self + .id + .strip_prefix("file_") + .ok_or("invalid managed file ID")?; + if self.marker != "file" || self.version != 1 { + return Err("unsupported managed File reference version".into()); + } + if suffix.len() != 64 + || !suffix + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)) + || !valid_name(&self.name) + || !matches!(self.mime_type.as_str(), "image/png" | "image/jpeg") + || self.size_bytes == 0 + || self.size_bytes > MAX_FILE_BYTES + || self.image.width == 0 + || self.image.height == 0 + || self.image.width > MAX_DIMENSION + || self.image.height > MAX_DIMENSION + || self.pixels() > MAX_PIXELS + { + return Err("invalid managed File reference metadata".into()); + } + Ok(()) + } +} + +#[derive(Default)] +struct Selection { + nodes: usize, + files: Vec<(String, FileReference)>, +} + +impl Selection { + fn walk(&mut self, value: &Value, pointer: String, depth: usize) -> Result<()> { + self.nodes += 1; + if self.nodes > MAX_NODES || depth > MAX_DEPTH || pointer.len() > 2048 { + return Err("compose return exceeds file-selection traversal budget".into()); + } + match value { + Value::Object(object) if object.get("$kit").and_then(Value::as_str) == Some("file") => { + if self.files.len() >= MAX_OCCURRENCES { + return Err("compose return exceeds 64 File reference occurrences".into()); + } + self.files + .push((pointer, FileReference::from_value(value)?)); + } + Value::Object(object) => { + // Check before allocating a sorted traversal frontier. + if object.len() > MAX_NODES - self.nodes { + return Err("compose return exceeds file-selection traversal budget".into()); + } + let mut entries = object.iter().collect::>(); + entries.sort_unstable_by(|(a, _), (b, _)| a.cmp(b)); + for (key, child) in entries { + if key.len() > MAX_HEADER_BYTES { + return Err("compose return key exceeds file-selection label budget".into()); + } + self.walk( + child, + format!("{pointer}/{}", key.replace('~', "~0").replace('/', "~1")), + depth + 1, + )?; + } + } + Value::Array(array) => { + if array.len() > MAX_NODES - self.nodes { + return Err("compose return exceeds file-selection traversal budget".into()); + } + for (index, child) in array.iter().enumerate() { + self.walk(child, format!("{pointer}/{index}"), depth + 1)?; + } + } + _ => {} + } + Ok(()) + } +} + +fn inspect_image(bytes: &[u8]) -> Result<(String, ImageDimensions)> { + let format = image::guess_format(bytes).map_err(display)?; + if !matches!(format, ImageFormat::Png | ImageFormat::Jpeg) { + return Err("read_file supports only nonanimated PNG and JPEG images".into()); + } + let mut limits = Limits::default(); + limits.max_image_width = Some(MAX_DIMENSION); + limits.max_image_height = Some(MAX_DIMENSION); + limits.max_alloc = Some(MAX_DECODE_BYTES); + let decoder: Box = if format == ImageFormat::Png { + let decoder = image::codecs::png::PngDecoder::with_limits(Cursor::new(bytes), limits) + .map_err(display)?; + if decoder.is_apng().map_err(display)? { + return Err("animated PNG images are not supported by read_file".into()); + } + Box::new(decoder) + } else { + let mut reader = ImageReader::with_format(Cursor::new(bytes), format); + reader.limits(limits); + Box::new(reader.into_decoder().map_err(display)?) + }; + let (width, height) = decoder.dimensions(); + if width == 0 + || height == 0 + || u64::from(width) * u64::from(height) > MAX_PIXELS + || decoder.total_bytes() > MAX_DECODE_BYTES + { + return Err("image exceeds decoded pixel or allocation budget".into()); + } + // Fully decode to reject corrupt payloads before publication, then discard + // pixels. Preserve source bytes, EXIF orientation and metadata unchanged. + DynamicImage::from_decoder(decoder).map_err(display)?; + Ok(( + format.to_mime_type().into(), + ImageDimensions { width, height }, + )) +} + +fn valid_name(name: &str) -> bool { + !name.is_empty() && name.len() <= 255 && !name.chars().any(char::is_control) +} + +fn check_cancelled(cancellation: Option<&TurnCancellation>) -> Result<()> { + if cancellation.is_some_and(TurnCancellation::is_cancelled) { + Err("managed file operation cancelled".into()) + } else { + Ok(()) + } +} + +fn display(error: impl std::fmt::Display) -> String { + error.to_string() +} + +#[cfg(test)] +mod tests; diff --git a/src/managed_files/tests.rs b/src/managed_files/tests.rs new file mode 100644 index 0000000..8394c93 --- /dev/null +++ b/src/managed_files/tests.rs @@ -0,0 +1,558 @@ +#![allow( + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::disallowed_methods, + clippy::disallowed_macros +)] +use super::*; +use serde_json::json; +use std::fs as disk; +use tempfile::TempDir; + +struct Fixture { + dir: TempDir, + store: FileStore, +} +impl Fixture { + fn new() -> Self { + let dir = tempfile::tempdir().unwrap(); + let store = FileStore { + base: dir.path().canonicalize().unwrap().join("store"), + }; + Self { dir, store } + } + fn source(&self, name: &str, format: ImageFormat, width: u32, height: u32) -> PathBuf { + let mut bytes = Cursor::new(Vec::new()); + DynamicImage::new_luma8(width, height) + .write_to(&mut bytes, format) + .unwrap(); + let path = self.dir.path().join(name); + disk::write(&path, bytes.into_inner()).unwrap(); + path + } + fn import(&self, name: &str) -> FileReference { + let path = self.source(name, ImageFormat::Png, 3, 2); + self.store.import("session", &path, None).unwrap() + } + fn object(&self, reference: &FileReference) -> PathBuf { + self.store.session_directory("session").join(&reference.id) + } + fn select(&self, value: &Value) -> Result> { + self.store.selected_parts("session", value, None) + } +} +#[test] +fn png_and_jpeg_snapshots_survive_source_deletion_and_reopen() { + for (format, name, mime) in [ + (ImageFormat::Png, "image.png", "image/png"), + (ImageFormat::Jpeg, "image.jpg", "image/jpeg"), + ] { + let f = Fixture::new(); + let source = f.source(name, format, 3, 2); + let bytes = disk::read(&source).unwrap(); + let reference = f.store.import("session", &source, None).unwrap(); + assert_eq!(reference.mime_type, mime); + assert_eq!( + reference.image, + ImageDimensions { + width: 3, + height: 2 + } + ); + assert_eq!(reference.size_bytes, bytes.len() as u64); + disk::remove_file(source).unwrap(); + let reopened = FileStore { + base: f.store.base.clone(), + }; + let parts = reopened + .selected_parts("session", &json!(reference), None) + .unwrap(); + assert_eq!(parts.len(), 2); + match &parts[1] { + Part::Media(media) => match &media.data { + DataRef::InlineBytes(actual) => assert_eq!(actual, &bytes), + _ => panic!("expected inline snapshot"), + }, + _ => panic!("expected image"), + } + assert!( + reopened + .selected_parts("other-session", &json!(reference), None) + .is_err() + ); + } +} +#[test] +fn rejects_unknown_versions_fields_forged_stale_and_traversal_refs() { + let f = Fixture::new(); + let reference = f.import("image.png"); + let good = json!(reference); + for (key, value) in [ + ("version", json!(2)), + ("unexpected", json!(true)), + ("name", json!("forged.png")), + ("id", json!(format!("file_{}", "0".repeat(64)))), + ("id", json!("../image.png")), + ("id", json!("/tmp/image.png")), + ("mime_type", json!("image/gif")), + ("size_bytes", json!(0)), + ] { + let mut altered = good.clone(); + altered[key] = value; + assert!(f.select(&altered).is_err(), "accepted {altered}"); + } + let mut nested_unknown = good.clone(); + nested_unknown["image"]["unexpected"] = json!(1); + assert!(f.select(&nested_unknown).is_err()); + disk::remove_file(f.object(&reference)).unwrap(); + assert!(f.select(&good).is_err()); +} +#[test] +fn nested_selection_is_sorted_escaped_deduplicated_and_return_scoped() { + let f = Fixture::new(); + let a = f.import("a.png"); + let b = f.import("b.png"); + let parts = f.select(&json!({"z": a, "a/~": [a, b]})).unwrap(); + assert_eq!(parts.len(), 4); + for (index, expected) in [ + (0, "Image #1 at JSON Pointer \"/a~1~0/0\": a.png"), + (2, "Image #2 at JSON Pointer \"/a~1~0/1\": b.png"), + ] { + match &parts[index] { + Part::Text(text) => assert_eq!(text.text, expected), + _ => panic!("expected label"), + } + } + assert!( + f.select(&json!({"done": true, "intermediate": null})) + .unwrap() + .is_empty() + ); + let mut contradictory = json!(a); + contradictory["name"] = json!("different.png"); + assert!( + f.select(&json!([a, contradictory])) + .unwrap_err() + .contains("contradictory") + ); +} +#[test] +fn traversal_depth_label_and_occurrence_budgets() { + let f = Fixture::new(); + let reference = json!(f.import("image.png")); + assert_eq!( + f.select(&Value::Array(vec![reference.clone(); MAX_OCCURRENCES])) + .unwrap() + .len(), + 2 + ); + assert!( + f.select(&Value::Array(vec![reference; MAX_OCCURRENCES + 1])) + .unwrap_err() + .contains("occurrences") + ); + assert!( + f.select(&Value::Array(vec![Value::Null; MAX_NODES - 1])) + .unwrap() + .is_empty() + ); + assert!( + f.select(&Value::Array(vec![Value::Null; MAX_NODES])) + .unwrap_err() + .contains("traversal") + ); + let mut nested = Value::Null; + for _ in 0..MAX_DEPTH { + nested = json!([nested]); + } + assert!(f.select(&nested).unwrap().is_empty()); + assert!( + f.select(&json!([nested])) + .unwrap_err() + .contains("traversal") + ); + assert!( + f.select(&json!({"x".repeat(MAX_HEADER_BYTES + 1): null})) + .unwrap_err() + .contains("label") + ); +} +#[test] +fn aggregate_image_count_and_pixel_budgets() { + let f = Fixture::new(); + let references: Vec<_> = (0..=MAX_IMAGES) + .map(|i| f.import(&format!("{i}.png"))) + .collect(); + assert_eq!( + f.select(&json!(&references[..MAX_IMAGES])).unwrap().len(), + MAX_IMAGES * 2 + ); + assert!( + f.select(&json!(references)) + .unwrap_err() + .contains("delivery budget") + ); + let source = f.source("large.png", ImageFormat::Png, 4096, 4096); + let large: Vec<_> = (0..3) + .map(|_| f.store.import("session", &source, None).unwrap()) + .collect(); + assert_eq!(f.select(&json!(&large[..2])).unwrap().len(), 4); + assert!( + f.select(&json!(large)) + .unwrap_err() + .contains("delivery budget") + ); +} +#[test] +fn aggregate_encoded_byte_budget() { + let f = Fixture::new(); + let source = f.source("padded.png", ImageFormat::Png, 3, 2); + // The PNG decoder accepts trailing bytes; preserved bytes still count toward delivery. + let mut bytes = disk::read(&source).unwrap(); + bytes.resize(MAX_FILE_BYTES as usize, 0); + disk::write(&source, bytes).unwrap(); + let references: Vec<_> = (0..3) + .map(|_| f.store.import("session", &source, None).unwrap()) + .collect(); + assert_eq!(f.select(&json!(&references[..2])).unwrap().len(), 4); + assert!( + f.select(&json!(references)) + .unwrap_err() + .contains("delivery budget") + ); +} +#[test] +fn rejects_malformed_corrupt_truncated_and_trailing_envelopes() { + let f = Fixture::new(); + let reference = f.import("image.png"); + let path = f.object(&reference); + let original = disk::read(&path).unwrap(); + let mut cases = vec![ + vec![], + original[..7].to_vec(), + original[..12].to_vec(), + original[..original.len() - 1].to_vec(), + ]; + let mut magic = original.clone(); + magic[0] ^= 1; + cases.push(magic); + let mut corrupt = original.clone(); + *corrupt.last_mut().unwrap() ^= 1; + cases.push(corrupt); + let mut trailing = original.clone(); + trailing.push(0); + cases.push(trailing); + for length in [0, MAX_HEADER_BYTES as u32 + 1] { + let mut invalid = original.clone(); + invalid[8..12].copy_from_slice(&length.to_le_bytes()); + cases.push(invalid); + } + let mut malformed = original.clone(); + malformed[12] = b'!'; + cases.push(malformed); + let header_len = u32::from_le_bytes(original[8..12].try_into().unwrap()) as usize; + let mut header: Value = serde_json::from_slice(&original[12..12 + header_len]).unwrap(); + header["unknown"] = json!(true); + let encoded = serde_json::to_vec(&header).unwrap(); + let mut unknown = MAGIC.to_vec(); + unknown.extend_from_slice(&(encoded.len() as u32).to_le_bytes()); + unknown.extend(encoded); + unknown.extend_from_slice(&original[12 + header_len..]); + cases.push(unknown); + for (index, bytes) in cases.into_iter().enumerate() { + disk::write(&path, bytes).unwrap(); + assert!( + f.select(&json!(reference)).is_err(), + "accepted damaged envelope {index}" + ); + } +} +#[cfg(unix)] +#[test] +fn rejects_symlink_objects_and_session_directories() { + use std::os::unix::fs::symlink; + let f = Fixture::new(); + let reference = f.import("image.png"); + let path = f.object(&reference); + let outside = f.dir.path().join("outside"); + disk::rename(&path, &outside).unwrap(); + symlink(&outside, &path).unwrap(); + assert!(f.select(&json!(reference)).is_err()); + disk::remove_file(&path).unwrap(); + disk::rename(&outside, &path).unwrap(); + let session = f.store.session_directory("session"); + let moved = f.dir.path().join("moved-session"); + disk::rename(&session, &moved).unwrap(); + symlink(&moved, &session).unwrap(); + assert!(f.select(&json!(reference)).is_err()); +} +#[test] +fn import_rejects_invalid_destination_and_cancelled_work() { + let f = Fixture::new(); + let source = f.source("image.png", ImageFormat::Png, 3, 2); + disk::write(&f.store.base, b"not a directory").unwrap(); + assert!(f.store.import("session", &source, None).is_err()); + let f = Fixture::new(); + let source = f.source("image.png", ImageFormat::Png, 3, 2); + let controller = agentkit_core::CancellationController::new(); + let cancellation = controller.handle().checkpoint(); + controller.interrupt(); + assert!( + f.store + .import("session", &source, Some(&cancellation)) + .unwrap_err() + .contains("cancelled") + ); + assert!(!f.store.base.exists(), "cancelled import published storage"); +} +#[test] +fn import_rejects_directory_nonregular_oversize_corrupt_and_dimensions() { + let f = Fixture::new(); + assert!(f.store.import("session", f.dir.path(), None).is_err()); + #[cfg(unix)] + { + let socket_path = f.dir.path().join("socket"); + let _socket = std::os::unix::net::UnixListener::bind(&socket_path).unwrap(); + assert!(f.store.import("session", &socket_path, None).is_err()); + } + let huge = f.dir.path().join("huge.png"); + disk::File::create(&huge) + .unwrap() + .set_len(MAX_FILE_BYTES + 1) + .unwrap(); + assert!(f.store.import("session", &huge, None).is_err()); + for (name, bytes) in [ + ("empty.png", vec![]), + ("junk.png", b"not an image".to_vec()), + ("truncated.png", vec![137, 80, 78, 71, 13, 10, 26, 10]), + ] { + let path = f.dir.path().join(name); + disk::write(&path, bytes).unwrap(); + assert!(f.store.import("session", &path, None).is_err()); + } + let source = f.source("corrupt.png", ImageFormat::Png, 3, 2); + let mut bytes = disk::read(&source).unwrap(); + bytes.truncate(bytes.len() / 2); + disk::write(&source, bytes).unwrap(); + assert!(f.store.import("session", &source, None).is_err()); + for (width, height) in [(MAX_DIMENSION + 1, 1), (4097, 4096)] { + let source = f.source("dimensions.png", ImageFormat::Png, width, height); + assert!(f.store.import("session", &source, None).is_err()); + } + assert!(!f.store.base.exists()); +} + +// PNG chunks use CRC-32 over their type and data. Keep fixture construction local +// rather than adding a production dependency solely to encode a one-frame APNG. +fn png_chunk(kind: &[u8; 4], data: &[u8]) -> Vec { + let mut chunk = (data.len() as u32).to_be_bytes().to_vec(); + chunk.extend_from_slice(kind); + chunk.extend_from_slice(data); + let mut crc = u32::MAX; + for byte in &chunk[4..] { + crc ^= u32::from(*byte); + for _ in 0..8 { + crc = (crc >> 1) ^ (0xedb8_8320 & 0_u32.wrapping_sub(crc & 1)); + } + } + chunk.extend_from_slice(&(!crc).to_be_bytes()); + chunk +} + +#[test] +fn import_rejects_apng_even_with_one_frame() { + let f = Fixture::new(); + let source = f.source("animated.png", ImageFormat::Png, 3, 2); + let bytes = disk::read(&source).unwrap(); + // First frame uses the existing IDAT payload with matching canvas dimensions. + let mut animation = bytes[..33].to_vec(); // signature and IHDR + animation.extend(png_chunk(b"acTL", &[0, 0, 0, 1, 0, 0, 0, 0])); + let mut frame = Vec::new(); + for value in [0_u32, 3, 2, 0, 0] { + frame.extend_from_slice(&value.to_be_bytes()); + } + frame.extend_from_slice(&[0, 1, 0, 10, 0, 0]); // delay, dispose, blend + animation.extend(png_chunk(b"fcTL", &frame)); + animation.extend_from_slice(&bytes[33..]); + disk::write(&source, animation).unwrap(); + let error = f.store.import("session", &source, None).unwrap_err(); + assert!(error.contains("animated PNG"), "{error}"); + assert!(!f.store.base.exists()); +} + +#[test] +fn import_rejects_decoded_allocation_over_budget() { + let f = Fixture::new(); + // RGBA16 requires 72 MB despite fitting both dimension and pixel limits. + let image = DynamicImage::new_rgba16(3000, 3000); + let mut bytes = Cursor::new(Vec::new()); + image.write_to(&mut bytes, ImageFormat::Png).unwrap(); + drop(image); + assert!(bytes.get_ref().len() as u64 <= MAX_FILE_BYTES); + let source = f.dir.path().join("allocation.png"); + disk::write(&source, bytes.into_inner()).unwrap(); + assert!(f.store.import("session", &source, None).is_err()); + assert!(!f.store.base.exists()); +} + +#[test] +fn aggregate_label_budget_is_checked_before_resolving_missing_payloads() { + let f = Fixture::new(); + let mut selected = serde_json::Map::new(); + for prefix in ['a', 'b', 'c'] { + let reference = f.import(&format!("{prefix}.png")); + let key = prefix.to_string().repeat(1500); + // Every individual path and label is permitted; only their sum exceeds + // 4 KiB. Resolve once to establish genuine same-session references. + let individual = json!({key.clone(): reference}); + assert_eq!(f.select(&individual).unwrap().len(), 2); + disk::remove_file(f.object(&reference)).unwrap(); + assert!(f.select(&individual).unwrap_err().contains("inaccessible")); + selected.insert(key, json!(reference)); + } + let error = f.select(&Value::Object(selected)).unwrap_err(); + assert_eq!(error, "selected image labels exceed the 4 KiB label budget"); +} + +#[test] +fn current_writer_descriptor_has_strict_shape_and_roundtrips() { + let f = Fixture::new(); + let source = f.source("shape.png", ImageFormat::Png, 3, 2); + let reference = f.store.import("session", &source, None).unwrap(); + let encoded = serde_json::to_value(&reference).unwrap(); + assert_eq!( + encoded, + json!({ + "$kit": "file", + "version": 1, + "id": reference.id, + "name": "shape.png", + "mime_type": "image/png", + "size_bytes": disk::metadata(&source).unwrap().len(), + "image": {"width": 3, "height": 2} + }) + ); + let decoded: FileReference = serde_json::from_value(encoded.clone()).unwrap(); + assert_eq!(decoded, reference); + assert_eq!(f.select(&json!(decoded)).unwrap().len(), 2); + for key in encoded.as_object().unwrap().keys() { + let mut missing = encoded.clone(); + missing.as_object_mut().unwrap().remove(key); + assert!( + serde_json::from_value::(missing).is_err(), + "accepted missing {key}" + ); + } + let mut unknown = encoded.clone(); + unknown["path"] = json!(source); + assert!(serde_json::from_value::(unknown).is_err()); + let mut unknown_dimension = encoded; + unknown_dimension["image"]["channels"] = json!(1); + assert!(serde_json::from_value::(unknown_dimension).is_err()); +} + +#[test] +fn import_rejects_control_character_name_and_unsupported_image_format() { + let f = Fixture::new(); + let invalid_name = f.source("invalid\nname.png", ImageFormat::Png, 3, 2); + let error = f.store.import("session", &invalid_name, None).unwrap_err(); + assert!(error.contains("file name"), "{error}"); + let unsupported = f.dir.path().join("unsupported.gif"); + DynamicImage::new_rgba8(3, 2) + .save_with_format(&unsupported, ImageFormat::Gif) + .unwrap(); + let error = f.store.import("session", &unsupported, None).unwrap_err(); + assert!(error.contains("only nonanimated PNG and JPEG"), "{error}"); + assert!(!f.store.base.exists()); +} + +const RESTART_MANIFEST_ENV: &str = "KIT_MANAGED_FILES_RESTART_TEST_MANIFEST"; + +#[derive(Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +struct RestartManifest { + base: PathBuf, + source: PathBuf, + reference: FileReference, + bytes: Vec, + digest: String, + parent_pid: u32, +} + +#[test] +fn snapshots_resolve_in_a_fresh_process_after_source_deletion() { + for (format, name) in [ + (ImageFormat::Png, "restart.png"), + (ImageFormat::Jpeg, "restart.jpg"), + ] { + let f = Fixture::new(); + let source = f.source(name, format, 3, 2); + let bytes = disk::read(&source).unwrap(); + let reference = f.store.import("session", &source, None).unwrap(); + let manifest = RestartManifest { + base: f.store.base.clone(), + source: source.clone(), + reference, + digest: blake3::hash(&bytes).to_hex().to_string(), + bytes, + parent_pid: std::process::id(), + }; + let manifest_path = f.dir.path().join("restart-manifest.json"); + disk::write(&manifest_path, serde_json::to_vec(&manifest).unwrap()).unwrap(); + disk::remove_file(source).unwrap(); + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "managed_files::tests::fresh_process_resolve_child", + "--ignored", + "--nocapture", + ]) + .env(RESTART_MANIFEST_ENV, &manifest_path) + .output() + .unwrap(); + let stdout = String::from_utf8_lossy(&output.stdout); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!( + output.status.success(), + "restart child failed for {name}: {}\n{stdout}\n{stderr}", + output.status + ); + assert!( + stdout.contains("managed_files::tests::fresh_process_resolve_child ... ok"), + "child test did not succeed: {stdout}" + ); + assert!( + stdout.contains("1 passed; 0 failed"), + "child did not run exactly one successful test: {stdout}" + ); + } +} + +#[test] +#[ignore = "invoked only by the fresh-process restart parent with a fixture manifest"] +fn fresh_process_resolve_child() { + let path = std::env::var_os(RESTART_MANIFEST_ENV).expect("restart manifest path"); + let manifest: RestartManifest = serde_json::from_slice(&disk::read(path).unwrap()).unwrap(); + assert_ne!(std::process::id(), manifest.parent_pid); + assert!(!manifest.source.exists(), "original source must be deleted"); + let store = FileStore { + base: manifest.base, + }; + let parts = store + .selected_parts("session", &json!(manifest.reference), None) + .unwrap(); + assert_eq!(parts.len(), 2); + match &parts[1] { + Part::Media(media) => match &media.data { + DataRef::InlineBytes(actual) => { + assert_eq!(actual, &manifest.bytes); + assert_eq!(blake3::hash(actual).to_hex().as_str(), manifest.digest); + } + _ => panic!("expected inline snapshot"), + }, + _ => panic!("expected image"), + } +} + +mod faults; diff --git a/src/managed_files/tests/faults.rs b/src/managed_files/tests/faults.rs new file mode 100644 index 0000000..370d3f5 --- /dev/null +++ b/src/managed_files/tests/faults.rs @@ -0,0 +1,273 @@ +//! Faults at the existing filesystem backend boundary, isolated per process. +use super::*; +use crate::resilient_fs::{ + Backend, BackendFile, BackendLease, DiskBackend, DiskEntry, DiskOpenOptions, FileIdentity, Fs, + LeaseRequest, +}; +use std::{ + io::{self, Read, Seek, SeekFrom, Write}, + sync::Arc, +}; + +const MANIFEST_ENV: &str = "KIT_MANAGED_FILES_FAULT_TEST_MANIFEST"; + +#[derive(Clone, Copy, Debug, Serialize, Deserialize)] +enum Mode { + Write, + Sync, + NoSpace, +} +impl Mode { + fn expected_error(self) -> String { + match self { + Self::Write => "injected managed object write failure".into(), + Self::Sync => "injected managed object sync failure".into(), + Self::NoSpace => io::Error::from_raw_os_error(libc::ENOSPC).to_string(), + } + } +} + +struct FaultBackend { + mode: Mode, + object_directory: PathBuf, +} +struct FaultFile { + disk: Box, + mode: Mode, +} + +// No mutable fault switches or counters: every process owns one fixed policy. +impl Read for FaultFile { + fn read(&mut self, bytes: &mut [u8]) -> io::Result { + self.disk.read(bytes) + } +} +impl Seek for FaultFile { + fn seek(&mut self, from: SeekFrom) -> io::Result { + self.disk.seek(from) + } +} +impl Write for FaultFile { + fn write(&mut self, bytes: &[u8]) -> io::Result { + match self.mode { + Mode::Write => Err(io::Error::other(self.mode.expected_error())), + Mode::NoSpace => Err(io::Error::from_raw_os_error(libc::ENOSPC)), + Mode::Sync => self.disk.write(bytes), + } + } + fn flush(&mut self) -> io::Result<()> { + self.disk.flush() + } +} +impl FaultFile { + fn check_sync(&self) -> io::Result<()> { + // Permit empty-object creation; fail durability only after native data + // exists. Directory syncs and all non-object files remain real. + if matches!(self.mode, Mode::Sync) && self.disk.metadata()?.len() > 0 { + return Err(io::Error::other(self.mode.expected_error())); + } + Ok(()) + } +} +impl BackendFile for FaultFile { + fn metadata(&self) -> io::Result { + self.disk.metadata() + } + fn identity(&self) -> io::Result> { + self.disk.identity() + } + fn set_len(&self, size: u64) -> io::Result<()> { + self.disk.set_len(size) + } + fn sync_data(&self) -> io::Result<()> { + self.check_sync()?; + self.disk.sync_data() + } + fn sync_all(&self) -> io::Result<()> { + self.check_sync()?; + self.disk.sync_all() + } + fn set_permissions(&self, p: disk::Permissions) -> io::Result<()> { + self.disk.set_permissions(p) + } +} +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. + if path.parent() == Some(self.object_directory.as_path()) { + Box::new(FaultFile { + disk, + mode: self.mode, + }) + } else { + disk + } + } +} +impl Backend for FaultBackend { + fn open(&self, path: &Path, options: &DiskOpenOptions) -> io::Result> { + Ok(self.wrap(path, DiskBackend.open(path, options)?)) + } + fn metadata(&self, path: &Path, follow: bool) -> io::Result { + DiskBackend.metadata(path, follow) + } + fn identity(&self, path: &Path, follow: bool) -> io::Result> { + DiskBackend.identity(path, follow) + } + fn read_dir(&self, path: &Path) -> io::Result> { + DiskBackend.read_dir(path) + } + fn read_link(&self, path: &Path) -> io::Result { + DiskBackend.read_link(path) + } + fn canonicalize(&self, path: &Path) -> io::Result { + DiskBackend.canonicalize(path) + } + fn create_dir(&self, path: &Path, private: bool) -> io::Result<()> { + DiskBackend.create_dir(path, private) + } + fn remove_file(&self, path: &Path) -> io::Result<()> { + DiskBackend.remove_file(path) + } + fn remove_dir(&self, path: &Path) -> io::Result<()> { + DiskBackend.remove_dir(path) + } + fn rename(&self, from: &Path, to: &Path) -> io::Result<()> { + DiskBackend.rename(from, to) + } + fn set_permissions(&self, path: &Path, p: disk::Permissions) -> io::Result<()> { + DiskBackend.set_permissions(path, p) + } + fn sync_directory(&self, path: &Path) -> io::Result<()> { + DiskBackend.sync_directory(path) + } + fn acquire_lease(&self, request: &LeaseRequest) -> io::Result> { + DiskBackend.acquire_lease(request) + } + fn open_beneath(&self, root: &Path, relative: &Path) -> io::Result> { + Ok(self.wrap( + &root.join(relative), + DiskBackend.open_beneath(root, relative)?, + )) + } +} + +#[derive(Serialize, Deserialize)] +struct Manifest { + base: PathBuf, + source: PathBuf, + mode: Mode, + parent_pid: u32, +} + +#[test] +fn write_sync_and_pending_recovery_fail_without_publishing_a_descriptor() { + for mode in [Mode::Write, Mode::Sync, Mode::NoSpace] { + let f = Fixture::new(); + let source = f.source("fault.png", ImageFormat::Png, 3, 2); + let manifest = Manifest { + base: f.store.base.clone(), + source, + mode, + parent_pid: std::process::id(), + }; + let path = f.dir.path().join("fault-manifest.json"); + disk::write(&path, serde_json::to_vec(&manifest).unwrap()).unwrap(); + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "managed_files::tests::faults::fault_child", + "--ignored", + "--nocapture", + ]) + .env(MANIFEST_ENV, &path) + .output() + .unwrap(); + let stdout = String::from_utf8_lossy(&output.stdout); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!( + output.status.success(), + "{mode:?} child failed: {}\n{stdout}\n{stderr}", + output.status + ); + assert!( + stdout.contains("managed_files::tests::faults::fault_child ... ok"), + "{stdout}" + ); + assert!(stdout.contains("1 passed; 0 failed"), "{stdout}"); + let error: String = + serde_json::from_slice(&disk::read(path.with_extension("result.json")).unwrap()) + .unwrap(); + assert_eq!(error, mode.expected_error()); + } +} + +#[test] +#[ignore = "invoked by the parent in an isolated process with one immutable fault mode"] +fn fault_child() { + let path = PathBuf::from(std::env::var_os(MANIFEST_ENV).expect("fault manifest path")); + let manifest: Manifest = serde_json::from_slice(&disk::read(&path).unwrap()).unwrap(); + assert_ne!(std::process::id(), manifest.parent_pid); + let store = FileStore { + base: manifest.base, + }; + assert!( + fs::initialize_global(Fs::new(Arc::new(FaultBackend { + mode: manifest.mode, + object_directory: store.session_directory("session"), + }))) + .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. + 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() + .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()); + } + 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); + assert!(fs::global().status().pending_operations > 0); + assert_eq!( + fs::require_disk(&object).unwrap_err().raw_os_error(), + Some(libc::ENOSPC) + ); + } + } + disk::write( + path.with_extension("result.json"), + serde_json::to_vec(&error).unwrap(), + ) + .unwrap(); +} diff --git a/src/provider/adapter.rs b/src/provider/adapter.rs index 4faf0eb..04e6059 100644 --- a/src/provider/adapter.rs +++ b/src/provider/adapter.rs @@ -823,6 +823,249 @@ pub struct SpeakeasyKitSession { context_window: Option, } +/// 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, +/// but detached notification images always need ordinary user attachments. +pub(super) fn project_tool_output_images( + mut request: TurnRequest, + native: bool, +) -> Result { + // Validate before moving or recursively visiting parts. The iterator stack + // bounds both depth and work without allocating a sibling-sized frontier. + const MAX_NODES: usize = 100_000; + const MAX_DEPTH: usize = 64; + let mut visited = 0; + let mut pending = Vec::new(); + for item in &request.transcript { + visited += 1; + if visited > MAX_NODES { + return Err(tool_image_traversal_error()); + } + // Store only one iterator per nesting level, never a transcript-wide + // frontier or one entry per sibling. Bound both traversal and stack size. + pending.push(item.parts.iter()); + while let Some(parts) = pending.last_mut() { + let Some(part) = parts.next() else { + pending.pop(); + continue; + }; + visited += 1; + if visited > MAX_NODES { + return Err(tool_image_traversal_error()); + } + if let Part::ToolResult(result) = part + && let ToolOutput::Parts(parts) = &result.output + { + if pending.len() >= MAX_DEPTH { + return Err(tool_image_traversal_error()); + } + pending.push(parts.iter()); + } + } + } + + let mut transcript = Vec::with_capacity(request.transcript.len()); + let mut outstanding = std::collections::HashSet::new(); + let mut images = Vec::new(); + for mut item in request.transcript { + // The loop has already answered detached calls with placeholders. Its + // completion notification contains serialized ToolResultPart values, + // not another tool answer. Even native transports must lift these images. + if item.kind == agentkit_core::ItemKind::Notification + && matches!(item.parts.first(), Some(Part::Text(text)) if text.text.starts_with("Background tool results: ")) + { + for part in &mut item.parts { + if let Part::Structured(value) = part + && is_detached_result(&value.value) + { + project_detached_images(&mut value.value, &mut images, &mut visited, 1)?; + } + } + } + // Register the whole item before processing any answers. Calls may also + // span multiple assistant items, and results multiple tool items. + for part in &item.parts { + if let Part::ToolCall(call) = part { + outstanding.insert(call.id.clone()); + } + } + 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); + if outstanding.is_empty() && !images.is_empty() { + let mut attachment = agentkit_core::Item::new( + agentkit_core::ItemKind::User, + std::mem::take(&mut images), + ); + attachment + .metadata + .insert("kit.projected_tool_images".into(), Value::Bool(true)); + transcript.push(attachment); + } + } + 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(), + )); + } + request.transcript = transcript; + Ok(request) +} + +// maybe_convert_detached adds no dedicated metadata marker. Match its exact +// result envelope only inside its Background tool results notification, never +// reinterpret arbitrary Structured tool/user output as typed media. +fn is_detached_result(value: &Value) -> bool { + value.as_object().is_some_and(|object| { + object.len() == 4 + && value["call_id"].is_string() + && value["is_error"].is_boolean() + && value["metadata"].is_object() + && value["output"].is_object() + }) +} + +// Walk only the serialized Parts/ToolResult domain, not arbitrary JSON values +// or byte arrays. This keeps large images out of the node budget and never +// deserializes an unbounded recursive ToolResult tree. +fn project_detached_images( + result: &mut Value, + images: &mut Vec, + visited: &mut usize, + depth: usize, +) -> Result<(), LoopError> { + if depth >= 64 { + return Err(tool_image_traversal_error()); + } + let call_id = result["call_id"].as_str().unwrap_or_default().to_owned(); + let Some(parts) = result["output"] + .get_mut("Parts") + .and_then(Value::as_array_mut) + else { + return Ok(()); + }; + let mut previous_text = None; + for part in parts { + *visited += 1; + if *visited > 100_000 { + return Err(tool_image_traversal_error()); + } + if part + .get("Media") + .is_some_and(|media| media["modality"] == "Image") + { + let label = format!( + "Image from background tool call {call_id}: see the user image message immediately after this complete result batch." + ); + let placeholder = + serde_json::to_value(Part::text(&label)).map_err(tool_image_projection_error)?; + let image = std::mem::replace(part, placeholder); + let image: Part = serde_json::from_value(image).map_err(|error| LoopError::InvalidState(format!( + "selected-images-not-delivered: invalid detached image: {error}. Do not retry or rerun the program." + )))?; + images.push(Part::text(label)); + if let Some(text) = previous_text.take() { + images.push(Part::Text(text)); + } + images.push(image); + } else if let Some(nested) = part.get_mut("ToolResult") { + project_detached_images(nested, images, visited, depth + 1)?; + previous_text = None; + } else { + previous_text = part + .get("Text") + .and_then(|text| serde_json::from_value(text.clone()).ok()); + } + } + Ok(()) +} + +// Recursion is safe after the complete request passes the depth/node preflight. +fn project_result_images( + result: &mut agentkit_core::ToolResultPart, + images: &mut Vec, + native: bool, +) -> Result<(), LoopError> { + let ToolOutput::Parts(parts) = &mut result.output else { + return Ok(()); + }; + let mut previous_text: Option<&agentkit_core::TextPart> = None; + for part in parts.iter_mut() { + match part { + Part::Media(media) if !native && media.modality == Modality::Image => { + let label = format!( + "Image from tool call {}: see the user image message immediately after this complete tool-result batch.", + result.call_id.0 + ); + images.push(Part::text(&label)); + if let Some(text) = previous_text.take() { + images.push(Part::Text(text.clone())); + } + images.push(std::mem::replace(part, Part::text(label))); + } + Part::ToolResult(nested) => { + project_result_images(nested, images, native)?; + previous_text = None; + } + Part::Text(text) => previous_text = Some(text), + _ => previous_text = None, + } + } + if !parts.iter().any(|part| matches!(part, Part::ToolResult(_))) { + return Ok(()); + } + let mut flat = Vec::new(); + for part in std::mem::take(parts) { + if let Part::ToolResult(nested) = part { + // Nested calls are content, not new protocol calls. Preserve their + // provenance and diagnostics while exposing supported content parts. + let provenance = Value::Object(serde_json::Map::from_iter([ + ("call_id".into(), Value::String(nested.call_id.0)), + ("is_error".into(), Value::Bool(nested.is_error)), + ( + "metadata".into(), + serde_json::to_value(nested.metadata).map_err(tool_image_projection_error)?, + ), + ])); + flat.push(Part::structured(Value::Object(serde_json::Map::from_iter( + [("nested_tool_result".into(), provenance)], + )))); + match nested.output { + ToolOutput::Parts(parts) => flat.extend(parts), + ToolOutput::Text(text) => flat.push(Part::text(text)), + ToolOutput::Structured(value) => flat.push(Part::structured(value)), + ToolOutput::Files(files) => flat.push(Part::structured(Value::Object( + serde_json::Map::from_iter([( + "files".into(), + serde_json::to_value(files).map_err(tool_image_projection_error)?, + )]), + ))), + } + } else { + flat.push(part); + } + } + *parts = flat; + Ok(()) +} + +fn tool_image_projection_error(error: serde_json::Error) -> LoopError { + LoopError::InvalidState(format!( + "selected-images-not-delivered: image request projection failed: {error}. The program may already have completed; do not retry or rerun the program." + )) +} + +fn tool_image_traversal_error() -> LoopError { + LoopError::InvalidState( + "selected-images-not-delivered: tool-output image validation exceeds its traversal budget. The program may already have completed; do not retry or rerun the program.".into(), + ) +} + #[async_trait] impl ModelSession for KitSession { type Turn = KitTurn; @@ -832,6 +1075,11 @@ impl ModelSession for KitSession { request: TurnRequest, cancellation: Option, ) -> Result { + let request = if matches!(self, Self::OpenAiSubscription(_)) { + request + } else { + project_tool_output_images(request, false)? + }; match self { Self::OpenAiSubscription(session) => session .begin_turn(request, cancellation) @@ -1744,6 +1992,95 @@ mod tests { } } + fn selected_image_request(nested: bool) -> TurnRequest { + let image = Part::media( + Modality::Image, + "image/png", + DataRef::InlineBytes(vec![1, 2, 3]), + ); + let output = if nested { + vec![Part::ToolResult(ToolResultPart::success( + "nested", + ToolOutput::Parts(vec![image]), + ))] + } else { + vec![Part::text("Selected image"), image] + }; + TurnRequest { + session_id: SessionId::new("provider-identity-test"), + turn_id: TurnId::new("replay"), + transcript: vec![ + Item::new( + ItemKind::Tool, + vec![Part::ToolResult(ToolResultPart::success( + "completed-call", + ToolOutput::Parts(output), + ))], + ), + Item::text(ItemKind::User, "Continue after switching providers"), + ], + available_tools: Vec::new(), + cache: None, + metadata: MetadataMap::new(), + } + } + + #[test] + fn selected_image_projection_bounds_nested_and_wide_outputs() { + let mut request = selected_image_request(false); + let mut part = Part::text("deep"); + for _ in 0..65 { + part = Part::ToolResult(ToolResultPart::success( + "nested", + ToolOutput::Parts(vec![part]), + )); + } + request.transcript[0].parts = vec![part]; + let error = super::project_tool_output_images(request.clone(), false).unwrap_err(); + assert!(error.to_string().contains("traversal budget")); + request.transcript[0].parts = vec![Part::ToolResult(ToolResultPart::success( + "wide", + ToolOutput::Parts(vec![Part::text("text"); 100_001]), + ))]; + let error = super::project_tool_output_images(request.clone(), false).unwrap_err(); + assert!(error.to_string().contains("do not retry or rerun")); + } + + #[test] + fn selected_image_projection_preserves_user_images_and_text_tool_outputs() { + let mut request = selected_image_request(false); + request.transcript[0] = Item::new( + ItemKind::User, + vec![Part::media( + Modality::Image, + "image/png", + DataRef::InlineBytes(vec![1, 2, 3]), + )], + ); + for output in [ + ToolOutput::Text("done".into()), + ToolOutput::Structured(json!({"ok": true})), + ToolOutput::Parts(vec![Part::text("done")]), + ] { + request.transcript.push(Item::new( + ItemKind::Tool, + vec![Part::ToolResult(ToolResultPart::success( + "text-call", + output, + ))], + )); + } + let original = serde_json::to_value(&request.transcript).unwrap(); + for native in [false, true] { + let projected = super::project_tool_output_images(request.clone(), native).unwrap(); + assert_eq!( + serde_json::to_value(&projected.transcript).unwrap(), + original + ); + } + assert_eq!(serde_json::to_value(&request.transcript).unwrap(), original); + } + #[tokio::test] async fn kit_session_delegates_initial_provider_identity() { let session = openrouter_session("test/initial").await; @@ -1934,3 +2271,11 @@ mod tests { } } } + +#[cfg(test)] +#[path = "adapter_image_tests.rs"] +mod image_tests; + +#[cfg(test)] +#[path = "adapter_background_image_tests.rs"] +mod background_image_tests; diff --git a/src/provider/adapter_background_image_tests.rs b/src/provider/adapter_background_image_tests.rs new file mode 100644 index 0000000..9950c66 --- /dev/null +++ b/src/provider/adapter_background_image_tests.rs @@ -0,0 +1,201 @@ +#![allow(clippy::disallowed_methods, clippy::disallowed_macros)] +//! Regression through the real loop/task manager, not a fabricated notification. +#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] + +use std::{collections::VecDeque, time::Duration}; + +use agentkit_core::{ + DataRef, FinishReason, Item, ItemKind, MetadataMap, Modality, Part, ToolCallPart, ToolOutput, + ToolResultPart, TurnCancellation, +}; +use agentkit_http::Authentication; +use agentkit_loop::{ + Agent, LoopError, LoopInterrupt, LoopStep, ModelAdapter, ModelSession, ModelTurn, + ModelTurnEvent, ModelTurnResult, SessionConfig, TurnRequest, +}; +use agentkit_provider_openai::OpenAIResponsesConfig; +use agentkit_task_manager::{AsyncTaskManager, RoutingDecision, TaskEvent, TaskManager}; +use agentkit_tools_core::{ + AllowAllPermissions, Tool, ToolAnnotations, ToolContext, ToolError, ToolName, ToolRegistry, + ToolRequest, ToolResult, ToolSpec, +}; +use async_trait::async_trait; +use serde_json::json; +use tokio::sync::{mpsc, watch}; + +struct Model(mpsc::UnboundedSender); +struct Session(mpsc::UnboundedSender); +struct Turn(VecDeque); + +#[async_trait] +impl ModelAdapter for Model { + type Session = Session; + + async fn start_session(&self, _: SessionConfig) -> Result { + Ok(Session(self.0.clone())) + } +} + +#[async_trait] +impl ModelSession for Session { + type Turn = Turn; + + async fn begin_turn( + &mut self, + request: TurnRequest, + _: Option, + ) -> Result { + let answered = request + .transcript + .iter() + .any(|item| item.kind == ItemKind::Tool); + self.0.send(request).unwrap(); + let mut events = VecDeque::new(); + let (finish_reason, output) = if answered { + ( + FinishReason::Completed, + Item::text(ItemKind::Assistant, "done"), + ) + } else { + let call = ToolCallPart::new("background-image", "image", json!({})); + events.push_back(ModelTurnEvent::ToolCall(call.clone())); + ( + FinishReason::ToolCall, + Item::new(ItemKind::Assistant, vec![Part::ToolCall(call)]), + ) + }; + events.push_back(ModelTurnEvent::Finished(ModelTurnResult { + model: None, + response_id: None, + finish_reason, + output_items: vec![output], + usage: None, + metadata: MetadataMap::new(), + })); + Ok(Turn(events)) + } +} + +#[async_trait] +impl ModelTurn for Turn { + async fn next_event( + &mut self, + _: Option, + ) -> Result, LoopError> { + Ok(self.0.pop_front()) + } +} + +struct ImageTool { + spec: ToolSpec, + release: watch::Receiver, +} + +#[async_trait] +impl Tool for ImageTool { + fn spec(&self) -> &ToolSpec { + &self.spec + } + + async fn invoke( + &self, + request: ToolRequest, + _: &mut ToolContext<'_>, + ) -> Result { + // The external operation cannot finish until the test has observed the + // real loop's detach placeholder. No sleeps or scheduler races. + self.release + .clone() + .wait_for(|released| *released) + .await + .unwrap(); + Ok(ToolResult { + result: ToolResultPart::success( + request.call_id, + ToolOutput::Parts(vec![ + Part::text("Selected background screenshot: $.image"), + Part::media( + Modality::Image, + "image/png", + DataRef::InlineBytes(vec![1, 2, 3]), + ), + ]), + ) + .with_metadata(MetadataMap::from_iter([( + "diagnostic".into(), + json!("kept"), + )])), + duration: None, + metadata: MetadataMap::new(), + }) + } +} + +#[tokio::test] +async fn real_detached_completion_projects_notification_images_for_responses() { + // Timeout is a deadlock guard, not an assertion about performance. + tokio::time::timeout(Duration::from_secs(5), async { + let (requests, mut received) = mpsc::unbounded_channel(); + let (release, gate) = watch::channel(false); + let manager = AsyncTaskManager::new().routing(|_: &ToolRequest| { + RoutingDecision::ForegroundThenDetachAfter(Duration::ZERO) + }); + let handle = manager.handle(); + let agent = Agent::builder() + .model(Model(requests)) + .add_tool_source(ToolRegistry::new().with(ImageTool { + spec: ToolSpec { + name: ToolName::new("image"), + description: "Return a screenshot after release".into(), + input_schema: json!({"type": "object", "properties": {}, "additionalProperties": false}), + output_schema: None, + annotations: ToolAnnotations::default(), + metadata: MetadataMap::new(), + }, + release: gate, + })) + .permissions(AllowAllPermissions) + .task_manager(manager) + .build().unwrap(); + let mut driver = agent.start(SessionConfig::new("real-background-image")).await.unwrap(); + driver.submit_input(vec![Item::text(ItemKind::User, "Get the screenshot")]).unwrap(); + assert!(matches!(driver.next().await.unwrap(), LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_)))); + let detached = driver.snapshot().transcript; + assert!(detached.iter().flat_map(|item| &item.parts).any(|part| { + matches!(part, Part::ToolResult(result) if matches!(&result.output, ToolOutput::Text(text) if text.contains("running in the background"))) + })); + release.send(true).unwrap(); + // Task events and the ready-item queue are independent public APIs. + // Await completion without draining the result the loop must consume. + loop { + if matches!(handle.next_event().await.unwrap(), TaskEvent::Completed(_, _)) { break; } + } + loop { + match driver.next().await.unwrap() { + LoopStep::Finished(_) => break, + LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_)) => {}, + other => panic!("unexpected loop step: {other:?}"), + } + } + let request = std::iter::from_fn(|| received.try_recv().ok()) + .find(|request| request.transcript.iter().any(|item| item.kind == ItemKind::Notification)) + .expect("real loop must send its detached completion notification to the model"); + let original = serde_json::to_value(&request.transcript).unwrap(); + assert!(original.to_string().contains("InlineBytes")); + assert_eq!(request.transcript.iter().filter(|item| item.kind == ItemKind::Tool).count(), 1); + for native in [false, true] { + let projected = super::project_tool_output_images(request.clone(), native).unwrap(); + let attachment = projected.transcript.iter().find(|item| item.metadata.get("kit.projected_tool_images") == Some(&json!(true))).unwrap(); + assert_eq!(attachment.kind, ItemKind::User); + let wire = OpenAIResponsesConfig::chatgpt_private("gpt-5.4", Authentication::bearer("test")) + .encode_request(&projected).unwrap(); + let text = wire.to_string(); + assert_eq!(text.matches("data:image/png;base64,AQID").count(), 1); + assert!(!text.contains("InlineBytes")); + assert!(text.contains("Selected background screenshot")); + assert!(text.contains("diagnostic")); + assert_eq!(wire["input"].as_array().unwrap().iter().filter(|item| item["type"] == "function_call_output").count(), 1); + assert_eq!(serde_json::to_value(&request.transcript).unwrap(), original); + } + }).await.expect("detached completion test deadlocked"); +} diff --git a/src/provider/adapter_image_tests.rs b/src/provider/adapter_image_tests.rs new file mode 100644 index 0000000..7dd9b1f --- /dev/null +++ b/src/provider/adapter_image_tests.rs @@ -0,0 +1,383 @@ +#![allow(clippy::disallowed_methods, clippy::disallowed_macros)] +use super::*; +use agentkit_core::{Item, ItemKind, MetadataMap, SessionId, ToolCallPart, ToolResultPart, TurnId}; +use agentkit_http::{ + HeaderMap, Http, HttpClient, HttpError, HttpRequest, HttpResponse, StatusCode, +}; +use serde_json::json; + +fn request() -> TurnRequest { + let image = Part::Media( + agentkit_core::MediaPart::new( + Modality::Image, + "image/png", + DataRef::InlineBytes(vec![1, 2, 3]), + ) + .with_metadata(MetadataMap::from_iter([( + "position".into(), + json!("$.image"), + )])), + ); + TurnRequest { + session_id: SessionId::new("projection"), + turn_id: TurnId::new("replay"), + transcript: vec![ + Item::new( + ItemKind::Assistant, + vec![Part::ToolCall(ToolCallPart::new("a", "compose", json!({})))], + ), + Item::new( + ItemKind::Assistant, + vec![Part::ToolCall(ToolCallPart::new("b", "shell", json!({})))], + ), + Item::new( + ItemKind::Tool, + vec![Part::ToolResult(ToolResultPart::success( + "a", + ToolOutput::Parts(vec![ + Part::structured(json!({"answer": 42})), + Part::text("original output"), + Part::ToolResult(ToolResultPart::success( + "nested", + ToolOutput::Parts(vec![Part::text("$.image"), image]), + )), + ]), + ))], + ), + Item::new( + ItemKind::Tool, + vec![Part::ToolResult(ToolResultPart::success( + "b", + ToolOutput::Text("parallel result".into()), + ))], + ), + Item::text(ItemKind::User, "continue"), + ], + available_tools: vec![], + cache: None, + metadata: MetadataMap::new(), + } +} + +#[test] +fn projection_preserves_canonical_native_metadata_and_replay() { + let canonical = request(); + let original = serde_json::to_value(&canonical.transcript).unwrap(); + let projected = project_tool_output_images(canonical.clone(), false).unwrap(); + assert_eq!( + serde_json::to_value(&canonical.transcript).unwrap(), + original + ); + let native = project_tool_output_images(canonical, true).unwrap(); + assert_eq!(native.transcript.len(), 5); + assert!( + serde_json::to_string(&native.transcript) + .unwrap() + .contains("InlineBytes") + ); + assert_eq!(projected.transcript.len(), 6); + assert_eq!(projected.transcript[3].kind, ItemKind::Tool); + let image = projected.transcript[4] + .parts + .iter() + .find_map(|p| match p { + Part::Media(m) => Some(m), + _ => None, + }) + .unwrap(); + assert_eq!(image.metadata["position"], "$.image"); + let first = serde_json::to_value(&projected.transcript).unwrap(); + let replay = project_tool_output_images(projected, false).unwrap(); + assert_eq!(serde_json::to_value(replay.transcript).unwrap(), first); +} + +#[test] +fn whole_item_batch_keeps_all_results_before_multiple_direct_images() { + let mut request = request(); + let second_call = request.transcript.remove(1); + request.transcript[0].parts.extend(second_call.parts); + let second_result = request.transcript.remove(2); + request.transcript[1].parts.extend(second_result.parts); + let Part::ToolResult(result) = &mut request.transcript[1].parts[1] else { + panic!() + }; + result.output = ToolOutput::Parts(vec![ + Part::text("second image label"), + Part::media( + Modality::Image, + "image/jpeg", + DataRef::InlineBytes(vec![4, 5, 6]), + ), + Part::text("trailing text"), + ]); + let projected = project_tool_output_images(request, false).unwrap(); + assert_eq!(projected.transcript.len(), 4); + assert_eq!(projected.transcript[1].parts.len(), 2); + assert_eq!(projected.transcript[1].kind, ItemKind::Tool); + assert_eq!(projected.transcript[2].kind, ItemKind::User); + let images: Vec<_> = projected.transcript[2] + .parts + .iter() + .filter_map(|p| match p { + Part::Media(m) => Some(m), + _ => None, + }) + .collect(); + assert_eq!(images.len(), 2); + assert_eq!(images[0].data, DataRef::InlineBytes(vec![1, 2, 3])); + assert_eq!(images[1].data, DataRef::InlineBytes(vec![4, 5, 6])); + let tool = serde_json::to_string(&projected.transcript[1]).unwrap(); + assert!(tool.contains("trailing text")); + assert!(tool.contains("second image label")); +} + +#[test] +fn unfinished_batch_does_not_deliver_images_between_results() { + let mut request = request(); + request.transcript.remove(3); + assert!( + project_tool_output_images(request, false) + .unwrap_err() + .to_string() + .contains("all outstanding tool calls") + ); +} + +// Genuine HTTP boundary: inspect the bytes emitted by the pinned public +// Completions encoder, also used by OpenRouter (a CompletionsSession alias). +struct CaptureClient(tokio::sync::mpsc::UnboundedSender); + +#[async_trait] +impl HttpClient for CaptureClient { + async fn execute(&self, request: HttpRequest) -> Result { + self.0 + .send(serde_json::from_slice(&request.body.unwrap()).unwrap()) + .unwrap(); + Ok(HttpResponse::new( + StatusCode::BAD_REQUEST, + HeaderMap::new(), + request.url, + Box::pin(futures_util::stream::empty()), + )) + } +} + +async fn capture_wire( + provider: P, + request: TurnRequest, +) -> Value { + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); + let adapter = CompletionsAdapter::with_client(provider, Http::new(CaptureClient(tx))); + let mut session = adapter + .start_session(SessionConfig::new("projection")) + .await + .unwrap(); + assert!(session.begin_turn(request, None).await.is_err()); + rx.try_recv().expect("request must reach the HTTP encoder") +} + +async fn assert_wire(provider: P, request: TurnRequest) { + let body = capture_wire(provider, request).await; + let messages = body["messages"].as_array().unwrap(); + let roles: Vec<_> = messages + .iter() + .map(|m| m["role"].as_str().unwrap()) + .collect(); + assert_eq!( + roles, + ["assistant", "assistant", "tool", "tool", "user", "user"] + ); + assert_eq!(messages[2]["tool_call_id"], "a"); + assert_eq!(messages[3]["tool_call_id"], "b"); + let output = messages[2]["content"].as_str().unwrap(); + assert!(output.contains("original output")); + assert!(output.contains("answer")); + assert!(!output.contains("inline_bytes")); + let content = messages[4]["content"].as_array().unwrap(); + assert!(content.iter().any(|p| p["text"] == "$.image")); + let images: Vec<_> = content + .iter() + .filter(|p| p["type"] == "image_url") + .collect(); + assert_eq!(images.len(), 1); + assert_eq!(images[0]["image_url"]["url"], "data:image/png;base64,AQID"); + assert_eq!(messages[5]["content"], "continue"); +} + +#[tokio::test] +async fn completions_and_openrouter_encode_parallel_mixed_replay_as_user_images() { + let projected = project_tool_output_images(request(), false).unwrap(); + let speakeasy = SpeakeasyProvider { + openrouter: OpenRouterProvider::from(OpenRouterConfig::new("test", "test/model")), + api_key: "test".into(), + project: "test".into(), + chat_id: None, + }; + assert_wire(speakeasy, projected.clone()).await; + assert_wire( + OpenRouterProvider::from(OpenRouterConfig::new("test", "test/model")), + projected, + ) + .await; +} + +#[test] +fn nested_results_encode_with_private_responses_in_both_modes() { + use agentkit_http::Authentication; + use agentkit_provider_openai::OpenAIResponsesConfig; + for native in [false, true] { + let projected = project_tool_output_images(request(), native).unwrap(); + let config = + OpenAIResponsesConfig::chatgpt_private("gpt-5.4", Authentication::bearer("test")); + let wire = config.encode_request(&projected).unwrap(); + let text = serde_json::to_string(&wire).unwrap(); + assert!(text.contains("nested_tool_result")); + assert!(text.contains("original output")); + assert!(text.contains("is_error")); + assert_eq!(text.matches("data:image/png;base64,AQID").count(), 1); + assert!(!text.contains("InlineBytes")); + } +} + +// Exact maybe_convert_detached representation from patched agentkit-loop +// 8e4ee26. That private method retains arbitrary item/result metadata and adds +// no dedicated marker. It emits a summary followed by serialized results. +fn detached_request() -> TurnRequest { + let mut request = request(); + let Part::ToolResult(result) = &request.transcript[2].parts[0] else { + panic!() + }; + let mut result = result.clone(); + result + .metadata + .insert("diagnostic".into(), json!("preserve me")); + result.is_error = true; + request.transcript[2].parts = vec![Part::ToolResult(ToolResultPart::success("a", ToolOutput::Text( + "Tool compose is now running in the background. The result will be delivered when it completes.".into() + )))]; + let mut notification = Item::new( + ItemKind::Notification, + vec![ + Part::text( + "Background tool results: 1 total, 1 failed, 1 with metadata. a failed: parts payload (3 parts)", + ), + Part::structured(serde_json::to_value(result).unwrap()), + ], + ); + notification + .metadata + .insert("delivery".into(), json!("deferred")); + // Another foreground call is still outstanding when notification arrives. + request.transcript.insert(3, notification); + request +} + +#[tokio::test] +async fn detached_notifications_deliver_once_even_on_native_transport() { + use agentkit_http::Authentication; + use agentkit_provider_openai::OpenAIResponsesConfig; + for native in [false, true] { + let request = detached_request(); + let original = serde_json::to_value(&request.transcript).unwrap(); + let projected = project_tool_output_images(request.clone(), native).unwrap(); + assert_eq!(serde_json::to_value(&request.transcript).unwrap(), original); + assert_eq!(projected.transcript[3].kind, ItemKind::Notification); + assert_eq!(projected.transcript[3].metadata["delivery"], "deferred"); + assert_eq!(projected.transcript[4].kind, ItemKind::Tool); + assert_eq!( + projected.transcript[5].metadata["kit.projected_tool_images"], + true + ); + let note = serde_json::to_string(&projected.transcript[3]).unwrap(); + assert!(!note.contains("InlineBytes")); + assert!(note.contains("preserve me")); + assert!(note.contains("$.image")); + let wire = + OpenAIResponsesConfig::chatgpt_private("gpt-5.4", Authentication::bearer("test")) + .encode_request(&projected) + .unwrap(); + let text = serde_json::to_string(&wire).unwrap(); + assert_eq!(text.matches("data:image/png;base64,AQID").count(), 1); + assert_eq!( + wire["input"] + .as_array() + .unwrap() + .iter() + .filter(|i| i["type"] == "function_call_output") + .count(), + 2 + ); + assert!(!text.contains("InlineBytes")); + let body = capture_wire( + OpenRouterProvider::from(OpenRouterConfig::new("test", "test/model")), + projected.clone(), + ) + .await; + assert_eq!( + body["messages"] + .as_array() + .unwrap() + .iter() + .filter(|m| m["role"] == "tool") + .count(), + 2 + ); + assert_eq!( + body.to_string() + .matches("data:image/png;base64,AQID") + .count(), + 1 + ); + assert!(!body.to_string().contains("InlineBytes")); + let replay = project_tool_output_images(projected.clone(), native).unwrap(); + assert_eq!( + serde_json::to_value(replay.transcript).unwrap(), + serde_json::to_value(projected.transcript).unwrap() + ); + } +} + +#[test] +fn arbitrary_structured_results_are_not_detached_notifications() { + let mut request = detached_request(); + request.transcript[3].parts[0] = Part::text("Unrelated notification"); + let expected = request.transcript[3].clone(); + let projected = project_tool_output_images(request, false).unwrap(); + assert_eq!(projected.transcript[3], expected); +} + +#[test] +fn detached_output_traversal_is_bounded_in_both_modes() { + for native in [false, true] { + let mut request = detached_request(); + let Part::Structured(result) = &mut request.transcript[3].parts[1] else { + panic!() + }; + let mut nested = json!({"Text": {"text": "deep", "metadata": {}}}); + for _ in 0..65 { + nested = json!({"ToolResult": { + "call_id": "nested", "is_error": false, "metadata": {}, + "output": {"Parts": [nested]}, + }}); + } + result.value["output"] = json!({"Parts": [nested]}); + assert!( + project_tool_output_images(request, native) + .unwrap_err() + .to_string() + .contains("traversal budget") + ); + let mut request = detached_request(); + let Part::Structured(result) = &mut request.transcript[3].parts[1] else { + panic!() + }; + result.value["output"] = + json!({"Parts": vec![json!({"Text": {"text": "wide", "metadata": {}}}); 100_001]}); + assert!( + project_tool_output_images(request, native) + .unwrap_err() + .to_string() + .contains("traversal budget") + ); + } +} diff --git a/src/provider/chatgpt.rs b/src/provider/chatgpt.rs index 06fe275..348f7ac 100644 --- a/src/provider/chatgpt.rs +++ b/src/provider/chatgpt.rs @@ -46,7 +46,8 @@ const MAX_ITEMS: usize = 10_000; const MAX_FIELD_BYTES: usize = 1024 * 1024; const MAX_SOURCE_IMAGE_BYTES: usize = 10 * 1024 * 1024; const MAX_DECODED_IMAGE_BYTES: u64 = 64 * 1024 * 1024; -const MAX_IMAGE_PIXELS: u64 = 10_000_000; +// Match managed-file import limits; retain the independent decoded-byte bound. +const MAX_IMAGE_PIXELS: u64 = 16 * 1024 * 1024; const MAX_IMAGE_DIMENSION: u32 = 8_192; const MAX_TOOL_RESULT_DEPTH: usize = 8; const JPEG_DATA_URL_PREFIX: &str = "data:image/jpeg;base64,"; @@ -314,13 +315,20 @@ impl ModelSession for OpenAiSubscriptionSession { mut request: TurnRequest, cancellation: Option, ) -> Result { + let model = self.inner.model_name().unwrap_or("unknown"); + request = super::adapter::project_tool_output_images( + request, + supports_tool_output_images(model), + )?; migrate_legacy_continuations(&mut request, &self.authentication_binding)?; let normalization_cancellation = cancellation.clone(); let request = tokio::task::spawn_blocking(move || { normalize_openai_images(request, normalization_cancellation.as_ref()) }) .await - .map_err(|_| protocol("Responses image normalization task failed"))??; + .map_err(|_| { + tool_image_normalization_error(protocol("Responses image normalization task failed")) + })??; self.inner .begin_turn(request, cancellation) .await @@ -337,6 +345,13 @@ impl ModelSession for OpenAiSubscriptionSession { } } +// Keep the verified native tool-output route. Other models use the ordinary +// user-image request path; this is a transport choice, not a vision allowlist. +// The pinned private Responses encoder supplies native image tool-output blocks. +fn supports_tool_output_images(model: &str) -> bool { + model == "gpt-5.4" +} + fn normalize_openai_images( mut request: TurnRequest, cancellation: Option<&agentkit_core::TurnCancellation>, @@ -347,7 +362,12 @@ fn normalize_openai_images( item.kind, ItemKind::User | ItemKind::Context | ItemKind::Tool ) { - normalize_openai_parts(&mut item.parts, cancellation, 0)?; + let result = normalize_openai_parts(&mut item.parts, cancellation, 0); + if item.metadata.get("kit.projected_tool_images") == Some(&Value::Bool(true)) { + result.map_err(tool_image_normalization_error)?; + } else { + result?; + } } } Ok(request) @@ -377,7 +397,8 @@ fn normalize_openai_parts( } Part::ToolResult(result) => { if let ToolOutput::Parts(parts) = &mut result.output { - normalize_openai_parts(parts, cancellation, depth + 1)?; + normalize_openai_parts(parts, cancellation, depth + 1) + .map_err(tool_image_normalization_error)?; } } _ => {} @@ -386,6 +407,16 @@ fn normalize_openai_parts( Ok(()) } +fn tool_image_normalization_error(error: LoopError) -> LoopError { + match error { + // Preserve cancellation and avoid wrapping an already contextualized error. + LoopError::Cancelled | LoopError::InvalidState(_) => error, + _ => LoopError::InvalidState(format!( + "selected-images-not-delivered: {error}. The program may already have completed; do not retry or rerun the program. The retained images could not be prepared for delivery." + )), + } +} + fn check_image_cancellation( cancellation: Option<&agentkit_core::TurnCancellation>, ) -> Result<(), LoopError> { @@ -989,6 +1020,139 @@ mod tests { use super::*; use serde_json::json; + fn image_tool_request() -> TurnRequest { + use agentkit_core::{Item, SessionId, ToolResultPart, TurnId}; + TurnRequest { + session_id: SessionId::new("image-session"), + turn_id: TurnId::new("image-turn"), + transcript: vec![Item::new( + ItemKind::Tool, + vec![Part::ToolResult(ToolResultPart::success( + "image-call", + ToolOutput::Parts(vec![ + Part::text("Selected screenshot"), + Part::media( + Modality::Image, + "image/png", + DataRef::InlineBytes(vec![1, 2, 3]), + ), + ]), + ))], + )], + available_tools: Vec::new(), + cache: None, + metadata: MetadataMap::new(), + } + } + + #[test] + fn selected_images_use_native_private_responses_wire_content() { + let request = image_tool_request(); + let original = serde_json::to_value(&request.transcript).unwrap(); + assert!(supports_tool_output_images("gpt-5.4")); + let request = super::super::adapter::project_tool_output_images( + request, + supports_tool_output_images("gpt-5.4"), + ) + .unwrap(); + let config = + OpenAIResponsesConfig::chatgpt_private("gpt-5.4", Authentication::bearer("test-key")); + let wire = config.encode_request(&request).unwrap(); + assert_eq!( + wire["input"][0], + json!({ + "type": "function_call_output", + "call_id": "image-call", + "output": [ + {"type": "input_text", "text": "Selected screenshot"}, + {"type": "input_image", "image_url": "data:image/png;base64,AQID", "detail": "high"} + ] + }) + ); + assert_eq!(serde_json::to_value(&request.transcript).unwrap(), original); + } + + #[test] + fn selected_image_guard_leaves_user_images_and_text_wire_unchanged() { + use agentkit_core::{Item, ToolResultPart}; + let mut request = image_tool_request(); + request.transcript = vec![ + Item::new( + ItemKind::User, + vec![ + Part::text("Look at this"), + Part::media( + Modality::Image, + "image/png", + DataRef::InlineBytes(vec![1, 2, 3]), + ), + ], + ), + Item::new( + ItemKind::Tool, + vec![Part::ToolResult(ToolResultPart::success( + "text-call", + ToolOutput::Text("done".into()), + ))], + ), + ]; + // An unverified model is not blocked for ordinary user images or text. + let request = super::super::adapter::project_tool_output_images(request, false).unwrap(); + let config = OpenAIResponsesConfig::chatgpt_private( + "gpt-future", + Authentication::bearer("test-key"), + ); + let wire = config.encode_request(&request).unwrap(); + assert_eq!( + wire["input"][0]["content"][0], + json!({"type": "input_text", "text": "Look at this"}) + ); + assert_eq!(wire["input"][0]["content"][1]["type"], "input_image"); + assert_eq!(wire["input"][1]["output"], "done"); + } + + #[test] + fn unlisted_subscription_model_projects_images_on_each_request_only() { + let request = image_tool_request(); + let original = serde_json::to_value(&request.transcript).unwrap(); + let config = OpenAIResponsesConfig::chatgpt_private( + "gpt-6-astra", + Authentication::bearer("test-key"), + ); + assert!(!supports_tool_output_images("gpt-6-astra")); + let mut previous_wire = None; + // A resumed or continued turn starts from the same canonical typed data. + for _ in 0..2 { + let projected = super::super::adapter::project_tool_output_images( + request.clone(), + supports_tool_output_images("gpt-6-astra"), + ) + .unwrap(); + let projected = normalize_openai_images(projected, None).unwrap(); + let wire = config.encode_request(&projected).unwrap(); + assert_eq!(wire["input"].as_array().unwrap().len(), 2); + assert_eq!(wire["input"][0]["type"], "function_call_output"); + assert_eq!(wire["input"][0]["call_id"], "image-call"); + let output = wire["input"][0]["output"].to_string(); + assert!(output.contains("Selected screenshot")); + assert!(!output.contains("AQID")); + assert_eq!(wire["input"][1]["role"], "user"); + let images: Vec<_> = wire["input"][1]["content"] + .as_array() + .unwrap() + .iter() + .filter(|part| part["type"] == "input_image") + .collect(); + assert_eq!(images.len(), 1); + assert_eq!(images[0]["image_url"], "data:image/png;base64,AQID"); + if let Some(previous) = previous_wire { + assert_eq!(wire, previous); + } + previous_wire = Some(wire); + assert_eq!(serde_json::to_value(&request.transcript).unwrap(), original); + } + } + #[test] fn subscription_config_accepts_models_without_a_client_release() { assert!(SubscriptionConfig::new("gpt-future".into()).is_ok()); @@ -1035,6 +1199,162 @@ mod tests { assert_eq!((decoded.width(), decoded.height()), (600, 600)); } + #[test] + fn projected_image_failure_retains_no_rerun_guidance() { + let mut request = image_tool_request(); + let Part::ToolResult(result) = &mut request.transcript[0].parts[0] else { + panic!("expected tool result"); + }; + result.output = ToolOutput::Parts(vec![Part::media( + Modality::Image, + "image/png", + DataRef::InlineBytes(vec![0; MAX_FIELD_BYTES]), + )]); + let projected = super::super::adapter::project_tool_output_images(request, false).unwrap(); + let error = normalize_openai_images(projected, None).unwrap_err(); + assert!(error.to_string().contains("selected-images-not-delivered")); + assert!(error.to_string().contains("do not retry or rerun")); + } + + #[test] + fn projected_tool_images_use_user_image_normalization_limits() { + let png = noisy_png(600, 600); + assert!(png.len() > MAX_NORMALIZED_IMAGE_BYTES); + let mut request = image_tool_request(); + let Part::ToolResult(result) = &mut request.transcript[0].parts[0] else { + panic!("expected tool result"); + }; + result.output = ToolOutput::Parts(vec![ + Part::text("Retained diagnostic and label"), + Part::media(Modality::Image, "image/png", DataRef::InlineBytes(png)), + ]); + let projected = super::super::adapter::project_tool_output_images(request, false).unwrap(); + let normalized = normalize_openai_images(projected, None).unwrap(); + let config = OpenAIResponsesConfig::chatgpt_private( + "gpt-6-astra", + Authentication::bearer("test-key"), + ) + .with_limits(OpenAIResponsesLimits { + max_request_bytes: MAX_REQUEST_BYTES, + max_attempt_bytes: MAX_ATTEMPT_BYTES, + max_wire_bytes: MAX_WIRE_BYTES, + max_items: MAX_ITEMS, + max_text_bytes: MAX_FIELD_BYTES, + }); + let wire = config.encode_request(&normalized).unwrap(); + let image = wire["input"][1]["content"] + .as_array() + .unwrap() + .iter() + .find(|part| part["type"] == "input_image") + .unwrap(); + let url = image["image_url"].as_str().unwrap(); + assert!(url.starts_with("data:image/jpeg;base64,")); + assert!(url.len() <= MAX_FIELD_BYTES); + assert!( + wire["input"][0]["output"] + .to_string() + .contains("Retained diagnostic and label") + ); + } + + #[test] + fn selected_twelve_megapixel_jpeg_normalizes_to_native_wire_budget() { + // Keep a high-detail region in an otherwise flat 12MP image: the valid + // source exceeds the wire field limit but fits the managed-file limits. + let image = RgbImage::from_fn(4000, 3000, |x, y| { + if x >= 800 || y >= 800 { + return Rgb([255, 255, 255]); + } + let mut value = x + .wrapping_mul(747_796_405) + .wrapping_add(y.wrapping_mul(2_891_336_453)); + value = (value ^ (value >> 16)).wrapping_mul(2_246_822_519); + value ^= value >> 13; + Rgb([value as u8, (value >> 8) as u8, (value >> 16) as u8]) + }); + let mut jpeg = Vec::new(); + image::codecs::jpeg::JpegEncoder::new_with_quality(&mut jpeg, 100) + .encode_image(&DynamicImage::ImageRgb8(image)) + .unwrap(); + assert!(jpeg.len() > MAX_NORMALIZED_IMAGE_BYTES); + assert!(jpeg.len() <= 8 * 1024 * 1024); + let mut request = image_tool_request(); + let Part::ToolResult(result) = &mut request.transcript[0].parts[0] else { + panic!("expected tool result"); + }; + result.output = ToolOutput::Parts(vec![ + Part::text("Selected 12MP JPEG"), + Part::media(Modality::Image, "image/jpeg", DataRef::InlineBytes(jpeg)), + ]); + // Exercise the full normalizer used by begin_turn, not just byte decoding. + let normalized = normalize_openai_images(request, None).unwrap(); + let config = + OpenAIResponsesConfig::chatgpt_private("gpt-5.4", Authentication::bearer("test-key")) + .with_limits(OpenAIResponsesLimits { + max_request_bytes: MAX_REQUEST_BYTES, + max_attempt_bytes: MAX_ATTEMPT_BYTES, + max_wire_bytes: MAX_WIRE_BYTES, + max_items: MAX_ITEMS, + max_text_bytes: MAX_FIELD_BYTES, + }); + let wire = config.encode_request(&normalized).unwrap(); + let output = &wire["input"][0]["output"]; + assert_eq!(wire["input"][0]["call_id"], "image-call"); + assert_eq!(output[0]["type"], "input_text"); + assert_eq!(output[1]["type"], "input_image"); + let url = output[1]["image_url"].as_str().unwrap(); + assert!(url.len() <= MAX_FIELD_BYTES); + let bytes = BASE64 + .decode(url.strip_prefix(JPEG_DATA_URL_PREFIX).unwrap()) + .unwrap(); + assert!( + ImageReader::new(Cursor::new(bytes)) + .with_guessed_format() + .unwrap() + .decode() + .is_ok() + ); + } + + #[tokio::test] + async fn selected_image_normalization_failure_does_not_invite_program_retry() { + let config = + OpenAIResponsesConfig::chatgpt_private("gpt-5.4", Authentication::bearer("test-key")); + let inner = OpenAIResponsesAdapter::new(config) + .unwrap() + .start_session(SessionConfig::new("image-session")) + .await + .unwrap(); + let mut session = OpenAiSubscriptionSession { + inner, + context_window: None, + authentication_binding: "unused".into(), + }; + let mut request = image_tool_request(); + let Part::ToolResult(result) = &mut request.transcript[0].parts[0] else { + panic!("expected tool result"); + }; + result.output = ToolOutput::Parts(vec![Part::media( + Modality::Image, + "image/jpeg", + DataRef::InlineBytes(vec![0; MAX_FIELD_BYTES]), + )]); + let error = match session.begin_turn(request, None).await { + Err(error) => error, + Ok(_) => panic!("invalid image reached provider"), + }; + assert!(matches!(error, LoopError::InvalidState(_))); + let message = error.to_string(); + assert!(message.contains("selected-images-not-delivered")); + assert!(message.contains("program may already have completed")); + assert!(message.contains("do not retry or rerun the program")); + assert!(matches!( + tool_image_normalization_error(LoopError::Cancelled), + LoopError::Cancelled + )); + } + #[test] fn image_normalization_observes_turn_cancellation() { let controller = agentkit_core::CancellationController::new(); @@ -1470,3 +1790,7 @@ mod tests { assert_eq!(catalog.context_windows.len(), 2); } } + +#[cfg(test)] +#[path = "chatgpt_image_tests.rs"] +mod image_tests; diff --git a/src/provider/chatgpt_image_tests.rs b/src/provider/chatgpt_image_tests.rs new file mode 100644 index 0000000..f86fbaf --- /dev/null +++ b/src/provider/chatgpt_image_tests.rs @@ -0,0 +1,236 @@ +#![allow(clippy::disallowed_methods, clippy::disallowed_macros)] +//! Wire-level coverage for fallback images with authentication-bound replay. +#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] + +use super::*; +use agentkit_core::{Item, SessionId, ToolResultPart, TurnId}; +use agentkit_http::{Bytes, Http, StatusCode}; +use serde_json::json; +use tokio::sync::mpsc; + +// Fake only the external HTTP boundary; decoding, authentication binding, +// projection, normalization, and request encoding remain production code. +struct ContinuationHttp { + requests: mpsc::UnboundedSender, +} + +#[async_trait] +impl HttpClient for ContinuationHttp { + async fn execute(&self, request: HttpRequest) -> Result { + let url = request.url.clone(); + self.requests.send(request).unwrap(); + let items = [ + json!({"id":"reason-1","type":"reasoning","summary":[{"type":"summary_text","text":"Inspect both results"}],"encrypted_content":"opaque-reasoning"}), + json!({"id":"call-item-image","type":"function_call","call_id":"call-image","name":"compose","arguments":"{}"}), + json!({"id":"call-item-text","type":"function_call","call_id":"call-text","name":"compose","arguments":"{}"}), + ]; + let mut events = vec![ + json!({"type":"response.created","response":{"id":"response-1","model":"gpt-6-astra"}}), + ]; + for (index, item) in items.into_iter().enumerate() { + events.push(json!({"type":"response.output_item.added","output_index":index,"item":{"id":item["id"],"type":item["type"]}})); + if item["type"] == "reasoning" { + events.push(json!({"type":"response.reasoning_summary_part.added","item_id":item["id"],"output_index":index,"summary_index":0,"part":{"type":"summary_text"}})); + events.push(json!({"type":"response.reasoning_summary_text.delta","item_id":item["id"],"output_index":index,"summary_index":0,"delta":"Inspect both results"})); + events.push(json!({"type":"response.reasoning_summary_text.done","item_id":item["id"],"output_index":index,"summary_index":0,"text":"Inspect both results"})); + events.push(json!({"type":"response.reasoning_summary_part.done","item_id":item["id"],"output_index":index,"summary_index":0,"part":item["summary"][0]})); + } else { + events.push(json!({"type":"response.function_call_arguments.delta","item_id":item["id"],"output_index":index,"delta":"{}"})); + events.push(json!({"type":"response.function_call_arguments.done","item_id":item["id"],"output_index":index,"arguments":"{}"})); + } + events + .push(json!({"type":"response.output_item.done","output_index":index,"item":item})); + } + events.push(json!({"type":"response.completed","response":{"id":"response-1","model":"gpt-6-astra","usage":{"input_tokens":1,"output_tokens":1}}})); + let body = events + .into_iter() + .enumerate() + .map(|(index, mut event)| { + event["sequence_number"] = json!(index + 1); + format!( + "event: {}\ndata: {event}\n\n", + event["type"].as_str().unwrap() + ) + }) + .collect::(); + Ok(HttpResponse::new( + StatusCode::OK, + HeaderMap::from_iter([( + agentkit_http::header::CONTENT_TYPE, + HeaderValue::from_static("text/event-stream"), + )]), + url, + Box::pin(futures_util::stream::once( + async move { Ok(Bytes::from(body)) }, + )), + )) + } +} + +async fn collect_output( + session: &mut OpenAiSubscriptionSession, + request: TurnRequest, +) -> Vec { + let mut turn = session.begin_turn(request, None).await.unwrap(); + while let Some(event) = turn.next_event(None).await.unwrap() { + if let ModelTurnEvent::Finished(result) = event { + return result.output_items; + } + } + panic!("response must finish"); +} + +#[tokio::test] +async fn authenticated_continuation_replays_parallel_results_with_fallback_images() { + let (sender, mut requests) = mpsc::unbounded_channel(); + let config = OpenAIResponsesConfig::chatgpt_private( + "gpt-6-astra", + Authentication::bearer("test-image-replay-key"), + ); + let adapter = OpenAIResponsesAdapter::with_client( + config, + Http::new(ContinuationHttp { requests: sender }), + ); + let mut session = OpenAiSubscriptionSession { + inner: adapter + .start_session(SessionConfig::new("image-replay")) + .await + .unwrap(), + context_window: None, + // No legacy metadata is injected: the real adapter creates and validates + // its current authentication-bound continuation metadata. + authentication_binding: "unused-legacy-binding".into(), + }; + let mut canonical = TurnRequest { + session_id: SessionId::new("image-replay"), + turn_id: TurnId::new("initial"), + transcript: vec![Item::text(ItemKind::User, "Inspect both results")], + available_tools: Vec::new(), + cache: None, + metadata: MetadataMap::new(), + }; + let output = collect_output(&mut session, canonical.clone()).await; + let initial = requests.try_recv().unwrap(); + assert_eq!( + initial.headers[agentkit_http::header::AUTHORIZATION], + "Bearer test-image-replay-key" + ); + let metadata: Vec<_> = output + .iter() + .flat_map(|item| &item.parts) + .filter_map(|part| match part { + Part::Reasoning(reasoning) => reasoning.metadata.get(CONTINUATION_METADATA), + Part::ToolCall(call) => call.metadata.get(CONTINUATION_METADATA), + _ => None, + }) + .collect(); + assert_eq!(metadata.len(), 3); + for value in metadata { + assert!(!value["authentication_binding"].as_str().unwrap().is_empty()); + assert_eq!(value["session_id"], "image-replay"); + } + canonical.transcript.extend(output); + let mut png = Cursor::new(Vec::new()); + DynamicImage::ImageRgb8(RgbImage::from_pixel(1, 1, Rgb([12, 34, 56]))) + .write_to(&mut png, ImageFormat::Png) + .unwrap(); + let expected_url = format!("data:image/png;base64,{}", BASE64.encode(png.get_ref())); + canonical.transcript.push(Item::new( + ItemKind::Tool, + vec![Part::ToolResult(ToolResultPart::success( + "call-image", + ToolOutput::Parts(vec![ + Part::text("Selected screenshot: $.image"), + Part::media( + Modality::Image, + "image/png", + DataRef::InlineBytes(png.into_inner()), + ), + ]), + ))], + )); + canonical.transcript.push(Item::new( + ItemKind::Tool, + vec![Part::ToolResult(ToolResultPart::success( + "call-text", + ToolOutput::Text("Second parallel result".into()), + ))], + )); + let original = serde_json::to_value(&canonical.transcript).unwrap(); + let mut encoded = Vec::new(); + for turn_id in ["continuation", "replay"] { + // A fresh provider session proves replay does not depend on hidden + // per-session request state or a persisted synthetic user turn. + session.inner = adapter + .start_session(SessionConfig::new("image-replay")) + .await + .unwrap(); + let mut request = canonical.clone(); + request.turn_id = TurnId::new(turn_id); + collect_output(&mut session, request).await; + let captured = requests.try_recv().unwrap(); + let wire: Value = serde_json::from_slice(captured.body.as_ref().unwrap()).unwrap(); + assert_eq!(wire["model"], "gpt-6-astra"); + let input = wire["input"].as_array().unwrap(); + let reasoning = input + .iter() + .find(|item| item["type"] == "reasoning") + .unwrap(); + assert_eq!(reasoning["id"], "reason-1"); + assert_eq!(reasoning["encrypted_content"], "opaque-reasoning"); + let calls: Vec<_> = input + .iter() + .enumerate() + .filter(|(_, item)| item["type"] == "function_call") + .collect(); + let results: Vec<_> = input + .iter() + .enumerate() + .filter(|(_, item)| item["type"] == "function_call_output") + .collect(); + assert_eq!(calls.len(), 2); + assert_eq!(results.len(), 2); + for (index, (call_id, item_id)) in [ + ("call-image", "call-item-image"), + ("call-text", "call-item-text"), + ] + .into_iter() + .enumerate() + { + assert_eq!(calls[index].1["call_id"], call_id); + assert_eq!(calls[index].1["id"], item_id); + assert_eq!(results[index].1["call_id"], call_id); + assert!(calls[index].0 < results[0].0); + } + assert!( + results[0].1["output"] + .as_str() + .unwrap() + .contains("Selected screenshot: $.image") + ); + assert_eq!(results[1].1["output"], "Second parallel result"); + let images: Vec<_> = input + .iter() + .enumerate() + .flat_map(|(index, item)| { + item["content"] + .as_array() + .into_iter() + .flatten() + .filter(|part| part["type"] == "input_image") + .map(move |part| (index, item, part)) + }) + .collect(); + assert_eq!(images.len(), 1); + assert_eq!(images[0].0, results[1].0 + 1); + assert_eq!(images[0].1["role"], "user"); + assert_eq!(images[0].2["image_url"], expected_url); + assert_eq!( + serde_json::to_value(&canonical.transcript).unwrap(), + original + ); + encoded.push(wire["input"].clone()); + } + assert_eq!(encoded[0], encoded[1]); + assert!(requests.try_recv().is_err()); +} diff --git a/src/runtime.rs b/src/runtime.rs index 68cbd8d..287e965 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -40,8 +40,8 @@ use crate::{ }, tools::{ A2aTool, ArtifactTool, AuthTool, CloseTool, DocsTool, EditTool, ForkTool, McpTool, - Observed, PromptTool, ShellTool, SubagentTool, Subagents, SubagentsTool, ToolSearch, - observe_shared, + Observed, PromptTool, ReadFileTool, ShellTool, SubagentTool, Subagents, SubagentsTool, + ToolSearch, observe_shared, }, }; @@ -1137,6 +1137,7 @@ impl Runtime { .with(Observed::new(ArtifactTool::new(crate::artifacts::base( &self.root, )))) + .with(Observed::new(ReadFileTool::new(self.root.clone()))) .with(Observed::new(DocsTool::new())) .with(Observed::new(ShellTool::new(self.root.clone()))) .with(Observed::new(EditTool::new(self.root.clone()))); @@ -2264,11 +2265,19 @@ impl Tool for BackgroundableCompose { let call_id = request.call_id.clone(); let artifact_directory = crate::artifacts::directory(&self.root, &request.session_id.0, &call_id.0); + let session = request.session_id.0.clone(); let request = Self::sanitized(request)?; let _job = self.begin_background(background, &call_id, ctx)?; match self.inner.invoke(request, ctx).await { Ok(mut result) => { - match crate::compose_output::guard(&artifact_directory, result.result.output).await + match crate::compose_output::finalize( + &self.root, + &session, + &artifact_directory, + result.result.output, + ctx.cancellation.clone(), + ) + .await { Ok(output) => { result.result.output = output; @@ -2290,6 +2299,7 @@ impl Tool for BackgroundableCompose { let call_id = request.call_id.clone(); let artifact_directory = crate::artifacts::directory(&self.root, &request.session_id.0, &call_id.0); + let session = request.session_id.0.clone(); let request = match Self::sanitized(request) { Ok(request) => request, Err(error) => return ToolExecutionOutcome::Failed(error), @@ -2300,7 +2310,14 @@ impl Tool for BackgroundableCompose { }; match self.inner.invoke_outcome(request, ctx).await { ToolExecutionOutcome::Completed(mut result) => { - match crate::compose_output::guard(&artifact_directory, result.result.output).await + match crate::compose_output::finalize( + &self.root, + &session, + &artifact_directory, + result.result.output, + ctx.cancellation.clone(), + ) + .await { Ok(output) => { result.result.output = output; diff --git a/src/runtime/tests.rs b/src/runtime/tests.rs index b792c9b..8334afb 100644 --- a/src/runtime/tests.rs +++ b/src/runtime/tests.rs @@ -1,3 +1,6 @@ +mod managed_background; +mod managed_files; + use std::{sync::Arc, time::Duration}; use agentkit_core::{ diff --git a/src/runtime/tests/managed_background.rs b/src/runtime/tests/managed_background.rs new file mode 100644 index 0000000..76509e2 --- /dev/null +++ b/src/runtime/tests/managed_background.rs @@ -0,0 +1,279 @@ +use super::*; +use agentkit_core::{DataRef, Modality, ToolResultPart}; +use agentkit_task_manager::{ + TaskKind, TaskLaunchRequest, TaskManager, TaskResolution, TaskStartContext, TaskStartOutcome, + TurnTaskUpdate, +}; +use agentkit_tools_core::{ToolContext, ToolError, ToolRegistry, ToolResult, ToolSpec}; +use tokio::sync::{Notify, mpsc}; + +// A genuine external-tool boundary. The receiver owns incoming requests and the +// test owns the single release permit. notify_one retains that permit even when +// release precedes the wait. No lock or production instrumentation is involved. +struct UploadImage { + spec: ToolSpec, + requests: mpsc::UnboundedSender, + release: Arc, +} + +#[async_trait::async_trait] +impl Tool for UploadImage { + fn spec(&self) -> &ToolSpec { + &self.spec + } + + async fn invoke( + &self, + request: ToolRequest, + context: &mut ToolContext<'_>, + ) -> Result { + self.requests.send(request.input["file"].clone()).unwrap(); + let cancellation = context.cancellation.as_ref().expect("compose cancellation"); + tokio::select! { + biased; + _ = cancellation.cancelled() => return Err(ToolError::Cancelled), + _ = self.release.notified() => {} + } + Ok(ToolResult::new(ToolResultPart::success( + request.call_id, + ToolOutput::Structured(json!({"accepted": true})), + ))) + } +} + +struct Fixture { + root: tempfile::TempDir, + session: String, + bytes: Vec, +} + +impl Fixture { + fn new() -> Self { + let root = tempfile::tempdir().unwrap(); + let session = format!("managed-background-{}", uuid::Uuid::new_v4()); + let mut png = std::io::Cursor::new(Vec::new()); + image::DynamicImage::new_rgb8(2, 3) + .write_to(&mut png, image::ImageFormat::Png) + .unwrap(); + let bytes = png.into_inner(); + std::fs::write(root.path().join("image.png"), &bytes).unwrap(); + Self { + root, + session, + bytes, + } + } + + async fn detached_completion(&self, input: Value, resumed: bool) -> Value { + // Rebuild the complete compose source on each invocation. In particular, + // resume cannot rely on a prior resolver object retaining the source. + let (requests, mut received) = mpsc::unbounded_channel(); + let release = Arc::new(Notify::new()); + let children = ToolRegistry::new() + .with(crate::tools::ReadFileTool::new( + self.root.path().to_path_buf(), + )) + .with(UploadImage { + spec: ToolSpec::new( + ToolName::new("upload_image"), + "Upload an image reference to an external service", + json!({"type":"object", "properties":{"file":{"type":"object"}}, + "required":["file"], "additionalProperties":false}), + ) + .with_output_schema( + json!({"type":"object", "properties":{"accepted":{"type":"boolean"}}, + "required":["accepted"], "additionalProperties":false}), + ), + requests, + release: release.clone(), + }); + let inner = agentkit_tool_compose::ComposeTool::wrap(children.clone()) + .with_backend(super::super::HiddenRunletBackend(children.clone())); + let compose = super::super::ComposeOnly { + compose: inner.clone(), + backgroundable: BackgroundableCompose::new( + inner, + BackgroundJobs::default(), + self.root.path().to_path_buf(), + children, + ), + }; + let executor: Arc = Arc::new(BasicToolExecutor::new([ + Arc::new(compose) as Arc + ])); + let manager = super::super::background_task_manager(); + let tasks = manager.handle(); + let controller = CancellationController::new(); + let cancellation = controller.handle().checkpoint(); + let session_id = SessionId::new(self.session.clone()); + let turn_id = TurnId::new(if resumed { + "resumed-turn" + } else { + "originating-turn" + }); + let call_id = ToolCallId::new(if resumed { + "resumed-call" + } else { + "image-call" + }); + let permissions = Arc::new(AllowAllPermissions); + let resources: Arc = Arc::new(()); + let context = OwnedToolContext { + session_id: session_id.clone(), + turn_id: turn_id.clone(), + metadata: MetadataMap::new(), + permissions: permissions.clone(), + resources: resources.clone(), + cancellation: Some(cancellation.clone()), + execution_scope: Some(ToolExecutionScope { + executor: executor.clone(), + session_id: session_id.clone(), + turn_id: turn_id.clone(), + permissions, + resources, + cancellation: Some(cancellation.clone()), + }), + approved_request: None, + }; + let script = if resumed { + "receipt = upload_image({file: input})\nreturn {file: input, accepted: receipt.accepted}" + } else { + "file = read_file({path: \"image.png\"})\nreceipt = upload_image({file})\nreturn {file, accepted: receipt.accepted}" + }; + let start = manager + .start_task( + TaskLaunchRequest::plain( + None, + ToolRequest::new( + call_id.clone(), + ToolName::new("compose"), + json!({"script":script, "input":input, "background":true}), + session_id, + turn_id.clone(), + ), + ), + TaskStartContext { + executor, + tool_context: context, + }, + ) + .await + .unwrap(); + let TaskStartOutcome::Pending { + task_id, + kind: TaskKind::Foreground, + } = start + else { + panic!("expected the real foreground-then-detach route: {start:?}"); + }; + let reference = received.recv().await.expect("external upload request"); + assert_eq!(reference["$kit"], "file"); + let update = manager + .wait_for_turn(&turn_id, Some(cancellation.clone())) + .await + .unwrap(); + let Some(TurnTaskUpdate::Detached(snapshot)) = update else { + panic!("expected actual runner detach: {update:?}"); + }; + assert_eq!(snapshot.id, task_id); + assert_eq!(snapshot.kind, TaskKind::Background); + // The foreground runner is finished while the upload remains blocked. + assert!( + manager + .wait_for_turn(&turn_id, None) + .await + .unwrap() + .is_none() + ); + assert!( + manager + .take_pending_loop_updates() + .await + .unwrap() + .resolutions + .is_empty() + ); + if !resumed { + std::fs::remove_file(self.root.path().join("image.png")).unwrap(); + } + assert!(!self.root.path().join("image.png").exists()); + controller.interrupt(); + assert!(cancellation.is_cancelled()); + manager.on_turn_interrupted(&turn_id).await.unwrap(); + assert_eq!(tasks.list_running().await[0].kind, TaskKind::Background); + release.notify_one(); + tasks.wait_for_idle().await; + let mut updates = manager + .take_pending_loop_updates() + .await + .unwrap() + .resolutions; + assert_eq!( + updates.len(), + 1, + "one deferred completion must reach the loop" + ); + let TaskResolution::Item(item) = updates.pop_front().unwrap() else { + panic!("expected completion item"); + }; + let Part::ToolResult(result) = &item.parts[0] else { + panic!("expected tool result") + }; + assert_eq!(result.call_id, call_id); + let ToolOutput::Parts(parts) = &result.output else { + panic!( + "detached image was lost or stringified: {:?}", + result.output + ); + }; + let images: Vec<_> = parts + .iter() + .filter_map(|part| match part { + Part::Media(media) if media.modality == Modality::Image => Some(media), + _ => None, + }) + .collect(); + assert_eq!(images.len(), 1); + assert_eq!(images[0].mime_type, "image/png"); + assert_eq!(images[0].data, DataRef::InlineBytes(self.bytes.clone())); + assert!(parts.iter().any(|part| matches!(part, Part::Structured(_)))); + assert!( + manager + .take_pending_loop_updates() + .await + .unwrap() + .resolutions + .is_empty() + ); + assert!(tasks.list_running().await.is_empty()); + assert_eq!(tasks.list_completed().await[0].id, task_id); + reference + } +} + +impl Drop for Fixture { + fn drop(&mut self) { + let base = crate::artifacts::base(self.root.path()); + let _ = crate::resilient_fs::remove_dir_all( + base.with_file_name("files") + .join(blake3::hash(self.session.as_bytes()).to_hex().as_str()), + ); + let _ = crate::resilient_fs::remove_dir_all(crate::artifacts::session_directory( + &base, + &self.session, + )); + } +} + +#[tokio::test] +async fn managed_image_survives_runner_detach_turn_cancellation_and_resume() { + tokio::time::timeout(Duration::from_secs(15), async { + let fixture = Fixture::new(); + let reference = fixture.detached_completion(Value::Null, false).await; + // Cross a serialization boundary, as a stored descriptor does on resume. + let restored = serde_json::from_str(&serde_json::to_string(&reference).unwrap()).unwrap(); + assert_eq!(fixture.detached_completion(restored, true).await, reference); + }) + .await + .expect("real background task runner did not deliver managed image"); +} diff --git a/src/runtime/tests/managed_files.rs b/src/runtime/tests/managed_files.rs new file mode 100644 index 0000000..5f74054 --- /dev/null +++ b/src/runtime/tests/managed_files.rs @@ -0,0 +1,254 @@ +use super::*; +use agentkit_core::{DataRef, Item, Modality, ToolResultPart}; +use agentkit_http::Authentication; +use agentkit_loop::TurnRequest; +use agentkit_provider_openai::OpenAIResponsesConfig; +use base64::Engine as _; + +struct Fixture { + root: tempfile::TempDir, + session: String, + bytes: Vec, +} + +impl Fixture { + fn new() -> Self { + let root = tempfile::tempdir().unwrap(); + let session = format!( + "managed-runtime-{}", + blake3::hash(root.path().to_string_lossy().as_bytes()).to_hex() + ); + let mut encoded = std::io::Cursor::new(Vec::new()); + image::DynamicImage::new_rgb8(2, 3) + .write_to(&mut encoded, image::ImageFormat::Png) + .unwrap(); + let bytes = encoded.into_inner(); + std::fs::write(root.path().join("image.png"), &bytes).unwrap(); + Self { + root, + session, + bytes, + } + } + + async fn execute( + &self, + script: &str, + input: Value, + background: bool, + outcome: bool, + ) -> Result { + // 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(); + let compose = runtime.compose(0); + assert_eq!(compose.specs().len(), 1); + assert_eq!(compose.specs()[0].name.0, "compose"); + assert!(compose.specs()[0].description.contains("read_file")); + let source: Arc = Arc::new(compose.compose.clone()); + let executor: Arc = Arc::new(BasicToolExecutor::new([source])); + let permissions = Arc::new(AllowAllPermissions); + let resources: Arc = Arc::new(()); + let session_id = SessionId::new(self.session.clone()); + let turn_id = TurnId::new("turn"); + let owned = OwnedToolContext { + session_id: session_id.clone(), + turn_id: turn_id.clone(), + metadata: MetadataMap::new(), + permissions: permissions.clone(), + resources: resources.clone(), + cancellation: None, + execution_scope: Some(ToolExecutionScope { + executor, + session_id: session_id.clone(), + turn_id: turn_id.clone(), + permissions, + resources, + cancellation: None, + }), + approved_request: None, + }; + let request = ToolRequest::new( + ToolCallId::new("image-call"), + ToolName::new("compose"), + json!({"script":script,"input":input,"background":background}), + session_id, + turn_id, + ); + if outcome { + match compose + .backgroundable + .invoke_outcome(request, &mut owned.borrowed()) + .await + { + ToolExecutionOutcome::Completed(result) => Ok(result.result.output), + other => Err(format!("{other:?}")), + } + } else { + compose + .backgroundable + .invoke(request, &mut owned.borrowed()) + .await + .map(|result| result.result.output) + .map_err(|error| error.to_string()) + } + } +} + +impl Drop for Fixture { + fn drop(&mut self) { + let base = crate::artifacts::base(self.root.path()); + let directory = base + .with_file_name("files") + .join(blake3::hash(self.session.as_bytes()).to_hex().as_str()); + let _ = crate::resilient_fs::remove_dir_all(directory); + let _ = crate::resilient_fs::remove_dir_all(crate::artifacts::session_directory( + &base, + &self.session, + )); + } +} + +fn assert_image(output: &ToolOutput, bytes: &[u8]) { + let ToolOutput::Parts(parts) = output else { + panic!("expected selected image 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); + assert_eq!(images[0].data, DataRef::InlineBytes(bytes.to_vec())); +} + +#[tokio::test] +async fn managed_files_compose_reaches_native_wire_and_survives_restart() { + let fixture = Fixture::new(); + for (background, outcome) in [(false, false), (false, true), (true, false), (true, true)] { + let output = fixture + .execute( + "return read_file({ path: \"image.png\" })", + Value::Null, + background, + outcome, + ) + .await + .unwrap(); + assert_image(&output, &fixture.bytes); + let request = TurnRequest { + session_id: SessionId::new(fixture.session.clone()), + turn_id: TurnId::new("wire"), + transcript: vec![Item::new( + ItemKind::Tool, + vec![Part::ToolResult(ToolResultPart::success( + "image-call", + output, + ))], + )], + available_tools: Vec::new(), + cache: None, + metadata: MetadataMap::new(), + }; + let config = + OpenAIResponsesConfig::chatgpt_private("gpt-5.4", Authentication::bearer("test-key")); + let wire = config.encode_request(&request).unwrap(); + let blocks = wire["input"][0]["output"].as_array().unwrap(); + assert!(blocks.iter().any(|block| block["type"] == "input_text")); + let image = blocks + .iter() + .find(|block| block["type"] == "input_image") + .unwrap(); + let expected = format!( + "data:image/png;base64,{}", + base64::engine::general_purpose::STANDARD.encode(&fixture.bytes) + ); + assert_eq!(image["image_url"], expected); + } + let output = fixture + .execute( + "return read_file({ path: \"image.png\" })", + Value::Null, + false, + false, + ) + .await + .unwrap(); + let ToolOutput::Parts(parts) = output else { + panic!() + }; + let Part::Structured(reference) = &parts[0] else { + panic!() + }; + std::fs::remove_file(fixture.root.path().join("image.png")).unwrap(); + let resumed = fixture + .execute( + "return { nested: [input, input] }", + reference.value.clone(), + true, + true, + ) + .await + .unwrap(); + assert_image(&resumed, &fixture.bytes); +} + +#[tokio::test] +async fn managed_files_select_only_returned_references_and_preserve_spilled_images() { + let fixture = Fixture::new(); + let output = fixture + .execute( + "image = read_file({ path: \"image.png\" })\nreturn { ok: true }", + Value::Null, + false, + true, + ) + .await + .unwrap(); + assert_eq!(output, ToolOutput::structured(json!({"ok":true}))); + let output = fixture.execute("image = read_file({ path: \"image.png\" })\nreturn { text: input, nested: [image, image] }", json!("large text ".repeat(3000)), true, true).await.unwrap(); + assert_image(&output, &fixture.bytes); + let ToolOutput::Parts(parts) = output else { + panic!() + }; + let Part::Structured(spill) = &parts[0] else { + panic!() + }; + assert!(spill.value["artifact"].is_string()); + let text_bytes = serde_json::to_vec(&spill.value).unwrap().len() + + parts + .iter() + .filter_map(|part| match part { + Part::Text(text) => Some(text.text.len()), + _ => None, + }) + .sum::(); + assert!(text_bytes <= 8192); + assert!(matches!(&parts[1], Part::Text(text) if text.text.contains("/nested/0"))); + assert!(matches!(&parts[2], Part::Media(_))); + let artifact = std::fs::read_to_string(spill.value["artifact"].as_str().unwrap()).unwrap(); + assert!(!artifact.contains("InlineBytes")); + assert!(!artifact.contains("inline_bytes")); +} + +#[tokio::test] +async fn managed_files_delivery_failure_does_not_claim_rollback() { + let fixture = Fixture::new(); + let error = fixture + .execute( + "_ = edit({ op: \"add\", path: \"effect.txt\", content: \"done\" })\nreturn input", + json!({"$kit":"file","version":999}), + false, + true, + ) + .await + .unwrap_err(); + assert!(error.contains("compose program completed")); + assert!(error.contains("do not rerun blindly")); + assert_eq!( + std::fs::read_to_string(fixture.root.path().join("effect.txt")).unwrap(), + "done" + ); +} diff --git a/src/tools/mod.rs b/src/tools/mod.rs index e776eff..df85a1f 100644 --- a/src/tools/mod.rs +++ b/src/tools/mod.rs @@ -4,6 +4,7 @@ mod docs; mod edit; pub(crate) mod mcp; mod observed; +mod read_file; mod shell; mod subagent; @@ -15,5 +16,6 @@ pub use edit::EditTool; pub use mcp::{AuthTool, McpTool, ToolSearch}; pub use observed::Observed; pub(crate) use observed::shared as observe_shared; +pub use read_file::ReadFileTool; pub use shell::ShellTool; pub use subagent::{CloseTool, ForkTool, PromptTool, SubagentTool, Subagents, SubagentsTool}; diff --git a/src/tools/read_file.rs b/src/tools/read_file.rs new file mode 100644 index 0000000..264c624 --- /dev/null +++ b/src/tools/read_file.rs @@ -0,0 +1,111 @@ +use std::path::PathBuf; + +use agentkit_core::{ToolOutput, ToolResultPart}; +use agentkit_tools_core::{ + Tool, ToolAnnotations, ToolContext, ToolError, ToolName, ToolRequest, ToolResult, ToolSpec, +}; +use async_trait::async_trait; +use serde::Deserialize; +use serde_json::{Map, Value}; + +use crate::managed_files::FileStore; + +#[derive(Clone)] +pub struct ReadFileTool { + root: PathBuf, + store: FileStore, + spec: ToolSpec, +} + +impl ReadFileTool { + pub fn new(root: PathBuf) -> Self { + Self { + store: FileStore::new(&root), + root, + spec: ToolSpec::new( + ToolName::new("read_file"), + "Import a regular local PNG or JPEG image as a durable immutable File reference. Only File references reachable from the final compose return deliver pixels; intermediate references remain private. Maximum 8 MiB, 8192 pixels per dimension and 16 megapixels; animation is unsupported. Source bytes, orientation and metadata are preserved. Files are session-scoped and survive restart and source deletion; other sessions have no implicit access. This is not a text-file reader.", + object_schema([ ("path", object([("type", Value::from("string")), ("minLength", Value::from(1)), ("maxLength", Value::from(4096))])) ]), + ) + .with_output_schema(object_schema([ + ("$kit", object([("type", Value::from("string")), ("enum", Value::Array(vec![Value::from("file")]))])), + ("version", object([("type", Value::from("integer")), ("enum", Value::Array(vec![Value::from(1)]))])), + ("id", object([("type", Value::from("string")), ("pattern", Value::from("^file_[0-9a-f]{64}$"))])), + ("name", object([("type", Value::from("string")), ("minLength", Value::from(1)), ("maxLength", Value::from(255))])), + ("mime_type", object([("type", Value::from("string")), ("enum", Value::Array(vec![Value::from("image/png"), Value::from("image/jpeg")]))])), + ("size_bytes", positive_integer(8_388_608)), + ("image", object_schema([("width", positive_integer(8192)), ("height", positive_integer(8192))])), + ])) + .with_annotations(ToolAnnotations::read_only()), + } + } +} + +fn object(fields: [(&str, Value); N]) -> Value { + Value::Object(Map::from_iter( + fields + .into_iter() + .map(|(key, value)| (key.to_owned(), value)), + )) +} + +fn object_schema(fields: [(&str, Value); N]) -> Value { + let required = Value::Array(fields.iter().map(|(key, _)| Value::from(*key)).collect()); + object([ + ("type", Value::from("object")), + ("properties", object(fields)), + ("required", required), + ("additionalProperties", Value::from(false)), + ]) +} + +fn positive_integer(maximum: u64) -> Value { + object([ + ("type", Value::from("integer")), + ("minimum", Value::from(1)), + ("maximum", Value::from(maximum)), + ]) +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct Input { + path: String, +} + +#[async_trait] +impl Tool for ReadFileTool { + fn spec(&self) -> &ToolSpec { + &self.spec + } + + async fn invoke( + &self, + request: ToolRequest, + context: &mut ToolContext<'_>, + ) -> Result { + let input: Input = serde_json::from_value(request.input) + .map_err(|error| ToolError::InvalidInput(error.to_string()))?; + if input.path.is_empty() || input.path.len() > 4096 { + return Err(ToolError::InvalidInput( + "path must contain 1 to 4096 UTF-8 bytes".into(), + )); + } + let path = self.root.join(input.path); + let store = self.store.clone(); + let cancellation = context.cancellation.clone(); + let session = request.session_id.0; + let reference = tokio::task::spawn_blocking(move || { + store.import(&session, &path, cancellation.as_ref()) + }) + .await + .map_err(|error| ToolError::Internal(error.to_string()))? + .map_err(ToolError::ExecutionFailed)?; + let value = serde_json::to_value(reference) + .map_err(|error| ToolError::Internal(error.to_string()))?; + Ok(ToolResult::new(ToolResultPart::success( + request.call_id, + ToolOutput::structured(value), + ))) + } +} diff --git a/src/tui/app.rs b/src/tui/app.rs index 60835c1..f0d4f30 100644 --- a/src/tui/app.rs +++ b/src/tui/app.rs @@ -26,9 +26,10 @@ use crate::events::{GenerationOutcome, RuntimeEvent, SubagentStatus}; use crate::file_search::FileMatch; const MAX_TOOL_OUTPUT_LINES: usize = 5_000; +pub(super) const MAX_TOOL_IMAGES: usize = 32; const MAX_IMAGE_BASE64_BYTES: usize = 14 * 1024 * 1024; const MAX_IMAGE_SOURCE_BYTES: usize = 10 * 1024 * 1024; -const MAX_RETAINED_IMAGE_SOURCE_BYTES: usize = 32 * 1024 * 1024; +pub(super) const MAX_RETAINED_IMAGE_SOURCE_BYTES: usize = 32 * 1024 * 1024; use super::{ command::{self, Command as SlashCommand, Parsed, known_token, parse}, @@ -99,6 +100,7 @@ pub enum Update { status: Option, script: Option, output: Option>, + images: Option>, append_output: bool, intent: Option>, backgrounded: bool, @@ -362,6 +364,8 @@ pub struct ToolCall { pub children: Vec, /// Raw tool output, kept whole but folded away until asked for. pub output: Vec, + /// Typed tool-result images sharing the user-image retention and decode budgets. + pub images: Vec, /// User-facing summary supplied by a compose caller. pub intent: Option, pub expanded: bool, @@ -521,7 +525,7 @@ pub enum Block { started: Instant, millis: Option, }, - Tool(ToolCall), + Tool(Box), TurnDuration(u64), Notice(String), Error(String), @@ -1193,14 +1197,14 @@ impl App { if self.transcript_revisions.len() != self.blocks.len() || self.focused_call_id.is_some() { if let Some(id) = &self.focused_call_id && let Some(call) = self.blocks.iter().rev().find_map(|block| match block { - Block::Tool(call) if &call.id == id => Some(call), + Block::Tool(call) if &call.id == id => Some(call.as_ref()), _ => None, }) { return Some(call); } return self.blocks.iter().rev().find_map(|block| match block { - Block::Tool(call) => Some(call), + Block::Tool(call) => Some(call.as_ref()), _ => None, }); } @@ -1214,7 +1218,7 @@ impl App { && call.running() && !call.backgrounded => { - Some(call) + Some(call.as_ref()) } _ => None, }) @@ -1707,7 +1711,7 @@ impl App { self.close_thought(); self.prepare_focused_call(id.clone()); let expanded = title == agentkit_tool_compose::COMPOSE_TOOL_NAME; - self.push_block(Block::Tool(ToolCall { + self.push_block(Block::Tool(Box::new(ToolCall { id, title, kind, @@ -1718,12 +1722,13 @@ impl App { progress: Box::default(), children: Vec::new(), output: Vec::new(), + images: Vec::new(), intent: None, expanded, compose_view: ComposeView::Output, expansion_explicit: false, backgrounded, - })); + }))); } Update::ToolPatched { id, @@ -1732,6 +1737,7 @@ impl App { status, script, output, + images, append_output, intent, backgrounded, @@ -1745,6 +1751,36 @@ impl App { backgrounded, }); } + if let Some(images) = images + && let Some(index) = self.call_index(&id) + && let Block::Tool(call) = &mut self.blocks[index] + { + if !append_output { + let replaced = call + .images + .iter() + .map(|image| image.data.len()) + .sum::(); + self.retained_image_source_bytes = + self.retained_image_source_bytes.saturating_sub(replaced); + call.images.clear(); + } + for image in images { + if call.images.len() >= MAX_TOOL_IMAGES { + break; + } + if call.images.iter().any(|existing| existing.key == image.key) { + continue; + } + let retained = self + .retained_image_source_bytes + .saturating_add(image.data.len()); + if retained <= MAX_RETAINED_IMAGE_SOURCE_BYTES { + self.retained_image_source_bytes = retained; + call.images.push(image); + } + } + } let Some(call) = self.call_mut(&id) else { return; }; @@ -4203,6 +4239,45 @@ mod tests { assert!(UserImage::new(decoded_too_large, "image/png".into(), 0).is_none()); } + #[test] + fn tool_images_share_user_budget_and_release_replaced_sources() { + let mut app = app(); + let bytes = 9 * 1024 * 1024; + let image = |byte: char| { + UserImage::new(byte.to_string().repeat(bytes), "image/png".into(), 0).unwrap() + }; + app.apply(Update::UserMessage { + id: "user".into(), + text: "[Image]".into(), + images: vec![image('A')], + append: false, + }); + let patch = |images, append_output| Update::ToolPatched { + id: "tool".into(), + title: None, + kind: None, + status: None, + script: None, + output: Some(vec!["[Image]".into()]), + images: Some(images), + append_output, + intent: None, + backgrounded: false, + }; + app.apply(patch(vec![image('B'), image('C'), image('D')], false)); + assert_eq!(app.retained_image_source_bytes, 3 * bytes); + let Block::Tool(call) = &app.blocks[1] else { + panic!("expected tool") + }; + assert_eq!(call.images.len(), 2); + app.apply(patch(vec![image('B')], true)); + assert_eq!(app.retained_image_source_bytes, 3 * bytes); + app.apply(patch(vec![image('D')], false)); + assert_eq!(app.retained_image_source_bytes, 2 * bytes); + app.apply(patch(Vec::new(), false)); + assert_eq!(app.retained_image_source_bytes, bytes); + } + #[test] fn retained_user_image_sources_have_an_aggregate_bound() { let source_bytes = 9 * 1024 * 1024; @@ -4398,6 +4473,7 @@ mod tests { status: Some(ToolCallStatus::Failed), script: None, output: None, + images: None, append_output: false, intent: None, backgrounded: false, @@ -4561,6 +4637,7 @@ mod tests { status: Some(ToolCallStatus::Completed), script: None, output: None, + images: None, append_output: false, intent: None, backgrounded: false, @@ -6043,6 +6120,7 @@ mod tests { status: Some(ToolCallStatus::Completed), script: None, output: None, + images: None, append_output: false, intent: None, backgrounded: false, @@ -6481,6 +6559,7 @@ mod tests { status: Some(ToolCallStatus::Completed), script: None, output: None, + images: None, append_output: false, intent: None, backgrounded: true, diff --git a/src/tui/mod.rs b/src/tui/mod.rs index ca5553a..6d76308 100644 --- a/src/tui/mod.rs +++ b/src/tui/mod.rs @@ -2579,6 +2579,13 @@ fn translate(notification: UpdateSessionNotification) -> (String, Vec) { MessageKind::Thought, ), SessionUpdate::ToolCallUpdate(update) => { + let images = match &update.content { + MaybeUndefined::Value(content) => Some(tool_images_of(content)), + MaybeUndefined::Null => Some(Vec::new()), + // Raw output is a parallel representation, not an authoritative + // content replacement. Only explicit content patches clear pixels. + MaybeUndefined::Undefined => None, + }; let output = match &update.content { MaybeUndefined::Value(content) => Some(output_of(Some(content))), MaybeUndefined::Null => Some(Vec::new()), @@ -2628,12 +2635,14 @@ fn translate(notification: UpdateSessionNotification) -> (String, Vec) { }, script, output, + images, append_output: false, intent, backgrounded, }] } SessionUpdate::ToolCallContentChunk(chunk) => { + let images = tool_images_of(std::slice::from_ref(&chunk.content)); let output = output_of(Some(std::slice::from_ref(&chunk.content))); let backgrounded = output .iter() @@ -2645,6 +2654,7 @@ fn translate(notification: UpdateSessionNotification) -> (String, Vec) { status: None, script: None, output: Some(output), + images: Some(images), append_output: true, intent: None, backgrounded, @@ -2820,19 +2830,152 @@ fn intent_of(input: &Value) -> Option { /// A tool call's output as readable lines, kept whole for the folded card. fn raw_output_lines(output: &Value) -> Vec { + let output = raw_output_without_media(output, 0); if let Some(text) = output .as_str() .or_else(|| output.get("text").and_then(Value::as_str)) { return readable(text); } - serde_json::to_string_pretty(output) + serde_json::to_string_pretty(&output) .unwrap_or_else(|_| output.to_string()) .lines() .map(str::to_string) .collect() } +// ACP raw output can serialize ToolOutput::Parts (Media/DataRef), ACP content, +// or a JSON-encoded version of either. Do not print their pixels in text-only +// cards. This fallback deliberately does not decode, fetch, or replace images. +fn raw_output_without_media(output: &Value, depth: usize) -> Value { + if depth >= 64 { + return Value::String("[Truncated output]".into()); + } + match output { + Value::Object(object) => { + let is_image = (object.contains_key("data") + && ["mime_type", "mimeType"].iter().any(|key| { + object + .get(*key) + .and_then(Value::as_str) + .is_some_and(|mime| mime.starts_with("image/")) + })) + || ["type", "modality"].iter().any(|key| { + object + .get(*key) + .and_then(Value::as_str) + .is_some_and(|kind| kind.eq_ignore_ascii_case("image")) + }); + if is_image { + return Value::String("[Image]".into()); + } + if object.contains_key("InlineBytes") || object.contains_key("InlineText") { + return Value::String("[Media]".into()); + } + Value::Object( + object + .iter() + .take(MAX_OUTPUT_LINES) + .map(|(key, value)| (key.clone(), raw_output_without_media(value, depth + 1))) + .collect(), + ) + } + Value::Array(values) => Value::Array( + values + .iter() + .take(MAX_OUTPUT_LINES) + .map(|value| raw_output_without_media(value, depth + 1)) + .collect(), + ), + Value::String(text) => { + if let Ok(value) = serde_json::from_str::(text) { + raw_output_without_media(&value, depth + 1) + } else { + Value::String(redact_image_data_urls(text)) + } + } + _ => output.clone(), + } +} + +// Recognize bounded data-URL headers and consume only their payload spans. +// A literal scheme mention is not an image, and diagnostics outside a URL +// must survive redaction. No payload decoding or network access is needed. +fn redact_image_data_urls(text: &str) -> String { + const PREFIX: &str = "data:image/"; + let mut result = String::new(); + let mut copied = 0; + for (start, _) in text.match_indices(PREFIX) { + if start < copied { + continue; + } + let tail = &text[start + PREFIX.len()..]; + let Some(header_len) = tail + .bytes() + .take(513) + .position(|byte| !byte.is_ascii_alphanumeric() && !b"+.-;=_%".contains(&byte)) + else { + continue; + }; + let header = &tail[..header_len]; + if header.is_empty() || header.starts_with(';') || tail.as_bytes()[header_len] != b',' { + continue; + } + let payload = &tail[header_len + 1..]; + let base64 = header + .rsplit(';') + .next() + .is_some_and(|part| part.eq_ignore_ascii_case("base64")); + let payload_len = payload + .bytes() + .take_while(|byte| { + byte.is_ascii_alphanumeric() + || if base64 { + b"+/=%".contains(byte) + } else { + b"%:@!$&*+-./;=?_~,".contains(byte) + } + }) + .count(); + if payload_len == 0 { + continue; + } + result.push_str(&text[copied..start]); + result.push_str("[Image]"); + copied = start + PREFIX.len() + header_len + 1 + payload_len; + } + result.push_str(&text[copied..]); + result +} + +// Keep pixels separate from text previews. Live chunks and replay snapshots +// use this path; never fetch a model-supplied URI. +fn tool_images_of(content: &[ToolCallContent]) -> Vec { + let mut images: Vec = Vec::new(); + let mut retained = 0usize; + for entry in content { + if images.len() >= app::MAX_TOOL_IMAGES { + break; + } + let ToolCallContent::Content(content) = entry else { + continue; + }; + let ContentBlock::Image(image) = &content.content else { + continue; + }; + if retained.saturating_add(image.data.len()) > app::MAX_RETAINED_IMAGE_SOURCE_BYTES { + continue; + } + if let Some(image) = UserImage::new(image.data.clone(), image.mime_type.to_string(), 0) + && !images.iter().any(|existing| existing.key == image.key) + { + retained += image.data.len(); + images.push(image); + } + } + images +} + fn output_of(content: Option<&[ToolCallContent]>) -> Vec { let mut output = Vec::new(); for entry in content.unwrap_or_default() { @@ -3921,6 +4064,8 @@ mod tests { assert_eq!(user.images.len(), 1); assert_eq!(user.images[0].data, "c2VjcmV0"); assert_eq!(summary, "summary"); + assert_eq!(tool.images.len(), 1); + assert_eq!(tool.images[0].data, "c2VjcmV0"); assert_eq!(tool.status, wire::ToolCallStatus::Completed); assert_eq!( tool.output, @@ -3969,6 +4114,210 @@ mod tests { } } + #[test] + fn tool_images_survive_live_chunks_and_replay_replacement() { + use super::app::Block; + let mut app = App::new( + PathBuf::from("/tmp"), + "provider".into(), + "model".into(), + "a2a".into(), + ); + let image = |data: &str| { + wire::ToolCallContent::Content(Box::new(wire::Content::new(ContentBlock::Image( + wire::ImageContent::new(data, "image/png"), + )))) + }; + let snapshots = [ + SessionUpdate::ToolCallContentChunk(wire::ToolCallContentChunk::new( + "tool", + image("AQID"), + )), + SessionUpdate::ToolCallContentChunk(wire::ToolCallContentChunk::new( + "tool", + image("AQID"), + )), + SessionUpdate::ToolCallContentChunk(wire::ToolCallContentChunk::new( + "tool", + image("BAUG"), + )), + SessionUpdate::ToolCallUpdate( + wire::ToolCallUpdate::new("tool").content(vec![image("AQID"), image("BAUG")]), + ), + SessionUpdate::ToolCallUpdate( + wire::ToolCallUpdate::new("tool").status(wire::ToolCallStatus::Completed), + ), + SessionUpdate::ToolCallUpdate( + wire::ToolCallUpdate::new("tool").content(vec![image("BAUG")]), + ), + SessionUpdate::ToolCallUpdate( + wire::ToolCallUpdate::new("tool").content(Vec::::new()), + ), + ]; + for (notification, expected) in snapshots.into_iter().zip([1, 1, 2, 2, 2, 1, 0]) { + for update in translate_for_session( + UpdateSessionNotification::new("session", notification), + "session", + ) { + app.apply(update); + } + let Block::Tool(tool) = &app.blocks[0] else { + panic!("expected tool") + }; + assert_eq!(tool.images.len(), expected); + assert!( + tool.output + .iter() + .all(|line| !line.contains("AQID") && !line.contains("BAUG")) + ); + } + } + + #[test] + fn raw_tool_output_patches_preserve_images_and_never_print_pixels() { + use super::app::Block; + use agentkit_core::{DataRef, Modality, Part, ToolOutput}; + use serde_json::Value; + + let mut app = App::new( + PathBuf::from("/tmp"), + "provider".into(), + "model".into(), + "a2a".into(), + ); + let content = wire::ToolCallContent::Content(Box::new(wire::Content::new( + ContentBlock::Image(wire::ImageContent::new("c2VjcmV0", "image/png")), + ))); + let initial = wire::ToolCallUpdate::new("tool").content(vec![content]); + for update in translate_for_session( + UpdateSessionNotification::new("session", SessionUpdate::ToolCallUpdate(initial)), + "session", + ) { + app.apply(update); + } + let inline = serde_json::to_value(ToolOutput::parts(vec![ + Part::text("image result"), + Part::media( + Modality::Image, + "image/png", + DataRef::inline_text("c2VjcmV0"), + ), + Part::media( + Modality::Image, + "image/png", + DataRef::inline_bytes(vec![231, 232, 233]), + ), + ])) + .unwrap(); + let raw_outputs = [ + inline.clone(), + Value::String(serde_json::to_string(&inline).unwrap()), + json!({"output": inline}), + json!({"content": [{"type": "image", "data": "c2VjcmV0", "mimeType": "image/png"}]}), + json!({"uri": "data:image/png;base64,c2VjcmV0"}), + json!({"text": "data:image/png;base64,c2VjcmV0"}), + json!({"text": "ordinary result"}), + Value::Null, + ]; + for raw in raw_outputs { + let notification = UpdateSessionNotification::new( + "session", + SessionUpdate::ToolCallUpdate( + wire::ToolCallUpdate::new("tool").raw_output(raw.clone()), + ), + ); + for update in translate_for_session(notification, "session") { + app.apply(update); + } + let Block::Tool(tool) = &app.blocks[0] else { + panic!("expected tool") + }; + assert_eq!(tool.images.len(), 1); + assert_eq!(tool.images[0].data, "c2VjcmV0"); + let text = tool.output.join("\n"); + assert!(!text.contains("c2VjcmV0"), "{text}"); + assert!(!text.contains("231"), "{text}"); + assert!(!text.contains("InlineBytes"), "{text}"); + assert!(!text.contains("data:image/"), "{text}"); + + // Raw-only replay uses the same safe fallback even without a prior + // typed content notification. It must not pretend to reconstruct pixels. + let notification = UpdateSessionNotification::new( + "session", + SessionUpdate::ToolCallUpdate( + wire::ToolCallUpdate::new("raw-only").raw_output(raw), + ), + ); + assert!( + matches!(translate_for_session(notification, "session").as_slice(), + [Update::ToolPatched { images: None, output: Some(output), .. }] + if output.iter().all(|line| !line.contains("c2VjcmV0") && !line.contains("231")) + ) + ); + } + } + + #[test] + fn raw_tool_output_redacts_only_image_spans_and_preserves_diagnostics() { + let cases = [ + ( + "Screenshot: data:image/png;base64,AQID\nUpload failed: permission denied", + "Screenshot: [Image]\nUpload failed: permission denied", + ), + ( + "The literal data:image/ is a URI prefix, not an image.", + "The literal data:image/ is a URI prefix, not an image.", + ), + ( + "Incomplete data:image/png;base64,\nUpload failed", + "Incomplete data:image/png;base64,\nUpload failed", + ), + ( + "First (data:image/png;base64,AQID), second \"data:image/jpeg;base64,BAUG\". Failed.", + "First ([Image]), second \"[Image]\". Failed.", + ), + ( + "Encoded: data:image/svg+xml,%3Csvg%3E\nUpload failed", + "Encoded: [Image]\nUpload failed", + ), + ( + "Image: data:image/svg+xml,%3Csvg%3E,%3C/svg%3E\nUpload failed", + "Image: [Image]\nUpload failed", + ), + ("data:image/png;base64,AQID", "[Image]"), + ]; + for (text, expected) in cases { + // Both native raw objects and JSON-encoded raw output occur in live + // and replayed updates. Neither may discard non-image diagnostics. + for raw in [ + json!({"text": text}), + json!(json!({"text": text}).to_string()), + ] { + let notification = UpdateSessionNotification::new( + "session", + SessionUpdate::ToolCallUpdate( + wire::ToolCallUpdate::new("raw-only").raw_output(raw), + ), + ); + let updates = translate_for_session(notification, "session"); + let [ + Update::ToolPatched { + images: None, + output: Some(lines), + .. + }, + ] = updates.as_slice() + else { + panic!("expected raw-only tool patch"); + }; + assert_eq!(lines.join("\n"), expected, "source: {text}"); + assert!(!lines.iter().any(|line| line.contains("AQID") + || line.contains("BAUG") + || line.contains("%3Csvg"))); + } + } + } + #[test] fn translates_replayed_raw_tool_output() { let update = UpdateSessionNotification::new( diff --git a/src/tui/progress_tests.rs b/src/tui/progress_tests.rs index a536e1b..874dfac 100644 --- a/src/tui/progress_tests.rs +++ b/src/tui/progress_tests.rs @@ -161,6 +161,7 @@ fn authoritative_progress_healed_mismatch_and_unicode_spans_stay_neutral() { status: None, script: Some(script.into()), output: None, + images: None, append_output: false, intent: None, backgrounded: false, @@ -201,6 +202,7 @@ fn authoritative_progress_incomplete_and_parent_completion_never_fabricate_succe status: Some(agent_client_protocol::schema::v2::ToolCallStatus::Completed), script: None, output: None, + images: None, append_output: false, intent: None, backgrounded: false, @@ -416,6 +418,7 @@ async fn authoritative_progress_real_compose_bridge_to_diagnostic_app_and_render status: None, script: Some(source.into()), output: None, + images: None, append_output: false, intent: None, backgrounded: false, @@ -449,6 +452,7 @@ async fn authoritative_progress_real_bridge_lag_invalidates_all_observations() { status: None, script: Some(source.into()), output: None, + images: None, append_output: false, intent: None, backgrounded: false, @@ -470,6 +474,7 @@ fn authoritative_progress_two_subagents_dependency_is_not_descendant_completion( status: None, script: Some(source.into()), output: None, + images: None, append_output: false, intent: None, backgrounded: false, @@ -621,6 +626,7 @@ async fn authoritative_progress_real_bridge_cancellation_invalidates() { status: None, script: Some(source.into()), output: None, + images: None, append_output: false, intent: None, backgrounded: false, @@ -652,6 +658,7 @@ async fn authoritative_progress_real_bridge_healing_stays_neutral() { status: None, script: Some(source.into()), output: None, + images: None, append_output: false, intent: None, backgrounded: false, @@ -674,6 +681,7 @@ async fn authoritative_progress_real_iterations_remain_distinct() { status: None, script: Some(source.into()), output: None, + images: None, append_output: false, intent: None, backgrounded: false, @@ -696,6 +704,7 @@ fn authoritative_progress_inline_nested_multiline_and_unicode_ranges() { status: None, script: Some(source.into()), output: None, + images: None, append_output: false, intent: None, backgrounded: false, @@ -952,6 +961,7 @@ fn authoritative_progress_terminal_conflicts_invalidate_completed_display() { status: Some(agent_client_protocol::schema::v2::ToolCallStatus::Completed), script: None, output: None, + images: None, append_output: false, intent: None, backgrounded: false, diff --git a/src/tui/ui.rs b/src/tui/ui.rs index 2a97dfe..64117d5 100644 --- a/src/tui/ui.rs +++ b/src/tui/ui.rs @@ -999,10 +999,12 @@ fn draw_transcript(frame: &mut Frame<'_>, app: &mut App, images: &mut ImageRunti frame.render_widget(Paragraph::new(visible), inner); draw_selection(frame, app, inner, offset, &row_widths); for (block_index, source_index, y) in visible_images { - let Some(Block::User(message)) = app.blocks.get(block_index) else { - continue; + let sources = match app.blocks.get(block_index) { + Some(Block::User(message)) => &message.images, + Some(Block::Tool(call)) => &call.images, + _ => continue, }; - let Some(source) = message.images.get(source_index) else { + let Some(source) = sources.get(source_index) else { continue; }; if let Some(image) = images.prepare(source, inner.width.max(1)) { @@ -1260,7 +1262,39 @@ fn transcript_block_rows( (line, (call.clone(), code, Some(line_index))) }) .collect::>(); - (wrap_linked_tagged(&lines, width), Vec::new()) + let mut rows = wrap_linked_tagged(&lines, width); + let mut placements = Vec::new(); + if let Block::Tool(call) = block + && call.expanded + && (!call.is_compose() || call.compose_view == ComposeView::Output) + { + for (source, image) in call.images.iter().enumerate() { + rows.extend(wrap_linked_tagged( + &[( + LinkedLine::plain(Line::from(Span::styled( + format!(" [Image: {}]", image.mime_type), + theme::dim(), + ))), + (Some(call.id.clone()), None, None), + )], + width, + )); + if !reserve_images { + continue; + } + let row = rows.len(); + rows.extend((0..RESERVED_ROWS).map(|_| { + ( + Line::default(), + (Some(call.id.clone()), None, None), + Vec::new(), + String::new(), + ) + })); + placements.push(CachedTranscriptImage { source, row }); + } + } + (rows, placements) } fn uncopyable(lines: Vec) -> Vec<(LinkedLine, Option>)> { @@ -3312,6 +3346,7 @@ mod tests { status: None, script: None, output: None, + images: None, append_output: false, intent: Some(Some(" Check every source file. ".into())), backgrounded: false, @@ -3338,6 +3373,7 @@ mod tests { app.apply(Update::ToolPatched { title: None, kind: None, + images: None, append_output: false, intent: None, id: "call-1".into(), @@ -3372,6 +3408,7 @@ mod tests { status: Some(agent_client_protocol::schema::v2::ToolCallStatus::Completed), script: Some("return 1".into()), output: Some(vec!["1".into()]), + images: None, append_output: false, intent: None, backgrounded: false, @@ -3383,6 +3420,7 @@ mod tests { status: None, script: None, output: None, + images: None, append_output: false, intent: None, backgrounded: false, @@ -3399,6 +3437,7 @@ mod tests { app.apply(Update::ToolPatched { title: None, kind: None, + images: None, append_output: false, intent: None, id: "call-1".into(), @@ -3441,6 +3480,7 @@ mod tests { app.apply(Update::ToolPatched { title: None, kind: None, + images: None, append_output: false, intent: None, id: "call-1".into(), @@ -3479,6 +3519,7 @@ mod tests { app.apply(Update::ToolPatched { title: None, kind: None, + images: None, append_output: false, intent: None, id: "call-1".into(), @@ -3591,6 +3632,7 @@ mod tests { app.apply(Update::ToolPatched { title: None, kind: None, + images: None, append_output: false, intent: None, id: "call-1".into(), @@ -3751,6 +3793,7 @@ mod tests { app.apply(Update::ToolPatched { title: None, kind: None, + images: None, append_output: false, intent: None, id: "call-1".into(), @@ -4281,6 +4324,64 @@ mod tests { ); } + #[test] + fn tool_images_render_with_shared_runtime_and_text_fallback() { + let mut png = std::io::Cursor::new(Vec::new()); + image::DynamicImage::new_rgb8(4, 2) + .write_to(&mut png, image::ImageFormat::Png) + .unwrap(); + let source = UserImage::new( + base64::engine::general_purpose::STANDARD.encode(png.into_inner()), + "image/png".into(), + 0, + ) + .unwrap(); + let mut app = App::new( + PathBuf::from("/tmp"), + "provider".into(), + "model".into(), + "a2a".into(), + ); + app.apply(Update::ToolPatched { + id: "tool".into(), + title: Some("compose".into()), + kind: None, + status: None, + script: None, + output: Some(vec!["[Image]".into()]), + images: Some(vec![source]), + append_output: false, + intent: None, + backgrounded: false, + }); + let mut images = ImageRuntime::with_picker(Picker::halfblocks()); + refresh_transcript_cache_with_images(&mut app, &mut images, 40); + assert_eq!(app.transcript_cache[0].as_ref().unwrap().images.len(), 1); + assert_eq!(images.cached_entries(), 0); + let mut terminal = Terminal::new(TestBackend::new(60, 40)).unwrap(); + terminal + .draw(|frame| draw(frame, &mut app, &mut images)) + .unwrap(); + assert_eq!(images.cached_entries(), 1); + let mut disabled = ImageRuntime::disabled(); + terminal + .draw(|frame| draw(frame, &mut app, &mut disabled)) + .unwrap(); + let cached = app.transcript_cache[0].as_ref().unwrap(); + assert!(cached.images.is_empty()); + assert!( + cached + .rows + .iter() + .any(|row| line_text(&row.0).contains("[Image: image/png]")) + ); + if let Block::Tool(call) = &mut app.blocks[0] { + call.expanded = false; + } + let (_, placements) = super::transcript_block_rows(&app, 0, 40, true); + assert!(placements.is_empty()); + } + #[test] fn image_rows_are_fixed_and_decoding_is_lazy() { let mut png = std::io::Cursor::new(Vec::new());