diff --git a/Cargo.lock b/Cargo.lock index 6520df03f..2d2628c6b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -32,6 +32,15 @@ version = "1.0.104" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" +[[package]] +name = "arbitrary" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3d036a3c4ab069c7b410a2ce876bd74808d2d0888a82667669f8e783a898bf1" +dependencies = [ + "derive_arbitrary", +] + [[package]] name = "async-trait" version = "0.1.92" @@ -67,12 +76,30 @@ version = "0.23.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5" +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + [[package]] name = "bitflags" version = "2.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" +[[package]] +name = "bitvec" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddcec3d12c579d40898fe0a9a358a803c23e9c52ca3c425707f81c9436211837" +dependencies = [ + "funty", + "radium", + "tap", + "wyz", +] + [[package]] name = "block-buffer" version = "0.12.1" @@ -88,6 +115,12 @@ version = "3.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" +[[package]] +name = "bytemuck" +version = "1.25.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95832e849adfb21180ccb6826a99da14e5d266ae5c2e668e1602cf234f153797" + [[package]] name = "bytes" version = "1.12.1" @@ -170,6 +203,12 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "crossbeam-utils" +version = "0.8.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a31eee39dddec8330830986fcd7625edb5a24ec90ea038215273bbc3adb08ac6" + [[package]] name = "crypto-common" version = "0.2.2" @@ -179,6 +218,17 @@ dependencies = [ "hybrid-array", ] +[[package]] +name = "derive_arbitrary" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e567bd82dcff979e4b03460c307b3cdc9e96fde3d73bed1496d2bc75d9dd62a" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + [[package]] name = "digest" version = "0.11.3" @@ -262,6 +312,25 @@ version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" +[[package]] +name = "fdeflate" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e6853b52649d4ac5c0bd02320cddc5ba956bdb407c4b75a2c6b75bf51500f8c" +dependencies = [ + "simd-adler32", +] + +[[package]] +name = "filetime" +version = "0.2.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c287a33c7f0a620c38e641e7f60827713987b3c0f26e8ddc9462cc69cf75759" +dependencies = [ + "cfg-if", + "libc", +] + [[package]] name = "find-msvc-tools" version = "0.1.9" @@ -309,6 +378,12 @@ dependencies = [ "winapi", ] +[[package]] +name = "funty" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6d5a32815ae3f33302d95fdcb2ce17862f8c65363dcfd29360480ba1001fc9c" + [[package]] name = "futures" version = "0.3.34" @@ -773,6 +848,24 @@ version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +[[package]] +name = "libdeflate-sys" +version = "1.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7870e5fbd2766179a937c725fb11f4ca0ef025d982beb61bd3ce755425bd19c" +dependencies = [ + "cc", +] + +[[package]] +name = "libdeflater" +version = "1.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cbb8285d2e2bd4b8d4155840ae81929354b4ba0d7b3c272c84761ffa2141e1f" +dependencies = [ + "libdeflate-sys", +] + [[package]] name = "libredox" version = "0.1.24" @@ -874,6 +967,20 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" +[[package]] +name = "oxipng" +version = "9.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26c613f0f566526a647c7473f6a8556dbce22c91b13485ee4b4ec7ab648e4973" +dependencies = [ + "bitvec", + "indexmap", + "libdeflater", + "log", + "rgb", + "rustc-hash", +] + [[package]] name = "parking_lot" version = "0.12.5" @@ -933,6 +1040,19 @@ version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" +[[package]] +name = "png" +version = "0.17.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82151a2fc869e011c153adc57cf2789ccb8d9906ce52c0b39a6b5697749d7526" +dependencies = [ + "bitflags 1.3.2", + "crc32fast", + "fdeflate", + "flate2", + "miniz_oxide", +] + [[package]] name = "potential_utf" version = "0.1.5" @@ -1036,6 +1156,12 @@ version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" +[[package]] +name = "radium" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09" + [[package]] name = "rand" version = "0.9.4" @@ -1071,7 +1197,7 @@ version = "0.5.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" dependencies = [ - "bitflags", + "bitflags 2.13.0", ] [[package]] @@ -1155,6 +1281,15 @@ dependencies = [ "webpki-roots", ] +[[package]] +name = "rgb" +version = "0.8.53" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47b34b781b31e5d73e9fbc8689c70551fd1ade9a19e3e28cfec8580a79290cc4" +dependencies = [ + "bytemuck", +] + [[package]] name = "ring" version = "0.17.14" @@ -1185,7 +1320,7 @@ version = "0.40.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "23f2a97da3e3873c73cb2a2e71b35c40ff95e0b1eefa8d72d8499a6928c3b5b3" dependencies = [ - "bitflags", + "bitflags 2.13.0", "fallible-iterator", "fallible-streaming-iterator", "hashlink", @@ -1206,7 +1341,7 @@ version = "1.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" dependencies = [ - "bitflags", + "bitflags 2.13.0", "errno", "libc", "linux-raw-sys", @@ -1448,6 +1583,22 @@ dependencies = [ "syn 2.0.118", ] +[[package]] +name = "tap" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369" + +[[package]] +name = "tar" +version = "0.4.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f6221d9a6003c78398e3b239969f352578258df48c8eb051caadae0015bc840" +dependencies = [ + "filetime", + "libc", +] + [[package]] name = "tempfile" version = "3.27.0" @@ -1521,6 +1672,8 @@ dependencies = [ "dirs", "flate2", "futures", + "oxipng", + "png", "regex", "reqwest", "rusqlite", @@ -1528,6 +1681,7 @@ dependencies = [ "serde", "serde_json", "sha2", + "tar", "tempfile", "thiserror", "tinyagents-definition", @@ -1542,6 +1696,7 @@ dependencies = [ "tracing", "uuid", "wait-timeout", + "zip", ] [[package]] @@ -1704,6 +1859,7 @@ name = "tinyinference-llm" version = "0.3.0" dependencies = [ "async-trait", + "base64 0.22.1", "bytes", "futures", "reqwest", @@ -1847,7 +2003,7 @@ version = "0.6.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840" dependencies = [ - "bitflags", + "bitflags 2.13.0", "bytes", "futures-util", "http", @@ -2340,6 +2496,15 @@ version = "0.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" +[[package]] +name = "wyz" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05f360fc0b24296329c78fda852a1e9ae82de9cf7b27dae4b7f62f118f77b9ed" +dependencies = [ + "tap", +] + [[package]] name = "yoke" version = "0.8.3" @@ -2443,6 +2608,22 @@ dependencies = [ "syn 2.0.118", ] +[[package]] +name = "zip" +version = "2.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fabe6324e908f85a1c52063ce7aa26b68dcb7eb6dbc83a2d148403c9bc3eba50" +dependencies = [ + "arbitrary", + "crc32fast", + "crossbeam-utils", + "displaydoc", + "flate2", + "indexmap", + "memchr", + "thiserror", +] + [[package]] name = "zmij" version = "1.0.21" diff --git a/crates/tinyagents-harness/Cargo.toml b/crates/tinyagents-harness/Cargo.toml index a597ff5f1..0f4a77db0 100644 --- a/crates/tinyagents-harness/Cargo.toml +++ b/crates/tinyagents-harness/Cargo.toml @@ -17,7 +17,10 @@ base64 = "0.22" chrono = { workspace = true } chrono-tz = { version = "0.10", optional = true } dirs = { version = "6", optional = true } +tar = { version = "0.4", default-features = false, optional = true } +zip = { version = "2", default-features = false, features = ["deflate-flate2", "flate2"], optional = true } flate2 = { version = "1", optional = true } +oxipng = { version = "9", default-features = false, optional = true } futures = { workspace = true } regex = "1" rustix = { version = "1", features = ["fs"] } @@ -48,7 +51,7 @@ sqlite = ["dep:rusqlite"] # still spell out the old name keep compiling. tools = ["builtin-tools"] builtin-tools = ["dep:chrono-tz"] -multimodal = ["dep:flate2", "dep:reqwest"] +multimodal = ["dep:flate2", "dep:reqwest", "dep:tar", "dep:zip", "dep:oxipng"] # Image and video generation tools over tinyinference-image / -video. media = ["dep:tinyinference-image", "dep:tinyinference-video"] # Gates the Claude Code CLI and Claude Agent SDK provider adapters, which are @@ -63,6 +66,7 @@ langfuse = ["dep:reqwest"] tracing = ["tinytools-agent/tracing"] [dev-dependencies] +png = "0.17" tokio = { workspace = true, default-features = true, features = ["macros", "rt-multi-thread", "time", "test-util"] } [lints] diff --git a/crates/tinyagents-harness/src/error.rs b/crates/tinyagents-harness/src/error.rs index 609ca93bf..d86ec7966 100644 --- a/crates/tinyagents-harness/src/error.rs +++ b/crates/tinyagents-harness/src/error.rs @@ -366,6 +366,7 @@ impl From for TinyAgentsError { tinyinference_embeddings::Error::Serialization(error) => Self::Serialization(error), tinyinference_embeddings::Error::Embedding(message) => Self::Embedding(message), tinyinference_embeddings::Error::Cancelled => Self::Cancelled, + tinyinference_embeddings::Error::Rerank(error) => Self::Embedding(error.to_string()), } } } @@ -457,3 +458,7 @@ impl From for TinyAgentsError { Self::Storage(err.to_string()) } } + +#[cfg(test)] +#[path = "error_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness/src/error_tests.rs b/crates/tinyagents-harness/src/error_tests.rs new file mode 100644 index 000000000..2e8beaebc --- /dev/null +++ b/crates/tinyagents-harness/src/error_tests.rs @@ -0,0 +1,15 @@ +use super::*; + +#[test] +fn rerank_errors_convert_to_embedding_errors_with_context() { + let source = tinyinference_embeddings::Error::Rerank( + tinyinference_embeddings::rerank::RerankError::ResponseTooLarge { limit: 17 }, + ); + + let converted = TinyAgentsError::from(source); + + assert!(matches!( + converted, + TinyAgentsError::Embedding(message) if message.contains("17 bytes") + )); +} diff --git a/crates/tinyagents-harness/src/multimodal/README.md b/crates/tinyagents-harness/src/multimodal/README.md index 87423f859..704d9637f 100644 --- a/crates/tinyagents-harness/src/multimodal/README.md +++ b/crates/tinyagents-harness/src/multimodal/README.md @@ -33,10 +33,14 @@ format contributes a header naming it plus a content hash. - [`data_uri`] — `data:` URI parsing (including gzip-compressed attachments with a required `original_mime` parameter), percent-decoding, and encoding. - [`resolve`] — [`TextExtractor`] (host-pluggable document text extraction), - [`NoTextExtractor`], and the two resolution entry points + [`NoTextExtractor`], generic [`resolve_attachment`], and legacy entry points [`resolve_image`]/[`resolve_file`] that turn one marker reference into a payload, trying `data:` → `http(s)` (gated by `allow_remote_fetch`) → local path in that order. +- [`types`] — [`ResolvedAttachment`] and [`UnknownMimePolicy`] for generic intake. +- [`archive`] — [`inspect_archive`] for bounded ZIP/TAR/TAR.GZ listing; + `archive/types.rs` holds formats, budgets, entries, truncation and errors. +- [`png`] — [`optimize_png_lossless`] for optional faithful PNG derivatives. - [`payload`] — [`FilePayload`] (`Extracted` / `Reference`), [`compose_multimodal_message`] (renders the final provider-bound message), and supporting helpers (`truncate_chars`, `sha256_prefix`, `format_size`, @@ -54,10 +58,16 @@ format contributes a header naming it plus a content hash. | `markers.rs` | Marker prefixes and pure string transforms over them. | | `mime.rs` | MIME detection: header, extension, and magic-byte sniffing. | | `data_uri.rs` | `data:` URI parsing, gzip decompression, percent-decoding. | -| `resolve.rs` | `TextExtractor` trait, `resolve_image`/`resolve_file` pipelines. | +| `resolve.rs` | `TextExtractor`, `resolve_attachment`, `resolve_image`/`resolve_file`. | +| `types.rs` | Generic resolved-attachment metadata and unknown-MIME policy. | +| `archive/mod.rs` | Archive module wiring and public exports. | +| `archive/ops.rs` | Bounded archive listing, without filesystem extraction. | +| `archive/types.rs` | Archive format, budgets, entries, errors and truncation. | +| `archive/zip_admission.rs` | Allocation-free ZIP preflight and admitted reader. | +| `png.rs` | Optional lossless PNG optimization. | | `payload.rs` | `FilePayload`, message composition, truncation, hashing. | | `error.rs` | `MultimodalError`, `Result`. | -| `test.rs` | Unit tests for the marker/MIME/size/path edge cases across both pipelines. | +| `mod_tests.rs`, `resolve_tests.rs`, `archive/ops_tests.rs`, `png_tests.rs`, `stash_tests.rs` | Marker, MIME, size, path, intake, archive, PNG and stash behavior. | ## Operational constraints @@ -74,6 +84,69 @@ format contributes a header naming it plus a content hash. message type). This module never decides which local paths may be read — that is `FileLimits::files_disabled`'s lever, not a filesystem allowlist here. -- Every extraction/fetch/read failure degrades to a `FilePayload::Reference` - or a skipped attachment rather than failing the whole turn — a damaged PDF - should cost the model its text, not the conversation. +- Text extraction failures degrade to a `FilePayload::Reference`. Resolution + errors (read/fetch/MIME/size) remain typed errors for the host to present or + skip according to its own policy. + +## Generic intake and archives + +`resolve_attachment(source, &FileLimits, max_bytes, &Client, UnknownMimePolicy)` +returns `ResolvedAttachment { bytes, name, mime, size_bytes }`. It neither +extracts text nor writes files. `Reject` requires the existing MIME allowlist +and preserves rejection of undetected local bytes; `Accept` lets a host retain +arbitrary media or unknown formats without disabling byte limits, remote +fetch gates, or `max_files == 0`. The host authorizes local paths and chooses +storage locations. Names are untrusted display metadata. In `Accept` mode, +explicit HTTP media MIME types precede content sniffing, and common audio/video +extensions and signatures are recognized before the UTF-8 fallback. `Reject` +continues to use the legacy file detector. + +Transport gzip is identified by `application/gzip` plus `original_mime` in a +data URI and decoded exactly once under the byte cap. A `.tar.gz` attachment +without that parameter retains its compressed bytes. Remote URL display names +come from the decoded final path segment, without query parameters. Local reads +and remote streams stop at the configured byte budget, including files that +change after their metadata is read. `resolve_file` uses the generic resolver +with `Reject`, then retains its existing extractor/degrade-to-reference flow. + +`ArchiveFormat::detect(name, mime, bytes)` identifies ZIP, TAR, and TAR.GZ; +Office document MIME types/extensions are excluded. `inspect_archive(bytes, +format, &ArchiveLimits)` returns `ArchiveListing` entries with `name`, `kind`, +and `declared_size`. It performs no disk extraction and follows no links. +Traversal paths are displayed as recorded. ZIP members are streamed to a sink +for CRC validation under a cumulative expansion cap; TAR.GZ is inflated once +under a stream-byte cap. GNU/PAX extension records are listed as `Other` +rather than interpreted into potentially unbounded names. + +Limits cap input bytes, entry count, individual and aggregate UTF-8 name bytes, +and decompressed bytes. Exhaustion returns `truncation` with the applicable +reason; malformed headers/CRC/streams return `ArchiveError`. A TAR/TAR.GZ +stream beyond its decoded budget produces an empty truncated listing; a ZIP +member whose declared size exceeds the remaining budget appears in the +partial listing without being inflated. A partial listing does not validate +members beyond the stopping point. ZIP listing and Office probing share an allocation-free central-directory +preflight before the eager ZIP parser. It checks ordinary and ZIP64 counts +against available directory bytes and caps eager metadata at 10,000 entries, +4096 bytes per name, 1 MiB across names, and 4 MiB of directory metadata. +Archives above those safety ceilings produce an empty truncated listing; +Office probing falls back to the usual extension/header detection. Smaller +caller listing budgets retain a prefix as usual. Invalid directory counts and +records are errors. The reader pins the admitted directory offset and hides +member footer signatures during metadata indexing so parser retries cannot +select unchecked embedded ZIPs. Member data is restored for normal CRC and +expansion validation. Office probing reads metadata and local headers only. + +`types.rs` owns generic intake types; `archive/types.rs` owns listing types +and budgets; `archive/ops.rs` owns inspection; sibling `*_tests.rs` cover fixtures +and error/limit behavior. No transcript or inference content-block variants +are introduced by these APIs. + +PNG derivatives use `optimize_png_lossless`: it changes only IDAT compression +and filtering, preserving pixel format, hidden RGB, interlacing and every other +chunk byte-for-byte. Animated, malformed and oversized inputs are skipped; the +caller retains the original and decides where to store the optional smaller copy. + +ZIP admission conservatively rejects footer signature bytes in central-directory +metadata or archive comments, even where ZIP permits those bytes. This prevents +the eager reader from falling back to a different allocation-heavy footer; hosts +should retain the original and report unavailable listing rather than extracting. diff --git a/crates/tinyagents-harness/src/multimodal/archive/README.md b/crates/tinyagents-harness/src/multimodal/archive/README.md new file mode 100644 index 000000000..d8644734b --- /dev/null +++ b/crates/tinyagents-harness/src/multimodal/archive/README.md @@ -0,0 +1,22 @@ +# Archive inspection + +`inspect_archive` lists ZIP, TAR and TAR.GZ entries from supplied bytes. Entry +names are untrusted metadata; no filesystem objects are created and links are +never followed. The public API is re-exported through `multimodal::archive` and +`multimodal` for compatibility. + +| File | Responsibility | +| --- | --- | +| `mod.rs` | Module wiring and deliberate public exports | +| `types.rs` | Formats, caller budgets, entries, errors and truncation | +| `ops.rs` | Listing, CRC validation and bounded decompression | +| `zip_admission.rs` | Shared allocation-free ZIP preflight and admitted reader | +| `ops_tests.rs` | Synthetic malformed, hostile and bounded archive fixtures | + +ZIP metadata preflight occurs before the eager parser allocates its index. +Office MIME probing uses the same internal admission path. Input, metadata, +entry, name and expansion ceilings remain enforced; oversized archives report +explicit truncation. TAR.GZ inflation is bounded once, and GNU/PAX records are +listed rather than followed into unbounded names. A partial listing validates +only the processed prefix. See the parent README for detailed limits and the +conservative rejection of ZIP footer signatures in metadata/comments. diff --git a/crates/tinyagents-harness/src/multimodal/archive/mod.rs b/crates/tinyagents-harness/src/multimodal/archive/mod.rs new file mode 100644 index 000000000..b01930329 --- /dev/null +++ b/crates/tinyagents-harness/src/multimodal/archive/mod.rs @@ -0,0 +1,8 @@ +//! Bounded in-memory ZIP, TAR and TAR.GZ inspection without filesystem extraction. + +mod ops; +mod types; +pub(super) mod zip_admission; + +pub use ops::inspect_archive; +pub use types::*; diff --git a/crates/tinyagents-harness/src/multimodal/archive/ops.rs b/crates/tinyagents-harness/src/multimodal/archive/ops.rs new file mode 100644 index 000000000..a8be3b6ed --- /dev/null +++ b/crates/tinyagents-harness/src/multimodal/archive/ops.rs @@ -0,0 +1,175 @@ +//! Bounded archive inspection. Entry names are data, including absolute and +//! traversal paths; this module never extracts or creates filesystem objects. + +use std::io::{self, Cursor, Read}; + +use super::types::*; +use super::zip_admission::{open_admitted_zip, zip_preflight}; + +/// List an archive in memory. Malformed headers and compressed streams are +/// errors; resource caps return a partial listing with a truncation reason. +/// ZIP members are streamed to a sink to validate CRC and expansion bounds. +/// TAR.GZ is expanded once under the decoded-byte cap. GNU/PAX metadata is +/// listed as `Other`, rather than followed or expanded into unbounded names. +pub fn inspect_archive( + bytes: &[u8], + format: ArchiveFormat, + limits: &ArchiveLimits, +) -> Result { + let limits = limits.effective(); + if bytes.len() > limits.max_input_bytes { + return Err(ArchiveError::InputTooLarge); + } + match format { + ArchiveFormat::Zip => inspect_zip(bytes, &limits), + ArchiveFormat::Tar => inspect_tar(bytes, &limits), + ArchiveFormat::TarGzip => { + let mut decoded = Vec::new(); + flate2::read::MultiGzDecoder::new(bytes) + .take(limits.max_decompressed_bytes.saturating_add(1) as u64) + .read_to_end(&mut decoded)?; + if decoded.len() > limits.max_decompressed_bytes { + return Ok(ArchiveListing { + entries: Vec::new(), + truncation: Some(ArchiveTruncation::DecompressedBytes), + }); + } + inspect_tar(&decoded, &limits) + } + } +} + +fn add_entry( + list: &mut ArchiveListing, + entry: ArchiveEntry, + limits: &ArchiveLimits, + names: &mut usize, +) -> bool { + if list.entries.len() >= limits.max_entries { + list.truncation = Some(ArchiveTruncation::Entries); + return false; + } + if entry.name.len() > limits.max_name_bytes + || names.saturating_add(entry.name.len()) > limits.max_total_name_bytes + { + list.truncation = Some(ArchiveTruncation::NameBytes); + return false; + } + *names += entry.name.len(); + list.entries.push(entry); + true +} + +fn inspect_zip(bytes: &[u8], limits: &ArchiveLimits) -> Result { + if let Some(reason) = zip_preflight(bytes)? { + return Ok(ArchiveListing { + entries: Vec::new(), + truncation: Some(reason), + }); + } + let mut archive = open_admitted_zip(bytes)?; + let mut list = ArchiveListing::default(); + let mut names = 0; + let mut expanded = 0u64; + for index in 0..archive.len() { + if list.entries.len() >= limits.max_entries { + list.truncation = Some(ArchiveTruncation::Entries); + break; + } + let mut file = archive + .by_index(index) + .map_err(|e| ArchiveError::Invalid(e.to_string()))?; + let kind = if file.is_dir() { + ArchiveEntryKind::Directory + } else if file + .unix_mode() + .is_some_and(|mode| mode & 0o170000 == 0o120000) + { + ArchiveEntryKind::Symlink + } else { + ArchiveEntryKind::File + }; + let size = file.size(); + if !add_entry( + &mut list, + ArchiveEntry { + name: file.name().to_string(), + kind, + declared_size: size, + }, + limits, + &mut names, + ) { + break; + } + let remaining = (limits.max_decompressed_bytes as u64).saturating_sub(expanded); + if size > remaining { + list.truncation = Some(ArchiveTruncation::DecompressedBytes); + break; + } + let read = io::copy( + &mut file.by_ref().take(remaining.saturating_add(1)), + &mut io::sink(), + )?; + if read > remaining { + list.truncation = Some(ArchiveTruncation::DecompressedBytes); + break; + } + expanded += read; + } + Ok(list) +} + +fn inspect_tar(bytes: &[u8], limits: &ArchiveLimits) -> Result { + if bytes.len() > limits.max_decompressed_bytes { + return Ok(ArchiveListing { + entries: Vec::new(), + truncation: Some(ArchiveTruncation::DecompressedBytes), + }); + } + let mut archive = tar::Archive::new(Cursor::new(bytes)); + let mut list = ArchiveListing::default(); + let mut names = 0; + for entry in archive.entries()?.raw(true) { + if list.entries.len() >= limits.max_entries { + list.truncation = Some(ArchiveTruncation::Entries); + break; + } + let entry = entry?; + let header = entry.header(); + let entry_type = header.entry_type(); + let kind = if entry_type.is_file() { + ArchiveEntryKind::File + } else if entry_type.is_dir() { + ArchiveEntryKind::Directory + } else if entry_type.is_symlink() { + ArchiveEntryKind::Symlink + } else if entry_type.is_hard_link() { + ArchiveEntryKind::HardLink + } else { + ArchiveEntryKind::Other + }; + let name = String::from_utf8_lossy(&entry.path_bytes()).into_owned(); + let size = header.size()?; + if !add_entry( + &mut list, + ArchiveEntry { + name, + kind, + declared_size: size, + }, + limits, + &mut names, + ) { + break; + } + if entry.raw_file_position().saturating_add(size) > bytes.len() as u64 { + return Err(ArchiveError::Invalid("truncated TAR member".to_string())); + } + } + Ok(list) +} + +#[cfg(test)] +#[path = "ops_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness/src/multimodal/archive/ops_tests.rs b/crates/tinyagents-harness/src/multimodal/archive/ops_tests.rs new file mode 100644 index 000000000..37d2c2417 --- /dev/null +++ b/crates/tinyagents-harness/src/multimodal/archive/ops_tests.rs @@ -0,0 +1,469 @@ +use super::*; + +#[test] +fn malformed_archives_are_errors() { + for format in [ + ArchiveFormat::Zip, + ArchiveFormat::Tar, + ArchiveFormat::TarGzip, + ] { + assert!(inspect_archive(b"not an archive", format, &ArchiveLimits::default()).is_err()); + } +} + +use std::io::{Cursor, Write}; + +fn zip_fixture(entries: &[(&str, &[u8])]) -> Vec { + let mut writer = zip::ZipWriter::new(Cursor::new(Vec::new())); + for (name, bytes) in entries { + writer + .start_file( + *name, + zip::write::SimpleFileOptions::default() + .compression_method(zip::CompressionMethod::Deflated), + ) + .unwrap(); + writer.write_all(bytes).unwrap(); + } + writer.finish().unwrap().into_inner() +} + +fn tar_fixture() -> Vec { + let mut builder = tar::Builder::new(Vec::new()); + let mut header = tar::Header::new_gnu(); + header.set_mode(0o644); + header.set_size(5); + // Set path bytes directly: the writer correctly refuses traversal paths, + // but the inspector must display a hostile archive without extracting it. + header.as_mut_bytes()[..14].copy_from_slice(b"../outside.txt"); + header.set_cksum(); + builder.append(&header, &b"hello"[..]).unwrap(); + let mut link = tar::Header::new_gnu(); + link.set_path("link").unwrap(); + link.set_entry_type(tar::EntryType::Symlink); + link.set_link_name("/etc/passwd").unwrap(); + link.set_size(0); + link.set_cksum(); + builder.append(&link, &b""[..]).unwrap(); + builder.into_inner().unwrap() +} + +#[test] +fn zip_listing_preserves_names_sizes_and_applies_each_budget() { + let bytes = zip_fixture(&[("../outside.txt", b"hello"), ("normal.txt", b"world")]); + let list = inspect_archive(&bytes, ArchiveFormat::Zip, &ArchiveLimits::default()).unwrap(); + assert_eq!( + list.entries[0], + ArchiveEntry { + name: "../outside.txt".into(), + kind: ArchiveEntryKind::File, + declared_size: 5 + } + ); + assert_eq!(list.entries.len(), 2); + assert_eq!(list.truncation, None); + let limits = ArchiveLimits { + max_entries: 1, + ..ArchiveLimits::default() + }; + let list = inspect_archive(&bytes, ArchiveFormat::Zip, &limits).unwrap(); + assert_eq!(list.entries.len(), 1); + assert_eq!(list.truncation, Some(ArchiveTruncation::Entries)); + for limits in [ + ArchiveLimits { + max_name_bytes: 3, + ..ArchiveLimits::default() + }, + ArchiveLimits { + max_total_name_bytes: 3, + ..ArchiveLimits::default() + }, + ] { + let list = inspect_archive(&bytes, ArchiveFormat::Zip, &limits).unwrap(); + assert!(list.entries.is_empty()); + assert_eq!(list.truncation, Some(ArchiveTruncation::NameBytes)); + } + assert!(matches!( + inspect_archive( + &bytes, + ArchiveFormat::Zip, + &ArchiveLimits { + max_input_bytes: 2, + ..ArchiveLimits::default() + } + ), + Err(ArchiveError::InputTooLarge) + )); +} + +#[test] +fn tar_and_tar_gzip_list_traversal_and_links_without_extracting() { + let bytes = tar_fixture(); + let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + encoder.write_all(&bytes).unwrap(); + let gzip = encoder.finish().unwrap(); + for (format, bytes) in [ + (ArchiveFormat::Tar, bytes.as_slice()), + (ArchiveFormat::TarGzip, gzip.as_slice()), + ] { + let list = inspect_archive(bytes, format, &ArchiveLimits::default()).unwrap(); + assert_eq!(list.entries[0].name, "../outside.txt"); + assert_eq!(list.entries[0].declared_size, 5); + assert_eq!(list.entries[1].kind, ArchiveEntryKind::Symlink); + assert_eq!(list.truncation, None); + let list = inspect_archive( + bytes, + format, + &ArchiveLimits { + max_decompressed_bytes: 100, + ..ArchiveLimits::default() + }, + ) + .unwrap(); + assert_eq!(list.truncation, Some(ArchiveTruncation::DecompressedBytes)); + } +} + +#[test] +fn compressed_zip_expansion_is_capped_and_truncated_tar_is_an_error() { + let bytes = zip_fixture(&[("bomb", &vec![0; 100_000])]); + assert!(bytes.len() < 1000); + let list = inspect_archive( + &bytes, + ArchiveFormat::Zip, + &ArchiveLimits { + max_decompressed_bytes: 1000, + ..ArchiveLimits::default() + }, + ) + .unwrap(); + assert_eq!(list.entries[0].declared_size, 100_000); + assert_eq!(list.truncation, Some(ArchiveTruncation::DecompressedBytes)); + let tar = tar_fixture(); + assert!(inspect_archive(&tar[..514], ArchiveFormat::Tar, &ArchiveLimits::default()).is_err()); +} + +#[test] +fn archive_detection_does_not_mistake_office_documents_for_archives() { + assert_eq!( + ArchiveFormat::detect("table.xlsx", "application/zip", b"PK\x03\x04"), + None + ); + assert_eq!( + ArchiveFormat::detect("archive.zip", "application/zip", b"PK\x03\x04"), + Some(ArchiveFormat::Zip) + ); + let mut tar_gzip = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + tar_gzip.write_all(&tar_fixture()).unwrap(); + let tar_gzip = tar_gzip.finish().unwrap(); + assert_eq!( + ArchiveFormat::detect("attachment", "application/gzip", &tar_gzip), + Some(ArchiveFormat::TarGzip) + ); + assert_eq!( + ArchiveFormat::detect("stale.zip", "application/gzip", &tar_gzip), + Some(ArchiveFormat::TarGzip) + ); + let tar = tar_fixture(); + assert_eq!( + ArchiveFormat::detect("stale.zip", "application/x-tar", &tar), + Some(ArchiveFormat::Tar) + ); + assert_eq!( + ArchiveFormat::detect("attachment.tgz", "application/octet-stream", &[]), + Some(ArchiveFormat::TarGzip) + ); + let mut text_gzip = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + text_gzip.write_all(b"compressed text").unwrap(); + assert_eq!( + ArchiveFormat::detect( + "attachment", + "application/gzip", + &text_gzip.finish().unwrap() + ), + None + ); + let mut empty_tar_gzip = + flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + empty_tar_gzip.write_all(&[0; 1024]).unwrap(); + assert_eq!( + ArchiveFormat::detect( + "empty.tar.gz", + "application/gzip", + &empty_tar_gzip.finish().unwrap() + ), + Some(ArchiveFormat::TarGzip) + ); + assert_eq!( + ArchiveFormat::detect("archive.tar", "application/x-tar", b""), + Some(ArchiveFormat::Tar) + ); +} + +#[test] +fn office_container_contents_override_generic_zip_headers() { + for (member, expected) in [ + ( + "word/document.xml", + "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + ), + ( + "xl/workbook.xml", + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + ), + ( + "ppt/presentation.xml", + "application/vnd.openxmlformats-officedocument.presentationml.presentation", + ), + ] { + let bytes = zip_fixture(&[ + ("[Content_Types].xml", b""), + (member, b""), + ]); + assert_eq!( + crate::multimodal::mime::detect_file_mime(None, &bytes, Some("application/zip")) + .as_deref(), + Some(expected) + ); + } +} + +#[test] +fn corrupt_zip_payloads_and_gzip_trailers_are_errors() { + let mut bytes = zip_fixture(&[("file.txt", b"hello world")]); + let start = zip::ZipArchive::new(Cursor::new(&bytes)) + .unwrap() + .by_index(0) + .unwrap() + .data_start() as usize; + bytes[start] ^= 0xff; + assert!(inspect_archive(&bytes, ArchiveFormat::Zip, &ArchiveLimits::default()).is_err()); + let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + encoder.write_all(&tar_fixture()).unwrap(); + let mut gzip = encoder.finish().unwrap(); + let checksum = gzip.len() - 8; + gzip[checksum] ^= 0xff; + assert!(inspect_archive(&gzip, ArchiveFormat::TarGzip, &ArchiveLimits::default()).is_err()); +} + +#[test] +fn exact_and_zero_listing_budgets_have_no_off_by_one() { + let bytes = zip_fixture(&[("a", b"hello")]); + let limits = ArchiveLimits { + max_entries: 1, + max_name_bytes: 1, + max_total_name_bytes: 1, + max_decompressed_bytes: 5, + ..ArchiveLimits::default() + }; + let list = inspect_archive(&bytes, ArchiveFormat::Zip, &limits).unwrap(); + assert_eq!(list.truncation, None); + assert_eq!(list.entries.len(), 1); + let limits = ArchiveLimits { + max_entries: 0, + ..ArchiveLimits::default() + }; + let list = inspect_archive(&bytes, ArchiveFormat::Zip, &limits).unwrap(); + assert!(list.entries.is_empty()); + assert_eq!(list.truncation, Some(ArchiveTruncation::Entries)); +} + +#[test] +fn zip_admission_checks_counts_and_names_before_constructing_parser() { + let bytes = zip_fixture(&[("one", b"")]); + let end = bytes.len() - 22; + let mut hostile = bytes.clone(); + hostile[end + 8..end + 12].copy_from_slice(&[255; 4]); + assert!(zip_preflight(&hostile).is_err()); + let mut inconsistent = bytes.clone(); + inconsistent[end + 8..end + 10].copy_from_slice(&u16::MAX.to_le_bytes()); + assert!(zip_preflight(&inconsistent).is_err()); + let mut zip64 = bytes[..end].to_vec(); + let zip64_offset = zip64.len() as u64; + zip64.extend_from_slice(b"PK\x06\x06"); + zip64.extend_from_slice(&44u64.to_le_bytes()); + zip64.extend_from_slice(&[0; 12]); + zip64.extend_from_slice(&u64::MAX.to_le_bytes()); + zip64.extend_from_slice(&u64::MAX.to_le_bytes()); + zip64.extend_from_slice(&(end as u64).to_le_bytes()); + zip64.extend_from_slice(&0u64.to_le_bytes()); + zip64.extend_from_slice(b"PK\x06\x07"); + zip64.extend_from_slice(&0u32.to_le_bytes()); + zip64.extend_from_slice(&zip64_offset.to_le_bytes()); + zip64.extend_from_slice(&1u32.to_le_bytes()); + zip64.extend_from_slice(&hostile[end..]); + assert!(zip_preflight(&zip64).is_err()); + let mut valid64 = zip64.clone(); + let offset = zip64_offset as usize; + valid64[offset + 24..offset + 32].copy_from_slice(&1u64.to_le_bytes()); + valid64[offset + 32..offset + 40].copy_from_slice(&1u64.to_le_bytes()); + valid64[offset + 40..offset + 48].copy_from_slice( + &(u32::from_le_bytes(bytes[end + 12..end + 16].try_into().unwrap()) as u64).to_le_bytes(), + ); + valid64[offset + 48..offset + 56].copy_from_slice( + &(u32::from_le_bytes(bytes[end + 16..end + 20].try_into().unwrap()) as u64).to_le_bytes(), + ); + assert_eq!(zip_preflight(&valid64).unwrap(), None); + assert_eq!( + inspect_archive(&valid64, ArchiveFormat::Zip, &ArchiveLimits::default()) + .unwrap() + .entries + .len(), + 1 + ); + let mut prefixed64 = b"executable preamble".to_vec(); + prefixed64.extend_from_slice(&valid64); + assert_eq!(zip_preflight(&prefixed64).unwrap(), None); + assert_eq!( + inspect_archive(&prefixed64, ArchiveFormat::Zip, &ArchiveLimits::default()) + .unwrap() + .entries + .len(), + 1 + ); + assert!(inspect_archive(&zip64, ArchiveFormat::Zip, &ArchiveLimits::default()).is_err()); + assert_eq!( + crate::multimodal::mime::detect_file_mime(None, &zip64, None).as_deref(), + Some("application/zip") + ); + let entries: Vec<_> = (0..10_001) + .map(|i| (format!("entry-{i}"), b"".as_slice())) + .collect(); + let refs: Vec<_> = entries + .iter() + .map(|(name, data)| (name.as_str(), *data)) + .collect(); + let bytes = zip_fixture(&refs); + assert_eq!( + zip_preflight(&bytes).unwrap(), + Some(ArchiveTruncation::Entries) + ); + let name = "a".repeat(4097); + let bytes = zip_fixture(&[(&name, b"")]); + assert_eq!( + zip_preflight(&bytes).unwrap(), + Some(ArchiveTruncation::NameBytes) + ); + let entries: Vec<_> = (0..300) + .map(|i| (format!("{i:04}{}", "x".repeat(4090)), b"".as_slice())) + .collect(); + let refs: Vec<_> = entries + .iter() + .map(|(name, data)| (name.as_str(), *data)) + .collect(); + let bytes = zip_fixture(&refs); + assert_eq!( + zip_preflight(&bytes).unwrap(), + Some(ArchiveTruncation::NameBytes) + ); + let list = inspect_archive(&bytes, ArchiveFormat::Zip, &ArchiveLimits::default()).unwrap(); + assert!(list.entries.is_empty()); + assert_eq!(list.truncation, Some(ArchiveTruncation::NameBytes)); +} + +#[test] +fn zip_admission_rejects_malformed_records_and_bounds_unicode_names() { + let bytes = zip_fixture(&[("a", b"")]); + let end = bytes.len() - 22; + let size = u32::from_le_bytes(bytes[end + 12..end + 16].try_into().unwrap()) as usize; + let central = end - size; + for (offset, value) in [ + (central, 0), + (central + 28, 255), + (central + 29, 255), + (end + 12, 255), + (end + 13, 255), + (end + 10, 0), + (end + 20, 1), + ] { + let mut corrupt = bytes.clone(); + corrupt[offset] = value; + assert!(zip_preflight(&corrupt).is_err()); + } + // Unicode path extras can replace a short raw filename during full parse. + let mut extra = Vec::new(); + extra.extend_from_slice(&0x7075u16.to_le_bytes()); + extra.extend_from_slice(&4102u16.to_le_bytes()); + extra.extend_from_slice(&[1, 0, 0, 0, 0]); + extra.extend_from_slice(&vec![b'x'; 4097]); + let mut unicode = bytes.clone(); + unicode.splice(central + 47..central + 47, extra.iter().copied()); + unicode[central + 30..central + 32].copy_from_slice(&(extra.len() as u16).to_le_bytes()); + let new_end = end + extra.len(); + unicode[new_end + 12..new_end + 16] + .copy_from_slice(&((size + extra.len()) as u32).to_le_bytes()); + assert_eq!( + zip_preflight(&unicode).unwrap(), + Some(ArchiveTruncation::NameBytes) + ); + // Incomplete and overrun extra-field records never reach the eager parser. + for declared in [1u16, 4] { + let mut corrupt = unicode.clone(); + corrupt[central + 30..central + 32].copy_from_slice(&declared.to_le_bytes()); + assert!(zip_preflight(&corrupt).is_err()); + } +} + +#[test] +fn zip_metadata_index_cannot_fall_back_to_a_footer_inside_member_data() { + let nested = zip_fixture(&[("inside", b"nested data")]); + let mut writer = zip::ZipWriter::new(Cursor::new(Vec::new())); + writer + .start_file( + "nested.zip", + zip::write::SimpleFileOptions::default() + .compression_method(zip::CompressionMethod::Stored), + ) + .unwrap(); + writer.write_all(&nested).unwrap(); + let bytes = writer.finish().unwrap().into_inner(); + assert_eq!( + inspect_archive(&bytes, ArchiveFormat::Zip, &ArchiveLimits::default()) + .unwrap() + .entries[0] + .declared_size, + nested.len() as u64 + ); + // Self-extracting preambles still work with the pinned directory offset. + let mut prefixed = b"executable preamble".to_vec(); + prefixed.extend_from_slice(&bytes); + assert_eq!( + inspect_archive(&prefixed, ArchiveFormat::Zip, &ArchiveLimits::default()) + .unwrap() + .entries[0] + .name, + "nested.zip" + ); + let end = bytes.len() - 22; + let size = u32::from_le_bytes(bytes[end + 12..end + 16].try_into().unwrap()) as usize; + let central = end - size; + let name_len = + u16::from_le_bytes(bytes[central + 28..central + 30].try_into().unwrap()) as usize; + let mut corrupt = bytes.clone(); + // A structurally bounded ZIP64 field that is semantically incomplete makes + // ZipArchive retry earlier footers. The nested member must remain hidden. + let extra = [1, 0, 1, 0, 0]; + corrupt.splice(central + 46 + name_len..central + 46 + name_len, extra); + corrupt[central + 30..central + 32].copy_from_slice(&5u16.to_le_bytes()); + corrupt[end + 5 + 12..end + 5 + 16].copy_from_slice(&((size + 5) as u32).to_le_bytes()); + // AES compression without its mandatory AES extra field is rejected by + // the full parser, triggering its earlier-footer fallback path. + corrupt[central + 10..central + 12].copy_from_slice(&99u16.to_le_bytes()); + assert_eq!(zip_preflight(&corrupt).unwrap(), None); + assert!(inspect_archive(&corrupt, ArchiveFormat::Zip, &ArchiveLimits::default()).is_err()); +} + +#[test] +fn footer_signature_in_valid_zip_comment_is_conservatively_rejected() { + let mut writer = zip::ZipWriter::new(Cursor::new(Vec::new())); + writer.set_comment("valid comment with PK\u{5}\u{6} signature"); + writer + .start_file("file.txt", zip::write::SimpleFileOptions::default()) + .unwrap(); + writer.write_all(b"payload").unwrap(); + let bytes = writer.finish().unwrap().into_inner(); + // The ZIP format permits these comment bytes. Our bounded admission + // deliberately rejects them rather than exposing eager footer fallback. + assert!(zip::ZipArchive::new(Cursor::new(&bytes)).is_ok()); + assert!(inspect_archive(&bytes, ArchiveFormat::Zip, &ArchiveLimits::default()).is_err()); +} diff --git a/crates/tinyagents-harness/src/multimodal/archive/types.rs b/crates/tinyagents-harness/src/multimodal/archive/types.rs new file mode 100644 index 000000000..6f93d0647 --- /dev/null +++ b/crates/tinyagents-harness/src/multimodal/archive/types.rs @@ -0,0 +1,167 @@ +//! Archive inspection types and caller-selected resource caps. + +use std::io::Read; + +/// Supported archive wire formats; Office ZIP containers are documents. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ArchiveFormat { + /// ZIP with stored or deflated members. + Zip, + /// Uncompressed TAR. + Tar, + /// Gzip-compressed TAR. + TarGzip, +} + +impl ArchiveFormat { + /// Identify archive formats by MIME, signatures, and a display filename. + /// Office MIME types and extensions take precedence over ZIP container magic. + pub fn detect(name: &str, mime: &str, bytes: &[u8]) -> Option { + let name = name.to_ascii_lowercase(); + let mime = mime + .split(';') + .next() + .unwrap_or_default() + .trim() + .to_ascii_lowercase(); + if mime.starts_with("application/vnd.openxmlformats-officedocument.") + || [".docx", ".xlsx", ".pptx"] + .iter() + .any(|ext| name.ends_with(ext)) + { + return None; + } + if mime == "application/gzip" || bytes.starts_with(b"\x1f\x8b") { + if is_tar_gzip(bytes) || name.ends_with(".tar.gz") || name.ends_with(".tgz") { + return Some(Self::TarGzip); + } + return None; + } + if bytes.get(257..262) == Some(b"ustar") { + return Some(Self::Tar); + } + if bytes.starts_with(b"PK\x03\x04") || bytes.starts_with(b"PK\x05\x06") { + return Some(Self::Zip); + } + if mime == "application/x-tar" { + return Some(Self::Tar); + } + if mime == "application/zip" { + return Some(Self::Zip); + } + if name.ends_with(".tar.gz") || name.ends_with(".tgz") { + return Some(Self::TarGzip); + } + if name.ends_with(".tar") { + return Some(Self::Tar); + } + if name.ends_with(".zip") { + return Some(Self::Zip); + } + None + } +} + +fn is_tar_gzip(bytes: &[u8]) -> bool { + let mut header = [0; 512]; + flate2::read::GzDecoder::new(bytes) + .take(header.len() as u64) + .read_exact(&mut header) + .is_ok() + && &header[257..262] == b"ustar" +} + +/// Inspection resource budget. Zero values are honored; values above the +/// safety ceilings are clamped before parsing. Limits concern listing and +/// validation, and never grant filesystem access. +#[derive(Debug, Clone)] +pub struct ArchiveLimits { + /// Maximum compressed/input bytes (ceiling 50 MiB). + pub max_input_bytes: usize, + /// Maximum entries retained (ceiling 10,000). + pub max_entries: usize, + /// Maximum UTF-8 bytes per displayed name (ceiling 4096). + pub max_name_bytes: usize, + /// Maximum UTF-8 bytes across names (ceiling 1 MiB). + pub max_total_name_bytes: usize, + /// Maximum inflated ZIP member bytes or TAR stream bytes (ceiling 50 MiB). + pub max_decompressed_bytes: usize, +} +impl Default for ArchiveLimits { + fn default() -> Self { + Self { + max_input_bytes: 50 * 1024 * 1024, + max_entries: 200, + max_name_bytes: 1024, + max_total_name_bytes: 64 * 1024, + max_decompressed_bytes: 16 * 1024 * 1024, + } + } +} +impl ArchiveLimits { + pub(crate) fn effective(&self) -> Self { + Self { + max_input_bytes: self.max_input_bytes.min(50 * 1024 * 1024), + max_entries: self.max_entries.min(10_000), + max_name_bytes: self.max_name_bytes.min(4096), + max_total_name_bytes: self.max_total_name_bytes.min(1024 * 1024), + max_decompressed_bytes: self.max_decompressed_bytes.min(50 * 1024 * 1024), + } + } +} + +/// An entry's declared filesystem type. No links are followed. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ArchiveEntryKind { + /// Regular file. + File, + /// Directory. + Directory, + /// Symbolic link. + Symlink, + /// Hard link. + HardLink, + /// Metadata, device, or another unsupported TAR kind. + Other, +} +/// Metadata for one entry. Names are untrusted data, never extraction paths. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ArchiveEntry { + /// Name recorded in the archive, including unsafe paths. + pub name: String, + /// Entry's declared kind. + pub kind: ArchiveEntryKind, + /// Uncompressed size declared by the header. + pub declared_size: u64, +} +/// A complete or explicitly truncated listing. +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub struct ArchiveListing { + /// Entries inspected in archive order. + pub entries: Vec, + /// Which budget stopped inspection, if any. + pub truncation: Option, +} +/// Resource ceiling that stopped inspection. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ArchiveTruncation { + /// Entry count cap reached. + Entries, + /// Individual or aggregate displayed name byte cap reached. + NameBytes, + /// Inflated content/stream byte cap reached. + DecompressedBytes, +} +/// Invalid or oversized archive input. +#[derive(Debug, thiserror::Error)] +pub enum ArchiveError { + /// Input exceeded the compressed/input byte budget. + #[error("archive input exceeds byte limit")] + InputTooLarge, + /// Invalid ZIP headers or unsupported/encrypted member. + #[error("invalid archive: {0}")] + Invalid(String), + /// Invalid TAR headers, compressed data, or ZIP CRC. + #[error("invalid archive: {0}")] + Io(#[from] std::io::Error), +} diff --git a/crates/tinyagents-harness/src/multimodal/archive/zip_admission.rs b/crates/tinyagents-harness/src/multimodal/archive/zip_admission.rs new file mode 100644 index 000000000..6c231c1ea --- /dev/null +++ b/crates/tinyagents-harness/src/multimodal/archive/zip_admission.rs @@ -0,0 +1,276 @@ +//! Allocation-free ZIP central-directory admission before the eager ZIP parser. + +use super::types::{ArchiveError, ArchiveTruncation}; + +fn invalid() -> ArchiveError { + ArchiveError::Invalid("invalid ZIP central directory".into()) +} + +fn number(bytes: &[u8], offset: usize, width: usize) -> Result { + let slice = bytes + .get(offset..offset.checked_add(width).ok_or_else(invalid)?) + .ok_or_else(invalid)?; + Ok(slice + .iter() + .enumerate() + .fold(0, |value, (i, byte)| value | ((*byte as u64) << (8 * i)))) +} + +fn zip64_eocd_offset(bytes: &[u8], locator: usize) -> Result { + let archive_relative = number(bytes, locator + 8, 8)?; + bytes + .get(..locator) + .ok_or_else(invalid)? + .windows(4) + .enumerate() + .rev() + .find_map(|(offset, signature)| { + if signature != b"PK\x06\x06" || (offset as u64) < archive_relative { + return None; + } + let record_size = number(bytes, offset + 4, 8).ok()?; + ((offset as u64).checked_add(12)?.checked_add(record_size)? == locator as u64) + .then_some(offset) + }) + .ok_or_else(invalid) +} + +fn footer_offset(bytes: &[u8]) -> Result { + let end = (bytes.len().saturating_sub(65_557)..bytes.len().saturating_sub(21)) + .rev() + .find(|&i| { + bytes.get(i..i + 4) == Some(b"PK\x05\x06") + && number(bytes, i + 20, 2).is_ok_and(|len| i + 22 + len as usize == bytes.len()) + }) + .ok_or_else(invalid)?; + Ok(end) +} + +/// Verify declared counts against available records and bound all metadata +/// before `ZipArchive::new` allocates. The safety ceilings are independent of +/// the caller's smaller listing budget so small listings retain their prefix. +/// Archives above a safety ceiling return an empty, truncated listing. +pub(in crate::multimodal) fn zip_preflight( + bytes: &[u8], +) -> Result, ArchiveError> { + let end = footer_offset(bytes)?; + let mut count = number(bytes, end + 10, 2)?; + if number(bytes, end + 8, 2)? != count + || number(bytes, end + 4, 2)? != 0 + || number(bytes, end + 6, 2)? != 0 + { + return Err(invalid()); + } + let mut size = number(bytes, end + 12, 4)?; + let mut directory_end = end; + let mut relative = number(bytes, end + 16, 4)?; + let mut zip64_locator_offset = None; + let mut zip64_eocd_absolute_offset = None; + if end >= 20 && bytes.get(end - 20..end - 16) == Some(b"PK\x06\x07") { + let offset = zip64_eocd_offset(bytes, end - 20)?; + zip64_locator_offset = Some(number(bytes, end - 12, 8)?); + zip64_eocd_absolute_offset = Some(offset as u64); + let record_size = number(bytes, offset + 4, 8)?; + if record_size < 44 + || (offset as u64) + .checked_add(12) + .and_then(|v| v.checked_add(record_size)) + != Some((end - 20) as u64) + { + return Err(invalid()); + } + if record_size > 4 * 1024 * 1024 { + return Ok(Some(ArchiveTruncation::NameBytes)); + } + count = number(bytes, offset + 32, 8)?; + if number(bytes, offset + 24, 8)? != count + || number(bytes, offset + 16, 4)? != 0 + || number(bytes, offset + 20, 4)? != 0 + || number(bytes, end - 16, 4)? != 0 + || number(bytes, end - 4, 4)? != 1 + { + return Err(invalid()); + } + size = number(bytes, offset + 40, 8)?; + directory_end = offset; + relative = number(bytes, offset + 48, 8)?; + for (legacy, sentinel, actual) in [ + (number(bytes, end + 10, 2)?, u16::MAX as u64, count), + (number(bytes, end + 12, 4)?, u32::MAX as u64, size), + (number(bytes, end + 16, 4)?, u32::MAX as u64, relative), + ] { + if legacy != sentinel && legacy != actual { + return Err(invalid()); + } + } + } else if count == 65535 || size == u32::MAX as u64 { + return Err(invalid()); + } + let size = usize::try_from(size).map_err(|_| invalid())?; + let start = directory_end.checked_sub(size).ok_or_else(invalid)?; + if relative > start as u64 { + return Err(invalid()); + } + let archive_offset = start as u64 - relative; + if zip64_locator_offset.is_some_and(|declared| { + zip64_eocd_absolute_offset.and_then(|offset| offset.checked_sub(archive_offset)) + != Some(declared) + }) { + return Err(invalid()); + } + // The eager reader retries earlier footer candidates. Embedded footers in + // admitted metadata are rejected; payload footers are hidden by the reader + // below while it constructs its metadata index. This deliberately excludes + // otherwise-valid names, extras and comments containing those byte sequences: + // listing availability is sacrificed rather than allowing an eager fallback. + for (i, signature) in bytes[start..].windows(4).enumerate() { + let position = start + i; + if (signature == b"PK\x05\x06" && position != end) + || (signature == b"PK\x06\x06" && (directory_end == end || position != directory_end)) + { + return Err(invalid()); + } + } + // Each directory entry requires at least 46 bytes. Check before comparing + // the count to a cap: impossible hostile counts are malformed, not truncated. + if count > (size / 46) as u64 { + return Err(invalid()); + } + if count > 10_000 { + return Ok(Some(ArchiveTruncation::Entries)); + } + // Extra fields and comments are also eagerly allocated by the ZIP reader. + if size > 4 * 1024 * 1024 { + return Ok(Some(ArchiveTruncation::NameBytes)); + } + let mut cursor = start; + let mut names = 0usize; + for _ in 0..count { + if bytes.get(cursor..cursor.saturating_add(4)) != Some(b"PK\x01\x02") { + return Err(invalid()); + } + let name_len = number(bytes, cursor + 28, 2)? as usize; + let extra_len = number(bytes, cursor + 30, 2)? as usize; + let comment_len = number(bytes, cursor + 32, 2)? as usize; + let name_start = cursor + 46; + let next = name_start + name_len + extra_len + comment_len; + if next > directory_end { + return Err(invalid()); + } + let name = &bytes[name_start..name_start + name_len]; + // CP437 characters can expand to three UTF-8 bytes. UTF-8 names use + // the exact lossy-decoding length (without allocating a String). + let mut displayed = if number(bytes, cursor + 8, 2)? & 0x800 != 0 { + std::str::from_utf8(name).map_or(name.len() * 3, str::len) + } else { + name.iter() + .map(|byte| if *byte < 128 { 1 } else { 3 }) + .sum() + }; + let mut field = name_start + name_len; + let extra_end = field + extra_len; + while field < extra_end { + if extra_end - field < 4 { + return Err(invalid()); + } + let tag = number(bytes, field, 2)?; + let len = number(bytes, field + 2, 2)? as usize; + field += 4; + if len > extra_end - field { + return Err(invalid()); + } + // The Unicode path extra field can replace the raw name. Bound + // it even when its CRC/version would later cause it to be ignored. + if tag == 0x7075 && len >= 5 { + displayed = displayed.max(len - 5); + } + field += len; + } + names += displayed; + if displayed > 4096 || names > 1024 * 1024 { + return Ok(Some(ArchiveTruncation::NameBytes)); + } + cursor = next; + } + if cursor != directory_end { + return Err(invalid()); + } + Ok(None) +} + +/// A reader that hides payload footers while the ZIP parser indexes metadata. +/// This prevents fallback to unchecked footers embedded in a member's bytes. +pub(in crate::multimodal) struct AdmittedReader<'a> { + cursor: std::io::Cursor<&'a [u8]>, + start: u64, + indexing: std::rc::Rc>, +} +impl std::io::Read for AdmittedReader<'_> { + fn read(&mut self, buffer: &mut [u8]) -> std::io::Result { + let position = self.cursor.position(); + let read = std::io::Read::read(&mut self.cursor, buffer)?; + if self.indexing.get() && position < self.start { + let hidden = ((self.start - position).min(read as u64)) as usize; + // Local header reads are needed during indexing. Hide only footer + // magic, including signatures crossing a read-buffer boundary. + for (index, byte) in buffer[..hidden].iter_mut().enumerate() { + let offset = position as usize + index; + if self + .cursor + .get_ref() + .get(offset..offset + 4) + .is_some_and(|magic| magic == b"PK\x05\x06" || magic == b"PK\x06\x06") + { + *byte = 0; + } + } + } + Ok(read) + } +} +impl std::io::Seek for AdmittedReader<'_> { + fn seek(&mut self, position: std::io::SeekFrom) -> std::io::Result { + std::io::Seek::seek(&mut self.cursor, position) + } +} + +/// Construct only after successful preflight. Pin the metadata offset and hide +/// all earlier member footer signatures until indexing finishes; payload reads then work +/// normally, including CRC and decompression validation. +pub(in crate::multimodal) fn open_admitted_zip( + bytes: &[u8], +) -> Result>, ArchiveError> { + let end = footer_offset(bytes)?; + let (directory_end, size, relative) = + if end >= 20 && bytes.get(end - 20..end - 16) == Some(b"PK\x06\x07") { + let offset = zip64_eocd_offset(bytes, end - 20)?; + ( + offset, + number(bytes, offset + 40, 8)?, + number(bytes, offset + 48, 8)?, + ) + } else { + ( + end, + number(bytes, end + 12, 4)?, + number(bytes, end + 16, 4)?, + ) + }; + let start = (directory_end as u64) + .checked_sub(size) + .ok_or_else(invalid)?; + let archive_offset = start.checked_sub(relative).ok_or_else(invalid)?; + let indexing = std::rc::Rc::new(std::cell::Cell::new(true)); + let reader = AdmittedReader { + cursor: std::io::Cursor::new(bytes), + start, + indexing: indexing.clone(), + }; + let config = zip::read::Config { + archive_offset: zip::read::ArchiveOffset::Known(archive_offset), + }; + let archive = zip::ZipArchive::with_config(config, reader) + .map_err(|error| ArchiveError::Invalid(error.to_string()))?; + indexing.set(false); + Ok(archive) +} diff --git a/crates/tinyagents-harness/src/multimodal/mime.rs b/crates/tinyagents-harness/src/multimodal/mime.rs index 0a39bd802..d3b176dc4 100644 --- a/crates/tinyagents-harness/src/multimodal/mime.rs +++ b/crates/tinyagents-harness/src/multimodal/mime.rs @@ -1,22 +1,14 @@ //! MIME detection for image and file attachments. //! -//! Three signals, consulted in a deliberate order: +//! Images and files use different MIME signals. For images the header wins +//! outright, then the extension, then magic bytes. For files a recognized +//! content type or file signature wins before extensions; unknown or generic +//! headers defer to signatures and extensions because many servers answer +//! `application/octet-stream` for everything. //! -//! 1. **The `Content-Type` header**, when one came from a fetch. -//! 2. **The file extension**, when the reference was a path. -//! 3. **Magic bytes.** -//! -//! Images and files order these differently, and the difference is not an -//! oversight. For images the header wins outright: an image server is -//! authoritative about what it served. For files the header wins only if it -//! names a format the allowlist could contain — a great many servers answer -//! `application/octet-stream` for everything, and taking that at face value -//! would degrade every fetched PDF to a metadata-only reference. -//! -//! [`file_mime_from_magic`] cannot separate the OOXML formats from a plain -//! zip: `.xlsx`, `.docx`, `.pptx` and `.zip` all begin `PK\x03\x04`, and -//! telling them apart means parsing the central directory. The extension is -//! what discriminates, which is why it is consulted before magic. +//! Office containers are recognized from standard parts in the ZIP central +//! directory before generic ZIP headers are consulted. An extension remains a +//! fallback for incomplete containers; no member contents are extracted. use std::path::Path; @@ -47,19 +39,50 @@ pub fn detect_image_mime( image_mime_from_magic(bytes).map(ToString::to_string) } -/// Detect a file's MIME type: header (only if recognised), then extension, -/// then magic bytes, then a UTF-8 sniff. +/// Detect a file's MIME type: Office container parts, known header, extension, +/// magic bytes, then a UTF-8 sniff. Generic ZIP/octet-stream headers defer to +/// Office extensions. pub fn detect_file_mime( path: Option<&Path>, bytes: &[u8], header_content_type: Option<&str>, ) -> Option { + if let Some(mime) = office_container_mime(bytes) { + return Some(mime.to_string()); + } + // A generic ZIP/octet-stream header must not hide an Office extension. + if header_content_type + .and_then(normalize_content_type) + .is_some_and(|mime| mime == "application/zip" || mime == "application/octet-stream") + && let Some(ext) = path + .and_then(|path| path.extension()) + .and_then(|ext| ext.to_str()) + && let Some(mime) = file_mime_from_extension(ext) + && mime.starts_with("application/vnd.openxmlformats-officedocument.") + { + return Some(mime.to_string()); + } if let Some(header_mime) = header_content_type.and_then(normalize_content_type) && file_mime_known(&header_mime) { return Some(header_mime); } + // A recognized signature outranks misleading media extensions. ZIP magic + // remains ambiguous, so an OOXML extension still supplies its subtype. + if let Some(mime) = file_mime_from_magic(bytes) { + let extension_is_ooxml = path + .and_then(|path| path.extension()) + .and_then(|ext| ext.to_str()) + .and_then(file_mime_from_extension) + .is_some_and(|extension_mime| { + extension_mime.starts_with("application/vnd.openxmlformats-officedocument.") + }); + if mime != "application/zip" || !extension_is_ooxml { + return Some(mime.to_string()); + } + } + if let Some(path) = path && let Some(ext) = path.extension().and_then(|value| value.to_str()) && let Some(mime) = file_mime_from_extension(ext) @@ -71,6 +94,9 @@ pub fn detect_file_mime( return Some(mime.to_string()); } + if let Some(mime) = image_mime_from_magic(bytes) { + return Some(mime.to_string()); + } if looks_like_utf8_text(bytes) { return Some("text/plain".to_string()); } @@ -78,6 +104,81 @@ pub fn detect_file_mime( None } +/// Generic attachments preserve explicit media types and recognize local media +/// before the legacy UTF-8 fallback. ZIP/octet-stream still allow Office probing. +pub(super) fn detect_attachment_mime( + path: &Path, + bytes: &[u8], + header: Option<&str>, +) -> Option { + if let Some(mime) = header.and_then(normalize_content_type) + && mime != "application/zip" + && mime != "application/octet-stream" + { + return Some(mime); + } + let extension = path + .extension() + .and_then(|value| value.to_str()) + .map(str::to_ascii_lowercase); + let media = if extension.as_deref() == Some("ogv") && bytes.starts_with(b"OggS") { + Some("video/ogg") + } else { + attachment_media_magic(bytes) + .or_else(|| image_mime_from_magic(bytes)) + .or(match extension.as_deref() { + Some("wav" | "wave") => Some("audio/wav"), + Some("mp3") => Some("audio/mpeg"), + Some("flac") => Some("audio/flac"), + Some("ogg" | "oga") => Some("audio/ogg"), + Some("ogv") => Some("video/ogg"), + Some("opus") => Some("audio/opus"), + Some("m4a" | "m4b" | "m4p") => Some("audio/mp4"), + Some("avif") => Some("image/avif"), + Some("heic") => Some("image/heic"), + Some("heif") => Some("image/heif"), + Some("aac") => Some("audio/aac"), + Some("mp4" | "m4v") => Some("video/mp4"), + Some("mov") => Some("video/quicktime"), + Some("webm") => Some("video/webm"), + Some("mkv") => Some("video/x-matroska"), + Some("avi") => Some("video/x-msvideo"), + _ => None, + }) + }; + media + .map(ToString::to_string) + .or_else(|| detect_file_mime(Some(path), bytes, header)) +} + +fn attachment_media_magic(bytes: &[u8]) -> Option<&'static str> { + if bytes.starts_with(b"RIFF") && bytes.get(8..12) == Some(b"WAVE") { + Some("audio/wav") + } else if bytes.starts_with(b"RIFF") && bytes.get(8..12) == Some(b"AVI ") { + Some("video/x-msvideo") + } else if bytes.starts_with(b"ID3") { + Some("audio/mpeg") + } else if bytes.starts_with(b"fLaC") { + Some("audio/flac") + } else if bytes.starts_with(b"OggS") { + Some("audio/ogg") + } else if bytes.get(4..8) == Some(b"ftyp") { + match bytes.get(8..12) { + Some(b"M4A " | b"M4B " | b"M4P ") => Some("audio/mp4"), + Some(b"qt ") => Some("video/quicktime"), + Some(b"avif" | b"avis") => Some("image/avif"), + Some(b"heic" | b"heix") => Some("image/heic"), + Some(b"mif1" | b"msf1") => Some("image/heif"), + Some(b"isom" | b"iso2" | b"mp41" | b"mp42" | b"M4V " | b"avc1" | b"dash") => { + Some("video/mp4") + } + _ => None, + } + } else { + None + } +} + /// Strip parameters from a `Content-Type` header and lower-case it. pub fn normalize_content_type(content_type: &str) -> Option { let mime = content_type.split(';').next()?.trim().to_ascii_lowercase(); @@ -145,6 +246,8 @@ pub fn file_mime_known(mime: &str) -> bool { | "text/csv" | "text/markdown" | "application/zip" + | "application/x-tar" + | "application/gzip" | "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" | "application/vnd.openxmlformats-officedocument.wordprocessingml.document" | "application/vnd.openxmlformats-officedocument.presentationml.presentation" @@ -160,6 +263,8 @@ pub fn file_mime_from_extension(ext: &str) -> Option<&'static str> { "md" | "markdown" => Some("text/markdown"), "csv" => Some("text/csv"), "zip" => Some("application/zip"), + "tar" => Some("application/x-tar"), + "gz" | "tgz" => Some("application/gzip"), "xlsx" => Some("application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"), "docx" => Some("application/vnd.openxmlformats-officedocument.wordprocessingml.document"), "pptx" => Some("application/vnd.openxmlformats-officedocument.presentationml.presentation"), @@ -184,6 +289,50 @@ pub fn file_mime_from_magic(bytes: &[u8]) -> Option<&'static str> { return Some("application/zip"); } + if bytes.starts_with(b"\x1f\x8b") { + return Some("application/gzip"); + } + if bytes.get(257..262) == Some(b"ustar") { + return Some("application/x-tar"); + } + None +} + +/// Identify standard Office parts using bounded container metadata, without +/// reading document contents. Invalid/incomplete ZIPs retain extension-based +/// detection for compatibility. No archive member is extracted. +fn office_container_mime(bytes: &[u8]) -> Option<&'static str> { + if bytes.len() > 50 * 1024 * 1024 || !bytes.starts_with(b"PK\x03\x04") { + return None; + } + if super::archive::zip_admission::zip_preflight(bytes) + .ok()? + .is_some() + { + return None; + } + let archive = super::archive::zip_admission::open_admitted_zip(bytes).ok()?; + if archive.len() > 10_000 || archive.index_for_name("[Content_Types].xml").is_none() { + return None; + } + for (part, mime) in [ + ( + "word/document.xml", + "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + ), + ( + "xl/workbook.xml", + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + ), + ( + "ppt/presentation.xml", + "application/vnd.openxmlformats-officedocument.presentationml.presentation", + ), + ] { + if archive.index_for_name(part).is_some() { + return Some(mime); + } + } None } diff --git a/crates/tinyagents-harness/src/multimodal/mod.rs b/crates/tinyagents-harness/src/multimodal/mod.rs index fe9901c31..ac1afdc36 100644 --- a/crates/tinyagents-harness/src/multimodal/mod.rs +++ b/crates/tinyagents-harness/src/multimodal/mod.rs @@ -1,4 +1,10 @@ -//! Attachment resolution for `[IMAGE:…]` and `[FILE:…]` markers. +//! Generic attachment intake, bounded archive inspection, and legacy +//! resolution for `[IMAGE:…]` and `[FILE:…]` markers. +//! +//! [`resolve_attachment`] returns original decoded bytes and metadata before +//! extraction. [`inspect_archive`] lists ZIP/TAR/TAR.GZ contents without disk +//! extraction. Hosts authorize paths and select storage and provider media +//! blocks; the APIs here introduce no new transcript variants. //! //! A user attaches a picture or a document. Somewhere between the text box and //! the provider, that attachment has to become bytes the model can read — @@ -64,19 +70,27 @@ //! clamped value admits exactly one attachment from a source that asked for //! zero. +pub mod archive; pub mod config; pub mod data_uri; pub mod error; pub mod markers; pub mod mime; pub mod payload; +pub mod png; +pub use png::optimize_png_lossless; pub mod resolve; pub mod stash; +pub mod types; #[cfg(test)] #[path = "mod_tests.rs"] mod test; +pub use archive::{ + ArchiveEntry, ArchiveEntryKind, ArchiveError, ArchiveFormat, ArchiveLimits, ArchiveListing, + ArchiveTruncation, inspect_archive, +}; pub use config::{ALLOWED_IMAGE_MIME_TYPES, FileLimits, ImageLimits}; pub use error::{MultimodalError, Result}; pub use markers::{ @@ -86,5 +100,9 @@ pub use markers::{ text_has_image_placeholders, }; pub use payload::{FilePayload, compose_multimodal_message, sha256_prefix, truncate_chars}; -pub use resolve::{NoTextExtractor, TextExtractor, resolve_file, resolve_image}; +pub use resolve::{ + NoTextExtractor, TextExtractor, resolve_attachment, resolve_file, resolve_image, +}; pub use stash::AttachmentStash; + +pub use types::{ResolvedAttachment, UnknownMimePolicy}; diff --git a/crates/tinyagents-harness/src/multimodal/mod_tests.rs b/crates/tinyagents-harness/src/multimodal/mod_tests.rs index a0890a149..1b8a7624c 100644 --- a/crates/tinyagents-harness/src/multimodal/mod_tests.rs +++ b/crates/tinyagents-harness/src/multimodal/mod_tests.rs @@ -265,6 +265,27 @@ fn the_extension_is_what_separates_ooxml_from_a_plain_zip() { ); } +#[test] +fn recognized_file_signatures_precede_misleading_extensions() { + for (name, bytes, expected) in [ + ( + "report.mp4", + b"%PDF-1.7 document".as_slice(), + "application/pdf", + ), + ( + "archive.txt", + b"PK\x03\x04archive".as_slice(), + "application/zip", + ), + ] { + assert_eq!( + detect_file_mime(Some(&PathBuf::from(name)), bytes, None).as_deref(), + Some(expected) + ); + } +} + #[test] fn an_unlabeled_text_file_falls_back_to_the_utf8_sniff() { assert!(looks_like_utf8_text(b"log line\n")); diff --git a/crates/tinyagents-harness/src/multimodal/png.rs b/crates/tinyagents-harness/src/multimodal/png.rs new file mode 100644 index 000000000..a66f92b9f --- /dev/null +++ b/crates/tinyagents-harness/src/multimodal/png.rs @@ -0,0 +1,70 @@ +//! Optional faithful PNG derivatives. Originals stay under host ownership. + +/// Return a smaller losslessly recompressed PNG, or `None` when optimization +/// is unsupported, malformed, oversized, animated, or offers no saving. +/// +/// This only changes IDAT compression/filtering. Pixel format, hidden RGB under +/// transparent pixels, interlacing, and every other chunk (including EXIF and +/// color metadata) remain byte-for-byte unchanged. The input is never modified. +/// Work is bounded to 32 MiB of input and 16 million pixels; hosts may enforce +/// tighter upload limits. This helper neither decodes nor changes other formats. +pub fn optimize_png_lossless(bytes: &[u8]) -> Option> { + if bytes.len() > 32 * 1024 * 1024 || bytes.get(..8)? != b"\x89PNG\r\n\x1a\n" { + return None; + } + if bytes.get(8..16)? != b"\0\0\0\rIHDR" { + return None; + } + let width = u32::from_be_bytes(bytes.get(16..20)?.try_into().ok()?); + let height = u32::from_be_bytes(bytes.get(20..24)?.try_into().ok()?); + if width == 0 || height == 0 || u64::from(width) * u64::from(height) > 16_000_000 { + return None; + } + let original_chunks = unchanged_chunks(bytes)?; + let options = oxipng::Options { + optimize_alpha: false, + bit_depth_reduction: false, + color_type_reduction: false, + palette_reduction: false, + grayscale_reduction: false, + scale_16: false, + interlace: None, + strip: oxipng::StripChunks::None, + timeout: Some(std::time::Duration::from_secs(5)), + ..oxipng::Options::from_preset(2) + }; + let optimized = oxipng::optimize_from_memory(bytes, &options).ok()?; + if optimized.len() >= bytes.len() || unchanged_chunks(&optimized)? != original_chunks { + return None; + } + Some(optimized) +} + +// Even if an optimizer starts rewriting a metadata chunk in a future release, +// the derivative is admitted only when every non-IDAT chunk stays identical. +fn unchanged_chunks(bytes: &[u8]) -> Option> { + let mut position = 8usize; + let mut chunks = Vec::new(); + while position < bytes.len() { + let header = bytes.get(position..position.checked_add(8)?)?; + let size = u32::from_be_bytes(header[..4].try_into().ok()?) as usize; + let end = position.checked_add(size)?.checked_add(12)?; + let chunk = bytes.get(position..end)?; + let name = &header[4..8]; + if (name == b"IHDR" && position != 8) || matches!(name, b"acTL" | b"fcTL" | b"fdAT") { + return None; + } + if name != b"IDAT" { + chunks.push(chunk); + } + position = end; + if name == b"IEND" { + return (position == bytes.len()).then_some(chunks); + } + } + None +} + +#[cfg(test)] +#[path = "png_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness/src/multimodal/png_tests.rs b/crates/tinyagents-harness/src/multimodal/png_tests.rs new file mode 100644 index 000000000..79d3ef9d8 --- /dev/null +++ b/crates/tinyagents-harness/src/multimodal/png_tests.rs @@ -0,0 +1,79 @@ +use super::*; + +fn fixture() -> Vec { + let mut out = Vec::new(); + { + let mut encoder = ::png::Encoder::new(&mut out, 64, 64); + encoder.set_color(::png::ColorType::Rgba); + encoder.set_depth(::png::BitDepth::Eight); + encoder.set_compression(::png::Compression::Fast); + encoder + .add_text_chunk("Author".into(), "Original metadata".into()) + .unwrap(); + let mut writer = encoder.write_header().unwrap(); + writer + .write_chunk( + ::png::chunk::ChunkType(*b"eXIf"), + b"opaque orientation metadata", + ) + .unwrap(); + let pixels = [123, 45, 67, 0].repeat(64 * 64); + writer.write_image_data(&pixels).unwrap(); + } + out +} + +fn decoded(bytes: &[u8]) -> Vec { + let mut reader = ::png::Decoder::new(std::io::Cursor::new(bytes)) + .read_info() + .unwrap(); + let mut pixels = vec![0; reader.output_buffer_size()]; + let info = reader.next_frame(&mut pixels).unwrap(); + pixels.truncate(info.buffer_size()); + pixels +} + +#[test] +fn png_derivative_is_smaller_preserves_hidden_rgb_and_metadata_and_original() { + let original = fixture(); + let snapshot = original.clone(); + let optimized = optimize_png_lossless(&original).unwrap(); + assert!(optimized.len() < original.len()); + assert_eq!(decoded(&optimized), decoded(&original)); + assert!( + optimized + .windows(b"Original metadata".len()) + .any(|w| w == b"Original metadata") + ); + assert!( + optimized + .windows(b"opaque orientation metadata".len()) + .any(|w| w == b"opaque orientation metadata") + ); + assert_eq!(unchanged_chunks(&optimized), unchanged_chunks(&original)); + assert_eq!(original, snapshot); + assert!(optimize_png_lossless(&optimized).is_none()); +} + +#[test] +fn animation_malformed_non_png_and_excessive_dimensions_are_skipped() { + assert!(optimize_png_lossless(b"jpeg").is_none()); + assert!(optimize_png_lossless(b"\x89PNG\r\n\x1a\n").is_none()); + let mut animated = fixture(); + // A structurally complete animation-control chunk need not be decoded: + // the optional derivative must skip it before optimization. + animated.splice( + 33..33, + [ + 0, 0, 0, 8, b'a', b'c', b'T', b'L', 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, + ], + ); + assert!(optimize_png_lossless(&animated).is_none()); + let mut duplicate_header = fixture(); + let second_header = duplicate_header[8..33].to_vec(); + duplicate_header.splice(33..33, second_header); + assert!(optimize_png_lossless(&duplicate_header).is_none()); + let mut giant = fixture(); + giant[16..20].copy_from_slice(&u32::MAX.to_be_bytes()); + assert!(optimize_png_lossless(&giant).is_none()); +} diff --git a/crates/tinyagents-harness/src/multimodal/resolve.rs b/crates/tinyagents-harness/src/multimodal/resolve.rs index 7b4ae3c39..8a1c441a2 100644 --- a/crates/tinyagents-harness/src/multimodal/resolve.rs +++ b/crates/tinyagents-harness/src/multimodal/resolve.rs @@ -40,6 +40,7 @@ use super::data_uri::{data_uri_param, encode_data_uri, gunzip, parse_data_uri}; use super::error::{MultimodalError, Result}; use super::mime::{detect_file_mime, detect_image_mime, is_allowed_image_mime}; use super::payload::FilePayload; +use super::types::{ResolvedAttachment, UnknownMimePolicy}; /// Host-supplied text extraction for formats this crate cannot decode itself. /// @@ -287,71 +288,127 @@ pub async fn resolve_file( remote_client: &Client, extractor: &dyn TextExtractor, ) -> Result { - if source.starts_with("data:") { - let (bytes, name, mime) = resolve_file_data_uri(source, max_bytes)?; - return build_file_payload( - source, - bytes, - name, - mime, - limits, - max_extracted_text_chars, - extractor, - ) - .await; - } + let resolved = resolve_attachment( + source, + limits, + max_bytes, + remote_client, + UnknownMimePolicy::Reject, + ) + .await?; + build_file_payload( + source, + resolved.bytes, + resolved.name, + resolved.mime, + limits, + max_extracted_text_chars, + extractor, + ) + .await +} - let (bytes, path_hint, name, header_content_type) = - if source.starts_with("http://") || source.starts_with("https://") { - if !limits.allow_remote_fetch { - return Err(MultimodalError::RemoteFileFetchDisabled { - input: source.to_string(), - }); - } - let (bytes, name, content_type) = - fetch_remote_file(source, max_bytes, remote_client).await?; - (bytes, None, name, content_type) +/// Resolve bytes and metadata without extracting text or writing to disk. +/// +/// The host must authorize local paths before calling. Unknown MIME acceptance +/// is explicit and does not affect legacy [`resolve_file`] callers. Only a gzip +/// data URI with `original_mime` is a transport envelope: a gzip attachment +/// without that parameter remains compressed for archive inspection. +pub async fn resolve_attachment( + source: &str, + limits: &FileLimits, + max_bytes: usize, + remote_client: &Client, + unknown_mime: UnknownMimePolicy, +) -> Result { + if limits.files_disabled() { + return Err(MultimodalError::TooManyFiles { + max_files: 0, + found: 1, + }); + } + let max_bytes = max_bytes.min(limits.max_file_bytes()); + let (bytes, name, mime) = if source.starts_with("data:") { + resolve_file_data_uri(source, max_bytes, unknown_mime)? + } else { + let (bytes, name, header) = + if source.starts_with("http://") || source.starts_with("https://") { + if !limits.allow_remote_fetch { + return Err(MultimodalError::RemoteFileFetchDisabled { + input: source.to_string(), + }); + } + fetch_remote_file(source, max_bytes, remote_client).await? + } else { + let (bytes, _path, name) = read_local_file(source, max_bytes).await?; + (bytes, name, None) + }; + let detected = if unknown_mime == UnknownMimePolicy::Accept { + super::mime::detect_attachment_mime(Path::new(&name), &bytes, header.as_deref()) } else { - let (bytes, path, name) = read_local_file(source, max_bytes).await?; - (bytes, Some(path), name, None) + detect_file_mime(Some(Path::new(&name)), &bytes, header.as_deref()) }; - - let mime = detect_file_mime(path_hint.as_deref(), &bytes, header_content_type.as_deref()) - .ok_or_else(|| MultimodalError::UnsupportedFileMime { + if detected.is_none() && unknown_mime == UnknownMimePolicy::Reject { + return Err(MultimodalError::UnsupportedFileMime { + input: source.to_string(), + mime: "unknown".to_string(), + supported: limits.supported_rendered(), + }); + } + let mime = detected + .or_else(|| { + header + .as_deref() + .and_then(super::mime::normalize_content_type) + }) + .unwrap_or_else(|| "application/octet-stream".to_string()); + (bytes, name, mime) + }; + if !limits.is_mime_allowed(&mime) && unknown_mime == UnknownMimePolicy::Reject { + return Err(MultimodalError::UnsupportedFileMime { input: source.to_string(), - mime: "unknown".to_string(), + mime, supported: limits.supported_rendered(), - })?; - - build_file_payload( - source, + }); + } + Ok(ResolvedAttachment { + size_bytes: bytes.len(), bytes, name, mime, - limits, - max_extracted_text_chars, - extractor, - ) - .await + }) } /// Resolve a `data:` file source: decompresses a gzip-wrapped payload when /// present, extracts the `name` parameter, and checks size. MIME allowlisting /// happens later in [`build_file_payload`]. -fn resolve_file_data_uri(source: &str, max_bytes: usize) -> Result<(Vec, String, String)> { +fn resolve_file_data_uri( + source: &str, + max_bytes: usize, + unknown_mime: UnknownMimePolicy, +) -> Result<(Vec, String, String)> { + // Reject payloads whose minimum decoded size is over cap before allocating. + if let Some((header, encoded)) = source.split_once(',') { + let minimum_size = if header + .split(';') + .any(|part| part.trim().eq_ignore_ascii_case("base64")) + { + (encoded.len() / 4).saturating_mul(3).saturating_sub(2) + } else { + encoded.len() / 3 + }; + check_file_size(source, minimum_size, max_bytes)?; + } let parsed = parse_data_uri(source).map_err(|reason| MultimodalError::InvalidFileMarker { input: source.to_string(), reason, })?; let name = data_uri_param(&parsed.params, "name").unwrap_or_else(|| "attachment".to_string()); - let (mime, bytes) = if parsed.mime == "application/gzip" { - let original_mime = data_uri_param(&parsed.params, "original_mime").ok_or_else(|| { - MultimodalError::InvalidFileMarker { - input: source.to_string(), - reason: "compressed file data URI missing original_mime parameter".to_string(), - } - })?; + check_file_size(source, parsed.bytes.len(), max_bytes)?; + let (mime, bytes) = if parsed.mime == "application/gzip" + && let Some(original_mime) = data_uri_param(&parsed.params, "original_mime") + { let bytes = gunzip(&parsed.bytes, max_bytes).map_err(|reason| { MultimodalError::InvalidFileMarker { input: source.to_string(), @@ -365,6 +422,19 @@ fn resolve_file_data_uri(source: &str, max_bytes: usize) -> Result<(Vec, Str check_file_size(source, bytes.len(), max_bytes)?; + let mime = if mime.is_empty() { + // RFC 2397 defaults an omitted media type to text/plain. + "text/plain".to_string() + } else if mime == "application/zip" || mime == "application/octet-stream" { + if unknown_mime == UnknownMimePolicy::Accept { + super::mime::detect_attachment_mime(Path::new(&name), &bytes, Some(&mime)) + } else { + detect_file_mime(Some(Path::new(&name)), &bytes, Some(&mime)) + } + .unwrap_or(mime) + } else { + mime + }; Ok((bytes, name, mime)) } @@ -461,13 +531,21 @@ async fn read_local_file(source: &str, max_bytes: usize) -> Result<(Vec, Pat check_file_size(source, metadata.len() as usize, max_bytes)?; - let bytes = - tokio::fs::read(&path) - .await - .map_err(|error| MultimodalError::LocalFileReadFailed { - input: source.to_string(), - reason: error.to_string(), - })?; + use tokio::io::AsyncReadExt; + let file = tokio::fs::File::open(&path).await.map_err(|error| { + MultimodalError::LocalFileReadFailed { + input: source.to_string(), + reason: error.to_string(), + } + })?; + let mut bytes = Vec::new(); + file.take(max_bytes.saturating_add(1) as u64) + .read_to_end(&mut bytes) + .await + .map_err(|error| MultimodalError::LocalFileReadFailed { + input: source.to_string(), + reason: error.to_string(), + })?; check_file_size(source, bytes.len(), max_bytes)?; @@ -514,25 +592,34 @@ async fn fetch_remote_file( .and_then(|value| value.to_str().ok()) .map(ToString::to_string); - let bytes = response - .bytes() - .await - .map_err(|error| MultimodalError::RemoteFileFetchFailed { - input: source.to_string(), - reason: error.to_string(), - })?; - - check_file_size(source, bytes.len(), max_bytes)?; + let mut response = response; + let mut bytes = Vec::new(); + while let Some(chunk) = + response + .chunk() + .await + .map_err(|error| MultimodalError::RemoteFileFetchFailed { + input: source.to_string(), + reason: error.to_string(), + })? + { + check_file_size(source, bytes.len().saturating_add(chunk.len()), max_bytes)?; + bytes.extend_from_slice(&chunk); + } // The last path segment, not a `Content-Disposition` filename: the header // is attacker-controlled on a fetched URL and has its own escaping rules, // and the payload header escapes whatever lands here anyway. - let name = source - .rsplit('/') - .next() + let name = reqwest::Url::parse(source) + .ok() + .and_then(|url| { + url.path_segments() + .and_then(|mut segments| segments.next_back()) + .map(str::to_owned) + }) .filter(|segment| !segment.is_empty()) - .unwrap_or(source) - .to_string(); + .and_then(|segment| super::data_uri::percent_decode(&segment)) + .unwrap_or_else(|| "attachment".to_string()); Ok((bytes.to_vec(), name, content_type)) } @@ -548,3 +635,7 @@ fn check_file_size(source: &str, size_bytes: usize, max_bytes: usize) -> Result< } Ok(()) } + +#[cfg(test)] +#[path = "resolve_tests.rs"] +mod tests; diff --git a/crates/tinyagents-harness/src/multimodal/resolve_tests.rs b/crates/tinyagents-harness/src/multimodal/resolve_tests.rs new file mode 100644 index 000000000..f638503ae --- /dev/null +++ b/crates/tinyagents-harness/src/multimodal/resolve_tests.rs @@ -0,0 +1,449 @@ +use super::*; +use base64::{Engine as _, engine::general_purpose::STANDARD}; +use std::io::Write; + +#[tokio::test] +async fn generic_resolution_retains_binary_bytes_and_decodes_only_transport_gzip() { + let bytes = [0, 255, 7, 0]; + let mut gzip = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + gzip.write_all(&bytes).unwrap(); + let uri = format!( + "data:application/gzip;original_mime=application/octet-stream;name=raw%20bytes.bin;base64,{}", + STANDARD.encode(gzip.finish().unwrap()) + ); + let resolved = resolve_attachment( + &uri, + &FileLimits::default(), + 1024, + &Client::new(), + UnknownMimePolicy::Reject, + ) + .await + .unwrap(); + assert_eq!(resolved.bytes, bytes); + assert_eq!(resolved.name, "raw bytes.bin"); + assert_eq!(resolved.mime, "application/octet-stream"); + assert_eq!(resolved.size_bytes, 4); +} + +#[tokio::test] +async fn unknown_types_need_explicit_host_opt_in() { + let uri = "data:application/x-private;name=a.custom;base64,AP8="; + let limits = FileLimits::default(); + assert!( + resolve_attachment( + uri, + &limits, + 1024, + &Client::new(), + UnknownMimePolicy::Reject + ) + .await + .is_err() + ); + let resolved = resolve_attachment( + uri, + &limits, + 1024, + &Client::new(), + UnknownMimePolicy::Accept, + ) + .await + .unwrap(); + assert_eq!(resolved.mime, "application/x-private"); + assert_eq!(resolved.bytes, [0, 255]); + assert!( + resolve_file(uri, &limits, 1024, 1000, &Client::new(), &NoTextExtractor) + .await + .is_err() + ); +} + +#[tokio::test] +async fn disabled_files_and_transport_expansion_are_rejected() { + let limits = FileLimits { + max_files: 0, + ..FileLimits::default() + }; + assert!( + resolve_attachment( + "data:text/plain,hello", + &limits, + 1024, + &Client::new(), + UnknownMimePolicy::Accept + ) + .await + .is_err() + ); + let mut gzip = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + gzip.write_all(&vec![b'a'; 4096]).unwrap(); + let uri = format!( + "data:application/gzip;original_mime=text/plain;base64,{}", + STANDARD.encode(gzip.finish().unwrap()) + ); + assert!( + resolve_attachment( + &uri, + &FileLimits::default(), + 100, + &Client::new(), + UnknownMimePolicy::Accept + ) + .await + .is_err() + ); +} + +#[tokio::test] +async fn archive_gzip_is_retained_and_nested_gzip_is_not_decoded_twice() { + let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + encoder.write_all(b"tar-like payload").unwrap(); + let original = encoder.finish().unwrap(); + let mut outer = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + outer.write_all(&original).unwrap(); + for uri in [ + format!( + "data:application/gzip;name=a.tar.gz;base64,{}", + STANDARD.encode(&original) + ), + format!( + "data:application/gzip;original_mime=application/gzip;name=a.tar.gz;base64,{}", + STANDARD.encode(outer.finish().unwrap()) + ), + ] { + let resolved = resolve_attachment( + &uri, + &FileLimits::default(), + 1024, + &Client::new(), + UnknownMimePolicy::Accept, + ) + .await + .unwrap(); + assert_eq!(resolved.bytes, original); + assert_eq!(resolved.name, "a.tar.gz"); + assert_eq!(resolved.mime, "application/gzip"); + } +} + +#[tokio::test] +async fn local_unknown_bytes_preserve_legacy_rejection_and_generic_metadata() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("opaque.custom"); + tokio::fs::write(&path, [0, 255, 7]).await.unwrap(); + let source = path.to_str().unwrap(); + assert!( + resolve_attachment( + source, + &FileLimits::default(), + 1024, + &Client::new(), + UnknownMimePolicy::Reject + ) + .await + .is_err() + ); + let resolved = resolve_attachment( + source, + &FileLimits::default(), + 1024, + &Client::new(), + UnknownMimePolicy::Accept, + ) + .await + .unwrap(); + assert_eq!(resolved.bytes, [0, 255, 7]); + assert_eq!(resolved.mime, "application/octet-stream"); + assert_eq!(resolved.name, "opaque.custom"); + assert!(matches!( + resolve_attachment( + source, + &FileLimits::default(), + 2, + &Client::new(), + UnknownMimePolicy::Accept + ) + .await, + Err(MultimodalError::FileTooLarge { .. }) + )); +} + +#[tokio::test] +async fn opt_in_retains_images_audio_video_and_native_documents_verbatim() { + for (mime, bytes) in [ + ("image/png", &b"\x89PNG\r\n\x1a\n"[..]), + ("audio/wav", &b"RIFF\0\0\0\0WAVE"[..]), + ("video/mp4", &b"\0\0\0\x18ftypmp42"[..]), + ("application/pdf", &b"%PDF-1.7"[..]), + ] { + let uri = format!("data:{mime};name=media;base64,{}", STANDARD.encode(bytes)); + let resolved = resolve_attachment( + &uri, + &FileLimits::default(), + 1024, + &Client::new(), + UnknownMimePolicy::Accept, + ) + .await + .unwrap(); + assert_eq!(resolved.mime, mime); + assert_eq!(resolved.bytes, bytes); + } +} + +#[tokio::test] +async fn malformed_and_oversized_data_uri_payloads_fail_before_extraction() { + for uri in [ + "data:text/plain;base64,%%%", + "data:text/plain,bad%Q0", + "data:application/gzip;original_mime=text/plain;base64,AP8=", + ] { + assert!(matches!( + resolve_attachment( + uri, + &FileLimits::default(), + 1024, + &Client::new(), + UnknownMimePolicy::Accept + ) + .await, + Err(MultimodalError::InvalidFileMarker { .. }) + )); + } + assert!(matches!( + resolve_attachment( + "data:text/plain;base64,YWFhYWFhYWFhYWFh", + &FileLimits::default(), + 2, + &Client::new(), + UnknownMimePolicy::Accept + ) + .await, + Err(MultimodalError::FileTooLarge { .. }) + )); +} + +#[tokio::test] +async fn generic_http_mime_precedes_utf8_sniff_and_legacy_stays_narrow() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}/media", listener.local_addr().unwrap()); + let server = tokio::spawn(async move { + for _ in 0..2 { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut request = [0; 2048]; + let mut received = 0; + loop { + let amount = stream.read(&mut request[received..]).await.unwrap(); + assert!(amount > 0); + received += amount; + if request[..received] + .windows(4) + .any(|part| part == b"\r\n\r\n") + { + break; + } + assert!(received < request.len()); + } + stream.write_all(b"HTTP/1.1 200 OK\r\nContent-Type: audio/wav; charset=binary\r\nContent-Length: 12\r\nConnection: close\r\n\r\nRIFF\0\0\0\0WAVE").await.unwrap(); + } + }); + let limits = FileLimits { + allow_remote_fetch: true, + ..FileLimits::default() + }; + let resolved = resolve_attachment( + &url, + &limits, + 1024, + &Client::new(), + UnknownMimePolicy::Accept, + ) + .await + .unwrap(); + assert_eq!(resolved.mime, "audio/wav"); + assert_eq!(resolved.bytes, b"RIFF\0\0\0\0WAVE"); + let legacy = resolve_attachment( + &url, + &limits, + 1024, + &Client::new(), + UnknownMimePolicy::Reject, + ) + .await + .unwrap(); + assert_eq!(legacy.mime, "text/plain"); + server.await.unwrap(); +} + +#[tokio::test] +async fn generic_local_media_uses_extensions_and_magic() { + let dir = tempfile::tempdir().unwrap(); + for (name, bytes, mime) in [ + ("sound.wav", &b"RIFF\0\0\0\0WAVE"[..], "audio/wav"), + ("movie.mp4", &b"\0\0\0\x18ftypmp42"[..], "video/mp4"), + ("recording.mp3", &b"ID3"[..], "audio/mpeg"), + ("sound.flac", &b"fLaC"[..], "audio/flac"), + ("sound.ogg", &b"OggS"[..], "audio/ogg"), + ("clip.ogv", &b"OggS"[..], "video/ogg"), + ("sound.opus", &b"opaque"[..], "audio/opus"), + ("sound.m4a", &b"opaque"[..], "audio/mp4"), + ("sound.aac", &b"opaque"[..], "audio/aac"), + ("movie.mov", &b"opaque"[..], "video/quicktime"), + ("movie.webm", &b"opaque"[..], "video/webm"), + ("movie.mkv", &b"opaque"[..], "video/x-matroska"), + ("movie.avi", &b"opaque"[..], "video/x-msvideo"), + ("magic-mp3", &b"ID3"[..], "audio/mpeg"), + ("magic-flac", &b"fLaC"[..], "audio/flac"), + ("magic-ogg", &b"OggS"[..], "audio/ogg"), + ("magic-avi", &b"RIFF\0\0\0\0AVI "[..], "video/x-msvideo"), + ("magic-mp4", &b"\0\0\0\x18ftypmp42"[..], "video/mp4"), + ("magic-m4a", &b"\0\0\0\x18ftypM4A "[..], "audio/mp4"), + ("magic-mov", &b"\0\0\0\x18ftypqt "[..], "video/quicktime"), + ("unlabelled", &b"RIFF\0\0\0\0WAVE"[..], "audio/wav"), + ] { + let path = dir.path().join(name); + tokio::fs::write(&path, bytes).await.unwrap(); + let resolved = resolve_attachment( + path.to_str().unwrap(), + &FileLimits::default(), + 1024, + &Client::new(), + UnknownMimePolicy::Accept, + ) + .await + .unwrap(); + assert_eq!(resolved.mime, mime); + assert_eq!(resolved.bytes, bytes); + } +} + +#[test] +fn iso_bmff_brands_distinguish_images_audio_and_known_video() { + for (brand, expected) in [ + ("avif", "image/avif"), + ("avis", "image/avif"), + ("heic", "image/heic"), + ("heix", "image/heic"), + ("mif1", "image/heif"), + ("msf1", "image/heif"), + ("M4A ", "audio/mp4"), + ("M4B ", "audio/mp4"), + ("M4P ", "audio/mp4"), + ("qt ", "video/quicktime"), + ("mp42", "video/mp4"), + ("isom", "video/mp4"), + ] { + let bytes = [b"\0\0\0\x18ftyp".as_slice(), brand.as_bytes()].concat(); + assert_eq!( + crate::multimodal::mime::detect_attachment_mime( + std::path::Path::new("attachment"), + &bytes, + None + ) + .as_deref(), + Some(expected), + "{brand}" + ); + } + let unknown = b"\0\0\0\x18ftypzzzz"; + assert_ne!( + crate::multimodal::mime::detect_attachment_mime( + std::path::Path::new("attachment"), + unknown, + None + ) + .as_deref(), + Some("video/mp4") + ); + for (name, expected) in [ + ("photo.avif", "image/avif"), + ("photo.heic", "image/heic"), + ("photo.heif", "image/heif"), + ("book.m4b", "audio/mp4"), + ] { + assert_eq!( + crate::multimodal::mime::detect_attachment_mime( + std::path::Path::new(name), + unknown, + None + ) + .as_deref(), + Some(expected) + ); + } +} + +#[tokio::test] +async fn omitted_data_uri_media_type_uses_rfc_text_default() { + for (uri, expected) in [ + ("data:,hello", b"hello".as_slice()), + ("data:;base64,aGVsbG8=", b"hello".as_slice()), + ] { + for policy in [UnknownMimePolicy::Accept, UnknownMimePolicy::Reject] { + let resolved = + resolve_attachment(uri, &FileLimits::default(), 1024, &Client::new(), policy) + .await + .unwrap(); + assert_eq!(resolved.mime, "text/plain"); + assert_eq!(resolved.bytes, expected); + } + } +} + +#[tokio::test] +async fn generic_data_uris_detect_media_like_local_sources_without_changing_legacy() { + for (name, bytes, expected) in [ + ("clip.mp4", b"\0\0\0\x18ftypmp42".as_slice(), "video/mp4"), + ("photo.avif", b"\0\0\0\x18ftypavif".as_slice(), "image/avif"), + ("sound.wav", b"RIFF\0\0\0\0WAVE".as_slice(), "audio/wav"), + ] { + let uri = format!( + "data:application/octet-stream;name={name};base64,{}", + STANDARD.encode(bytes) + ); + let accepted = resolve_attachment( + &uri, + &FileLimits::default(), + 1024, + &Client::new(), + UnknownMimePolicy::Accept, + ) + .await + .unwrap(); + assert_eq!(accepted.mime, expected); + assert_eq!(accepted.bytes, bytes); + let legacy = resolve_attachment( + &uri, + &FileLimits::default(), + 1024, + &Client::new(), + UnknownMimePolicy::Reject, + ) + .await + .unwrap(); + assert_eq!(legacy.mime, "application/octet-stream"); + assert_eq!(legacy.bytes, bytes); + } +} + +#[test] +fn generic_media_signatures_precede_misleading_extensions() { + for (name, bytes, expected) in [ + ("photo.avif", b"\0\0\0\x18ftypmp42".as_slice(), "video/mp4"), + ("movie.mp4", b"\0\0\0\x18ftypavif".as_slice(), "image/avif"), + ("movie.mp4", b"RIFF\0\0\0\0WAVE".as_slice(), "audio/wav"), + ("movie.mp4", b"\xff\xd8\xffbinary".as_slice(), "image/jpeg"), + ] { + assert_eq!( + crate::multimodal::mime::detect_attachment_mime( + std::path::Path::new(name), + bytes, + None + ) + .as_deref(), + Some(expected) + ); + } +} diff --git a/crates/tinyagents-harness/src/multimodal/types.rs b/crates/tinyagents-harness/src/multimodal/types.rs new file mode 100644 index 000000000..c34970347 --- /dev/null +++ b/crates/tinyagents-harness/src/multimodal/types.rs @@ -0,0 +1,25 @@ +//! Resolved, unconverted attachment data for host-controlled intake. + +/// Validated attachment bytes and display metadata. Names are untrusted +/// display strings, never authorized filesystem destinations. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ResolvedAttachment { + /// Decoded payload; transport gzip has already been removed once. + pub bytes: Vec, + /// Original display name when provided, otherwise `attachment`. + pub name: String, + /// Normalized MIME type; undetected bytes use `application/octet-stream`. + pub mime: String, + /// Decoded byte count, equal to `bytes.len()`. + pub size_bytes: usize, +} + +/// Whether a generic intake host admits types outside its file allowlist. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum UnknownMimePolicy { + /// Require the configured allowlist, as legacy file resolution does. + #[default] + Reject, + /// Admit arbitrary MIME types; size, source, and disable gates still apply. + Accept, +} diff --git a/crates/tinyagents-registry/model-catalog.snapshot.json b/crates/tinyagents-registry/model-catalog.snapshot.json index dbed791cf..c8b5e5bb4 100644 --- a/crates/tinyagents-registry/model-catalog.snapshot.json +++ b/crates/tinyagents-registry/model-catalog.snapshot.json @@ -41,7 +41,8 @@ "audio_output": false, "pdf_input": false, "prompt_caching": true, - "reasoning": false + "reasoning": false, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -107,7 +108,8 @@ "audio_output": false, "pdf_input": false, "prompt_caching": false, - "reasoning": false + "reasoning": false, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -172,7 +174,8 @@ "audio_output": false, "pdf_input": false, "prompt_caching": false, - "reasoning": false + "reasoning": false, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -240,7 +243,8 @@ "audio_output": false, "pdf_input": true, "prompt_caching": true, - "reasoning": false + "reasoning": false, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -307,7 +311,8 @@ "audio_output": false, "pdf_input": true, "prompt_caching": true, - "reasoning": false + "reasoning": false, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -374,7 +379,8 @@ "audio_output": false, "pdf_input": false, "prompt_caching": true, - "reasoning": false + "reasoning": false, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -441,7 +447,8 @@ "audio_output": false, "pdf_input": true, "prompt_caching": true, - "reasoning": false + "reasoning": false, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -508,7 +515,8 @@ "audio_output": false, "pdf_input": false, "prompt_caching": false, - "reasoning": false + "reasoning": false, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -574,7 +582,8 @@ "audio_output": false, "pdf_input": false, "prompt_caching": true, - "reasoning": false + "reasoning": false, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -640,7 +649,8 @@ "audio_output": false, "pdf_input": false, "prompt_caching": true, - "reasoning": false + "reasoning": false, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -706,7 +716,8 @@ "audio_output": false, "pdf_input": true, "prompt_caching": true, - "reasoning": true + "reasoning": true, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -785,7 +796,8 @@ "audio_output": false, "pdf_input": true, "prompt_caching": true, - "reasoning": true + "reasoning": true, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -865,7 +877,8 @@ "audio_output": false, "pdf_input": true, "prompt_caching": true, - "reasoning": true + "reasoning": true, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -939,7 +952,8 @@ "audio_output": false, "pdf_input": true, "prompt_caching": true, - "reasoning": true + "reasoning": true, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -1013,7 +1027,8 @@ "audio_output": false, "pdf_input": true, "prompt_caching": true, - "reasoning": true + "reasoning": true, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -1095,7 +1110,8 @@ "audio_output": false, "pdf_input": true, "prompt_caching": true, - "reasoning": true + "reasoning": true, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -1177,7 +1193,8 @@ "audio_output": false, "pdf_input": true, "prompt_caching": true, - "reasoning": true + "reasoning": true, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -1260,7 +1277,8 @@ "audio_output": false, "pdf_input": true, "prompt_caching": true, - "reasoning": true + "reasoning": true, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -1356,7 +1374,8 @@ "audio_output": false, "pdf_input": true, "prompt_caching": true, - "reasoning": true + "reasoning": true, + "video_input": true }, "source": "models.dev", "source_url": null, @@ -1458,7 +1477,8 @@ "audio_output": false, "pdf_input": true, "prompt_caching": true, - "reasoning": true + "reasoning": true, + "video_input": true }, "source": "models.dev", "source_url": null, @@ -1560,7 +1580,8 @@ "audio_output": false, "pdf_input": false, "prompt_caching": false, - "reasoning": true + "reasoning": true, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -1639,7 +1660,8 @@ "audio_output": false, "pdf_input": true, "prompt_caching": true, - "reasoning": true + "reasoning": true, + "video_input": true }, "source": "models.dev", "source_url": null, @@ -1719,7 +1741,8 @@ "audio_output": false, "pdf_input": false, "prompt_caching": true, - "reasoning": true + "reasoning": true, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -1786,7 +1809,8 @@ "audio_output": false, "pdf_input": true, "prompt_caching": true, - "reasoning": true + "reasoning": true, + "video_input": true }, "source": "models.dev", "source_url": null, @@ -1882,7 +1906,8 @@ "audio_output": false, "pdf_input": true, "prompt_caching": true, - "reasoning": true + "reasoning": true, + "video_input": true }, "source": "models.dev", "source_url": null, @@ -1974,7 +1999,8 @@ "audio_output": false, "pdf_input": true, "prompt_caching": true, - "reasoning": true + "reasoning": true, + "video_input": true }, "source": "models.dev", "source_url": null, @@ -2055,7 +2081,8 @@ "audio_output": false, "pdf_input": false, "prompt_caching": false, - "reasoning": false + "reasoning": false, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -2118,7 +2145,8 @@ "audio_output": false, "pdf_input": false, "prompt_caching": false, - "reasoning": false + "reasoning": false, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -2182,7 +2210,8 @@ "audio_output": false, "pdf_input": false, "prompt_caching": false, - "reasoning": false + "reasoning": false, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -2246,7 +2275,8 @@ "audio_output": false, "pdf_input": false, "prompt_caching": false, - "reasoning": false + "reasoning": false, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -2310,7 +2340,8 @@ "audio_output": false, "pdf_input": false, "prompt_caching": false, - "reasoning": false + "reasoning": false, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -2374,7 +2405,8 @@ "audio_output": false, "pdf_input": false, "prompt_caching": false, - "reasoning": false + "reasoning": false, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -2438,7 +2470,8 @@ "audio_output": false, "pdf_input": false, "prompt_caching": false, - "reasoning": false + "reasoning": false, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -2499,7 +2532,8 @@ "audio_output": false, "pdf_input": false, "prompt_caching": false, - "reasoning": false + "reasoning": false, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -2562,7 +2596,8 @@ "audio_output": false, "pdf_input": false, "prompt_caching": false, - "reasoning": false + "reasoning": false, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -2625,7 +2660,8 @@ "audio_output": false, "pdf_input": false, "prompt_caching": false, - "reasoning": false + "reasoning": false, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -2688,7 +2724,8 @@ "audio_output": false, "pdf_input": false, "prompt_caching": true, - "reasoning": true + "reasoning": true, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -2771,7 +2808,8 @@ "audio_output": false, "pdf_input": false, "prompt_caching": true, - "reasoning": true + "reasoning": true, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -2854,7 +2892,8 @@ "audio_output": false, "pdf_input": false, "prompt_caching": true, - "reasoning": true + "reasoning": true, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -2937,7 +2976,8 @@ "audio_output": false, "pdf_input": false, "prompt_caching": true, - "reasoning": true + "reasoning": true, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -3033,7 +3073,8 @@ "audio_output": false, "pdf_input": true, "prompt_caching": true, - "reasoning": false + "reasoning": false, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -3131,7 +3172,8 @@ "audio_output": false, "pdf_input": true, "prompt_caching": true, - "reasoning": true + "reasoning": true, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -3230,7 +3272,8 @@ "audio_output": false, "pdf_input": true, "prompt_caching": true, - "reasoning": true + "reasoning": true, + "video_input": false }, "source": "models.dev", "source_url": null, @@ -3339,7 +3382,8 @@ "audio_output": false, "pdf_input": true, "prompt_caching": true, - "reasoning": true + "reasoning": true, + "video_input": false }, "source": "models.dev", "source_url": null, diff --git a/crates/tinyagents-registry/src/bin/catalog_gen.rs b/crates/tinyagents-registry/src/bin/catalog_gen.rs index 530305f84..ea31d8c37 100644 --- a/crates/tinyagents-registry/src/bin/catalog_gen.rs +++ b/crates/tinyagents-registry/src/bin/catalog_gen.rs @@ -241,6 +241,7 @@ fn to_entry(catalog_provider: &str, model_id: &str, model: &Value) -> Option tinyinference_llm::model::Mo image_in: caps.vision, audio_in: caps.audio_input, audio_out: caps.audio_output, + video_in: caps.video_input, + document_in: caps.pdf_input, ..Modalities::default() }, tool_calling: caps.tool_calling, diff --git a/crates/tinyagents-registry/src/catalog_tests.rs b/crates/tinyagents-registry/src/catalog_tests.rs index 6cce84f60..4be5631bb 100644 --- a/crates/tinyagents-registry/src/catalog_tests.rs +++ b/crates/tinyagents-registry/src/catalog_tests.rs @@ -208,3 +208,57 @@ fn from_json_rejects_an_invalid_snapshot() { let json = serde_json::to_string(&snapshot).unwrap(); assert!(ModelCatalog::from_json(&json).is_err()); } + +#[test] +fn media_profile_facts_are_explicit_and_unknown_models_stay_unknown() { + let mut entry = base_entry(); + let profile = profile_from_entry(&entry); + assert!(!profile.modalities.image_in); + assert!(!profile.modalities.audio_in); + assert!(!profile.modalities.video_in); + assert!(!profile.modalities.document_in); + entry.capabilities.pdf_input = true; + entry.capabilities.audio_input = true; + entry.capabilities.video_input = true; + let profile = profile_from_entry(&entry); + assert!(profile.modalities.audio_in); + assert!(profile.modalities.video_in); + assert!(profile.modalities.document_in); + assert!( + ModelCatalog::seed() + .unwrap() + .profile("unknown", "unknown") + .is_none() + ); +} + +#[test] +fn bundled_video_capabilities_match_raw_input_modalities() { + let catalog = ModelCatalog::seed().unwrap(); + let mut video_models = 0; + for entry in &catalog.snapshot().models { + let expected = entry.raw["modalities"]["input"] + .as_array() + .is_some_and(|inputs| inputs.iter().any(|input| input == "video")); + assert_eq!( + entry.capabilities.video_input, expected, + "{}", + entry.model_id + ); + assert_eq!( + profile_from_entry(entry).modalities.video_in, + expected, + "{}", + entry.model_id + ); + video_models += usize::from(expected); + } + assert_eq!(video_models, 6); + assert!( + catalog + .profile("gemini", "gemini-2.5-pro") + .unwrap() + .modalities + .video_in + ); +} diff --git a/crates/tinyagents-session/src/transcript.rs b/crates/tinyagents-session/src/transcript.rs index d5651bedc..d601b0552 100644 --- a/crates/tinyagents-session/src/transcript.rs +++ b/crates/tinyagents-session/src/transcript.rs @@ -163,8 +163,8 @@ pub use thread_lookup::{ }; pub use types::{ CompactionMarker, DisplayMessage, DisplayRecord, DisplaySessionTranscript, LegacyText, - MessageUsage, SessionTranscript, ToolFailure, TranscriptMessage, TranscriptMeta, - TranscriptPart, TranscriptToolCall, TurnUsage, + MessageUsage, SessionTranscript, ToolFailure, TranscriptMediaRef, TranscriptMessage, + TranscriptMeta, TranscriptPart, TranscriptToolCall, TurnUsage, }; pub use writer::{ append_interrupted_partial, append_tools_record, append_transcript_turn, diff --git a/crates/tinyagents-session/src/transcript/jsonl.rs b/crates/tinyagents-session/src/transcript/jsonl.rs index 3d49d67c1..b93d2bfc4 100644 --- a/crates/tinyagents-session/src/transcript/jsonl.rs +++ b/crates/tinyagents-session/src/transcript/jsonl.rs @@ -520,11 +520,27 @@ fn rebuild_content( } TypedShape::ToolResult => Some(encode_tool_envelope(tool_call_id?, content)), TypedShape::UserParts => { - let parts: Vec = parts? + let original = parts?; + if original.iter().any(|part| { + matches!( + part, + TranscriptPart::Audio { .. } + | TranscriptPart::Video { .. } + | TranscriptPart::Document { .. } + ) + }) { + return Some(serde_json::json!({"_tinyagents_media_parts": original}).to_string()); + } + let parts: Vec = original .iter() .map(|part| match part { TranscriptPart::Text { text } => ContentPart::Text(text.clone()), TranscriptPart::Image { url } => ContentPart::Image(url.clone()), + TranscriptPart::Audio { .. } + | TranscriptPart::Video { .. } + | TranscriptPart::Document { .. } => { + unreachable!("new media uses the typed compatibility envelope") + } }) .collect(); Some(join_image_parts(&parts)) diff --git a/crates/tinyagents-session/src/transcript/typed_rows_tests.rs b/crates/tinyagents-session/src/transcript/typed_rows_tests.rs index 9d6897da5..8575b57fb 100644 --- a/crates/tinyagents-session/src/transcript/typed_rows_tests.rs +++ b/crates/tinyagents-session/src/transcript/typed_rows_tests.rs @@ -525,3 +525,74 @@ fn rows_lifted_from_different_legacy_envelopes_are_different_rows() { ); assert!(redacted.same_row_as(&built)); } + +#[test] +fn durable_media_parts_round_trip_without_inline_bytes() { + let parts = vec![ + TranscriptPart::Text { + text: "inspect these".into(), + }, + TranscriptPart::Audio { + source: TranscriptMediaRef::Path { + path: "uploads/a/song.mp3".into(), + }, + mime_type: "audio/mpeg".into(), + }, + TranscriptPart::Video { + source: TranscriptMediaRef::Url { + url: "https://example.com/movie.mp4".into(), + }, + mime_type: "video/mp4".into(), + }, + TranscriptPart::Document { + source: TranscriptMediaRef::Path { + path: "uploads/a/report.pdf".into(), + }, + mime_type: "application/pdf".into(), + }, + ]; + let row = TranscriptMessage::user_with_parts(parts.clone()); + assert_eq!(row.content, "inspect these"); + assert_eq!(row.parts, Some(parts)); + assert_eq!( + TranscriptMessage::from_legacy("user", row.legacy_content()).parts, + row.parts + ); + let dir = tempdir().unwrap(); + let path = dir.path().join("media.jsonl"); + write_transcript(&path, std::slice::from_ref(&row), &meta(), None).unwrap(); + let read = read_transcript(&path).unwrap(); + assert_eq!(read.messages[0].parts, row.parts); + let stored = fs::read_to_string(&path).unwrap(); + assert!(!stored.contains("base64")); + assert!(!stored.contains("data:")); + assert!(row.display_content().contains("[AUDIO:uploads/a/song.mp3]")); +} + +#[test] +fn arbitrary_json_is_not_a_media_compatibility_envelope() { + let audio = TranscriptMessage::user_with_parts(vec![TranscriptPart::Audio { + source: TranscriptMediaRef::Path { + path: "uploads/song.mp3".into(), + }, + mime_type: "audio/mpeg".into(), + }]); + let canonical = audio.legacy_content(); + let mut extra: serde_json::Value = serde_json::from_str(&canonical).unwrap(); + extra["user_note"] = serde_json::json!("keep this ordinary JSON"); + let image_only = serde_json::json!({"_tinyagents_media_parts":[TranscriptPart::Image {url:"https://example.com/image.png".into()}]}).to_string(); + let pretty = serde_json::to_string_pretty( + &serde_json::from_str::(&canonical).unwrap(), + ) + .unwrap(); + for text in [extra.to_string(), image_only, pretty] { + let row = TranscriptMessage::from_legacy("user", &text); + assert_eq!(row.content, text); + assert!(row.parts.is_none()); + assert_eq!(row.legacy_content(), text); + } + assert_eq!( + TranscriptMessage::from_legacy("user", canonical).parts, + audio.parts + ); +} diff --git a/crates/tinyagents-session/src/transcript/types.rs b/crates/tinyagents-session/src/transcript/types.rs index db4558264..fa326a879 100644 --- a/crates/tinyagents-session/src/transcript/types.rs +++ b/crates/tinyagents-session/src/transcript/types.rs @@ -21,7 +21,27 @@ pub struct TranscriptToolCall { pub extra_content: Option, } -/// One ordered part of a user row that mixes text and images. +/// A durable media location. Inline bytes are deliberately absent: hosts resolve +/// these references only on ephemeral provider requests. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum TranscriptMediaRef { + /// A path relative to the agent's acting workspace, or a host-approved path. + Path { path: String }, + /// An external location. Hosts apply fetch policy before resolving it. + Url { url: String }, +} + +impl TranscriptMediaRef { + fn location(&self) -> &str { + match self { + Self::Path { path } => path, + Self::Url { url } => url, + } + } +} + +/// One ordered part of a user row that mixes text and media. /// /// Serialized as `{"type":"text","text":..}` / `{"type":"image","url":..}`, /// which is also the on-disk `parts` shape of a typed `user_parts` line. @@ -32,6 +52,21 @@ pub enum TranscriptPart { Text { text: String }, /// An image reference (a `data:` URI, an `http(s)` URL or a path). Image { url: String }, + /// Audio whose original bytes remain at the referenced location. + Audio { + source: TranscriptMediaRef, + mime_type: String, + }, + /// Video whose original bytes remain at the referenced location. + Video { + source: TranscriptMediaRef, + mime_type: String, + }, + /// A document whose original bytes remain at the referenced location. + Document { + source: TranscriptMediaRef, + mime_type: String, + }, } /// The legacy string a row was lifted from ([`TranscriptMessage::normalized`]), @@ -202,14 +237,14 @@ impl TranscriptMessage { row } - /// A `user` row from ordered text/image parts. `content` becomes the - /// concatenated text parts. A parts list with no image is just text. + /// A `user` row from ordered text/media parts. `content` becomes the + /// concatenated text parts. A parts list with only text is just text. pub fn user_with_parts(parts: Vec) -> Self { let text = parts_text(&parts); let mut row = Self::user(text); if parts .iter() - .any(|part| matches!(part, TranscriptPart::Image { .. })) + .any(|part| !matches!(part, TranscriptPart::Text { .. })) { row.parts = Some(parts); } @@ -259,6 +294,24 @@ impl TranscriptMessage { } } "user" => { + if let Ok(value) = serde_json::from_str::(&self.content) + && let Some(parts) = value.get("_tinyagents_media_parts") + && let Ok(parts) = serde_json::from_value::>(parts.clone()) + && parts.iter().any(|part| { + matches!( + part, + TranscriptPart::Audio { .. } + | TranscriptPart::Video { .. } + | TranscriptPart::Document { .. } + ) + }) + && encode_media_parts(&parts) == self.content + { + self.content = parts_text(&parts); + self.parts = Some(parts); + self.legacy = LegacyText(Some(original)); + return self; + } let parts = split_image_parts(&self.content); if parts .iter() @@ -286,7 +339,9 @@ impl TranscriptMessage { /// The string a flat `{role, content}` row would have held for this row: /// the native envelope for an assistant row with calls (`content` as its /// text), the tool-result envelope for a tool row with a call id, and - /// `[OH_IMAGE:]` markers for a user row with parts. Everything else is + /// `[OH_IMAGE:]` markers for image-only media rows. Audio/video/document + /// rows use a `_tinyagents_media_parts` JSON compatibility envelope; old + /// binaries cannot interpret those new media types. Everything else is /// `content` as is. /// /// For the compatibility adapters that must keep writing the old string @@ -314,11 +369,26 @@ impl TranscriptMessage { return encode_tool_envelope(id, &self.content); } if let Some(parts) = self.parts.as_deref() { + if parts.iter().any(|part| { + matches!( + part, + TranscriptPart::Audio { .. } + | TranscriptPart::Video { .. } + | TranscriptPart::Document { .. } + ) + }) { + return encode_media_parts(parts); + } let parts: Vec = parts .iter() .map(|part| match part { TranscriptPart::Text { text } => ContentPart::Text(text.clone()), TranscriptPart::Image { url } => ContentPart::Image(url.clone()), + TranscriptPart::Audio { .. } + | TranscriptPart::Video { .. } + | TranscriptPart::Document { .. } => { + unreachable!("new media uses the typed compatibility envelope") + } }) .collect(); return join_image_parts(&parts); @@ -327,7 +397,7 @@ impl TranscriptMessage { } /// The row's text as a person reads it: `content`, except that a user row's - /// image parts render in place as `[IMAGE:]`. + /// media parts render in place as `[IMAGE:]`, `[AUDIO:]`, etc. #[must_use] pub fn display_content(&self) -> String { let Some(parts) = self.parts.as_deref() else { @@ -342,6 +412,16 @@ impl TranscriptMessage { out.push_str(url); out.push(']'); } + TranscriptPart::Audio { source, .. } + | TranscriptPart::Video { source, .. } + | TranscriptPart::Document { source, .. } => { + let label = match part { + TranscriptPart::Audio { .. } => "AUDIO", + TranscriptPart::Video { .. } => "VIDEO", + _ => "DOCUMENT", + }; + out.push_str(&format!("[{label}:{}]", source.location())); + } } } out @@ -461,12 +541,19 @@ impl From for TranscriptMessageWire { } /// The concatenated text parts. +fn encode_media_parts(parts: &[TranscriptPart]) -> String { + serde_json::json!({"_tinyagents_media_parts": parts}).to_string() +} + fn parts_text(parts: &[TranscriptPart]) -> String { parts .iter() .filter_map(|part| match part { TranscriptPart::Text { text } => Some(text.as_str()), - TranscriptPart::Image { .. } => None, + TranscriptPart::Image { .. } + | TranscriptPart::Audio { .. } + | TranscriptPart::Video { .. } + | TranscriptPart::Document { .. } => None, }) .collect() } diff --git a/docs/modules/harness/multimodal.md b/docs/modules/harness/multimodal.md index 6b239121d..eca58263d 100644 --- a/docs/modules/harness/multimodal.md +++ b/docs/modules/harness/multimodal.md @@ -65,3 +65,20 @@ counter (alongside the existing `images: usize`) that counts `Audio`/`Video`/ `ContentBlock::estimated_char_weight` charges each of them the same flat `IMAGE_CHAR_WEIGHT` an image gets, so a transcript dominated by non-text attachments is not under-counted by token-budgeting/compaction gating. + +## Generic attachment intake + +For uploads that must retain their original media, use the optional +`multimodal` feature's `resolve_attachment` API. It returns decoded bytes plus +display name, MIME, and size without text extraction. `UnknownMimePolicy::Accept` +is the host's explicit choice to admit arbitrary formats; existing +`resolve_file` callers use `Reject` and keep their MIME allowlist. Only gzip +data URIs with `original_mime` are transport-decoded, exactly once. Hosts own +local path authorization, storage, and choice of native `ContentBlock` variants. + +ZIP/TAR/TAR.GZ inspection uses `ArchiveFormat::detect` and `inspect_archive`. +Typed listings report names, kinds, declared sizes, and resource truncation; +inspection never extracts files or follows archive paths/links. Office ZIP +containers are detected from standard document parts and remain documents. +See [the module README](../../../crates/tinyagents-harness/src/multimodal/README.md) +for byte/name/entry budgets and partial-validation semantics. diff --git a/docs/modules/session/README.md b/docs/modules/session/README.md index 06897206f..92ef3c60a 100644 --- a/docs/modules/session/README.md +++ b/docs/modules/session/README.md @@ -281,3 +281,19 @@ Like every other bridge method in this crate, each trait method is a single bounded synchronous SQLite statement — no `spawn_blocking`, matching the tradeoff `store::with_connection`/`with_transaction` already make throughout. + +## Durable media parts + +`TranscriptMessage::user_with_parts` preserves ordered text, image, audio, video +and document parts. Text remains in `content`; media lives in typed `parts` in +the JSONL row, including on replay. Existing image URL rows keep their wire shape. +Audio/video/document parts use `TranscriptMediaRef::Path { path }` or +`TranscriptMediaRef::Url { url }`, with a `mime_type`; there is deliberately no +inline-byte variant. Hosts store originals and resolve these references only on +ephemeral provider requests, applying their own workspace/fetch policy. + +Legacy image markers continue to round-trip unchanged. `legacy_content` for new +media uses a `_tinyagents_media_parts` JSON envelope that `normalized` can lift +back into ordered parts. Older binaries cannot read these new media types. +Display projection shows compact `[AUDIO:]`, `[VIDEO:]` and +`[DOCUMENT:]` markers without reading or encoding media bytes. diff --git a/vendor/tinyinference b/vendor/tinyinference index a5f129d0b..c71e6b50b 160000 --- a/vendor/tinyinference +++ b/vendor/tinyinference @@ -1 +1 @@ -Subproject commit a5f129d0b258b23dfa2d64d14de0e265ac2d0690 +Subproject commit c71e6b50b40a10e3bf5d97c11ebb8fa2caa7cd60