refactor(core)!: one type each for source, generated file, reasoning output and response metadata

CHANGELOG.md

@@ -23,6 +23,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   `GenerateContent::ToolCall` with a non-string `input` no longer
   deserializes; the standalone `tool::ToolResult` gains a required
   `tool_name` and optional `dynamic` / `provider_metadata`.
+- Source, generated-file and reasoning types are unified, one per concept
+  (`aimux-core`). `result::Source`, `result::GeneratedFile` and
+  `result::ReasoningOutput` are the payloads of `GenerateContent::{Source,
+  File, Reasoning}` and `StreamPart::{Source, File}` (newtype variants now) and
+  the elements of `sources`, `files` and `reasoning` on `GenerateTextResult`
+  and `StreamTextResultAggregated`; `StreamPart::ResponseMetadata` wraps
+  `types::ResponseMetadata`. `SourcePart`, `FilePart` and `ReasoningPart` are
+  removed (TypeScript: `Source`, `GeneratedFile`, `ReasoningOutput`). Wire
+  format: enum variants keep their JSON shape; the top-level `sources` /
+  `files` / `reasoning` arrays gain `provider_metadata` (so the reasoning
+  signature is no longer dropped), and a source's `url` / `title` serialize as
+  `null` when absent instead of being omitted.
 
 - Removed the generated `ProviderName` type in every binding (Rust enum, TS
   const object, Go/Java/Kotlin consts, Swift enum, Dart consts, Python

aimux-core/Cargo.toml

@@ -17,12 +17,10 @@ futures = { workspace = true }
 async-stream = { workspace = true }
 tracing = { workspace = true }
 ts-rs = { workspace = true }
-# Deliberately narrower than the workspace-wide `full`: core's production
-# code uses only timers (`time`) and `select!` (`macros`). Keeping the
-# declaration at the real API surface documents what Core actually asks
-# of the runtime (RFC-0016 had removed core's tokio dependency; RFC-0031
-# reintroduces it for operation deadlines and retry backoff).
-tokio = { version = "1", features = ["time", "macros", "rt", "sync"] }
+# Core also resolves generated file URLs before exposing file outputs.
+tokio = { version = "1", features = ["time", "macros", "rt", "sync", "net"] }
+reqwest = { workspace = true }
+url = { workspace = true }
 tokio-util = "0.7"
 httpdate = { workspace = true }
 rand = { workspace = true }

aimux-core/src/composite.rs

@@ -21,7 +21,7 @@ use std::sync::Arc;
 use crate::language_model::LanguageModel;
 use crate::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage};
 use crate::result::GenerateContent;
-use crate::types::{TokenUsage, Usage};
+use crate::types::{InputTokenUsage, OutputTokenUsage, Usage};
 
 /// Composite-model child handle. `Arc` (not `Box`) so children are `Clone`,
 /// match the FFI/Node registry shape, and stay alive after a child handle is
@@ -31,28 +31,26 @@ pub type ChildModel = Arc<dyn LanguageModel>;
 /// Add two `Usage` values field-by-field. `Usage::raw` (provider-opaque) is
 /// dropped — summing raw maps across providers is meaningless.
 ///
-/// Each `TokenUsage` field is `Option<u32>`; `None` is treated as zero so a
+/// Each usage detail field is `Option<u32>`; `None` is treated as zero so a
 /// child that doesn't report a breakdown doesn't erase the other child's data.
 #[must_use]
 pub fn add_usage(a: Usage, b: &Usage) -> Usage {
     Usage {
-        input_tokens: add_token_usage(a.input_tokens, &b.input_tokens),
-        output_tokens: add_token_usage(a.output_tokens, &b.output_tokens),
+        input_tokens: InputTokenUsage {
+            total: opt_add(a.input_tokens.total, b.input_tokens.total),
+            no_cache: opt_add(a.input_tokens.no_cache, b.input_tokens.no_cache),
+            cache_read: opt_add(a.input_tokens.cache_read, b.input_tokens.cache_read),
+            cache_write: opt_add(a.input_tokens.cache_write, b.input_tokens.cache_write),
+        },
+        output_tokens: OutputTokenUsage {
+            total: opt_add(a.output_tokens.total, b.output_tokens.total),
+            text: opt_add(a.output_tokens.text, b.output_tokens.text),
+            reasoning: opt_add(a.output_tokens.reasoning, b.output_tokens.reasoning),
+        },
         raw: None,
     }
 }
 
