diff --git a/Cargo.lock b/Cargo.lock index 94ed1f9fb..ac9b63a0b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -117,7 +117,7 @@ dependencies = [ [[package]] name = "aion-agent" version = "0.2.11" -source = "git+https://github.com/iOfficeAI/aionrs.git?tag=v0.2.11#8e61a90329fa9f67c4fdf7e97fe02c24dba33f75" +source = "git+https://github.com/aimlapi/aionrs-aimlapi.git?branch=main#ba23022731059ce4a45bc1ec03e4bb99c3345da6" dependencies = [ "aion-compact", "aion-config", @@ -148,7 +148,7 @@ dependencies = [ [[package]] name = "aion-compact" version = "0.2.11" -source = "git+https://github.com/iOfficeAI/aionrs.git?tag=v0.2.11#8e61a90329fa9f67c4fdf7e97fe02c24dba33f75" +source = "git+https://github.com/aimlapi/aionrs-aimlapi.git?branch=main#ba23022731059ce4a45bc1ec03e4bb99c3345da6" dependencies = [ "regex", "serde", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "aion-config" version = "0.2.11" -source = "git+https://github.com/iOfficeAI/aionrs.git?tag=v0.2.11#8e61a90329fa9f67c4fdf7e97fe02c24dba33f75" +source = "git+https://github.com/aimlapi/aionrs-aimlapi.git?branch=main#ba23022731059ce4a45bc1ec03e4bb99c3345da6" dependencies = [ "aion-compact", "aion-process", @@ -185,7 +185,7 @@ dependencies = [ [[package]] name = "aion-mcp" version = "0.2.11" -source = "git+https://github.com/iOfficeAI/aionrs.git?tag=v0.2.11#8e61a90329fa9f67c4fdf7e97fe02c24dba33f75" +source = "git+https://github.com/aimlapi/aionrs-aimlapi.git?branch=main#ba23022731059ce4a45bc1ec03e4bb99c3345da6" dependencies = [ "aion-config", "aion-protocol", @@ -206,7 +206,7 @@ dependencies = [ [[package]] name = "aion-memory" version = "0.2.11" -source = "git+https://github.com/iOfficeAI/aionrs.git?tag=v0.2.11#8e61a90329fa9f67c4fdf7e97fe02c24dba33f75" +source = "git+https://github.com/aimlapi/aionrs-aimlapi.git?branch=main#ba23022731059ce4a45bc1ec03e4bb99c3345da6" dependencies = [ "aion-config", "chrono", @@ -220,7 +220,7 @@ dependencies = [ [[package]] name = "aion-process" version = "0.2.11" -source = "git+https://github.com/iOfficeAI/aionrs.git?tag=v0.2.11#8e61a90329fa9f67c4fdf7e97fe02c24dba33f75" +source = "git+https://github.com/aimlapi/aionrs-aimlapi.git?branch=main#ba23022731059ce4a45bc1ec03e4bb99c3345da6" dependencies = [ "libc", "tokio", @@ -231,7 +231,7 @@ dependencies = [ [[package]] name = "aion-protocol" version = "0.2.11" -source = "git+https://github.com/iOfficeAI/aionrs.git?tag=v0.2.11#8e61a90329fa9f67c4fdf7e97fe02c24dba33f75" +source = "git+https://github.com/aimlapi/aionrs-aimlapi.git?branch=main#ba23022731059ce4a45bc1ec03e4bb99c3345da6" dependencies = [ "aion-types", "serde", @@ -244,7 +244,7 @@ dependencies = [ [[package]] name = "aion-providers" version = "0.2.11" -source = "git+https://github.com/iOfficeAI/aionrs.git?tag=v0.2.11#8e61a90329fa9f67c4fdf7e97fe02c24dba33f75" +source = "git+https://github.com/aimlapi/aionrs-aimlapi.git?branch=main#ba23022731059ce4a45bc1ec03e4bb99c3345da6" dependencies = [ "aion-config", "aion-types", @@ -273,7 +273,7 @@ dependencies = [ [[package]] name = "aion-skills" version = "0.2.11" -source = "git+https://github.com/iOfficeAI/aionrs.git?tag=v0.2.11#8e61a90329fa9f67c4fdf7e97fe02c24dba33f75" +source = "git+https://github.com/aimlapi/aionrs-aimlapi.git?branch=main#ba23022731059ce4a45bc1ec03e4bb99c3345da6" dependencies = [ "aion-config", "aion-mcp", @@ -300,7 +300,7 @@ dependencies = [ [[package]] name = "aion-tools" version = "0.2.11" -source = "git+https://github.com/iOfficeAI/aionrs.git?tag=v0.2.11#8e61a90329fa9f67c4fdf7e97fe02c24dba33f75" +source = "git+https://github.com/aimlapi/aionrs-aimlapi.git?branch=main#ba23022731059ce4a45bc1ec03e4bb99c3345da6" dependencies = [ "aion-config", "aion-process", @@ -322,7 +322,7 @@ dependencies = [ [[package]] name = "aion-types" version = "0.2.11" -source = "git+https://github.com/iOfficeAI/aionrs.git?tag=v0.2.11#8e61a90329fa9f67c4fdf7e97fe02c24dba33f75" +source = "git+https://github.com/aimlapi/aionrs-aimlapi.git?branch=main#ba23022731059ce4a45bc1ec03e4bb99c3345da6" dependencies = [ "async-trait", "base64", @@ -1101,7 +1101,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -1112,7 +1112,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -1143,7 +1143,7 @@ dependencies = [ "objc2-foundation", "parking_lot", "percent-encoding", - "windows-sys 0.60.2", + "windows-sys 0.59.0", "x11rb", ] @@ -2222,6 +2222,7 @@ dependencies = [ "crossterm_winapi", "derive_more", "document-features", + "futures-core", "mio", "parking_lot", "rustix", @@ -2517,7 +2518,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.60.2", + "windows-sys 0.59.0", ] [[package]] @@ -2662,7 +2663,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -3121,6 +3122,17 @@ dependencies = [ "foldhash 0.1.5", ] +[[package]] +name = "hashbrown" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash 0.2.0", +] + [[package]] name = "hashbrown" version = "0.17.1" @@ -3597,7 +3609,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -3692,7 +3704,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8cfc352a66ba903c23239ef51e809508b6fc2b0f90e3476ac7a9ff47e863ae95" dependencies = [ "scopeguard", - "windows-sys 0.60.2", + "windows-sys 0.59.0", ] [[package]] @@ -3977,7 +3989,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.59.0", ] [[package]] @@ -4562,7 +4574,7 @@ dependencies = [ "once_cell", "socket2", "tracing", - "windows-sys 0.60.2", + "windows-sys 0.59.0", ] [[package]] @@ -4964,7 +4976,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -5438,7 +5450,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -5777,7 +5789,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -6665,7 +6677,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.59.0", ] [[package]] @@ -7117,17 +7129,18 @@ dependencies = [ [[package]] name = "workspace-hack" version = "0.1.0" -source = "git+https://github.com/iOfficeAI/aionrs.git?tag=v0.2.11#8e61a90329fa9f67c4fdf7e97fe02c24dba33f75" +source = "git+https://github.com/aimlapi/aionrs-aimlapi.git?branch=main#ba23022731059ce4a45bc1ec03e4bb99c3345da6" dependencies = [ "bitflags", "cc", + "either", "futures-channel", "futures-core", "futures-sink", "futures-task", "futures-util", "getrandom 0.2.17", - "getrandom 0.3.4", + "hashbrown 0.16.1", "hyper", "hyper-rustls", "hyper-util", @@ -7148,6 +7161,7 @@ dependencies = [ "smallvec", "subtle", "syn 2.0.117", + "syn 3.0.3", "time", "tokio", "tokio-rustls", @@ -7157,7 +7171,6 @@ dependencies = [ "tracing", "tracing-core", "uuid", - "windows-sys 0.60.2", "windows-sys 0.61.2", "winnow 1.0.2", "zeroize", diff --git a/Cargo.toml b/Cargo.toml index 5846fbdde..71baf6e26 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -65,12 +65,12 @@ aionui-cron = { path = "crates/aionui-cron" } aionui-assistant = { path = "crates/aionui-assistant" } aionui-app = { path = "crates/aionui-app" } -aion-agent = { git = "https://github.com/iOfficeAI/aionrs.git", tag = "v0.2.11" } -aion-providers = { git = "https://github.com/iOfficeAI/aionrs.git", tag = "v0.2.11" } -aion-types = { git = "https://github.com/iOfficeAI/aionrs.git", tag = "v0.2.11" } -aion-protocol = { git = "https://github.com/iOfficeAI/aionrs.git", tag = "v0.2.11" } -aion-config = { git = "https://github.com/iOfficeAI/aionrs.git", tag = "v0.2.11" } -aion-mcp = { git = "https://github.com/iOfficeAI/aionrs.git", tag = "v0.2.11" } +aion-agent = { git = "https://github.com/aimlapi/aionrs-aimlapi.git", branch = "main" } +aion-providers = { git = "https://github.com/aimlapi/aionrs-aimlapi.git", branch = "main" } +aion-types = { git = "https://github.com/aimlapi/aionrs-aimlapi.git", branch = "main" } +aion-protocol = { git = "https://github.com/aimlapi/aionrs-aimlapi.git", branch = "main" } +aion-config = { git = "https://github.com/aimlapi/aionrs-aimlapi.git", branch = "main" } +aion-mcp = { git = "https://github.com/aimlapi/aionrs-aimlapi.git", branch = "main" } # Core framework tokio = { version = "1", features = ["full"] } diff --git a/crates/aionui-ai-agent/src/factory/aionrs.rs b/crates/aionui-ai-agent/src/factory/aionrs.rs index d4c09cde6..fa725c7f6 100644 --- a/crates/aionui-ai-agent/src/factory/aionrs.rs +++ b/crates/aionui-ai-agent/src/factory/aionrs.rs @@ -1,4 +1,4 @@ -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use std::path::Path; use std::sync::Arc; @@ -345,6 +345,44 @@ fn resolve_build_session( Ok(None) } +/// AI/ML API attributes traffic by header. Without them a request is served +/// normally and simply counts for nobody, so a missing or misspelled value +/// fails silently — which is why the host match below is exact and why there +/// are tests for the near-misses rather than only the hit. +/// +/// `X-AIMLAPI-Partner-ID` must match `^part_[A-Za-z0-9]{1,64}$` — alphanumerics +/// only after the prefix. `HTTP-Referer` and `X-Title` are the OpenRouter +/// convention, and mirror what the desktop client already sends on the paths +/// it owns; conversations do not go through those, which is the gap this fills. +const AIMLAPI_HOST: &str = "api.aimlapi.com"; +const AIMLAPI_PARTNER_ID: &str = "part_UJK4IAHBjvT9g4cPDrb7B7KT"; +const AIMLAPI_SOURCE: &str = "agent/aionui"; +const AIMLAPI_REFERER: &str = "https://aionui.com"; +const AIMLAPI_TITLE: &str = "AionUi"; + +/// Exact host match, mirroring [`is_openai_host`]. +/// +/// Substring matching would be wrong in the direction that matters: a URL such +/// as `https://api.aimlapi.com.example.test/v1` is a different origin, and +/// sending a partner id there hands our attribution to whoever owns it. +fn is_aimlapi_host(url: &str) -> bool { + let lower = url.to_lowercase(); + lower + .strip_prefix("https://") + .or_else(|| lower.strip_prefix("http://")) + .map(|rest| rest == AIMLAPI_HOST || rest.starts_with(&format!("{AIMLAPI_HOST}/"))) + .unwrap_or(false) +} + +fn aimlapi_attribution_headers() -> BTreeMap { + BTreeMap::from([ + ("X-AIMLAPI-Partner-ID".to_owned(), AIMLAPI_PARTNER_ID.to_owned()), + ("X-AIMLAPI-Source".to_owned(), AIMLAPI_SOURCE.to_owned()), + ("HTTP-Referer".to_owned(), AIMLAPI_REFERER.to_owned()), + ("X-Title".to_owned(), AIMLAPI_TITLE.to_owned()), + ]) +} + /// Map AionUi DB platform/protocol settings to the aionrs provider identifier. pub(crate) fn map_aionrs_provider( platform: &str, @@ -395,6 +433,11 @@ pub(crate) fn resolve_aionrs_url_and_compat_with_mode( openai_api_mode_override: Option, ) -> (Option, AionrsCompatOverrides) { let mut compat = AionrsCompatOverrides::default(); + // Set before every return below, including the `is_full_url` early exit: + // a user who pasted a complete endpoint is still talking to the same host. + if is_aimlapi_host(raw_base_url) { + compat.extra_headers = aimlapi_attribution_headers(); + } let openai_api_mode = resolve_openai_api_mode(platform, mapped_provider, model_id, openai_api_mode_override); let use_responses = openai_api_mode == Some(OpenAiApiMode::Responses); @@ -893,6 +936,82 @@ mod tests { const TEST_USER_ID: &str = "user-1"; + fn headers_for(base_url: &str, is_full_url: bool) -> BTreeMap { + let (_url, compat) = resolve_aionrs_url_and_compat("custom", base_url, "openai", "some-model", is_full_url); + compat.extra_headers + } + + #[test] + fn aimlapi_base_url_resolves_attribution_headers() { + let headers = headers_for("https://api.aimlapi.com/v1", false); + + assert_eq!( + headers.get("X-AIMLAPI-Partner-ID").map(String::as_str), + Some(AIMLAPI_PARTNER_ID) + ); + assert_eq!( + headers.get("X-AIMLAPI-Source").map(String::as_str), + Some("agent/aionui") + ); + assert_eq!( + headers.get("HTTP-Referer").map(String::as_str), + Some("https://aionui.com") + ); + assert_eq!(headers.get("X-Title").map(String::as_str), Some("AionUi")); + } + + #[test] + fn attribution_survives_a_user_supplied_complete_endpoint() { + // `is_full_url` returns early, before the rest of the compat work. A + // user who pasted the whole endpoint is still talking to the same host. + let headers = headers_for("https://api.aimlapi.com/v1/chat/completions", true); + + assert_eq!( + headers.get("X-AIMLAPI-Partner-ID").map(String::as_str), + Some(AIMLAPI_PARTNER_ID) + ); + } + + #[test] + fn other_providers_get_no_headers() { + for base_url in ["https://api.openai.com/v1", "https://openrouter.ai/api/v1", ""] { + assert!( + headers_for(base_url, false).is_empty(), + "attribution must not leak to {base_url}" + ); + } + } + + #[test] + fn a_lookalike_domain_gets_no_headers() { + // The failure this guards against is silent and expensive: a substring + // match would hand the partner id to whoever owns the lookalike. + for base_url in [ + "https://api.aimlapi.com.example.test/v1", + "https://not-api.aimlapi.com/v1", + "https://evil.test/?u=https://api.aimlapi.com/v1", + ] { + assert!( + headers_for(base_url, false).is_empty(), + "attribution must not be sent to {base_url}" + ); + } + } + + #[test] + fn the_partner_id_matches_the_gateway_contract() { + // `^part_[A-Za-z0-9]{1,64}$` — a malformed id is accepted by the API and + // simply attributed to nobody, so nothing at runtime would report this. + let rest = AIMLAPI_PARTNER_ID + .strip_prefix("part_") + .expect("partner id must carry the part_ prefix"); + assert!((1..=64).contains(&rest.len()), "length out of range: {}", rest.len()); + assert!( + rest.chars().all(|c| c.is_ascii_alphanumeric()), + "only alphanumerics are allowed after the prefix: {AIMLAPI_PARTNER_ID}" + ); + } + fn path_test_lock() -> &'static tokio::sync::Mutex<()> { static LOCK: OnceLock> = OnceLock::new(); LOCK.get_or_init(|| tokio::sync::Mutex::new(())) diff --git a/crates/aionui-ai-agent/src/manager/aionrs/agent.rs b/crates/aionui-ai-agent/src/manager/aionrs/agent.rs index 3291d92f5..a1c529962 100644 --- a/crates/aionui-ai-agent/src/manager/aionrs/agent.rs +++ b/crates/aionui-ai-agent/src/manager/aionrs/agent.rs @@ -211,6 +211,13 @@ impl AionrsAgentManager { if let Some(path) = config_extra.compat_overrides.api_path { config.compat.transport.api_path = Some(path); } + // Attribution headers resolved in `factory::aionrs`. Extend rather than + // assign so a preset that already carries headers keeps them. + config + .compat + .transport + .extra_headers + .extend(config_extra.compat_overrides.extra_headers.clone()); if !config_extra.extra_mcp_servers.is_empty() { config.mcp.servers.extend(config_extra.extra_mcp_servers.clone()); diff --git a/crates/aionui-ai-agent/src/services/provider_health.rs b/crates/aionui-ai-agent/src/services/provider_health.rs index 3453f1202..92fea495a 100644 --- a/crates/aionui-ai-agent/src/services/provider_health.rs +++ b/crates/aionui-ai-agent/src/services/provider_health.rs @@ -240,6 +240,13 @@ async fn build_probe_engine(config_extra: AionrsResolvedConfig) -> Result, pub max_tokens_field: Option, pub api_path: Option, + /// Headers added to every request to this provider. Empty for providers + /// that ask for none, which is all of them except the ones resolved in + /// `factory::aionrs`. + pub extra_headers: BTreeMap, } /// Fully resolved Aionrs configuration passed to the agent manager.