-fn add_token_usage(a: TokenUsage, b: &TokenUsage) -> TokenUsage {
-    TokenUsage {
-        total: opt_add(a.total, b.total),
-        no_cache: opt_add(a.no_cache, b.no_cache),
-        cache_read: opt_add(a.cache_read, b.cache_read),
-        cache_write: opt_add(a.cache_write, b.cache_write),
-        text: opt_add(a.text, b.text),
-        reasoning: opt_add(a.reasoning, b.reasoning),
-    }
-}
-
 fn opt_add(a: Option<u32>, b: Option<u32>) -> Option<u32> {
     match (a, b) {
         (None, None) => None,

aimux-core/src/content.rs

@@ -19,6 +19,35 @@ pub enum ContentPart {
         provider_options: Option<Value>,
     },
 
+    /// Provider-specific custom content.
+    Custom {
+        kind: String,
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        provider_options: Option<Value>,
+    },
+
+    /// A file generated during reasoning.
+    ReasoningFile {
+        data: crate::shared::GeneratedFileData,
+        media_type: String,
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        provider_options: Option<Value>,
+    },
+
+    /// A request to approve a tool call on a later turn.
+    ToolApprovalRequest {
+        approval_id: String,
+        tool_call_id: String,
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        reason: Option<String>,
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        is_automatic: Option<bool>,
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        signature: Option<String>,
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        input_schema_input: Option<Value>,
+    },
+
     /// An image (raw bytes + MIME type).
     Image {
         image: Vec<u8>,

aimux-core/src/download.rs

@@ -0,0 +1,130 @@
+//! Resolve generated file URLs before exposing their file data.
+
+use std::net::SocketAddr;
+
+use futures::StreamExt;
+
+use crate::abort_signal::AbortSignal;
+use crate::download_guard::{validate_download_target, validate_download_url};
+use crate::error::AiMuxError;
+use crate::shared::{FileBytes, GeneratedFileData};
+
+const DEFAULT_MAX_DOWNLOAD_SIZE: usize = 2 * 1024 * 1024 * 1024;
+const MAX_DOWNLOAD_REDIRECTS: usize = 10;
+
+pub(crate) async fn resolve_generated_file_data(
+    data: &GeneratedFileData,
+    abort_signal: Option<&AbortSignal>,
+) -> Result<FileBytes, AiMuxError> {
+    match data {
+        GeneratedFileData::Data { data } => Ok(data.clone()),
+        GeneratedFileData::Url { url, .. } => {
+            let (data, _) = download(url, None, abort_signal).await?;
+            Ok(FileBytes::Binary(data))
+        }
+    }
+}
+
+async fn download(
+    url: &str,
+    max_bytes: Option<usize>,
+    abort_signal: Option<&AbortSignal>,
+) -> Result<(Vec<u8>, Option<String>), AiMuxError> {
+    let operation = download_inner(url, max_bytes.unwrap_or(DEFAULT_MAX_DOWNLOAD_SIZE));
+    if let Some(signal) = abort_signal {
+        tokio::select! {
+            biased;
+            () = signal.cancelled() => Err(AiMuxError::from_abort_signal(signal)),
+            result = operation => result,
+        }
+    } else {
+        operation.await
+    }
+}
+
+async fn download_inner(
+    url: &str,
+    max_bytes: usize,
+) -> Result<(Vec<u8>, Option<String>), AiMuxError> {
+    let mut current = url.to_owned();
+    for redirects in 0..=MAX_DOWNLOAD_REDIRECTS {
+        let parsed = validate_download_url(&current)?;
+        let addresses = validate_download_target(&current, None).await?;
+        let host = parsed
+            .host_str()
+            .ok_or_else(|| AiMuxError::InvalidArgument("download URL has no host".into()))?;
+        let port = parsed
+            .port_or_known_default()
+            .ok_or_else(|| AiMuxError::InvalidArgument("download URL has no known port".into()))?;
+        let addresses: Vec<_> = addresses
+            .into_iter()
+            .map(|address| SocketAddr::new(address, port))
+            .collect();
+        // Direct connections use only the addresses validated above. An
+        // automatic proxy or redirect must not bypass the address policy.
+        let client = reqwest::Client::builder()
+            .no_proxy()
+            .redirect(reqwest::redirect::Policy::none())
+            .resolve_to_addrs(host, &addresses)
+            .build()
+            .map_err(download_error)?;
+        let response = client
+            .get(parsed.clone())
+            .header(
+                reqwest::header::USER_AGENT,
+                concat!("aimux/", env!("CARGO_PKG_VERSION")),
+            )
+            .send()
+            .await
+            .map_err(download_error)?;
+        let status = response.status();
+        if matches!(status.as_u16(), 301 | 302 | 303 | 307 | 308)
+            && let Some(location) = response.headers().get(reqwest::header::LOCATION)
+        {
+            if redirects == MAX_DOWNLOAD_REDIRECTS {
+                return Err(AiMuxError::Other("download exceeded 10 redirects".into()));
+            }
+            let location = location.to_str().map_err(download_error)?;
+            current = parsed.join(location).map_err(download_error)?.to_string();
+            drop(response);
+            continue;
+        }
+        if !status.is_success() {
+            return Err(AiMuxError::Other(format!(
+                "download failed with HTTP status {status}"
+            )));
+        }
+        if response
+            .content_length()
+            .is_some_and(|length| length > max_bytes as u64)
+        {
+            return Err(size_error(max_bytes));
+        }
+        let media_type = response
+            .headers()
+            .get(reqwest::header::CONTENT_TYPE)
+            .and_then(|value| value.to_str().ok())
+            .map(str::to_owned);
+        let mut stream = response.bytes_stream();
+        let mut data = Vec::new();
+        while let Some(chunk) = stream.next().await {
+            let chunk = chunk.map_err(download_error)?;
+            if chunk.len() > max_bytes.saturating_sub(data.len()) {
+                return Err(size_error(max_bytes));
+            }
+            data.extend_from_slice(&chunk);
+        }
+        return Ok((data, media_type));
+    }
+    unreachable!("download loop returns a response or redirect error")
+}
+
+fn download_error(error: impl std::fmt::Display) -> AiMuxError {
+    AiMuxError::Other(format!("download failed: {error}"))
+}
+
+fn size_error(max_bytes: usize) -> AiMuxError {
+    AiMuxError::Other(format!(
+        "download exceeded maximum size of {max_bytes} bytes"
+    ))
+}

aimux-provider-utils/src/download_guard.rs → aimux-core/src/download_guard.rs

@@ -10,7 +10,7 @@
 
 use std::net::IpAddr;
 
-use aimux_core::error::AiMuxError;
+use crate::error::AiMuxError;
 
 /// Compare scheme, host, and effective port. Unparseable inputs are never
 /// same-origin. Public so providers can build credential allowlists on top
@@ -28,7 +28,8 @@ pub fn same_origin(url: &str, origin: &str) -> bool {
 /// Trust flows only within the trusted origin: the exemption applies to a
 /// redirect target only when the redirecting URL is itself on the trusted
 /// origin, so a foreign hop cannot launder a request into it.
-pub(crate) fn hop_trusted_origin<'a>(
+#[must_use]
+pub fn hop_trusted_origin<'a>(
     trusted_origin: Option<&'a str>,
     current_url: &str,
 ) -> Option<&'a str> {
@@ -55,7 +56,7 @@ fn host_ip_literal(host: &str) -> Option<IpAddr> {
 ///
 /// Callers that resolve DNS afterwards rely on this having rejected every
 /// non-public literal, so keep this the single literal-check site.
-pub(crate) fn validate_download_url(url: &str) -> Result<url::Url, AiMuxError> {
+pub fn validate_download_url(url: &str) -> Result<url::Url, AiMuxError> {
     let parsed = url::Url::parse(url)
         .map_err(|error| AiMuxError::InvalidArgument(format!("invalid download URL: {error}")))?;
     if !matches!(parsed.scheme(), "http" | "https") {
@@ -93,7 +94,7 @@ pub(crate) fn validate_download_url(url: &str) -> Result<url::Url, AiMuxError> {
 /// A URL same-origin with `trusted_origin` (normally the configured
 /// `base_url`) is exempt and returns no addresses; self-hosted deployments
 /// legitimately serve assets from private space on their own origin.
-pub(crate) async fn validate_download_target(
+pub async fn validate_download_target(
     url: &str,
     trusted_origin: Option<&str>,
 ) -> Result<Vec<IpAddr>, AiMuxError> {
@@ -206,7 +207,7 @@ fn is_public_download_address(address: IpAddr) -> bool {
 /// request; they leak deployment topology or unlock metadata endpoints when
 /// forwarded to a provider-supplied host. Mirrors AI SDK's download header
 /// policy (auth headers survive; the redirect loop clears them cross-origin).
-pub(crate) fn sanitize_download_headers(headers: &mut Vec<(String, String)>) {
+pub fn sanitize_download_headers(headers: &mut Vec<(String, String)>) {
     const BLOCKED: &[&str] = &[
         "connection",
         "keep-alive",
@@ -237,7 +238,7 @@ pub(crate) fn sanitize_download_headers(headers: &mut Vec<(String, String)>) {
 }
 
 /// Drop caller headers except `User-Agent` when a redirect crosses origin.
-pub(crate) fn retain_user_agent(headers: &mut Vec<(String, String)>) {
+pub fn retain_user_agent(headers: &mut Vec<(String, String)>) {
     let user_agent = headers
         .iter()
         .find(|(name, _)| name.eq_ignore_ascii_case("user-agent"))

aimux-core/src/error.rs

@@ -178,6 +178,12 @@ pub enum AiMuxError {
     #[error("invalid response data: {0}")]
     InvalidResponseData(String),
 
+    #[error("Tool call \"{tool_call_id}\" not found for approval request \"{approval_id}\".")]
+    ToolCallNotFoundForApproval {
+        tool_call_id: String,
+        approval_id: String,
+    },
+
     /// A model requested a tool that was not present in the call's tool set.
     /// Message templates for the three tool variants match the AI SDK verbatim
     /// (`NoSuchToolError` / `InvalidToolInputError` / `ToolCallRepairError`):