diff --git a/.env.example b/.env.example index 090d0159ff..d168ca6553 100644 --- a/.env.example +++ b/.env.example @@ -146,6 +146,12 @@ ASTRA_RUNTIME_ROOT_SECRET=your-runtime-root-secret-min-32-characters-change-me # # MEMORIA_BASE_URL: HTTP endpoint of the Memoria API. MEMORIA_BASE_URL=http://127.0.0.1:8100 +# Stable identity authority (defaults to the normalized API URL). +# MEMORIA_ISSUER=https://memoria.example.com +# Server-owned login website; leave unset for self-hosted password login. +# MEMORIA_WEB_URL=https://thememoria.ai +# Only set after confirming which issuer owns pre-issuer identity mappings. +# MEMORIA_LEGACY_ISSUER=http://127.0.0.1:8100 # MEMORIA_MASTER_KEY: API authentication key. # Requirements: minimum 16 characters. @@ -175,6 +181,25 @@ MEMORIA_EMBEDDING_API_KEY= # http://localhost:11434/v1 — Ollama (local) MEMORIA_EMBEDDING_BASE_URL= +# ── Personal BYOK endpoint policy ──────────────────────────────────────── +# Personal OpenAI-compatible BYOK endpoints: public HTTPS by default. +# Opt into the existing administrator host/port registry with trusted-domains. +# Both modes reject private/metadata targets, pin public DNS, and disable redirects. +# ASTRA_BYOK_ENDPOINT_POLICY=public-https +# Cloud BYOK rejects deployment-owned model inference for all end users. +# Memoria identities are always personal BYOK, even in self-hosted mode. +# ASTRA_DEPLOYMENT_MODE=self-hosted +# Custom BYOK DNS queries go directly to system-configured DNS servers, avoiding +# OS Fake-IP caches. DNS priority is preserved rather than racing nameservers. +# Optionally choose reachable DNS IPs (comma-separated, ports optional). +# ASTRA_BYOK_DNS_SERVERS=10.0.0.53,10.0.0.54:53 +# If UDP DNS is intercepted, explicitly select TCP-only DNS (no UDP fallback). +# ASTRA_BYOK_DNS_SERVERS=tcp://10.0.0.53:53 +# Default is direct egress. Optional operator-owned HTTP(S) CONNECT or SOCKS5 proxy. +# Ambient HTTP_PROXY/HTTPS_PROXY/ALL_PROXY are not used for custom BYOK endpoints. +# The proxy must accept validated public IP targets; origin TLS/SNI remains unchanged. +# ASTRA_BYOK_PROXY_URL=http://127.0.0.1:7890 + # ── Skill Selector Rerank (optional) ────────────────────────────────────── # Optional LLM-based reranker for skill selection. Uses any OpenAI-compatible # chat completions endpoint. Leave all unset to disable (keyword matching only). diff --git a/Cargo.lock b/Cargo.lock index 03a49b6aa9..de4dd62f74 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -580,12 +580,16 @@ dependencies = [ "flate2", "fs2", "futures-util", + "hickory-resolver", "hmac", "jsonwebtoken", "libc", "openssl", + "percent-encoding", + "rcgen", "regex", "reqwest 0.12.28", + "rustls", "serde", "serde_json", "serde_yaml_ng", @@ -595,10 +599,12 @@ dependencies = [ "tempfile", "thiserror 1.0.69", "tokio", + "tokio-rustls", "tokio-util", "tracing", "unicode-segmentation", "uuid", + "webpki-roots 1.0.7", ] [[package]] @@ -1133,6 +1139,17 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" +[[package]] +name = "chacha20" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65c35e4b699c7e15ccbe7ee35c005e4fc0a278d22238a2857e6ce2dadeda1b06" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.1", + "rand_core 0.10.1", +] + [[package]] name = "chrono" version = "0.4.45" @@ -1334,6 +1351,16 @@ dependencies = [ "crossterm 0.23.2", ] +[[package]] +name = "core-foundation" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e195e091a93c46f7102ec7818a2aa394e1e1771c3ab4825963fa03e45afb8f" +dependencies = [ + "core-foundation-sys", + "libc", +] + [[package]] name = "core-foundation" version = "0.10.1" @@ -1359,6 +1386,15 @@ dependencies = [ "libc", ] +[[package]] +name = "cpufeatures" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ca28b0ae3115b884660db4118d803791fd6756b6e88f39c0f3f7859060d7566" +dependencies = [ + "libc", +] + [[package]] name = "crc" version = "3.4.0" @@ -1593,7 +1629,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "97fb8b7c4503de7d6ae7b42ab72a5a59857b4c937ec27a3d4539dba95b5ab2be" dependencies = [ "cfg-if", - "cpufeatures", + "cpufeatures 0.2.17", "curve25519-dalek-derive", "digest", "fiat-crypto", @@ -2395,6 +2431,7 @@ dependencies = [ "cfg-if", "libc", "r-efi 6.0.0", + "rand_core 0.10.1", "wasip2", "wasip3", ] @@ -3418,6 +3455,76 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" +[[package]] +name = "hickory-net" +version = "0.26.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "084e7bd6a377435d568f652153e571b50970d7ccc1d1eeec0519f834632287e1" +dependencies = [ + "async-trait", + "cfg-if", + "data-encoding", + "futures-channel", + "futures-io", + "futures-util", + "hickory-proto", + "idna", + "ipnet", + "jni", + "rand 0.10.2", + "thiserror 2.0.18", + "tinyvec", + "tokio", + "tracing", + "url", +] + +[[package]] +name = "hickory-proto" +version = "0.26.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e2da0694c15b44c6f68a6b05e0233617008c54080e31d6eb848d858a9c5b38d" +dependencies = [ + "data-encoding", + "idna", + "ipnet", + "jni", + "once_cell", + "prefix-trie", + "rand 0.10.2", + "ring", + "thiserror 2.0.18", + "tinyvec", + "tracing", + "url", +] + +[[package]] +name = "hickory-resolver" +version = "0.26.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e4f9f4603319422d482738f3f6fe5aac03157fdbfed1cd85a3ff45adb09072f" +dependencies = [ + "cfg-if", + "futures-util", + "hickory-net", + "hickory-proto", + "ipconfig", + "ipnet", + "jni", + "moka", + "ndk-context", + "once_cell", + "parking_lot", + "rand 0.10.2", + "resolv-conf", + "smallvec", + "system-configuration", + "thiserror 2.0.18", + "tokio", + "tracing", +] + [[package]] name = "hkdf" version = "0.12.4" @@ -3830,11 +3937,27 @@ dependencies = [ "winapi", ] +[[package]] +name = "ipconfig" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4d40460c0ce33d6ce4b0630ad68ff63d6661961c48b6dba35e5a4d81cfb48222" +dependencies = [ + "socket2", + "widestring", + "windows-registry", + "windows-result", + "windows-sys 0.61.2", +] + [[package]] name = "ipnet" version = "2.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" +dependencies = [ + "serde", +] [[package]] name = "is-terminal" @@ -4321,6 +4444,29 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "moka" +version = "0.12.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "957228ad12042ee839f93c8f257b62b4c0ab5eaae1d4fa60de53b27c9d7c5046" +dependencies = [ + "crossbeam-channel", + "crossbeam-epoch", + "crossbeam-utils", + "equivalent", + "parking_lot", + "portable-atomic", + "smallvec", + "tagptr", + "uuid", +] + +[[package]] +name = "ndk-context" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "27b02d87554356db9e9a873add8782d4ea6e3e58ea071a9adb9a2e8ddb884a8b" + [[package]] name = "new_debug_unreachable" version = "1.0.6" @@ -4527,6 +4673,10 @@ name = "once_cell" version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +dependencies = [ + "critical-section", + "portable-atomic", +] [[package]] name = "once_cell_polyfill" @@ -4752,6 +4902,16 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c5a797f0e07bdf071d15742978fc3128ec6c22891c31a3a931513263904c982a" +[[package]] +name = "pem" +version = "3.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d30c53c26bc5b31a98cd02d20f25a7c8567146caf63ed593a9d87b2775291be" +dependencies = [ + "base64", + "serde_core", +] + [[package]] name = "pem-rfc7468" version = "0.7.0" @@ -5018,6 +5178,17 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "925383efa346730478fb4838dbe9137d2a47675ad789c546d150a6e1dd4ab31c" +[[package]] +name = "prefix-trie" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cf6e3177f0684016a5c209b00882e15f8bdd3f3bb48f0491df10cd102d0c6e7" +dependencies = [ + "either", + "ipnet", + "num-traits", +] + [[package]] name = "pretty_assertions" version = "1.4.1" @@ -5265,6 +5436,17 @@ dependencies = [ "rand_core 0.9.5", ] +[[package]] +name = "rand" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7f5fa3a058cd35567ef9bfa5e75732bee0f9e4c55fa90477bef2dfcdbc4be80" +dependencies = [ + "chacha20", + "getrandom 0.4.2", + "rand_core 0.10.1", +] + [[package]] name = "rand_chacha" version = "0.3.1" @@ -5303,6 +5485,12 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rand_core" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" + [[package]] name = "rand_xorshift" version = "0.4.0" @@ -5425,6 +5613,19 @@ dependencies = [ "unicode-width 0.2.2", ] +[[package]] +name = "rcgen" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75e669e5202259b5314d1ea5397316ad400819437857b90861765f24c4cf80a2" +dependencies = [ + "pem", + "ring", + "rustls-pki-types", + "time", + "yasna", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -5584,6 +5785,12 @@ dependencies = [ "web-sys", ] +[[package]] +name = "resolv-conf" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e061d1b48cb8d38042de4ae0a7a6401009d6143dc80d2e2d6f31f0bdd6470c7" + [[package]] name = "rfc6979" version = "0.4.0" @@ -5760,7 +5967,7 @@ version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "26d1e2536ce4f35f4846aa13bff16bd0ff40157cdb14cc056c7b14ba41233ba0" dependencies = [ - "core-foundation", + "core-foundation 0.10.1", "core-foundation-sys", "jni", "log", @@ -5925,7 +6132,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" dependencies = [ "bitflags 2.13.0", - "core-foundation", + "core-foundation 0.10.1", "core-foundation-sys", "libc", "security-framework-sys", @@ -6107,7 +6314,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" dependencies = [ "cfg-if", - "cpufeatures", + "cpufeatures 0.2.17", "digest", ] @@ -6128,7 +6335,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" dependencies = [ "cfg-if", - "cpufeatures", + "cpufeatures 0.2.17", "digest", ] @@ -6650,6 +6857,33 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "system-configuration" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a13f3d0daba03132c0aa9767f98351b3488edc2c100cda2d2ec2b04f3d8d3c8b" +dependencies = [ + "bitflags 2.13.0", + "core-foundation 0.9.4", + "system-configuration-sys", +] + +[[package]] +name = "system-configuration-sys" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e1d1b10ced5ca923a1fcb8d03e96b8d3268065d724548c0211415ff6ac6bac4" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "tagptr" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417" + [[package]] name = "tar" version = "0.4.46" @@ -7923,6 +8157,12 @@ dependencies = [ "wasite", ] +[[package]] +name = "widestring" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72069c3113ab32ab29e5584db3c6ec55d416895e60715417b5b883a357c3e471" + [[package]] name = "winapi" version = "0.3.9" @@ -8037,6 +8277,17 @@ dependencies = [ "windows-link", ] +[[package]] +name = "windows-registry" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02752bf7fbdcce7f2a27a742f798510f3e5ad88dbe84871e5168e2120c3d5720" +dependencies = [ + "windows-link", + "windows-result", + "windows-strings", +] + [[package]] name = "windows-result" version = "0.4.1" @@ -8446,6 +8697,15 @@ version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cfe53a6657fd280eaa890a3bc59152892ffa3e30101319d168b781ed6529b049" +[[package]] +name = "yasna" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e17bb3549cc1321ae1296b9cdc2698e2b6cb1992adfa19a8c72e5b7a738f44cd" +dependencies = [ + "time", +] + [[package]] name = "yoke" version = "0.8.2" diff --git a/Makefile b/Makefile index 7265105a54..4d51ad1dc5 100644 --- a/Makefile +++ b/Makefile @@ -1296,6 +1296,17 @@ test-online: # test-online because pressure timings are operational evidence, not a normal # per-case correctness budget. .PHONY: test-memoria-online-contract +# Scoped auth contract requires a separately provisioned Memoria API with the +# scoped-key capability. Explicit selection fails if required inputs are absent. +.PHONY: test-memoria-auth-online-contract +test-memoria-auth-online-contract: + @test -n "$$ASTRA_TEST_MEMORIA_URL" || { echo "ASTRA_TEST_MEMORIA_URL is required"; exit 2; } + @test -n "$$ASTRA_TEST_MEMORIA_MASTER_KEY" || { echo "ASTRA_TEST_MEMORIA_MASTER_KEY is required"; exit 2; } + @test -n "$$ASTRA_TEST_DATABASE" || { echo "ASTRA_TEST_DATABASE must explicitly designate an isolated DB"; exit 2; } + ASTRA_TEST_DB_IT=1 ASTRA_DATABASE="$$ASTRA_TEST_DATABASE" ASTRA_DATABASE_PREFIX="" \ + CARGO_INCREMENTAL=0 cargo test --locked -p astra-services --features external-contract-tests \ + --test memoria_live_contract_it -- --ignored + test-memoria-online-contract: @if [ ! -f .env ]; then echo "❌ .env is required for the real Memoria contract"; exit 2; fi @set -a; . ./.env; set +a; \ diff --git a/crates/astra-cli/src/cli/auth_flow.rs b/crates/astra-cli/src/cli/auth_flow.rs index 7bddfa1b97..0c2e87e2af 100644 --- a/crates/astra-cli/src/cli/auth_flow.rs +++ b/crates/astra-cli/src/cli/auth_flow.rs @@ -5,6 +5,7 @@ use crate::cli::cli_config::cli_utils::{ use crate::cli::session::session_state::SessionState; use serde::Deserialize; use std::time::Duration; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; /// Session authentication failure that can be repaired by `/login`. /// @@ -138,6 +139,405 @@ pub(crate) async fn do_login( Ok(tokens.access_token) } +pub(crate) async fn do_memoria_login_with_key( + api: &astra_thin_client::ThinClient, + profile: Option<&str>, + connection_key: &str, +) -> Result { + let tokens = request_memoria_tokens(api, connection_key).await?; + save_profile_auth_tokens(profile, "memoria", &tokens)?; + credential_store() + .mutate(|creds| { + let name = + CredentialStore::resolve_profile_name(profile, creds.current_profile.as_deref()); + if let Some(entry) = creds.profiles.get_mut(&name) { + // The connection key belongs only on the Astra server. Older + // CLI versions stored a Memoria key here, so clear it during + // the migration login as well. + entry.memoria_api_key = None; + } + }) + .map_err(|error| error.to_string())?; + Ok(tokens.access_token) +} + +async fn request_memoria_tokens( + api: &astra_thin_client::ThinClient, + connection_key: &str, +) -> Result { + let body = api + .post_auth_memoria_json(&serde_json::json!({ "connection_key": connection_key })) + .await + .map_err(map_thin_err)?; + parse_auth_tokens(&body) +} + +#[derive(Deserialize)] +struct MemoriaConnectionCallback { + state: String, + memoria_connection_key: String, +} + +pub(crate) async fn do_memoria_browser_login( + api: &astra_thin_client::ThinClient, + profile: Option<&str>, + website_base: &str, +) -> Result { + use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .map_err(|error| format!("failed to start local login callback: {error}"))?; + let port = listener + .local_addr() + .map_err(|error| format!("failed to inspect local login callback: {error}"))? + .port(); + let mut state_bytes = [0_u8; 32]; + state_bytes[..16].copy_from_slice(uuid::Uuid::new_v4().as_bytes()); + state_bytes[16..].copy_from_slice(uuid::Uuid::new_v4().as_bytes()); + let expected_state = URL_SAFE_NO_PAD.encode(state_bytes); + let website = validate_login_website(website_base)?; + let allowed_origin = website.origin().ascii_serialization(); + let connect_url = format!( + "{}/connect/astra?port={port}&state={expected_state}&cli_version={}", + website_base.trim_end_matches('/'), + env!("CARGO_PKG_VERSION") + ); + eprintln!("Open this page to connect Astra:\n{connect_url}"); + open_login_url(&connect_url); + + let deadline = tokio::time::Instant::now() + Duration::from_secs(300); + let mut rejected = 0_u8; + loop { + let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); + if remaining.is_zero() { + return Err("browser login timed out; run `astra login` to try again".to_string()); + } + let (mut stream, _) = tokio::time::timeout(remaining, listener.accept()) + .await + .map_err(|_| "browser login timed out; run `astra login` to try again".to_string())? + .map_err(|error| format!("local login callback failed: {error}"))?; + let request = tokio::time::timeout_at(deadline, read_callback_request(&mut stream)) + .await + .map_err(|_| "browser login timed out; run `astra login` to try again")?; + let Ok(request) = request else { + rejected = rejected.saturating_add(1); + write_callback_response(&mut stream, "400 Bad Request", None, "invalid request").await; + if rejected >= 3 { + return Err("too many invalid browser login callbacks".to_string()); + } + continue; + }; + if request.method == "OPTIONS" { + let origin = (request.origin.as_deref() == Some(allowed_origin.as_str())) + .then_some(allowed_origin.as_str()); + write_callback_response(&mut stream, "204 No Content", origin, "").await; + continue; + } + if request.method != "POST" + || request.path != "/callback" + || request.origin.as_deref() != Some(allowed_origin.as_str()) + || request.content_type.as_deref() != Some("application/json") + { + rejected = rejected.saturating_add(1); + write_callback_response(&mut stream, "403 Forbidden", None, "callback rejected").await; + if rejected >= 3 { + return Err("too many invalid browser login callbacks".to_string()); + } + continue; + } + let callback: MemoriaConnectionCallback = match serde_json::from_slice(&request.body) { + Ok(callback) => callback, + Err(_) => { + rejected = rejected.saturating_add(1); + write_callback_response( + &mut stream, + "400 Bad Request", + Some(&allowed_origin), + "invalid callback", + ) + .await; + if rejected >= 3 { + return Err("too many invalid browser login callbacks".to_string()); + } + continue; + } + }; + if !constant_time_eq(callback.state.as_bytes(), expected_state.as_bytes()) + || callback.memoria_connection_key.is_empty() + || callback.memoria_connection_key.len() > 4096 + { + rejected = rejected.saturating_add(1); + write_callback_response( + &mut stream, + "403 Forbidden", + Some(&allowed_origin), + "callback rejected", + ) + .await; + if rejected >= 3 { + return Err("too many invalid browser login callbacks".to_string()); + } + continue; + } + match do_memoria_login_with_key(api, profile, &callback.memoria_connection_key).await { + Ok(token) => { + write_callback_response( + &mut stream, + "200 OK", + Some(&allowed_origin), + r#"{"status":"connected"}"#, + ) + .await; + return Ok(token); + } + Err(error) => { + write_callback_response( + &mut stream, + "502 Bad Gateway", + Some(&allowed_origin), + "Astra could not verify the connection key.", + ) + .await; + return Err(error); + } + } + } +} + +struct CallbackRequest { + method: String, + path: String, + origin: Option, + content_type: Option, + body: Vec, +} + +async fn read_callback_request( + stream: &mut tokio::net::TcpStream, +) -> Result { + tokio::time::timeout(Duration::from_secs(5), read_callback_request_inner(stream)) + .await + .map_err(|_| "callback read timed out".to_string())? +} + +async fn read_callback_request_inner( + stream: &mut tokio::net::TcpStream, +) -> Result { + let mut data = Vec::with_capacity(2048); + let mut chunk = [0_u8; 1024]; + loop { + let read = stream.read(&mut chunk).await.map_err(|e| e.to_string())?; + if read == 0 { + return Err("callback request is incomplete".into()); + } + data.extend_from_slice(&chunk[..read]); + if data.len() > 8192 { + return Err("callback request is too large".into()); + } + let Some(end) = data + .windows(4) + .position(|w| w == b"\r\n\r\n") + .map(|i| i + 4) + else { + continue; + }; + let (mut request, length) = parse_callback_headers(&data[..end])?; + if data.len() < end + length { + continue; + } + if data.len() != end + length { + return Err("callback has trailing data".into()); + } + request.body = data[end..].to_vec(); + return Ok(request); + } +} + +fn parse_callback_headers(data: &[u8]) -> Result<(CallbackRequest, usize), String> { + let text = std::str::from_utf8(data).map_err(|_| "invalid callback headers")?; + let mut lines = text.split("\r\n"); + let line = lines.next().ok_or("missing callback request line")?; + let parts: Vec<_> = line.split_whitespace().collect(); + if parts.len() != 3 || !matches!(parts[2], "HTTP/1.1" | "HTTP/1.0") { + return Err("invalid callback request line".into()); + } + let mut origin = None; + let mut content_type = None; + let mut length = None; + for line in lines.filter(|line| !line.is_empty()) { + let (name, value) = line.split_once(':').ok_or("invalid callback header")?; + let value = value.trim(); + match name.to_ascii_lowercase().as_str() { + "origin" => { + if origin.replace(value.to_string()).is_some() { + return Err("duplicate callback origin".into()); + } + } + "content-type" => { + if content_type + .replace( + value + .split(';') + .next() + .unwrap_or("") + .trim() + .to_ascii_lowercase(), + ) + .is_some() + { + return Err("duplicate callback content type".into()); + } + } + "content-length" => { + let value = value + .parse::() + .map_err(|_| "invalid callback content length")?; + if value > 4096 || length.replace(value).is_some() { + return Err("invalid or duplicate callback content length".into()); + } + } + "transfer-encoding" => return Err("callback transfer encoding is unsupported".into()), + _ => {} + } + } + if parts[0] == "POST" && length.is_none() { + return Err("callback content length is required".into()); + } + Ok(( + CallbackRequest { + method: parts[0].into(), + path: parts[1].into(), + origin, + content_type, + body: vec![], + }, + length.unwrap_or(0), + )) +} + +async fn write_callback_response( + stream: &mut tokio::net::TcpStream, + status: &str, + origin: Option<&str>, + body: &str, +) { + let cors = origin + .map(|origin| { + format!( + "Access-Control-Allow-Origin: {origin}\r\nAccess-Control-Allow-Methods: POST, OPTIONS\r\nAccess-Control-Allow-Headers: Content-Type\r\nAccess-Control-Allow-Private-Network: true\r\nVary: Origin\r\n" + ) + }) + .unwrap_or_default(); + let response = format!( + "HTTP/1.1 {status}\r\n{cors}Content-Type: application/json\r\nCache-Control: no-store\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ); + let _ = stream.write_all(response.as_bytes()).await; +} + +fn constant_time_eq(left: &[u8], right: &[u8]) -> bool { + if left.len() != right.len() { + return false; + } + left.iter() + .zip(right) + .fold(0_u8, |diff, (left, right)| diff | (left ^ right)) + == 0 +} + +fn open_login_url(url: &str) { + #[cfg(target_os = "macos")] + let result = std::process::Command::new("open").arg(url).spawn(); + #[cfg(target_os = "linux")] + let result = std::process::Command::new("xdg-open").arg(url).spawn(); + #[cfg(target_os = "windows")] + let result = windows_browser_command(url).spawn(); + #[cfg(not(any(target_os = "macos", target_os = "linux", target_os = "windows")))] + let result: std::io::Result = Err(std::io::Error::new( + std::io::ErrorKind::Unsupported, + "browser launch is unsupported", + )); + if let Err(error) = result { + eprintln!("Could not open a browser automatically: {error}"); + } +} + +pub(crate) async fn discover_login_website( + api: &astra_thin_client::ThinClient, +) -> Result, String> { + let methods = match api.get_auth_methods().await { + Ok(value) => value, + // Old/self-hosted servers keep the original password journey. Do not + // silently downgrade on network errors, denied access or malformed JSON. + Err(astra_thin_client::ThinClientError::Api { status, .. }) if status.as_u16() == 404 => { + return Ok(None); + } + Err(error) => return Err(map_thin_err(error)), + }; + #[derive(Deserialize)] + struct Methods { + password: bool, + memoria: Option, + } + #[derive(Deserialize)] + struct BrowserProvider { + issuer: String, + authorization_url: String, + } + let methods: Methods = + serde_json::from_value(methods).map_err(|_| "Invalid Server login configuration")?; + if let Some(provider) = methods.memoria { + if provider.issuer.trim().is_empty() { + return Err("Server login issuer is missing".into()); + } + validate_login_website(&provider.authorization_url)?; + return Ok(Some(provider.authorization_url)); + } + if methods.password { + Ok(None) + } else { + Err("Server has no available login method".into()) + } +} + +fn validate_login_website(value: &str) -> Result { + let url = url::Url::parse(value).map_err(|_| "Invalid Server login URL")?; + let loopback = url.host_str().is_some_and(|h| { + h == "localhost" + || h.trim_matches(['[', ']']) + .parse::() + .is_ok_and(|ip| ip.is_loopback()) + }); + if !matches!(url.scheme(), "http" | "https") + || url.host_str().is_none() + || !url.username().is_empty() + || url.password().is_some() + || url.query().is_some() + || url.fragment().is_some() + || (url.scheme() == "http" && !loopback) + { + return Err("Server login URL requires HTTPS (HTTP is allowed only on loopback)".into()); + } + Ok(url) +} + +#[cfg(any(target_os = "windows", test))] +fn windows_browser_command(url: &str) -> std::process::Command { + let mut command = std::process::Command::new("powershell.exe"); + // The URL is data in the child environment, never shell source. In + // particular the callback's &state= cannot become another command. + command + .args([ + "-NoProfile", + "-NonInteractive", + "-Command", + "Start-Process -FilePath $env:ASTRA_LOGIN_URL", + ]) + .env("ASTRA_LOGIN_URL", url); + command +} + async fn request_login_tokens( api: &astra_thin_client::ThinClient, username: &str, @@ -287,15 +687,164 @@ pub(crate) async fn do_register_for_session( #[cfg(test)] mod tests { + + #[test] + fn loopback_callback_rejects_ambiguous_and_oversized_headers() { + for headers in [ + "Content-Length: 1\r\nContent-Length: 2", + "Content-Length: invalid", + "Content-Length: 4097", + "Content-Length: 0\r\nOrigin: https://a.example\r\nOrigin: https://b.example", + "Transfer-Encoding: chunked", + ] { + let input = format!("POST /callback HTTP/1.1\r\n{headers}\r\n\r\n"); + assert!( + super::parse_callback_headers(input.as_bytes()).is_err(), + "{headers}" + ); + } + } + + #[tokio::test] + async fn loopback_callback_rejects_truncated_body() { + use tokio::io::AsyncWriteExt; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let sender = tokio::spawn(async move { + let mut stream = tokio::net::TcpStream::connect(addr).await.unwrap(); + stream + .write_all(b"POST /callback HTTP/1.1\r\nContent-Length: 10\r\n\r\n{}") + .await + .unwrap(); + }); + let (mut stream, _) = listener.accept().await.unwrap(); + assert!(super::read_callback_request(&mut stream).await.is_err()); + sender.await.unwrap(); + } + + #[test] + fn login_urls_require_https_except_loopback() { + for good in [ + "https://thememoria.ai", + "http://localhost", + "http://127.0.0.1:3000", + "http://[::1]:3000", + ] { + assert!(super::validate_login_website(good).is_ok(), "{good}"); + } + for bad in [ + "http://thememoria.ai", + "http://127.0.0.1.example", + "javascript:alert(1)", + "https://user:password@example.com", + "https://example.com?x=1", + "https://example.com#fragment", + ] { + assert!(super::validate_login_website(bad).is_err(), "{bad}"); + } + } + + #[test] + fn windows_browser_launch_keeps_callback_url_out_of_shell_source() { + let url = "https://thememoria.ai/connect/astra?port=1234&state=abc&cli_version=0.2.1"; + let command = super::windows_browser_command(url); + let args: Vec<_> = command + .get_args() + .map(|v| v.to_string_lossy().into_owned()) + .collect(); + assert!( + args.iter() + .all(|arg| !arg.contains(url) && !arg.contains("&state")) + ); + assert!( + command + .get_envs() + .any(|(name, value)| name == "ASTRA_LOGIN_URL" + && value == Some(std::ffi::OsStr::new(url))) + ); + } + + #[tokio::test] + async fn login_discovery_preserves_local_servers_and_rejects_bad_cloud_urls() { + use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{method, path}, + }; + for (status, payload, expected) in [ + (404, serde_json::json!({}), "password"), + ( + 200, + serde_json::json!({"password":true,"memoria":null}), + "password", + ), + ( + 200, + serde_json::json!({"password":true,"memoria":{"issuer":"https://mem.example","authorization_url":"https://thememoria.ai"}}), + "browser", + ), + ( + 200, + serde_json::json!({"password":true,"memoria":{"issuer":"https://mem.example","authorization_url":"http://remote.example"}}), + "error", + ), + (503, serde_json::json!({}), "error"), + ] { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/auth/methods")) + .respond_with(ResponseTemplate::new(status).set_body_json(payload)) + .mount(&server) + .await; + let api = astra_thin_client::ThinClient::new(&server.uri(), None).unwrap(); + let result = super::discover_login_website(&api).await; + match expected { + "password" => assert_eq!(result.unwrap(), None), + "browser" => assert_eq!(result.unwrap().as_deref(), Some("https://thememoria.ai")), + _ => assert!(result.is_err()), + } + } + } + use super::{ - AuthTokenPayload, clear_profile_auth, do_login, do_login_for_session, is_auth_error, - is_llm_provider_auth_error, parse_auth_tokens, save_refreshed_profile_tokens, + AuthTokenPayload, clear_profile_auth, do_login, do_login_for_session, + do_memoria_login_with_key, is_auth_error, is_llm_provider_auth_error, parse_auth_tokens, + read_callback_request, save_refreshed_profile_tokens, }; use crate::cli::cli_config::cli_utils::{Profile, load_credentials, save_credentials}; use serde_json::json; use wiremock::matchers::{body_json, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; + #[tokio::test] + async fn loopback_callback_parser_reads_origin_content_type_and_secret_body() { + use tokio::io::AsyncWriteExt; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let client = tokio::spawn(async move { + let mut stream = tokio::net::TcpStream::connect(address).await.unwrap(); + let body = r#"{"state":"abc","memoria_connection_key":"secret"}"#; + let request = format!( + "POST /callback HTTP/1.1\r\nOrigin: https://thememoria.ai\r\nContent-Type: application/json; charset=utf-8\r\nContent-Length: {}\r\n\r\n{}", + body.len(), + body + ); + stream.write_all(request.as_bytes()).await.unwrap(); + }); + let (mut stream, _) = listener.accept().await.unwrap(); + let request = read_callback_request(&mut stream).await.unwrap(); + client.await.unwrap(); + + assert_eq!(request.method, "POST"); + assert_eq!(request.path, "/callback"); + assert_eq!(request.origin.as_deref(), Some("https://thememoria.ai")); + assert_eq!(request.content_type.as_deref(), Some("application/json")); + assert_eq!( + serde_json::from_slice::(&request.body).unwrap(), + json!({"state": "abc", "memoria_connection_key": "secret"}) + ); + } + #[test] fn auth_token_payload_requires_server_issued_user_identity() { let Err(missing) = @@ -528,4 +1077,45 @@ mod tests { assert_eq!(profile.access_token.as_deref(), Some("internal-access")); assert_eq!(profile.refresh_token.as_deref(), Some("internal-refresh")); } + + #[serial_test::serial] + #[tokio::test] + async fn memoria_login_sends_key_once_and_does_not_persist_it() { + let _creds_guard = crate::tests::isolate_credentials(); + let mut creds = load_credentials(); + creds.profiles.insert( + "default".to_string(), + Profile { + memoria_api_key: Some("legacy-key".to_string()), + ..Default::default() + }, + ); + save_credentials(&creds).unwrap(); + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/auth/memoria")) + .and(body_json(json!({"connection_key": "scoped-key"}))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "user_id": "memoria-user-id", + "access_token": "astra-access", + "refresh_token": "astra-refresh", + "memory_access": "read_only", + "granted_scopes": ["identity:read", "memory:read"] + }))) + .expect(1) + .mount(&server) + .await; + + let api = astra_thin_client::ThinClient::new(&server.uri(), None).unwrap(); + let token = do_memoria_login_with_key(&api, None, "scoped-key") + .await + .unwrap(); + + assert_eq!(token, "astra-access"); + let creds = load_credentials(); + let profile = &creds.profiles["default"]; + assert_eq!(profile.account_id.as_deref(), Some("memoria-user-id")); + assert_eq!(profile.username.as_deref(), Some("memoria")); + assert_eq!(profile.memoria_api_key, None); + } } diff --git a/crates/astra-cli/src/cli/cli_config/cli_args.rs b/crates/astra-cli/src/cli/cli_config/cli_args.rs index 19060e1958..57b45670ca 100644 --- a/crates/astra-cli/src/cli/cli_config/cli_args.rs +++ b/crates/astra-cli/src/cli/cli_config/cli_args.rs @@ -330,7 +330,9 @@ pub(crate) struct RegisterArgs { } #[derive(Args, Debug)] -#[command(after_help = "Examples:\n astra login --username alice --password secret")] +#[command( + after_help = "Examples:\n astra login\n astra login --manual\n astra login --username alice --password secret" +)] pub(crate) struct LoginArgs { /// Username to log in with #[arg(long)] @@ -338,6 +340,9 @@ pub(crate) struct LoginArgs { /// Password to log in with #[arg(long)] pub password: Option, + /// Enter a Memoria connection key instead of opening the browser + #[arg(long)] + pub manual: bool, } #[derive(Args, Debug)] @@ -1009,12 +1014,45 @@ pub(crate) struct SessionShowArgs { } #[derive(Subcommand, Debug)] -#[command(after_help = "Examples:\n astra model list\n astra model show gpt-4o")] +#[command( + after_help = "Examples:\n astra model add\n astra model list\n astra model add deepseek --provider deepseek --model deepseek-v4-flash --context-window 1000000 --api-key-stdin --default\n astra model add gateway --provider openai-compatible --base-url https://gateway.example/v1 --model MODEL_ID --api-key-stdin\n astra model show deepseek\n astra model probe deepseek" +)] pub(crate) enum ModelCmd { /// List available models List, + /// Add a personal Cloud BYOK model + Add(ModelAddArgs), /// Show model details Show(ModelShowArgs), + /// Check a personal model credential and endpoint + Probe(ModelShowArgs), + /// Delete a personal Cloud BYOK model + Delete(ModelShowArgs), +} + +#[derive(Args, Debug)] +pub(crate) struct ModelAddArgs { + /// Configuration alias used with chat --model (wizard default: provider model ID) + #[arg(value_name = "ALIAS")] + pub name: Option, + /// Model provider (omit to choose interactively) + #[arg(long, value_parser = ["openai", "anthropic", "deepseek", "openai-compatible"])] + pub provider: Option, + /// Exact model ID from your provider's API documentation (not the Astra alias) + #[arg(long)] + pub model: Option, + /// HTTPS API base URL, required for openai-compatible + #[arg(long)] + pub base_url: Option, + /// Read the API key from standard input + #[arg(long)] + pub api_key_stdin: bool, + /// Make this the default model for the current user + #[arg(long)] + pub default: bool, + /// Model context-window size in tokens + #[arg(long, default_value_t = 128_000)] + pub context_window: i32, } #[derive(Args, Debug)] @@ -1370,7 +1408,7 @@ pub(crate) struct ConfigShowPolicyArgs { #[cfg(test)] mod tests { - use super::{Cli, Command, SessionCmd, WorkSubcommand}; + use super::{Cli, Command, ModelCmd, SessionCmd, WorkSubcommand}; use clap::Parser; #[test] @@ -1433,4 +1471,55 @@ mod tests { let message = Cli::try_parse_from(["astra", "explain", "history", "please"]).unwrap(); assert!(message.validate_external_message_shorthand().is_ok()); } + + #[test] + fn personal_byok_model_add_is_a_typed_user_command() { + let cli = Cli::try_parse_from([ + "astra", + "model", + "add", + "deepseek", + "--provider", + "deepseek", + "--model", + "deepseek-chat", + "--api-key-stdin", + "--default", + ]) + .expect("personal model add command"); + let Some(Command::Model(ModelCmd::Add(args))) = cli.command else { + panic!("expected ModelCmd::Add") + }; + assert_eq!(args.name.as_deref(), Some("deepseek")); + assert_eq!(args.provider.as_deref(), Some("deepseek")); + assert_eq!(args.model.as_deref(), Some("deepseek-chat")); + assert!(args.api_key_stdin); + assert!(args.default); + assert_eq!(args.context_window, 128_000); + } + + #[test] + fn personal_byok_supports_wizard_and_compatible_flags() { + assert!(Cli::try_parse_from(["astra", "model", "add"]).is_ok()); + let cli = Cli::try_parse_from([ + "astra", + "model", + "add", + "gateway", + "--provider", + "openai-compatible", + "--base-url", + "https://gateway.example/v1", + "--model", + "custom-model", + "--api-key-stdin", + ]) + .unwrap(); + let Some(Command::Model(ModelCmd::Add(args))) = cli.command else { + panic!("model add") + }; + assert_eq!(args.base_url.as_deref(), Some("https://gateway.example/v1")); + assert_eq!(args.provider.as_deref(), Some("openai-compatible")); + assert!(Cli::try_parse_from(["astra", "model", "add", "--provider", "unknown"]).is_err()); + } } diff --git a/crates/astra-cli/src/cli/command_router.rs b/crates/astra-cli/src/cli/command_router.rs index 0c7754456b..cb3f47dd81 100644 --- a/crates/astra-cli/src/cli/command_router.rs +++ b/crates/astra-cli/src/cli/command_router.rs @@ -4,8 +4,8 @@ use crate::cli::arg_render::{ render_review_args, render_team_args, }; use crate::cli::auth_flow::{ - clear_profile_auth, do_login, do_register, is_auth_error, parse_auth_tokens, - save_refreshed_profile_tokens, + clear_profile_auth, do_login, do_memoria_browser_login, do_memoria_login_with_key, do_register, + is_auth_error, parse_auth_tokens, save_refreshed_profile_tokens, }; use crate::cli::cli_config::cli_args::{ AuditCmd, Cli, Command, JournalCmd, ModelCmd, SessionCaptureCmd, SessionCmd, SkillCmd, @@ -775,6 +775,85 @@ pub(crate) fn execute_cli_command<'a>( )) } +async fn find_personal_model_id( + api: &astra_thin_client::ThinClient, + token: &str, + name: &str, +) -> Result, String> { + let body = api + .get_bearer_path_query_text(token, paths::ME_MODELS, &[]) + .await + .map_err(map_thin_err)?; + let value: serde_json::Value = serde_json::from_str(&body) + .map_err(|error| format!("Invalid /me/models response: {error}"))?; + let items = value + .get("items") + .and_then(serde_json::Value::as_array) + .ok_or_else(|| "Invalid /me/models response: items is missing".to_string())?; + Ok(items.iter().find_map(|item| { + (item.get("name").and_then(serde_json::Value::as_str) == Some(name)) + .then(|| item.get("model_id").and_then(serde_json::Value::as_str)) + .flatten() + .map(str::to_string) + })) +} + +fn personal_model_id_prompt() -> inquire::Text<'static, 'static> { + inquire::Text::new("Model ID from your provider:") + .with_help_message("Copy the exact model ID from your provider's API documentation; this is sent to the provider.") +} + +fn personal_model_alias_prompt(model: &str) -> inquire::Text<'_, '_> { + inquire::Text::new("Configuration alias in Astra:") + .with_default(model) + .with_help_message("Press Enter to use the model ID, or choose a unique alias (e.g. work-model). Use this alias with astra chat --model .") +} + +#[cfg(test)] +mod personal_model_prompt_tests { + use super::*; + + #[test] + fn provider_id_has_explanation_but_no_guessed_default() { + let prompt = personal_model_id_prompt(); + assert_eq!(prompt.message, "Model ID from your provider:"); + assert_eq!(prompt.default, None); + assert!( + prompt + .help_message + .unwrap() + .contains("sent to the provider") + ); + } + + #[test] + fn alias_defaults_to_exact_provider_model_id() { + for model in ["deepseek-v4-flash", "vendor/model-1", "kimi-k2.6"] { + let prompt = personal_model_alias_prompt(model); + assert_eq!(prompt.message, "Configuration alias in Astra:"); + assert_eq!(prompt.default, Some(model)); + assert!( + prompt + .help_message + .unwrap() + .contains("astra chat --model ") + ); + } + } +} + +fn read_model_api_key_from_stdin() -> Result { + let mut api_key = String::new(); + std::io::stdin() + .read_to_string(&mut api_key) + .map_err(|error| format!("Failed to read model API key from stdin: {error}"))?; + let api_key = api_key.trim().to_string(); + if api_key.is_empty() { + return Err("Model API key from stdin cannot be empty".to_string()); + } + Ok(api_key) +} + #[allow(clippy::too_many_arguments)] async fn execute_cli_command_impl( command: Option, @@ -1016,9 +1095,22 @@ async fn execute_cli_command_impl( .magenta() .bold() ); - let username = prompt_or("Username", args.username)?; - let password = prompt_password_masked("Password", args.password)?; - do_login(api, profile.as_deref(), &username, &password).await?; + if args.username.is_some() || args.password.is_some() { + let username = prompt_or("Username", args.username)?; + let password = prompt_password_masked("Password", args.password)?; + do_login(api, profile.as_deref(), &username, &password).await?; + } else if args.manual { + let connection_key = prompt_password_masked("Memoria connection key", None)?; + do_memoria_login_with_key(api, profile.as_deref(), &connection_key).await?; + } else { + if let Some(website) = crate::cli::auth_flow::discover_login_website(api).await? { + do_memoria_browser_login(api, profile.as_deref(), &website).await?; + } else { + let username = prompt_or("Username", None)?; + let password = prompt_password_masked("Password", None)?; + do_login(api, profile.as_deref(), &username, &password).await?; + } + } eprintln!( "{}", " ✓ Logged in. Run `astra` to start chatting.".green() @@ -2032,16 +2124,167 @@ async fn execute_cli_command_impl( Ok(ExitCode::Success) } + Some(Command::Model(ModelCmd::Add(args))) => { + let (_, _, _, token) = get_profile_and_token(profile.as_deref())?; + let interactive = + !args.api_key_stdin && std::io::IsTerminal::is_terminal(&std::io::stdin()); + let wizard = args.provider.is_none() || args.model.is_none() || args.name.is_none(); + let required_text = + |value: Option, label: &str, flag: &str| -> Result { + let value = match value { + Some(value) => value, + None if interactive => inquire::Text::new(&format!("{label}:")) + .prompt() + .map_err(|e| e.to_string())?, + None => return Err(format!("{flag} is required in non-interactive mode")), + }; + if value.trim().is_empty() { + return Err(format!("{label} must not be empty")); + } + Ok(value.trim().to_string()) + }; + let provider = match args.provider { + Some(provider) => provider, + None if interactive => { + let label = inquire::Select::new( + "Model provider:", + vec!["OpenAI", "Anthropic", "DeepSeek", "OpenAI-compatible"], + ) + .prompt() + .map_err(|e| e.to_string())?; + label.to_ascii_lowercase() + } + None => return Err("--provider is required in non-interactive mode".into()), + }; + let base_url = if provider == "openai-compatible" { + Some(required_text( + args.base_url, + "API base URL (HTTPS, including /v1 when required)", + "--base-url", + )?) + } else { + if args.base_url.is_some() { + return Err("Use --provider openai-compatible to configure --base-url".into()); + } + None + }; + if let Some(base_url) = &base_url { + api.post_bearer_path_json_text( + &token, + paths::ME_MODEL_VALIDATE_ENDPOINT, + &serde_json::json!({ "base_url": base_url }), + ) + .await + .map_err(map_thin_err)?; + } + let model = match args.model { + None if interactive => Some( + personal_model_id_prompt() + .prompt() + .map_err(|e| e.to_string())?, + ), + value => value, + }; + let model = required_text(model, "Model ID from your provider", "--model")?; + let name = match args.name { + None if interactive => Some( + personal_model_alias_prompt(&model) + .prompt() + .map_err(|e| e.to_string())?, + ), + value => value, + }; + let name = required_text(name, "Configuration alias in Astra", "configuration alias")?; + let context_window = if wizard && interactive { + inquire::CustomType::::new("Context window (tokens):") + .with_default(args.context_window) + .prompt() + .map_err(|e| e.to_string())? + } else { + args.context_window + }; + if context_window <= 0 { + return Err("Context window must be positive".into()); + } + let is_default = if wizard && interactive && !args.default { + inquire::Confirm::new("Use this as your default model?") + .with_default(true) + .prompt() + .map_err(|e| e.to_string())? + } else { + args.default + }; + let api_key = if args.api_key_stdin { + read_model_api_key_from_stdin()? + } else { + prompt_password_masked( + "Model API key (input hidden; press Enter to confirm)", + None, + )? + }; + let payload = serde_json::json!({ + "name": name, + "provider": provider, + "model": model, + "base_url": base_url, + "api_key": api_key, + "context_window": context_window, + "is_default": is_default, + }); + let body = api + .post_bearer_path_json_text(&token, paths::ME_MODELS, &payload) + .await + .map_err(map_thin_err)?; + print_json_or_raw(&body); + Ok(ExitCode::Success) + } + Some(Command::Model(ModelCmd::Show(args))) => { let (_, _, _, token) = get_profile_and_token(profile.as_deref())?; + let body = if let Some(model_id) = + find_personal_model_id(api, &token, &args.model_name).await? + { + api.get_bearer_path_query_text(&token, &paths::me_model(&model_id), &[]) + .await + .map_err(map_thin_err)? + } else { + api.get_model_text(&token, &args.model_name) + .await + .map_err(map_thin_err)? + }; + print_json_or_raw(&body); + Ok(ExitCode::Success) + } + + Some(Command::Model(ModelCmd::Probe(args))) => { + let (_, _, _, token) = get_profile_and_token(profile.as_deref())?; + let model_id = find_personal_model_id(api, &token, &args.model_name) + .await? + .ok_or_else(|| format!("Personal model '{}' not found", args.model_name))?; let body = api - .get_model_text(&token, &args.model_name) + .post_bearer_path_empty_text(&token, &paths::me_model_check(&model_id)) .await .map_err(map_thin_err)?; print_json_or_raw(&body); Ok(ExitCode::Success) } + Some(Command::Model(ModelCmd::Delete(args))) => { + let (_, _, _, token) = get_profile_and_token(profile.as_deref())?; + let model_id = find_personal_model_id(api, &token, &args.model_name) + .await? + .ok_or_else(|| format!("Personal model '{}' not found", args.model_name))?; + api.delete_bearer_path_text(&token, &paths::me_model(&model_id)) + .await + .map_err(map_thin_err)?; + stdout_println!( + "{} Deleted personal model {}", + theme::icon_ok(), + args.model_name + ); + Ok(ExitCode::Success) + } + Some(Command::Skill(SkillCmd::List(args))) => { let pipeline_modules = create_pipeline_modules_quiet(api, profile.as_deref()); let filter = SkillCatalogFilter { diff --git a/crates/astra-memoria/src/lib.rs b/crates/astra-memoria/src/lib.rs index 53bd175fab..0d1edae084 100644 --- a/crates/astra-memoria/src/lib.rs +++ b/crates/astra-memoria/src/lib.rs @@ -222,6 +222,13 @@ pub fn validate_strict_memories( /// Provider-neutral Memoria operations used by runtime orchestration. #[async_trait::async_trait] pub trait MemoriaPort: Send + Sync { + /// Admission precedes retrieval, inference, reflection and cleanup work. + /// Explicit transports default to admitted; user-bound transports resolve + /// current consent without contacting the external memory service. + async fn admits_operation(&self, _write: bool) -> Result { + Ok(true) + } + /// Return a transport bound to an authenticated owner. /// /// Server runtimes are multi-tenant and must scope the port before placing diff --git a/crates/astra-server-types/src/lib.rs b/crates/astra-server-types/src/lib.rs index 4ed3069de7..125fa72d30 100644 --- a/crates/astra-server-types/src/lib.rs +++ b/crates/astra-server-types/src/lib.rs @@ -680,6 +680,13 @@ pub struct AuthLoginRequest { pub password: String, } +#[cfg(feature = "server")] +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +pub struct AuthMemoriaRequest { + pub connection_key: String, +} + #[cfg(feature = "server")] #[derive(Deserialize)] #[serde(deny_unknown_fields)] @@ -691,7 +698,9 @@ pub struct AuthRefreshRequest { #[derive(Deserialize)] #[serde(deny_unknown_fields)] pub struct AuthReauthenticateRequest { + #[serde(default)] pub password: String, + pub memoria_proof: Option, pub purpose: ReauthenticationPurpose, } @@ -996,6 +1005,18 @@ pub struct AuthTokenResponse { pub expires_in: u32, } +#[cfg(feature = "server")] +#[derive(Serialize, PartialEq, Eq)] +pub struct AuthMemoriaResponse { + pub user_id: String, + pub access_token: String, + pub refresh_token: String, + pub token_type: String, + pub expires_in: u32, + pub memory_access: String, + pub granted_scopes: Vec, +} + #[cfg(feature = "server")] #[derive(Serialize, PartialEq, Eq)] pub struct AuthLogoutResponse { @@ -1850,6 +1871,7 @@ impl From for ReauthenticationRequestData { fn from(value: AuthReauthenticateRequest) -> Self { Self { password: value.password, + memoria_proof: value.memoria_proof, purpose: value.purpose, } } diff --git a/crates/astra-thin-client/src/client.rs b/crates/astra-thin-client/src/client.rs index 6fe75ebbd0..87d6d13ecf 100644 --- a/crates/astra-thin-client/src/client.rs +++ b/crates/astra-thin-client/src/client.rs @@ -455,6 +455,18 @@ impl ThinClient { Self::text_or_api(resp).await } + pub async fn post_auth_memoria_json(&self, body: &Value) -> Result { + let url = self.url(paths::AUTH_MEMORIA)?; + let resp = self + .http + .post(url) + .header(header::CONTENT_TYPE, "application/json") + .json(body) + .send() + .await?; + Self::text_or_api(resp).await + } + pub async fn get_auth_me_text(&self, token: &str) -> Result { let url = self.url(paths::AUTH_ME)?; let resp = self @@ -466,6 +478,17 @@ impl ThinClient { Self::text_or_api(resp).await } + pub async fn get_auth_methods(&self) -> Result { + let response = self + .http + .get(self.url(paths::AUTH_METHODS)?) + .timeout(Duration::from_secs(10)) + .send() + .await?; + let text = Self::text_or_api(response).await?; + Ok(serde_json::from_str(&text)?) + } + pub async fn get_auth_me_text_timeout( &self, token: &str, diff --git a/crates/astra-thin-client/src/paths.rs b/crates/astra-thin-client/src/paths.rs index c36d3d46b3..12d623fb72 100644 --- a/crates/astra-thin-client/src/paths.rs +++ b/crates/astra-thin-client/src/paths.rs @@ -301,6 +301,8 @@ pub fn chat_run_delegations_resume(run_id: &str) -> String { pub const AUTH_REGISTER: &str = "/auth/register"; pub const AUTH_LOGIN: &str = "/auth/login"; +pub const AUTH_MEMORIA: &str = "/auth/memoria"; +pub const AUTH_METHODS: &str = "/auth/methods"; pub const AUTH_REFRESH: &str = "/auth/refresh"; pub const AUTH_LOGOUT: &str = "/auth/logout"; pub const AUTH_REAUTHENTICATE: &str = "/auth/reauthenticate"; @@ -310,6 +312,18 @@ pub const HEALTH: &str = "/health"; pub const MODELS: &str = "/models"; pub const MODEL_ACCESS: &str = "/model-access"; +pub const ME_MODELS: &str = "/me/models"; +pub const ME_MODEL_VALIDATE_ENDPOINT: &str = "/me/models/validate-endpoint"; + +#[inline] +pub fn me_model(model_id: &str) -> String { + format!("{ME_MODELS}/{}", model_segment(model_id)) +} + +#[inline] +pub fn me_model_check(model_id: &str) -> String { + format!("{}/check", me_model(model_id)) +} #[inline] pub fn model(name: &str) -> String { @@ -635,6 +649,9 @@ mod tests { "/models/bedrock%2Fclaude%3Fvariant%231" ); assert_eq!(MODEL_ACCESS, "/model-access"); + assert_eq!(ME_MODELS, "/me/models"); + assert_eq!(me_model("model/id"), "/me/models/model%2Fid"); + assert_eq!(me_model_check("model/id"), "/me/models/model%2Fid/check"); } #[test] diff --git a/crates/core/src/config.rs b/crates/core/src/config.rs index 662f363d66..caec82249f 100644 --- a/crates/core/src/config.rs +++ b/crates/core/src/config.rs @@ -1133,6 +1133,9 @@ impl AppSettings { memoria: MemoriaSettings { base_url: value_or_default(&lookup, "MEMORIA_BASE_URL", DEFAULT_MEMORIA_URL), master_key: lookup("MEMORIA_MASTER_KEY"), + issuer: lookup("MEMORIA_ISSUER"), + web_url: lookup("MEMORIA_WEB_URL"), + legacy_issuer: lookup("MEMORIA_LEGACY_ISSUER"), }, runtime_root_secret: required_value( &lookup, @@ -1498,6 +1501,9 @@ impl fmt::Debug for ApiSettings { pub struct MemoriaSettings { pub base_url: String, pub master_key: Option, + pub issuer: Option, + pub web_url: Option, + pub legacy_issuer: Option, } impl MemoriaSettings { @@ -1507,6 +1513,9 @@ impl MemoriaSettings { base_url: env::var("MEMORIA_BASE_URL") .unwrap_or_else(|_| DEFAULT_MEMORIA_URL.to_string()), master_key: env::var("MEMORIA_MASTER_KEY").ok(), + issuer: env::var("MEMORIA_ISSUER").ok(), + web_url: env::var("MEMORIA_WEB_URL").ok(), + legacy_issuer: env::var("MEMORIA_LEGACY_ISSUER").ok(), } } diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs index 2a0c0ddbaa..5dfc6687d0 100644 --- a/crates/core/src/lib.rs +++ b/crates/core/src/lib.rs @@ -32,6 +32,7 @@ pub mod history_work; pub mod history_work_baseline; pub mod identity; pub mod local_state; +pub mod model_wire; pub mod process_runtime; pub mod work_unit; diff --git a/crates/core/src/model_wire.rs b/crates/core/src/model_wire.rs new file mode 100644 index 0000000000..52e59a5073 --- /dev/null +++ b/crates/core/src/model_wire.rs @@ -0,0 +1,42 @@ +//! Shared wire rules for chat inference and credential/connectivity probes. +//! Keep this below services/runtime so probes cannot invent a second contract. +use serde_json::{Value, json}; + +/// Apply a caller-bounded output budget to Anthropic Messages or OpenAI-style +/// chat completions. Bedrock Converse has a separate inferenceConfig shape. +/// Thinking-budget policy belongs to the caller, not this serialization helper. +/// DeepSeek shares the message shape, but its documented output bound remains +/// [`max_tokens`](https://api-docs.deepseek.com/api/create-chat-completion/). +pub fn apply_chat_output_token_limit(body: &mut Value, provider: &str, tokens: usize) { + let (field, obsolete) = if matches!(provider, "anthropic" | "deepseek") { + ("max_tokens", "max_completion_tokens") + } else { + ("max_completion_tokens", "max_tokens") + }; + if let Some(object) = body.as_object_mut() { + object.remove(obsolete); + object.insert(field.into(), json!(tokens)); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn output_limit_uses_the_protocol_not_the_model_name() { + // Expectations come from each provider's wire contract, not the + // production branch condition. The model name must not select it. + for (provider, field, absent) in [ + ("openai", "max_completion_tokens", "max_tokens"), + ("openai-compatible", "max_completion_tokens", "max_tokens"), + ("deepseek", "max_tokens", "max_completion_tokens"), + ("anthropic", "max_tokens", "max_completion_tokens"), + ] { + let mut body = json!({"model":"o3", "max_tokens":1, "max_completion_tokens":2}); + apply_chat_output_token_limit(&mut body, provider, 32); + assert_eq!(body[field], 32, "{provider}"); + assert!(body.get(absent).is_none(), "{provider}: {absent}"); + } + } +} diff --git a/crates/runtime/src/app_state.rs b/crates/runtime/src/app_state.rs index 39097cbcb4..6b71764767 100644 --- a/crates/runtime/src/app_state.rs +++ b/crates/runtime/src/app_state.rs @@ -208,6 +208,11 @@ pub struct AppState { pub memoria_base_url: String, pub memoria_master_key: Option, pub memoria_forwarder: Arc, + /// True only when the application composition explicitly injects a + /// forwarder (normally an in-process test double). Production BYOK calls + /// resolve a per-user credential instead of falling back to the + /// server-wide forwarder merely because one is configured. + pub(crate) memoria_forwarder_is_override: bool, memoria_health_cache: Arc>, memoria_health_refresh: Arc>, pub shared_pool: Option, @@ -327,6 +332,7 @@ impl AppState { memoria_base_url: default_memoria.base_url, memoria_master_key: default_memoria.master_key, memoria_forwarder: Arc::new(NoopMemoriaForwarder), + memoria_forwarder_is_override: false, memoria_health_cache: Arc::new(std::sync::RwLock::new(CachedMemoriaHealth::new( MemoriaHealth::Disabled, ))), @@ -419,6 +425,7 @@ impl AppState { } else { Arc::new(ReqwestMemoriaForwarder::new(base_url.clone(), key)) }; + self.memoria_forwarder_is_override = false; *astra_core::sync_poison::recover_rwlock_write(&self.memoria_health_cache) = CachedMemoriaHealth::new( if master_key.as_deref().is_some_and(|key| !key.is_empty()) { @@ -435,6 +442,7 @@ impl AppState { /// Inject a custom MemoriaForwarder (for testing). pub fn with_memoria_forwarder(mut self, forwarder: Arc) -> Self { self.memoria_forwarder = forwarder; + self.memoria_forwarder_is_override = true; *astra_core::sync_poison::recover_rwlock_write(&self.memoria_health_cache) = CachedMemoriaHealth::new(MemoriaHealth::Unavailable("probe pending".to_string())); self diff --git a/crates/runtime/src/data_layer/models.rs b/crates/runtime/src/data_layer/models.rs index e562077646..fbc3e8b80b 100644 --- a/crates/runtime/src/data_layer/models.rs +++ b/crates/runtime/src/data_layer/models.rs @@ -57,6 +57,120 @@ fn default_model_catalog_limit() -> u32 { DEFAULT_MODEL_CATALOG_LIMIT } +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +pub struct UserModelEndpointRequest { + base_url: String, +} + +pub async fn validate_user_model_endpoint_handler( + State(state): State, + headers: HeaderMap, + Json(request): Json, +) -> Result)> { + let user = state.auth_service.current_user(&headers).await?; + state + .model_service + .validate_user_model_endpoint(user.user_id, request.base_url) + .await?; + Ok(StatusCode::NO_CONTENT) +} + +pub async fn create_user_model_handler( + State(state): State, + headers: HeaderMap, + Json(request): Json, +) -> Result<(StatusCode, Json), (StatusCode, Json)> { + let user = state.auth_service.current_user(&headers).await?; + let model = state + .model_service + .create_user_model( + user.user_id, + UserModelCreateRequestData { + name: request.name, + provider: request.provider, + model: request.model, + base_url: request.base_url, + api_key: request.api_key, + context_window: request.context_window, + is_default: request.is_default, + }, + ) + .await?; + Ok((StatusCode::CREATED, Json(model))) +} + +pub async fn list_user_models_handler( + State(state): State, + headers: HeaderMap, +) -> Result, (StatusCode, Json)> { + let user = state.auth_service.current_user(&headers).await?; + let items = state.model_service.list_user_models(user.user_id).await?; + Ok(Json(UserModelListResponse { items })) +} + +pub async fn get_user_model_handler( + State(state): State, + Path(model_id): Path, + headers: HeaderMap, +) -> Result, (StatusCode, Json)> { + let user = state.auth_service.current_user(&headers).await?; + let model = state + .model_service + .get_user_model(user.user_id, model_id) + .await?; + Ok(Json(model)) +} + +pub async fn update_user_model_handler( + State(state): State, + Path(model_id): Path, + headers: HeaderMap, + Json(request): Json, +) -> Result, (StatusCode, Json)> { + let user = state.auth_service.current_user(&headers).await?; + let model = state + .model_service + .update_user_model( + user.user_id, + model_id, + UserModelUpdateRequestData { + api_key: request.api_key, + context_window: request.context_window, + is_default: request.is_default, + is_active: request.is_active, + }, + ) + .await?; + Ok(Json(model)) +} + +pub async fn delete_user_model_handler( + State(state): State, + Path(model_id): Path, + headers: HeaderMap, +) -> Result)> { + let user = state.auth_service.current_user(&headers).await?; + state + .model_service + .delete_user_model(user.user_id, model_id) + .await?; + Ok(StatusCode::NO_CONTENT) +} + +pub async fn check_user_model_handler( + State(state): State, + Path(model_id): Path, + headers: HeaderMap, +) -> Result, (StatusCode, Json)> { + let user = state.auth_service.current_user(&headers).await?; + let model = state + .model_service + .check_user_model(user.user_id, model_id) + .await?; + Ok(Json(model)) +} + pub async fn create_model_handler( State(state): State, headers: HeaderMap, @@ -183,25 +297,78 @@ async fn effective_model_catalog( } let user = principal.user; let is_admin = !active_only && state.admin.authorizer.require_admin(headers).await.is_ok(); + let user_id = user.user_id.clone(); let page = state .model_service - .list_models_page(user.user_id.clone(), is_admin, query.limit, cursor) + .list_models_page(user_id.clone(), is_admin, query.limit, cursor) .await?; let catalog_revision = state .model_service - .model_catalog_revision(user.user_id, is_admin) + .model_catalog_revision(user_id.clone(), is_admin) .await?; - Ok(EffectiveModelCatalog { - declared: vec![DeclaredModelAccess { + let mut declared = Vec::new(); + let allows_deployment = is_admin + || state + .model_service + .allows_deployment_models(user_id.clone()) + .await?; + if allows_deployment { + declared.push(DeclaredModelAccess { id: "self-hosted".to_string(), kind: ModelAccessKind::SelfHosted, label: "Self-hosted".to_string(), execution_placement: ModelExecutionPlacement::Server, availability: ModelAccessAvailability::Ready, - }], + }); + } + let provider_default = state + .model_service + .default_user_model_offering_id(user_id) + .await? + .map(|offering_id| ModelDefaultCandidate { + offering_id, + source: ModelDefaultSource::Astra, + scope: ModelDefaultScope::EffectiveCatalog, + }); + // Model Access needs the complete catalog both to resolve a default and + // to publish per-access counts that stay stable across pagination. + let default_catalog: Option> = + if !is_admin || active_only || provider_default.is_some() { + Some( + state + .model_service + .list_models(user.user_id, false) + .await? + .into_iter() + .map(ModelListItemResponse::from) + .collect(), + ) + } else { + None + }; + // Match run admission, and inspect the complete catalog rather than the + // current page so access sources cannot disappear across pagination. + if !is_admin + && (!allows_deployment + || default_catalog.as_ref().is_some_and(|catalog| { + catalog + .iter() + .any(|offering| offering.access_kind == ModelAccessKind::CloudByok) + })) + { + declared.push(DeclaredModelAccess { + id: "cloud-byok".to_string(), + kind: ModelAccessKind::CloudByok, + label: "Cloud BYOK".to_string(), + execution_placement: ModelExecutionPlacement::Server, + availability: ModelAccessAvailability::Ready, + }); + } + Ok(EffectiveModelCatalog { + declared, offerings: page.items, - provider_default: None, - default_catalog: None, + provider_default, + default_catalog, next_cursor: page.next_cursor, limit: page.limit, total: page.total, @@ -327,7 +494,7 @@ pub async fn get_memory_model_handler( State(state): State, headers: HeaderMap, ) -> Result, (StatusCode, Json)> { - let _user = state.auth_service.current_user(&headers).await?; + let user = state.auth_service.current_user(&headers).await?; let matrixone = crate::matrix_cloud_runtime::matrix_settings_from_env().map_err(|e| { error_response( StatusCode::SERVICE_UNAVAILABLE, @@ -335,14 +502,15 @@ pub async fn get_memory_model_handler( ) })?; let pool_ref = state.shared_pool.as_ref().map(|sp| sp.get()); - let resolved = resolve_memory_offerings(&matrixone, &state.fernet_encryptor, pool_ref) - .await - .map_err(|e| { - error_response( - StatusCode::SERVICE_UNAVAILABLE, - format!("Memory model resolution failed: {e}"), - ) - })?; + let resolved = + resolve_memory_offerings(&matrixone, &state.fernet_encryptor, &user.user_id, pool_ref) + .await + .map_err(|e| { + error_response( + StatusCode::SERVICE_UNAVAILABLE, + format!("Memory model resolution failed: {e}"), + ) + })?; let offerings = resolved .into_iter() .map(|offering| MemoryInferenceOfferingResponse { diff --git a/crates/runtime/src/matrix_cloud_runtime.rs b/crates/runtime/src/matrix_cloud_runtime.rs index 653506c4ac..162ea117f9 100644 --- a/crates/runtime/src/matrix_cloud_runtime.rs +++ b/crates/runtime/src/matrix_cloud_runtime.rs @@ -56,6 +56,7 @@ impl crate::session_memory::MemoryInferenceResolver for PoolMemoryInferenceResol let offerings = match astra_services::models::resolve_memory_offerings( settings, &self.encryptor, + user_id, Some(pool), ) .await @@ -72,29 +73,14 @@ impl crate::session_memory::MemoryInferenceResolver for PoolMemoryInferenceResol }; offerings .into_iter() - .filter_map(|offering| { - let offering_id = offering.offering_id.clone(); - let model_name = offering.model.model_name.clone(); - match crate::memory_hooks::DurableMemoryInferenceClient::from_offering( - offering, + .map(|offering| { + Arc::new(crate::memory_hooks::DurableMemoryInferenceClient::new( + offering.offering_id, + offering.model.model_name, self.pool.clone(), + self.encryptor.clone(), user_id, - ) { - Ok(client) => { - Some(std::sync::Arc::new(client) - as crate::memory_hooks::MemoryInferenceClient) - } - Err(error) => { - tracing::warn!( - target: "astra_runtime::memory_model", - %offering_id, - %model_name, - %error, - "memory model execution configuration is invalid" - ); - None - } - } + )) as crate::memory_hooks::MemoryInferenceClient }) .collect() } @@ -171,23 +157,23 @@ impl MatrixCloudRuntime { /// Also spins up the [`crate::session_memory::MemoryExtractionService`] /// here, because it needs all three of: encryptor (for selector /// resolve), ingestion sender (for events), and a [`MemoriaPort`] - /// (the sole persistence target for L1 session memory). If - /// [`HttpMemoriaPort::from_env`] returns `None` (no Memoria - /// endpoint configured / offline), the service is NOT built — - /// extraction is opt-in on connectivity, not silent fallback. - pub fn with_encryptor(mut self, enc: Arc) -> Self { + /// (the sole persistence target for L1 session memory). The port resolves + /// each owner's consented credential at operation time; users without + /// write access never cause a Memoria request. + pub fn with_encryptor( + mut self, + enc: Arc, + memoria_client: Option>, + ) -> Self { self.encryptor = Some(Arc::clone(&enc)); let ingestion = self.ingestion.lock().ok().and_then(|g| g.as_ref().cloned()); - let memoria = crate::turn::cloud::memoria_compact::HttpMemoriaPort::from_env(); - if let (Some(ingestion), Some(memoria)) = (ingestion, memoria) { + if let (Some(ingestion), Some(memoria_client)) = (ingestion, memoria_client) { let resolver: Arc = Arc::new(PoolMemoryInferenceResolver { pool: self.shared_pool.clone(), encryptor: Arc::clone(&enc), }); let broker = Arc::new(crate::session_memory::BackgroundActivityBroker::new()); - let memoria_client: Arc = - Arc::new(memoria); let svc = Arc::new( crate::session_memory::MemoryExtractionService::new_owner_scoped_template( resolver, diff --git a/crates/runtime/src/memory_hooks/inference.rs b/crates/runtime/src/memory_hooks/inference.rs index 9ac5282ffc..79cc5eea12 100644 --- a/crates/runtime/src/memory_hooks/inference.rs +++ b/crates/runtime/src/memory_hooks/inference.rs @@ -98,59 +98,83 @@ impl DirectMemoryInferenceClient { } pub(crate) struct DurableMemoryInferenceClient { - direct: DirectMemoryInferenceClient, - ledger: crate::turn::llm::durable::DurableInferenceLedger, + offering_id: String, + model_name: String, + shared_pool: astra_core::SharedPool, + encryptor: Arc, + user_id: String, } impl DurableMemoryInferenceClient { - pub(crate) fn from_offering( - offering: astra_services::ResolvedModelOffering, + pub(crate) fn new( + offering_id: String, + model_name: String, shared_pool: astra_core::SharedPool, + encryptor: Arc, user_id: impl Into, - ) -> Result { - let admitted_execution = - astra_services::AdmittedModelExecution::from_offering(offering.clone())?; - let model = offering.model; - let header_overrides = model.execution_header_overrides()?; - Ok(Self { - direct: DirectMemoryInferenceClient { - base_url: model.base_url, - api_key: model.api_key, - model_name: model.model_name, - wire_model_name: model.wire_model_name, - provider: model.provider, - header_overrides, - request_body_overrides: model.request_body_overrides, - completions_url_override: None, - request_timeout: None, - }, - ledger: crate::turn::llm::durable::DurableInferenceLedger::new( - shared_pool, - user_id, - admitted_execution, - ), - }) + ) -> Self { + Self { + offering_id, + model_name, + shared_pool, + encryptor, + user_id: user_id.into(), + } } } impl std::fmt::Debug for DurableMemoryInferenceClient { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - self.direct.fmt(f) + f.debug_struct("DurableMemoryInferenceClient") + .field("offering_id", &self.offering_id) + .field("model_name", &self.model_name) + .finish() } } #[async_trait] impl MemoryInferencePort for DurableMemoryInferenceClient { fn model_name(&self) -> &str { - &self.direct.model_name + &self.model_name } async fn complete( &self, request: MemoryInferenceRequest<'_>, ) -> Result { - let result = self - .ledger + // Reuse the run/completion admission contract for every background + // provider attempt. Never retain stale plaintext routes in the client. + let execution = astra_services::revalidate_admitted_model_execution( + self.shared_pool.settings(), + &self.encryptor, + &self.user_id, + &self.offering_id, + Some(self.shared_pool.get()), + ) + .await + .map_err(|error| { + astra_core::ClassifiedError::new( + astra_core::ErrorKind::PolicyDenied, + format!("Memory model admission failed: {error}"), + ) + })?; + let direct = DirectMemoryInferenceClient { + base_url: execution.base_url.clone(), + api_key: execution.api_key.clone(), + model_name: execution.model_name.clone(), + wire_model_name: execution.wire_model_name.clone(), + provider: execution.provider.clone(), + header_overrides: execution.header_overrides.clone(), + request_body_overrides: execution.request_body_overrides.clone(), + completions_url_override: None, + request_timeout: None, + }; + let ledger = crate::turn::llm::durable::DurableInferenceLedger::new( + self.shared_pool.clone(), + &self.user_id, + execution, + ); + let result = ledger .execute_nonstream( global_llm_client(), request.invocation_scope.clone(), @@ -159,7 +183,7 @@ impl MemoryInferencePort for DurableMemoryInferenceClient { messages: request.messages, tools: &[], cache_capability: None, - route: self.direct.execution_route(), + route: direct.execution_route(), max_output_tokens: Some(request.max_output_tokens), temperature: Some(request.temperature), has_fallback: false, diff --git a/crates/runtime/src/server/auth_handlers.rs b/crates/runtime/src/server/auth_handlers.rs index 10e21f6d84..6a92160867 100644 --- a/crates/runtime/src/server/auth_handlers.rs +++ b/crates/runtime/src/server/auth_handlers.rs @@ -1,7 +1,32 @@ use axum::extract::Extension; +use serde_json::Value; use super::*; +#[cfg(test)] +fn memory_access_for_scopes(scopes: &[String]) -> Option<&'static str> { + astra_services::auth::memoria::memory_access_for_scopes(scopes).map(|v| v.as_str()) +} + +pub(super) async fn auth_methods_handler(State(state): State) -> Json { + let provider = state.auth_service.memoria_credentials().map(|r| r.provider); + Json(serde_json::json!({ + "password": true, + "memoria": provider.and_then(|p| p.web_url.map(|url| serde_json::json!({ + "issuer": p.issuer, "authorization_url": url + }))) + })) +} + +pub(super) async fn auth_memoria_disconnect_handler( + State(state): State, + headers: HeaderMap, +) -> Result)> { + let user = state.auth_service.current_user(&headers).await?; + state.auth_service.disconnect_memoria(&user.user_id).await?; + Ok(StatusCode::NO_CONTENT) +} + pub(super) async fn auth_register_handler( Extension(trace): Extension, State(state): State, @@ -81,6 +106,27 @@ pub(super) async fn auth_login_handler( Ok(Json(AuthTokenResponse::from(tokens))) } +pub(super) async fn auth_memoria_handler( + Extension(trace): Extension, + State(state): State, + Json(request): Json, +) -> Result, (StatusCode, Json)> { + let login = state + .auth_service + .login_memoria(request.connection_key.trim()) + .await?; + tracing::info!(target: "astra_runtime::auth", request_id = %trace.request_id, "Memoria login succeeded"); + Ok(Json(AuthMemoriaResponse { + user_id: login.tokens.user_id, + access_token: login.tokens.access_token, + refresh_token: login.tokens.refresh_token, + token_type: login.tokens.token_type, + expires_in: login.tokens.expires_in, + memory_access: login.memory_access.as_str().to_string(), + granted_scopes: login.granted_scopes, + })) +} + pub(super) async fn auth_refresh_handler( Extension(trace): Extension, State(state): State, @@ -126,7 +172,7 @@ pub(super) async fn auth_reauthenticate_handler( State(state): State, headers: HeaderMap, Json(request): Json, -) -> Result, (StatusCode, Json)> { +) -> Result)> { let user = state.auth_service.current_user(&headers).await?; let proof = state .auth_service @@ -139,7 +185,26 @@ pub(super) async fn auth_reauthenticate_handler( purpose = proof.purpose.as_str(), "reauthentication proof issued" ); - Ok(Json(AuthReauthenticateResponse::from(proof))) + Ok(axum::response::IntoResponse::into_response(( + [(axum::http::header::CACHE_CONTROL, "no-store")], + Json(AuthReauthenticateResponse::from(proof)), + ))) +} + +pub(super) async fn auth_reauthentication_options_handler( + State(state): State, + headers: HeaderMap, +) -> Result)> { + let user = state.auth_service.current_user(&headers).await?; + Ok(axum::response::IntoResponse::into_response(( + [(axum::http::header::CACHE_CONTROL, "no-store")], + Json( + state + .auth_service + .reauthentication_options(&user.user_id) + .await?, + ), + ))) } pub(super) async fn auth_me_handler( @@ -168,6 +233,9 @@ async fn memory_proxy_call_for_user( endpoint: &str, body: serde_json::Value, ) -> Result, (StatusCode, Json)> { + let requires_write = method != reqwest::Method::GET + && !endpoint.ends_with("/retrieve") + && !endpoint.ends_with("/search"); let requested_scope = memory_proxy_scope(&body, user_id)?; if let Some(scope) = requested_scope.as_ref() { ensure_memory_proxy_session_owner(state, scope).await?; @@ -182,32 +250,45 @@ async fn memory_proxy_call_for_user( } else { None }; + let strict_validation_scope = match strict_recall_scope.as_ref() { + Some(scope) => Some( + astra_memoria::MemoryScope::new( + &memoria_owner_id_for_user(state, user_id).await?, + &scope.session_id, + ) + .map_err(|error| error_response(StatusCode::INTERNAL_SERVER_ERROR, error))?, + ), + None => None, + }; let body = apply_memory_proxy_identity(body, user_id, endpoint); let strict_recall_limit = strict_recall_scope .as_ref() .map(|_| strict_session_recall_limit(&body)); - let mut response = state - .memoria_forwarder - .forward(method, endpoint, body) - .await - .map_err(|error| { - tracing::warn!( - target: "astra_runtime::auth", - endpoint = endpoint, - error = %error, - "memory proxy forward failed" - ); - if error.contains("not configured") { - error_response(StatusCode::SERVICE_UNAVAILABLE, &error) - } else if let Some(status) = parse_memoria_forward_status(&error) { - error_response(status, &error) - } else { - internal_error(&error) - } - })?; + let mut response = + forward_memoria_for_user(state, user_id, requires_write, method, endpoint, body) + .await + .map_err(|error| { + tracing::warn!( + target: "astra_runtime::auth", + endpoint = endpoint, + error = %error, + "memory proxy forward failed" + ); + if error.contains("not configured") { + error_response(StatusCode::SERVICE_UNAVAILABLE, &error) + } else if error.contains("disabled by the user") + || error.contains("not enabled for this Astra account") + { + error_response(StatusCode::FORBIDDEN, &error) + } else if let Some(status) = parse_memoria_forward_status(&error) { + error_response(status, &error) + } else { + internal_error(&error) + } + })?; - if let Some(scope) = strict_recall_scope.as_ref() { + if let Some(scope) = strict_validation_scope.as_ref() { if let Err(error) = astra_memoria::validate_strict_recall_payload(&response, scope) { tracing::error!( target: "astra_runtime::auth", @@ -239,10 +320,15 @@ async fn memory_proxy_call_for_user( "memory_type": "working", "limit": limit, }); - match state - .memoria_forwarder - .forward(reqwest::Method::GET, "/v1/memories", list_request) - .await + match forward_memoria_for_user( + state, + user_id, + false, + reqwest::Method::GET, + "/v1/memories", + list_request, + ) + .await { Ok(working) => { if let Err(error) = @@ -297,6 +383,95 @@ async fn memory_proxy_call_for_user( Ok(Json(response)) } +async fn memoria_owner_id_for_user( + state: &AppState, + user_id: &str, +) -> Result)> { + if state.memoria_forwarder_is_override { + return Ok(user_id.to_string()); + } + let Some(resolver) = state.auth_service.memoria_credentials() else { + return Err(error_response( + StatusCode::FORBIDDEN, + "Memoria connection is not configured", + )); + }; + resolver + .resolve(user_id) + .await + .map_err(internal_error)? + .map(|credential| credential.owner) + .ok_or_else(|| error_response(StatusCode::FORBIDDEN, "memory access is not enabled")) +} + +async fn forward_memoria_for_user( + state: &AppState, + user_id: &str, + requires_write: bool, + method: reqwest::Method, + endpoint: &str, + mut body: serde_json::Value, +) -> Result { + // Explicit composition overrides are used by bounded in-process fixtures + // and custom deployments. They are never inferred from a configured + // server master key, so normal production requests remain BYOK-only. + if state.memoria_forwarder_is_override { + return state + .memoria_forwarder + .forward(method, endpoint, body) + .await; + } + let resolver = state + .auth_service + .memoria_credentials() + .ok_or("Memoria connection is not configured")?; + let credential = resolver + .resolve(user_id) + .await? + .ok_or("memory access is not enabled for this Astra account")?; + if !credential.access.allows(requires_write) { + return Err("memory access is disabled by the user".into()); + } + let connection_key = credential.key; + if let Some(object) = body.as_object_mut() { + object.remove("user_id"); + } + let url = format!( + "{}{}", + resolver.provider.base_url.trim_end_matches('/'), + endpoint + ); + let request = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .build() + .map_err(|_| "Memoria HTTP client unavailable")? + .request(method.clone(), url) + .bearer_auth(connection_key) + .header("X-Memoria-Tool", "astra") + .timeout(std::time::Duration::from_secs(30)); + let response = if method == reqwest::Method::GET { + request.query(&body) + } else { + request.json(&body) + } + .send() + .await + .map_err(|error| format!("Memoria request failed: {error}"))?; + let status = response.status(); + let text = response + .text() + .await + .map_err(|error| format!("Memoria response read error: {error}"))?; + if !status.is_success() { + let bounded: String = text.chars().take(4096).collect(); + return Err(format!("Memoria error {status}: {bounded}")); + } + if text.trim().is_empty() { + return Ok(serde_json::json!({})); + } + serde_json::from_str(&text).map_err(|error| format!("Memoria parse error: {error}")) +} + const MAX_STRICT_SESSION_RECALL_ITEMS: usize = 50; fn strict_session_recall_limit(body: &serde_json::Value) -> usize { @@ -685,9 +860,8 @@ async fn memoria_management_proxy_call( body.unwrap_or_else(|| serde_json::json!({})), &user.user_id, ); - state - .memoria_forwarder - .forward(method, endpoint, body) + let requires_write = method != reqwest::Method::GET; + forward_memoria_for_user(state, &user.user_id, requires_write, method, endpoint, body) .await .map(Json) .map_err(|error| { @@ -699,6 +873,12 @@ async fn memoria_management_proxy_call( ); if error.contains("not configured") { error_response(StatusCode::SERVICE_UNAVAILABLE, &error) + } else if error.contains("disabled by the user") + || error.contains("not enabled for this Astra account") + { + error_response(StatusCode::FORBIDDEN, &error) + } else if let Some(status) = parse_memoria_forward_status(&error) { + error_response(status, &error) } else { internal_error(&error) } @@ -885,12 +1065,39 @@ pub(super) async fn memoria_proxy_consolidate_handler( mod tests { use super::{ apply_memoria_management_identity, apply_memory_proxy_identity, encode_memoria_memory_id, - exact_memory_ids_for_user_purge, is_strict_session_recall, memory_proxy_scope, - normalize_exact_memory_purge_receipt, parse_memoria_forward_status, + exact_memory_ids_for_user_purge, is_strict_session_recall, memory_access_for_scopes, + memory_proxy_scope, normalize_exact_memory_purge_receipt, parse_memoria_forward_status, }; use axum::http::StatusCode; use serde_json::json; + #[test] + fn memoria_scope_sets_map_only_to_supported_access_modes() { + let strings = |values: &[&str]| { + values + .iter() + .map(|value| value.to_string()) + .collect::>() + }; + assert_eq!( + memory_access_for_scopes(&strings(&["identity:read"])), + Some("none") + ); + assert_eq!( + memory_access_for_scopes(&strings(&["memory:read", "identity:read"])), + Some("read_only") + ); + assert_eq!( + memory_access_for_scopes(&strings(&["memory:write", "identity:read", "memory:read",])), + Some("read_write") + ); + assert_eq!(memory_access_for_scopes(&strings(&["memory:read"])), None); + assert_eq!( + memory_access_for_scopes(&strings(&["identity:read", "admin"])), + None + ); + } + #[test] fn apply_memory_proxy_identity_overwrites_user_but_preserves_authorized_session() { let body = json!({ diff --git a/crates/runtime/src/server/completions.rs b/crates/runtime/src/server/completions.rs index b9b52d252b..c4b098efa6 100644 --- a/crates/runtime/src/server/completions.rs +++ b/crates/runtime/src/server/completions.rs @@ -89,6 +89,7 @@ pub(super) async fn completions_handler( }; super::model_execution_admission::admit_model_execution( &state.model_service, + &user.user_id, &selection, None, None, @@ -96,39 +97,52 @@ pub(super) async fn completions_handler( ) .await? } else { - let matrixone = - crate::matrix_cloud_runtime::matrix_settings_from_env().map_err(|error| { + let offering_id = if let Some(offering_id) = state + .model_service + .default_user_model_offering_id(user.user_id.clone()) + .await? + { + offering_id + } else { + if !state + .model_service + .allows_deployment_models(user.user_id.clone()) + .await? + { + return Err(crate::error_response_coded( + StatusCode::BAD_REQUEST, + "Configure and select your own BYOK model; deployment model fallback is disabled", + "missing_model_selection", + )); + } + let matrixone = + crate::matrix_cloud_runtime::matrix_settings_from_env().map_err(|error| { + crate::error_response_coded( + StatusCode::SERVICE_UNAVAILABLE, + format!("MatrixOne configuration unavailable: {error}"), + "model_catalog_unavailable", + ) + })?; + astra_services::resolve_reasoning_offering( + &matrixone, + &state.fernet_encryptor, + state.admin.config_service.as_ref(), + state.shared_pool.as_ref().map(|pool| pool.get()), + ) + .await + .map_err(|error| { crate::error_response_coded( StatusCode::SERVICE_UNAVAILABLE, - format!("MatrixOne configuration unavailable: {error}"), - "model_catalog_unavailable", + format!("Default Offering resolution failed: {error}"), + "model_default_unavailable", ) - })?; - let selected = astra_services::resolve_reasoning_offering( - &matrixone, - &state.fernet_encryptor, - state.admin.config_service.as_ref(), - state.shared_pool.as_ref().map(|pool| pool.get()), - ) - .await - .map_err(|error| { - crate::error_response_coded( - StatusCode::SERVICE_UNAVAILABLE, - format!("Default Offering resolution failed: {error}"), - "model_default_unavailable", - ) - })?; - let offering = state + })? + .offering_id + }; + state .model_service - .revalidate_model_offering(selected.offering_id) - .await?; - astra_services::AdmittedModelExecution::from_offering(offering).map_err(|error| { - crate::error_response_coded( - StatusCode::SERVICE_UNAVAILABLE, - format!("Default Offering execution configuration is invalid: {error}"), - "model_execution_configuration_invalid", - ) - })? + .admit_model_offering(user.user_id.clone(), offering_id) + .await? }; // 3. Durably admit the logical invocation before provider I/O. Auxiliary diff --git a/crates/runtime/src/server/edge/edge_callback_handlers.rs b/crates/runtime/src/server/edge/edge_callback_handlers.rs index 8b8d660fc4..038a9e744f 100644 --- a/crates/runtime/src/server/edge/edge_callback_handlers.rs +++ b/crates/runtime/src/server/edge/edge_callback_handlers.rs @@ -1342,7 +1342,8 @@ pub(crate) async fn post_provider_interaction_respond_handler( provider_scope_id: context.provider_scope_id.clone(), } } - astra_services::AuthPrincipalOrigin::Internal => { + astra_services::AuthPrincipalOrigin::Internal + | astra_services::AuthPrincipalOrigin::VerifiedProvider { .. } => { return Err(error_response( StatusCode::FORBIDDEN, "Provider interaction responses require provider authorization", diff --git a/crates/runtime/src/server/model_execution_admission.rs b/crates/runtime/src/server/model_execution_admission.rs index 271bb69c4e..bd767f4185 100644 --- a/crates/runtime/src/server/model_execution_admission.rs +++ b/crates/runtime/src/server/model_execution_admission.rs @@ -8,7 +8,7 @@ use astra_services::{ use astra_turn_types::ModelSelection; use axum::{Json, http::StatusCode}; -use crate::{error_response_coded, internal_error}; +use crate::error_response_coded; /// Admit one Offering into the single execution-material contract consumed by /// every agent and inference adapter. @@ -19,6 +19,7 @@ use crate::{error_response_coded, internal_error}; /// non-serializable value. pub(crate) async fn admit_model_execution( model_service: &Arc, + user_id: &str, selection: &ModelSelection, resolved: Option<&ResolvedModelSelection>, gateway: Option<&RuntimeCapabilityDescriptorRequest>, @@ -83,10 +84,9 @@ pub(crate) async fn admit_model_execution( "model_selection_invalid", )); } - let offering = model_service - .revalidate_model_offering(selection.offering_id.clone()) - .await?; - AdmittedModelExecution::from_offering(offering).map_err(internal_error) + model_service + .admit_model_offering(user_id.to_string(), selection.offering_id.clone()) + .await } fn is_exact_runtime_identity(value: &str) -> bool { @@ -198,6 +198,7 @@ mod tests { let service: Arc = Arc::new(StaticModelService); let catalog = admit_model_execution( &service, + "user-1", &ModelSelection { offering_id: "offer-server".into(), }, @@ -213,6 +214,7 @@ mod tests { let endpoint = admit_model_execution( &service, + "user-1", &ModelSelection { offering_id: "offer-edge".into(), }, @@ -250,6 +252,7 @@ mod tests { for context_window in [None, Some(0)] { let error = admit_model_execution( &service, + "user-1", &ModelSelection { offering_id: "offer-edge".into(), }, @@ -285,6 +288,7 @@ mod tests { let service: Arc = Arc::new(StaticModelService); let error = admit_model_execution( &service, + "user-1", &ModelSelection { offering_id: "offer-requested".into(), }, diff --git a/crates/runtime/src/server/router_builder/catalog.rs b/crates/runtime/src/server/router_builder/catalog.rs index 1e6ad064e4..32549f5539 100644 --- a/crates/runtime/src/server/router_builder/catalog.rs +++ b/crates/runtime/src/server/router_builder/catalog.rs @@ -2,6 +2,25 @@ use super::*; pub(super) fn add_routes(router: Router) -> Router { router + .route( + "/me/models", + get(data_layer::models::list_user_models_handler) + .post(data_layer::models::create_user_model_handler), + ) + .route( + "/me/models/validate-endpoint", + post(data_layer::models::validate_user_model_endpoint_handler), + ) + .route( + "/me/models/{model_id}", + get(data_layer::models::get_user_model_handler) + .put(data_layer::models::update_user_model_handler) + .delete(data_layer::models::delete_user_model_handler), + ) + .route( + "/me/models/{model_id}/check", + post(data_layer::models::check_user_model_handler), + ) .route( "/model-access", get(data_layer::models::get_model_access_handler), diff --git a/crates/runtime/src/server/router_builder/realtime.rs b/crates/runtime/src/server/router_builder/realtime.rs index 40c8df94da..8805411df3 100644 --- a/crates/runtime/src/server/router_builder/realtime.rs +++ b/crates/runtime/src/server/router_builder/realtime.rs @@ -11,11 +11,18 @@ pub(super) fn add_routes(router: Router) -> Router { .route("/metrics", get(meta_handlers::metrics_handler)) .route("/auth/register", post(auth_handlers::auth_register_handler)) .route("/auth/login", post(auth_handlers::auth_login_handler)) + .route("/auth/methods", get(auth_handlers::auth_methods_handler)) + .route( + "/auth/memoria", + post(auth_handlers::auth_memoria_handler) + .delete(auth_handlers::auth_memoria_disconnect_handler), + ) .route("/auth/refresh", post(auth_handlers::auth_refresh_handler)) .route("/auth/logout", post(auth_handlers::auth_logout_handler)) .route( "/auth/reauthenticate", - post(auth_handlers::auth_reauthenticate_handler), + post(auth_handlers::auth_reauthenticate_handler) + .get(auth_handlers::auth_reauthentication_options_handler), ) .route("/auth/me", get(auth_handlers::auth_me_handler)) .route( diff --git a/crates/runtime/src/server/run/lifecycle/mod.rs b/crates/runtime/src/server/run/lifecycle/mod.rs index f004c6ba33..32ac733323 100644 --- a/crates/runtime/src/server/run/lifecycle/mod.rs +++ b/crates/runtime/src/server/run/lifecycle/mod.rs @@ -8888,19 +8888,51 @@ impl AgenticRunLifecycleService { .model_service .list_models(user_id.to_string(), false) .await?; - let projection = astra_services::project_model_access( - vec![astra_services::DeclaredModelAccess { + let allows_deployment = self + .model_service + .allows_deployment_models(user_id.to_string()) + .await?; + let has_cloud_byok = !allows_deployment + || offerings.iter().any(|offering| { + offering.access_kind == astra_services::ModelAccessKind::CloudByok + }); + let mut declared = Vec::new(); + if allows_deployment { + declared.push(astra_services::DeclaredModelAccess { id: "self-hosted".to_string(), kind: astra_services::ModelAccessKind::SelfHosted, label: "Self-hosted".to_string(), execution_placement: astra_services::ModelExecutionPlacement::Server, availability: astra_services::ModelAccessAvailability::Ready, - }], - offerings - .into_iter() - .filter(|offering| offering.is_active) - .map(astra_services::ModelListItemResponse::from) - .collect(), + }); + } + if has_cloud_byok { + declared.push(astra_services::DeclaredModelAccess { + id: "cloud-byok".to_string(), + kind: astra_services::ModelAccessKind::CloudByok, + label: "Cloud BYOK".to_string(), + execution_placement: astra_services::ModelExecutionPlacement::Server, + availability: astra_services::ModelAccessAvailability::Ready, + }); + } + let user_default = self + .model_service + .default_user_model_offering_id(user_id.to_string()) + .await? + .map(|offering_id| astra_services::ModelDefaultCandidate { + offering_id, + source: astra_services::ModelDefaultSource::Astra, + scope: astra_services::ModelDefaultScope::EffectiveCatalog, + }); + let offering_views = offerings + .into_iter() + .filter(|offering| offering.is_active) + .map(astra_services::ModelListItemResponse::from) + .collect::>(); + let projection = astra_services::project_model_access_with_default( + declared, + offering_views, + user_default, chrono::Utc::now().to_rfc3339(), ) .map_err(|error| { @@ -8921,6 +8953,7 @@ impl AgenticRunLifecycleService { let selection = ModelSelection { offering_id }; let admitted = crate::server::model_execution_admission::admit_model_execution( &self.model_service, + user_id, &selection, None, None, @@ -8975,6 +9008,7 @@ impl AgenticRunLifecycleService { request.admitted_model_execution = Some( crate::server::model_execution_admission::admit_model_execution( &self.model_service, + user_id, selection, Some(resolved), Some(gateway), @@ -8994,6 +9028,7 @@ impl AgenticRunLifecycleService { } let admitted = crate::server::model_execution_admission::admit_model_execution( &self.model_service, + user_id, selection, None, None, @@ -10508,6 +10543,12 @@ impl AgenticRunLifecycleService { builder = builder.with_canonical_work_context_binding(binding.context_binding()); } + builder = builder.with_memoria_client( + self.memory_extraction_service + .as_ref() + .and_then(|svc| svc.memoria_client_for_owner(user_id).ok()), + ); + if let Some(pool) = &self.shared_pool { builder = builder.with_pool(pool.clone()); } @@ -13119,7 +13160,8 @@ impl RunLifecycleService for AgenticRunLifecycleService { // explicit workspace/executor binding and cannot silently fall back. let mut root_runtime_context_guard = None; if let Some(workspace) = server_tool_executor_workspace { - let memoria_base = Some(astra_core::MemoriaSettings::from_env().base_url); + // Memory access is provided by the composition-owned, scoped port. + let memoria_base = None; let mut executor = runtime_tool_executor::RuntimeToolExecutor::new( workspace.clone(), user_id.clone(), @@ -15458,7 +15500,7 @@ impl RunLifecycleService for AgenticRunLifecycleService { // tools to edge or blocks when edge is unavailable. let mut root_runtime_context_guard = None; if let Some(workspace) = server_tool_executor_workspace { - let memoria_base = Some(astra_core::MemoriaSettings::from_env().base_url); + let memoria_base = None; let mut executor = runtime_tool_executor::RuntimeToolExecutor::new( workspace.clone(), user_id.clone(), @@ -20033,21 +20075,22 @@ impl ServerSubRunExecutor { } return Ok(Some(execution.clone())); } - let offering = astra_services::resolve_active_llm_offering( + let execution = astra_services::revalidate_admitted_model_execution( &self.matrixone, self.encryptor.as_ref(), + &config.user_id, offering_id, self.shared_pool.as_ref().map(SharedPool::get), ) .await .map_err(|error| error.to_string())?; - if offering.model.model_name != expected_model_name { + if execution.model_name != expected_model_name { return Err( "durable sub-run Offering changed after admission; refusing route drift" .to_string(), ); } - astra_services::AdmittedModelExecution::from_offering(offering).map(Some) + Ok(Some(execution)) } async fn select_subrun_execution( @@ -20061,15 +20104,16 @@ impl ServerSubRunExecutor { .or(self.admitted_model_execution.as_ref()) .cloned()); }; - let offering = astra_services::revalidate_active_llm_offering( + let execution = astra_services::revalidate_admitted_model_execution( &self.matrixone, self.encryptor.as_ref(), + &config.user_id, &selection.offering_id, self.shared_pool.as_ref().map(SharedPool::get), ) .await .map_err(|error| error.to_string())?; - astra_services::AdmittedModelExecution::from_offering(offering).map(Some) + Ok(Some(execution)) } async fn exact_durable_subrun_control_authority( @@ -21005,6 +21049,9 @@ impl SubRunExecutor for ServerSubRunExecutor { .with_edge_callback_ledger(self.edge_callback_ledger.clone()) .with_interaction_mode(Some(config.interaction_mode)); + builder = builder.with_memoria_client(self.memory_extraction_service.as_ref() + .and_then(|svc| svc.memoria_client_for_owner(&config.user_id).ok())); + if let Some(pool) = &self.shared_pool { builder = builder.with_pool(pool.clone()); } @@ -21322,7 +21369,7 @@ impl SubRunExecutor for ServerSubRunExecutor { // Without this, the headless pipeline fallback cannot execute tools // server-side and sub-agents would get edge-protocol errors. { - let memoria_base = Some(astra_core::MemoriaSettings::from_env().base_url); + let memoria_base = None; let agent_working_dir = subrun_workspace.clone(); let mut executor = runtime_tool_executor::RuntimeToolExecutor::new( subrun_workspace, diff --git a/crates/runtime/src/server/server_loop_host.rs b/crates/runtime/src/server/server_loop_host.rs index 7cfe0eb76d..a85edfb1f6 100644 --- a/crates/runtime/src/server/server_loop_host.rs +++ b/crates/runtime/src/server/server_loop_host.rs @@ -3119,7 +3119,6 @@ impl ServerAgenticLoopHostBuilder { user_id: String, session_id: String, ) -> Self { - let memoria_owner_user_id = user_id.clone(); Self { matrixone, encryptor, @@ -3146,12 +3145,7 @@ impl ServerAgenticLoopHostBuilder { static_tool_catalog_admissible: true, plan_resume_hint: None, plan_authoring_active: false, - memoria_client: crate::turn::cloud::memoria_compact::HttpMemoriaPort::from_env().map( - |client| { - let client = client.with_owner_user_id(memoria_owner_user_id.clone()); - Arc::new(client) as Arc - }, - ), + memoria_client: None, server_service_tool_catalog_enabled: true, control_plane_tool_catalog_enabled: true, #[cfg(feature = "e2e-hooks")] @@ -6215,15 +6209,15 @@ impl ServerAgenticLoopHost { // no Server-owned route or secret to refresh at this boundary. return Ok(()); } - let offering = astra_services::revalidate_active_llm_offering( + let execution = astra_services::revalidate_admitted_model_execution( &self.matrixone, self.encryptor.as_ref(), + &self.user_id, &admitted.offering_id, self.shared_pool.as_ref().map(SharedPool::get), ) .await .map_err(|error| error.to_string())?; - let execution = astra_services::AdmittedModelExecution::from_offering(offering)?; self.admitted_model_execution = Some(execution); self.clear_resolved_llm_config(); Ok(()) diff --git a/crates/runtime/src/server/server_skill_subrun.rs b/crates/runtime/src/server/server_skill_subrun.rs index a38dccb5d5..fe5e91cf7f 100644 --- a/crates/runtime/src/server/server_skill_subrun.rs +++ b/crates/runtime/src/server/server_skill_subrun.rs @@ -500,7 +500,7 @@ impl ServerSkillSubRunExecutor { invocation_ledger: crate::server::tool_invocation_runtime::RuntimeToolInvocationLedger, ) -> Result { let workspace = self.provision_skill_workspace(skill_name, presentation_session_id)?; - let memoria_base = Some(astra_core::MemoriaSettings::from_env().base_url); + let memoria_base = None; let mut builder = ToolExecutionService::builder(); if let Some(pool) = &self.edge_connection_pool { builder = builder.edge_connection_pool(pool.clone()); @@ -955,6 +955,8 @@ impl SkillSubRunExecutor for ServerSkillSubRunExecutor { builder = builder.with_dedup_state(dedup.clone()); } + builder = builder.with_memoria_client(self.memory_extraction_service.as_ref() + .and_then(|svc| svc.memoria_client_for_owner(&self.user_id).ok())); let mut host = builder.build(); if let Some(sink) = &self.interaction_sink { host.set_interaction_sink(Arc::clone(sink)); diff --git a/crates/runtime/src/server/session/session_handlers.rs b/crates/runtime/src/server/session/session_handlers.rs index 3299359b4f..abe97b2518 100644 --- a/crates/runtime/src/server/session/session_handlers.rs +++ b/crates/runtime/src/server/session/session_handlers.rs @@ -1899,7 +1899,7 @@ pub(crate) async fn close_session_handler( ) .await?; astra_tools::memoria::MemoriaToolGateway::reset_session_process_state(&session_id); - schedule_session_end_governance(owner_id, session_id); + schedule_session_end_governance(&state, owner_id, session_id); Ok(Json(SessionResponse::from(session))) } @@ -1912,16 +1912,17 @@ pub(crate) async fn close_session_handler( /// Memoria latency cannot make the UI wait, while the existing per-owner, /// per-session debouncer prevents duplicate close requests from writing two /// episodes or purging twice. -fn schedule_session_end_governance(owner_id: String, session_id: String) { - let Some(memoria_client) = crate::turn::cloud::memoria_compact::HttpMemoriaPort::from_env() - else { +fn schedule_session_end_governance(state: &AppState, owner_id: String, session_id: String) { + let Some(resolver) = state.auth_service.memoria_credentials() else { tracing::debug!( owner_id = %owner_id, session_id = %session_id, - "session-end governance skipped because no memory provider is configured" + "session-end governance skipped because credential storage is unavailable" ); return; }; + let memoria_client = + crate::turn::cloud::memoria_compact::UserScopedMemoriaPort::new(resolver, owner_id.clone()); tokio::spawn(async move { let debouncer = crate::turn::session_end_debounce::global(); @@ -1934,7 +1935,6 @@ fn schedule_session_end_governance(owner_id: String, session_id: String) { return; }; - let memoria_client = memoria_client.with_owner_user_id(owner_id.clone()); let session_facts = astra_turn_types::session_facts::SessionFacts::default(); let governance = crate::turn::cloud::session_end_governance::run_session_end_governance( &session_facts, diff --git a/crates/runtime/src/server/state_builder/core.rs b/crates/runtime/src/server/state_builder/core.rs index 79bd09f021..47907b58bb 100644 --- a/crates/runtime/src/server/state_builder/core.rs +++ b/crates/runtime/src/server/state_builder/core.rs @@ -24,6 +24,7 @@ pub(super) fn build_auth_service( ) -> Result, Box> { let mut service = DatabaseAuthService::new(settings.matrixone.clone(), settings.jwt.clone()) .with_pool(shared_pool.clone()) + .with_memoria_settings(&settings.memoria)? .with_encryptor(shared_encryptor.as_ref().clone()); if let Some(control_pool) = control_pool { service = service.with_control_pool(control_pool.clone()); diff --git a/crates/runtime/src/server/state_builder/runtime.rs b/crates/runtime/src/server/state_builder/runtime.rs index 431840200c..d56f402a23 100644 --- a/crates/runtime/src/server/state_builder/runtime.rs +++ b/crates/runtime/src/server/state_builder/runtime.rs @@ -36,7 +36,16 @@ pub(super) async fn build_runtime_wiring( )); let matrix_rt = Arc::new( crate::matrix_cloud_runtime::MatrixCloudRuntime::attach(shared_pool.clone(), "default") - .with_encryptor(Arc::clone(run_encryptor)), + .with_encryptor( + Arc::clone(run_encryptor), + state.auth_service.memoria_credentials().map(|resolver| { + Arc::new( + crate::turn::cloud::memoria_compact::UserScopedMemoriaPort::template( + resolver, + ), + ) as Arc + }), + ), ); let memory_extraction_service = matrix_rt.clone_memory_extraction_service(); let workspace_record_store = Arc::new(astra_services::DatabaseWorkspaceRecordStore::new( diff --git a/crates/runtime/src/session_memory/service.rs b/crates/runtime/src/session_memory/service.rs index 9db40220d2..0bf04d961e 100644 --- a/crates/runtime/src/session_memory/service.rs +++ b/crates/runtime/src/session_memory/service.rs @@ -210,6 +210,10 @@ impl std::fmt::Debug for MemoryExtractionService { } impl MemoryExtractionService { + /// Share the composition-owned memory provider with the run and its children. + pub fn memoria_client_for_owner(&self, user: &str) -> Result, String> { + self.memoria_client.bind_owner(user) + } /// Build a service. `memoria_client` is required — callers that /// can't produce one (offline CLI, no Memoria configured) should /// simply skip constructing the service and leave @@ -798,6 +802,16 @@ impl MemoryExtractionService { // ── internals ───────────────────────────────────────────────────── async fn run_one(self: Arc, req: ExtractionRequest, content_fingerprint: u64) { + // This lightweight admission task must not load snapshots, resolve a + // model, spend tokens or schedule writes when consent forbids them. + match self.memoria_client.admits_operation(true).await { + Ok(true) => {} + Ok(false) => return, + Err(error) => { + tracing::warn!(%error, "memory extraction admission unavailable"); + return; + } + } let Some((session_id, turn)) = req.session_coordinates() else { tracing::error!( scope_kind = req.inference_scope.kind(), @@ -2701,6 +2715,97 @@ mod tests { } } + #[tokio::test] + async fn consent_admission_prevents_read_only_and_disabled_extraction_cost() { + use std::sync::atomic::{AtomicUsize, Ordering}; + struct ConsentPort { + read: bool, + calls: Arc, + } + #[async_trait] + impl MemoriaPort for ConsentPort { + async fn purge_working(&self, _: &str) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + Ok(0) + } + async fn admits_operation(&self, write: bool) -> Result { + Ok(self.read && !write) + } + async fn retrieve_ext( + &self, + _: &str, + _: Option<&str>, + _: usize, + _: bool, + ) -> Result, String> { + self.calls.fetch_add(1, Ordering::SeqCst); + Ok(vec![]) + } + async fn store( + &self, + _: &str, + _: &str, + _: Option<&str>, + _: Option<&str>, + ) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + Ok("unexpected".into()) + } + } + #[derive(Debug)] + struct Resolver(Arc); + #[async_trait] + impl MemoryInferenceResolver for Resolver { + async fn resolve_candidates(&self, _: &str) -> Vec { + self.0.fetch_add(1, Ordering::SeqCst); + vec![] + } + } + for read in [false, true] { + let calls = Arc::new(AtomicUsize::new(0)); + let port: Arc = Arc::new(ConsentPort { + read, + calls: calls.clone(), + }); + let (ingestion, _rx) = IngestionSender::for_tests(32); + let service = Arc::new(MemoryExtractionService::new( + Arc::new(Resolver(calls.clone())), + port.clone(), + ingestion, + "consent-user", + Arc::new(BackgroundActivityBroker::new()), + )); + service.maybe_spawn(sample_req("consent-session", 50_000, false)); + service.wait_for_pending(Duration::from_secs(2)).await; + crate::turn::cloud::session_end_governance::run_session_end_governance( + &Default::default(), + "consent-session", + port.as_ref(), + ) + .await + .unwrap(); + if !read { + let recalled = crate::turn::memory_prefetch::prefetch_memories_with_client( + port.as_ref(), + "recall", + "consent-user", + "consent-session", + 5, + ) + .await; + assert_eq!( + recalled.outcome, + astra_turn_types::MemoryRetrievalOutcome::NotAttempted + ); + } + assert_eq!( + calls.load(Ordering::SeqCst), + 0, + "no retrieval, model resolution or write before consent" + ); + } + } + fn build_ctx_with_memoria( selector: Option, memoria: Arc, diff --git a/crates/runtime/src/turn/cloud/memoria_compact.rs b/crates/runtime/src/turn/cloud/memoria_compact.rs index 866b8f4654..bb2cbd93e3 100644 --- a/crates/runtime/src/turn/cloud/memoria_compact.rs +++ b/crates/runtime/src/turn/cloud/memoria_compact.rs @@ -235,6 +235,7 @@ impl HttpMemoriaPort { api_key, http: astra_core::net::build_internal_http_client( reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) .connect_timeout(std::time::Duration::from_secs(10)) .timeout(std::time::Duration::from_secs(60)), "memoria compact client", @@ -295,7 +296,8 @@ impl HttpMemoriaPort { let request = self .http .request(method, url) - .header("Authorization", format!("Bearer {}", self.api_key)); + .header("Authorization", format!("Bearer {}", self.api_key)) + .header("X-Memoria-Tool", "astra"); Ok(match owner { Some(owner) => request.header("X-User-Id", owner), None => request, @@ -327,6 +329,211 @@ impl HttpMemoriaPort { } } +/// Resolves the current user's scoped Memoria credential for every operation. +/// This makes revocation and access-mode changes effective without restarting +/// an Astra runtime and prevents the server master key from becoming an +/// implicit end-user consent path. +#[derive(Clone)] +pub struct UserScopedMemoriaPort { + resolver: astra_services::auth::memoria::MemoriaCredentialResolver, + owner_user_id: Option, +} +impl UserScopedMemoriaPort { + pub fn new( + resolver: astra_services::auth::memoria::MemoriaCredentialResolver, + owner_user_id: String, + ) -> Self { + Self { + resolver, + owner_user_id: Some(owner_user_id), + } + } + pub fn template(resolver: astra_services::auth::memoria::MemoriaCredentialResolver) -> Self { + Self { + resolver, + owner_user_id: None, + } + } + fn owner_user_id(&self) -> Result<&str, String> { + self.owner_user_id + .as_deref() + .ok_or_else(|| "Memoria requires an authenticated owner".into()) + } + async fn client(&self, write: bool) -> Result<(HttpMemoriaPort, String), String> { + let credential = self + .resolver + .resolve(self.owner_user_id()?) + .await? + .ok_or("memory access is not enabled")?; + enforce_memory_access(credential.access.as_str(), write)?; + Ok(( + HttpMemoriaPort::new(self.resolver.provider.base_url.clone(), credential.key) + .with_owner_user_id(credential.owner.clone()), + credential.owner, + )) + } +} + +fn enforce_memory_access(access: &str, write: bool) -> Result<(), String> { + if access == "none" { + return Err("memory access is disabled by the user".to_string()); + } + if write && access != "read_write" { + return Err("memory write access is disabled by the user".to_string()); + } + if access != "read_only" && access != "read_write" { + return Err("Memoria credential has an invalid access mode".to_string()); + } + Ok(()) +} + +#[async_trait::async_trait] +impl MemoriaPort for UserScopedMemoriaPort { + async fn admits_operation(&self, write: bool) -> Result { + Ok(self + .resolver + .resolve(self.owner_user_id()?) + .await? + .is_some_and(|credential| credential.access.allows(write))) + } + fn bind_owner(&self, user_id: &str) -> Result, String> { + if self + .owner_user_id + .as_deref() + .is_some_and(|owner| owner != user_id) + { + return Err("memory_scope_violation: requested owner differs from bound owner".into()); + } + let mut bound = self.clone(); + bound.owner_user_id = Some(user_id.to_string()); + Ok(std::sync::Arc::new(bound)) + } + + async fn retrieve_for_prompt( + &self, + query: &str, + user_id: &str, + session_id: &str, + top_k: usize, + ) -> Result, String> { + if user_id != self.owner_user_id()? { + return Err("memory_scope_violation: requested owner differs from bound owner".into()); + } + let (client, memoria_user_id) = self.client(false).await?; + client + .retrieve_for_prompt(query, &memoria_user_id, session_id, top_k) + .await + } + + async fn retrieve_ext( + &self, + query: &str, + session_id: Option<&str>, + top_k: usize, + filter_session: bool, + ) -> Result, String> { + self.client(false) + .await? + .0 + .retrieve_ext(query, session_id, top_k, filter_session) + .await + } + + async fn retrieve_scoped_typed( + &self, + query: &str, + session_id: &str, + top_k: usize, + memory_types: &[&str], + ) -> Result, String> { + self.client(false) + .await? + .0 + .retrieve_scoped_typed(query, session_id, top_k, memory_types) + .await + } + + async fn store( + &self, + content: &str, + memory_type: &str, + session_id: Option<&str>, + trust_tier: Option<&str>, + ) -> Result { + self.client(true) + .await? + .0 + .store(content, memory_type, session_id, trust_tier) + .await + } + + async fn purge_working(&self, session_id: &str) -> Result { + self.client(true).await?.0.purge_working(session_id).await + } + + async fn purge_memory_types( + &self, + session_id: &str, + memory_types: &[&str], + ) -> Result { + self.client(true) + .await? + .0 + .purge_memory_types(session_id, memory_types) + .await + } + + async fn delete(&self, memory_id: &str) -> Result<(), String> { + self.client(true).await?.0.delete(memory_id).await + } + + async fn store_episode(&self, session_id: &str, overview: &str) -> Result { + self.client(true) + .await? + .0 + .store_episode(session_id, overview) + .await + } + + async fn store_scene( + &self, + session_id: &str, + signal: &str, + summary: &str, + ) -> Result { + self.client(true) + .await? + .0 + .store_scene(session_id, signal, summary) + .await + } + + async fn reflect_session( + &self, + session_id: &str, + force: bool, + ) -> Result { + self.client(true) + .await? + .0 + .reflect_session(session_id, force) + .await + } + + async fn feedback( + &self, + memory_id: &str, + signal: &str, + context: Option<&str>, + ) -> Result<(), String> { + self.client(true) + .await? + .0 + .feedback(memory_id, signal, context) + .await + } +} + fn parse_retrieved_memories(data: &Value) -> Vec { let Some(arr) = data .get("memories") @@ -1115,6 +1322,10 @@ pub async fn compact_with_memoria( compact_config: Option<&CompactConfig>, summary_client: Option<&dyn SummaryLlmClient>, ) -> CompactResult { + let client = match client { + Some(client) if client.admits_operation(false).await.unwrap_or(false) => Some(client), + _ => None, + }; // Check if we should attempt Memoria retrieval let should_retrieve = params.current_tokens >= config.min_tokens_for_retrieval && params.tier != CompactionTier::Normal @@ -1274,6 +1485,17 @@ mod tests { use std::sync::{Arc, Mutex}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; + #[test] + fn user_memory_access_is_enforced_before_transport_resolution() { + assert!(enforce_memory_access("none", false).is_err()); + assert!(enforce_memory_access("none", true).is_err()); + assert!(enforce_memory_access("read_only", false).is_ok()); + assert!(enforce_memory_access("read_only", true).is_err()); + assert!(enforce_memory_access("read_write", false).is_ok()); + assert!(enforce_memory_access("read_write", true).is_ok()); + assert!(enforce_memory_access("unexpected", false).is_err()); + } + async fn capture_one_http_request( status: &str, response_body: &'static [u8], diff --git a/crates/runtime/src/turn/cloud/session_end_governance.rs b/crates/runtime/src/turn/cloud/session_end_governance.rs index f49a6849c7..8844dbbb73 100644 --- a/crates/runtime/src/turn/cloud/session_end_governance.rs +++ b/crates/runtime/src/turn/cloud/session_end_governance.rs @@ -193,6 +193,9 @@ pub async fn run_session_end_governance( session_id: &str, client: &dyn super::memoria_compact::MemoriaPort, ) -> Result { + if !client.admits_operation(true).await? { + return Ok(SessionEndReport::default()); + } let mut report = SessionEndReport::default(); // Read before purge: the canonical narrative is the most valuable diff --git a/crates/runtime/src/turn/llm/client.rs b/crates/runtime/src/turn/llm/client.rs index 53b80a1f66..8e18461fde 100644 --- a/crates/runtime/src/turn/llm/client.rs +++ b/crates/runtime/src/turn/llm/client.rs @@ -3233,7 +3233,9 @@ fn build_provider_request_body_with_cache_capability( body["system"] = Value::Array(system); } if let Some(max_out) = max_output_tokens { - body["max_tokens"] = json!(max_out); + astra_core::model_wire::apply_chat_output_token_limit( + &mut body, provider, max_out, + ); } if let Some(temp) = temperature { body["temperature"] = json!(temp); @@ -3277,7 +3279,7 @@ fn build_provider_request_body_with_cache_capability( } if let Some(max_out) = max_output_tokens { // When thinking is active, providers like DeepSeek allocate a - // thinking_budget that must be LESS than max_completion_tokens. + // thinking_budget that must be LESS than the output token limit. // If max_out is too small, the request will 400. Bump to at // least thinking_budget + a headroom for the visible answer. // @@ -3295,7 +3297,7 @@ fn build_provider_request_body_with_cache_capability( tracing::debug!( user_max = max_out, bumped_to = required_floor, - "max_completion_tokens bumped to fit thinking budget" + "output token limit bumped to fit thinking budget" ); required_floor } else { @@ -3304,7 +3306,11 @@ fn build_provider_request_body_with_cache_capability( } else { max_out }; - body["max_completion_tokens"] = json!(effective_max); + astra_core::model_wire::apply_chat_output_token_limit( + &mut body, + provider, + effective_max, + ); } if let Some(temp) = temperature { body["temperature"] = json!(temp); @@ -4796,6 +4802,21 @@ async fn call_llm_and_collect_with_total_budget( // `ControlledProviderAttemptObserver` above. Keep deadline ownership in // that single layer so a durable-admission stall is classified as an // inference-ledger failure instead of racing an outer provider timer. + let compatible_client = if provider == astra_services::byok_endpoint::COMPATIBLE_PROVIDER { + Some( + astra_services::byok_endpoint::endpoint_client(&url) + .await + .map_err(|error| { + astra_core::ClassifiedError::new( + astra_core::ErrorKind::InvalidRequest, + error, + ) + })?, + ) + } else { + None + }; + let client = compatible_client.as_deref().unwrap_or(client); let observed_attempt = match attempt_observer { Some(observer) => Some(observer.begin_attempt(prepared_request.identity()).await?), None => None, @@ -6949,6 +6970,18 @@ async fn call_llm_nonstream_with_attempt_observer_and_tool_choice( ); let _registered_endpoint_permit = acquire_registered_endpoint_permit_for_override(&url, completions_url_override)?; + let compatible_client = if provider == astra_services::byok_endpoint::COMPATIBLE_PROVIDER { + Some( + astra_services::byok_endpoint::endpoint_client(&url) + .await + .map_err(|error| { + astra_core::ClassifiedError::new(astra_core::ErrorKind::InvalidRequest, error) + })?, + ) + } else { + None + }; + let client = compatible_client.as_deref().unwrap_or(client); let observed_attempt = match attempt_observer { Some(observer) => Some(observer.begin_attempt(prepared_request.identity()).await?), None => None, @@ -7526,6 +7559,82 @@ pub(crate) fn parse_openai_sse_json_stream( #[cfg(test)] mod tests { use super::*; + + #[test] + fn cloud_byok_compatible_reuses_openai_message_and_tool_wire_format() { + let messages = vec![json!({"role":"user","content":"Use the calculator"})]; + let tools = vec![json!({"type":"function", "function":{ + "name":"calculator", "description":"Add numbers", + "parameters":{"type":"object", "properties":{"a":{"type":"number"}}} + }})]; + for streaming in [false, true] { + let body = build_provider_request_body( + &messages, + &tools, + "upstream-model", + "openai-compatible", + Some(128), + None, + streaming, + &ThinkingConfig::Off, + ); + let expected = build_provider_request_body( + &messages, + &tools, + "upstream-model", + "openai", + Some(128), + None, + streaming, + &ThinkingConfig::Off, + ); + assert_eq!(body, expected); + assert_eq!(body["model"], "upstream-model"); + assert_eq!(body["tools"][0]["function"]["name"], "calculator"); + } + } + + #[tokio::test] + async fn cloud_byok_compatible_blocks_private_endpoint_in_both_transports() { + let messages = vec![json!({"role":"user","content":"hello"})]; + for streaming in [false, true] { + let call = LlmCall { + purpose: astra_turn_types::InferencePurpose::SubAgent, + messages: &messages, + tools: &[], + cache_capability: None, + route: LlmExecutionRoute { + model_name: "test-model", + wire_model_name: None, + api_key: "sk-private-secret", + base_url: "http://127.0.0.1:9/v1", + provider: "openai-compatible", + header_overrides: None, + request_body_overrides: None, + completions_url_override: None, + request_timeout: None, + }, + max_output_tokens: Some(128), + temperature: None, + has_fallback: false, + thinking: &ThinkingConfig::Off, + }; + let result = if streaming { + call_llm_and_collect(call, LlmCancel::None).await + } else { + call_llm_nonstream( + global_llm_client(), + call, + std::time::Duration::from_secs(10), + ) + .await + }; + let error = result.expect_err("private endpoint must fail before network I/O"); + assert_eq!(error.kind, astra_core::ErrorKind::InvalidRequest); + assert!(error.message.contains("HTTPS")); + assert!(!error.to_string().contains("sk-private-secret")); + } + } use axum::Router; use axum::body::Body; use axum::extract::State; @@ -16775,9 +16884,53 @@ mod tests { assert!(stop_slot["properties"].get("slots").is_none()); } - // --- Regression: max_completion_tokens bump respects user's ceiling --- #[test] - fn max_completion_tokens_honors_user_when_above_floor() { + fn build_provider_request_body_output_limits_follow_provider_contract() { + for (provider, model, limit, forbidden) in [ + ("openai", "o3", "max_completion_tokens", "max_tokens"), + ("openai", "gpt-4o", "max_completion_tokens", "max_tokens"), + ( + "openai-compatible", + "custom-model", + "max_completion_tokens", + "max_tokens", + ), + ( + "deepseek", + "deepseek-v4-flash", + "max_tokens", + "max_completion_tokens", + ), + ( + "anthropic", + "claude-sonnet-4-5", + "max_tokens", + "max_completion_tokens", + ), + ] { + for streaming in [false, true] { + let body = build_provider_request_body( + &[json!({"role": "user", "content": "hi"})], + &[], + model, + provider, + Some(4096), + None, + streaming, + &ThinkingConfig::Off, + ); + assert_eq!( + body[limit], 4096, + "{provider}, streaming={streaming}: {body}" + ); + assert!(body.get(forbidden).is_none(), "{provider}: {body}"); + } + } + } + + // --- Regression: output-limit bump respects user's ceiling --- + #[test] + fn deepseek_max_tokens_honors_user_when_above_floor() { use astra_turn_core::thinking_config::ThinkingConfig; // User sets 128K, thinking budget is 32K → floor = 40K → must keep 128K. let thinking = ThinkingConfig::Enabled { @@ -16794,14 +16947,15 @@ mod tests { &thinking, ); assert_eq!( - body["max_completion_tokens"].as_u64(), + body["max_tokens"].as_u64(), Some(128_000), "user ceiling above floor must not be bumped" ); + assert!(body.get("max_completion_tokens").is_none()); } #[test] - fn max_completion_tokens_bumps_when_user_below_floor() { + fn deepseek_max_tokens_bumps_when_user_below_floor() { use astra_turn_core::thinking_config::ThinkingConfig; // User sets 8K, thinking budget is 32K → floor = 32K + 8K = 40K → bump to 40K. let thinking = ThinkingConfig::Enabled { @@ -16818,14 +16972,15 @@ mod tests { &thinking, ); assert_eq!( - body["max_completion_tokens"].as_u64(), + body["max_tokens"].as_u64(), Some(40_192), "configured max below thinking_budget+headroom must be bumped to floor" ); + assert!(body.get("max_completion_tokens").is_none()); } #[test] - fn max_completion_tokens_unchanged_when_thinking_off() { + fn deepseek_max_tokens_unchanged_when_thinking_off() { use astra_turn_core::thinking_config::ThinkingConfig; let body = build_provider_request_body( &[json!({"role": "user", "content": "hi"})], @@ -16838,10 +16993,11 @@ mod tests { &ThinkingConfig::Off, ); assert_eq!( - body["max_completion_tokens"].as_u64(), + body["max_tokens"].as_u64(), Some(4_096), "thinking=off must never bump user's max" ); + assert!(body.get("max_completion_tokens").is_none()); } #[test] diff --git a/crates/runtime/src/turn/memory_prefetch.rs b/crates/runtime/src/turn/memory_prefetch.rs index 987d3285b3..37179b98dd 100644 --- a/crates/runtime/src/turn/memory_prefetch.rs +++ b/crates/runtime/src/turn/memory_prefetch.rs @@ -53,6 +53,16 @@ pub async fn prefetch_memories_with_client( session_id: &str, top_k: u32, ) -> MemoryPrefetchResult { + match client.admits_operation(false).await { + Ok(true) => {} + Ok(false) => return MemoryPrefetchResult::default(), + Err(_) => { + return MemoryPrefetchResult { + outcome: astra_turn_types::MemoryRetrievalOutcome::Unavailable, + ..Default::default() + }; + } + } if user_msg.trim().is_empty() { return MemoryPrefetchResult::default(); } @@ -169,6 +179,16 @@ pub async fn prefetch_session_start_memories_with_client( user_id: &str, session_id: &str, ) -> SessionStartPrefetchResult { + match client.admits_operation(false).await { + Ok(true) => {} + Ok(false) => return SessionStartPrefetchResult::default(), + Err(_) => { + return SessionStartPrefetchResult { + outcome: astra_turn_types::MemoryRetrievalOutcome::Unavailable, + ..Default::default() + }; + } + } let started = Instant::now(); // Two structured queries in parallel: diff --git a/crates/runtime/tests/memoria_auth_http.rs b/crates/runtime/tests/memoria_auth_http.rs new file mode 100644 index 0000000000..284e71d972 --- /dev/null +++ b/crates/runtime/tests/memoria_auth_http.rs @@ -0,0 +1,257 @@ +use astra_core::{JwtSettings, MatrixOneSettings, MemoriaSettings, SharedPool}; +use astra_runtime::{AppState, HealthChecker, ServiceInfo, build_app}; +use astra_services::{AuthService, DatabaseAuthService, FernetTokenEncryptor}; +use async_trait::async_trait; +use axum::{ + Json, Router, + body::{Body, to_bytes}, + http::{Request, StatusCode}, + routing::get, +}; +use serde_json::{Value, json}; +use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, +}; +use tower::ServiceExt; + +#[path = "../../services/tests/common/isolated_database.rs"] +mod isolated_database; + +struct Healthy; +#[async_trait] +impl HealthChecker for Healthy { + async fn database_healthy(&self) -> bool { + true + } +} + +async fn request( + app: Router, + method: &str, + path: &str, + token: Option<&str>, + body: Value, +) -> (StatusCode, Value) { + let mut builder = Request::builder() + .method(method) + .uri(path) + .header("content-type", "application/json"); + if let Some(token) = token { + builder = builder.header("authorization", format!("Bearer {token}")); + } + let response = app + .oneshot(builder.body(Body::from(body.to_string())).unwrap()) + .await + .unwrap(); + let status = response.status(); + let bytes = to_bytes(response.into_body(), 1_000_000).await.unwrap(); + ( + status, + if bytes.is_empty() { + Value::Null + } else { + serde_json::from_slice(&bytes).unwrap() + }, + ) +} + +#[tokio::test] +#[ignore = "requires isolated ASTRA_TEST_DATABASE and ASTRA_TEST_DB_IT=1"] +async fn public_memoria_auth_uses_one_provider_and_enforces_disconnect() { + assert_eq!(std::env::var("ASTRA_TEST_DB_IT").as_deref(), Ok("1")); + let db = MatrixOneSettings::from_env(); + isolated_database::require_isolated_database(&db.database); + astra_services::storage::ensure_core_schema(&db, "mysql") + .await + .unwrap(); + let pool = SharedPool::new(&db).await.unwrap(); + let calls = Arc::new(AtomicUsize::new(0)); + let owner = format!("http-{}", uuid::Uuid::new_v4()); + let read_calls = calls.clone(); + let app = Router::new() + .route("/auth/whoami", get(move |headers: axum::http::HeaderMap| { + let owner = owner.clone(); + async move { + let key = headers.get("authorization").and_then(|v| v.to_str().ok()).unwrap_or(""); + let scopes = if key == "Bearer readonly-key" { + vec!["identity:read", "memory:read"] + } else { vec!["identity:read"] }; + (if key == "Bearer invalid-key" { StatusCode::UNAUTHORIZED } else { StatusCode::OK }, + Json(json!({"user_id":owner, "key_id":key, "is_active":true, "is_master":false, + "scope":{"type":"personal","id":owner}, "api_version":"1", + "capabilities":["api_key_scopes","memory_filters_v1"], "granted_scopes":scopes}))) + } + })) + .route("/v1/profiles/me", get(move |headers: axum::http::HeaderMap| { + read_calls.fetch_add(1, Ordering::SeqCst); + assert_eq!(headers["authorization"], "Bearer readonly-key"); + async { Json(json!({"profile":"from-provider-a"})) } + })); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let server = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + let provider = MemoriaSettings { + base_url: base, + master_key: None, + issuer: None, + web_url: Some("http://localhost".into()), + legacy_issuer: None, + }; + let auth = Arc::new( + DatabaseAuthService::new( + db, + JwtSettings { + secret_key: "review-http-auth".into(), + algorithm: "HS256".into(), + access_token_expire_minutes: 30, + refresh_token_expire_days: 7, + }, + ) + .with_pool(pool.clone()) + .with_encryptor(FernetTokenEncryptor::new("review-http-encryption").unwrap()) + .with_memoria_settings(&provider) + .unwrap(), + ); + // An independent override must never reroute the scoped credential. + let app = build_app( + AppState::new(ServiceInfo::default(), Arc::new(Healthy)) + .with_shared_pool(pool) + .with_auth_service(auth.clone()) + .with_memoria_config("http://127.0.0.1:1", None), + ); + + let (status, methods) = request(app.clone(), "GET", "/auth/methods", None, json!({})).await; + assert_eq!(status, StatusCode::OK); + assert_eq!(methods["memoria"]["authorization_url"], "http://localhost"); + assert_eq!(methods["memoria"]["issuer"], provider.base_url); + assert_eq!( + request( + app.clone(), + "POST", + "/auth/memoria", + None, + json!({"connection_key":""}) + ) + .await + .0, + StatusCode::BAD_REQUEST + ); + assert_eq!( + request( + app.clone(), + "POST", + "/auth/memoria", + None, + json!({"connection_key":"invalid-key"}) + ) + .await + .0, + StatusCode::UNAUTHORIZED + ); + let (status, login) = request( + app.clone(), + "POST", + "/auth/memoria", + None, + json!({"connection_key":"identity-key"}), + ) + .await; + assert_eq!(status, StatusCode::OK, "{login}"); + assert_eq!(login["memory_access"], "none"); + assert!(!login.to_string().contains("identity-key")); + let token = login["access_token"].as_str().unwrap(); + assert_eq!( + request( + app.clone(), + "GET", + "/memory/profile", + Some(token), + json!({}) + ) + .await + .0, + StatusCode::FORBIDDEN + ); + assert_eq!(calls.load(Ordering::SeqCst), 0); + let (status, relink) = request( + app.clone(), + "POST", + "/auth/memoria", + None, + json!({"connection_key":"readonly-key"}), + ) + .await; + assert_eq!(status, StatusCode::OK); + assert_eq!(relink["user_id"], login["user_id"]); + let (status, profile) = request( + app.clone(), + "GET", + "/memory/profile", + Some(token), + json!({}), + ) + .await; + assert_eq!(status, StatusCode::OK, "{profile}"); + assert_eq!(profile["profile"], "from-provider-a"); + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!( + request( + app.clone(), + "POST", + "/memory/store", + Some(token), + json!({"content":"blocked","memory_type":"semantic"}) + ) + .await + .0, + StatusCode::FORBIDDEN + ); + assert_eq!( + request(app.clone(), "DELETE", "/auth/memoria", None, json!({})) + .await + .0, + StatusCode::UNAUTHORIZED + ); + assert_eq!( + request( + app.clone(), + "DELETE", + "/auth/memoria", + Some(token), + json!({}) + ) + .await + .0, + StatusCode::NO_CONTENT + ); + assert_eq!( + request( + app.clone(), + "POST", + "/auth/refresh", + None, + json!({"refresh_token":login["refresh_token"]}) + ) + .await + .0, + StatusCode::UNAUTHORIZED + ); + assert_eq!( + request(app, "GET", "/memory/profile", Some(token), json!({})) + .await + .0, + StatusCode::UNAUTHORIZED + ); + assert!( + auth.memoria_credentials() + .unwrap() + .resolve(login["user_id"].as_str().unwrap()) + .await + .unwrap() + .is_none() + ); + server.abort(); +} diff --git a/crates/runtime/tests/memoria_reauthentication_http.rs b/crates/runtime/tests/memoria_reauthentication_http.rs new file mode 100644 index 0000000000..f9eb165657 --- /dev/null +++ b/crates/runtime/tests/memoria_reauthentication_http.rs @@ -0,0 +1,401 @@ +//! Real Astra router/DB with a deterministic upstream fresh-auth attester. +//! The attester's actual email/one-time-consume contract is tested in the website. +use astra_core::{AppSettings, MatrixOneSettings, MemoriaSettings}; +use astra_runtime::{build_app, build_server_state}; +use astra_thin_client::device_proof::{DeviceProofPurpose, device_challenge_proof}; +use axum::{ + Json, Router, + body::{Body, to_bytes}, + http::{Request, StatusCode}, + routing::{get, post}, +}; +use serde_json::{Value, json}; +use std::{ + collections::HashMap, + sync::{ + Arc, Mutex, + atomic::{AtomicBool, Ordering}, + }, +}; +use tower::ServiceExt; +#[path = "../../services/tests/common/isolated_database.rs"] +mod isolated_database; + +async fn request( + app: Router, + method: &str, + path: &str, + token: &str, + body: Value, +) -> (StatusCode, Value) { + let res = app + .oneshot( + Request::builder() + .method(method) + .uri(path) + .header("content-type", "application/json") + .header("authorization", format!("Bearer {token}")) + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap(); + let status = res.status(); + if path == "/auth/reauthenticate" && status.is_success() { + assert_eq!(res.headers()["cache-control"], "no-store"); + } + let bytes = to_bytes(res.into_body(), 1_000_000).await.unwrap(); + ( + status, + serde_json::from_slice(&bytes) + .unwrap_or_else(|_| json!({"non_json":String::from_utf8_lossy(&bytes)})), + ) +} +fn evidence(proofs: &Mutex>, owner: &str, purpose: &str) -> String { + let proof = format!( + "msu_{}{}", + uuid::Uuid::new_v4().simple(), + uuid::Uuid::new_v4().simple() + ); + let now = chrono::Utc::now().timestamp(); + proofs.lock().unwrap().insert(proof.clone(),json!({"subject":owner,"key_id":"generation-1","purpose":purpose,"authenticated_at":now,"expires_at":now+120})); + proof +} +async fn reauth(app: &Router, token: &str, proof: &str, purpose: &str) -> (StatusCode, Value) { + request( + app.clone(), + "POST", + "/auth/reauthenticate", + token, + json!({"memoria_proof":proof,"purpose":purpose}), + ) + .await +} + +#[tokio::test] +#[ignore = "requires explicitly isolated ASTRA_TEST_DATABASE and ASTRA_TEST_DB_IT=1"] +async fn memoria_step_up_authorizes_real_device_and_takeover_routes_without_passwords() { + let db = MatrixOneSettings::from_env(); + isolated_database::require_isolated_database(&db.database); + let owner = uuid::Uuid::new_v4().to_string(); + let identity_owner = owner.clone(); + let active = Arc::new(AtomicBool::new(true)); + let identity_active = active.clone(); + let proofs = Arc::new(Mutex::new(HashMap::::new())); + let issued = proofs.clone(); + let pause_attestation = Arc::new(AtomicBool::new(false)); + let attestation_entered = Arc::new(tokio::sync::Notify::new()); + let attestation_release = Arc::new(tokio::sync::Notify::new()); + let pause = pause_attestation.clone(); + let entered = attestation_entered.clone(); + let release = attestation_release.clone(); + let upstream = Router::new() + .route("/auth/whoami",get(move |headers: axum::http::HeaderMap| { + let owner = identity_owner.clone(); let active = identity_active.load(Ordering::SeqCst); + async move { + let valid = active && headers.get("authorization").is_some_and(|v| v == "Bearer connection-key"); + (if valid {StatusCode::OK} else {StatusCode::UNAUTHORIZED},Json(json!({"user_id":owner,"key_id":"generation-1","is_active":true,"is_master":false,"scope":{"type":"personal","id":owner},"api_version":"1","capabilities":["api_key_scopes","memory_filters_v1"],"granted_scopes":["identity:read"]}))) + } + })) + .route("/api/auth/astra/reauthentication/consume",post(move |Json(body):Json| { + let result = issued.lock().unwrap().remove(body["proof"].as_str().unwrap_or("")); + let paused = pause.swap(false, Ordering::SeqCst); + let entered = entered.clone(); + let release = release.clone(); + async move { + if paused { + entered.notify_one(); + release.notified().await; + } + match result { Some(value) => (StatusCode::OK,Json(value)), None => (StatusCode::UNAUTHORIZED,Json(json!({}))) } + } + })); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let upstream_task = tokio::spawn(async move { + axum::serve(listener, upstream).await.unwrap(); + }); + let mut settings = AppSettings::from_map(&HashMap::from([ + ("MATRIXONE_PASSWORD".into(), db.password.clone()), + ( + "ASTRA_JWT_SECRET".into(), + "stepup-test-only-jwt-32-bytes-secret".into(), + ), + ( + "ASTRA_RUNTIME_ROOT_SECRET".into(), + "stepup-test-only-root-secret".into(), + ), + ( + "ASTRA_TOKEN_ENCRYPTION_KEY".into(), + "stepup-test-only-token-key".into(), + ), + ("ASTRA_AUTO_CREATE_DATABASE".into(), "1".into()), + ])) + .unwrap(); + settings.matrixone = db; + settings.memoria = MemoriaSettings { + base_url: base.clone(), + web_url: Some(base.clone()), + issuer: Some("https://stepup-issuer.test".into()), + master_key: None, + legacy_issuer: None, + }; + let state = build_server_state(settings).await.unwrap(); + let pool = state.shared_pool.as_ref().unwrap().clone(); + let app = build_app(state); + let (status, login) = request( + app.clone(), + "POST", + "/auth/memoria", + "", + json!({"connection_key":"connection-key"}), + ) + .await; + assert_eq!(status, StatusCode::OK, "{login}"); + let token = login["access_token"].as_str().unwrap(); + let user = login["user_id"].as_str().unwrap(); + let (status, options) = + request(app.clone(), "GET", "/auth/reauthenticate", token, json!({})).await; + assert_eq!(status, StatusCode::OK); + assert_eq!( + options["verification_url"], + format!("{base}/astra/reauthenticate") + ); + for old in ["connection-key", token, ""] { + assert_eq!( + reauth(&app, token, old, "device_trust").await.0, + StatusCode::UNAUTHORIZED + ); + } + for (subject, purpose) in [ + ("another-owner", "device_trust"), + (owner.as_str(), "device_reenroll"), + ] { + let proof = evidence(&proofs, subject, purpose); + assert_eq!( + reauth(&app, token, &proof, "device_trust").await.0, + StatusCode::UNAUTHORIZED + ); + } + let expired = evidence(&proofs, &owner, "device_trust"); + proofs.lock().unwrap().get_mut(&expired).unwrap()["authenticated_at"] = + json!(chrono::Utc::now().timestamp() - 180); + assert_eq!( + reauth(&app, token, &expired, "device_trust").await.0, + StatusCode::UNAUTHORIZED + ); + let session = uuid::Uuid::new_v4().to_string(); + sqlx::query("INSERT INTO agent_sessions (session_id,user_id) VALUES (?,?)") + .bind(&session) + .bind(user) + .execute(pool.get()) + .await + .unwrap(); + let path = format!("/sessions/{session}/device"); + let (status, enrolled) = request( + app.clone(), + "POST", + &format!("{path}/enroll"), + token, + json!({"device_id":"laptop","device_fingerprint":"fp-1"}), + ) + .await; + assert_eq!(status, StatusCode::CREATED, "{enrolled}"); + let (status, challenge) = request( + app.clone(), + "POST", + &format!("{path}/challenge"), + token, + json!({"device_id":"laptop","device_fingerprint":"fp-1","purpose":"trust"}), + ) + .await; + assert_eq!(status, StatusCode::OK, "{challenge}"); + let device_proof = device_challenge_proof( + enrolled["device_key"].as_str().unwrap(), + DeviceProofPurpose::Trust, + user, + &session, + "laptop", + "fp-1", + challenge["challenge_id"].as_str().unwrap(), + challenge["challenge"].as_str().unwrap(), + ); + let fresh = evidence(&proofs, &owner, "device_trust"); + let (status, trust) = reauth(&app, token, &fresh, "device_trust").await; + assert_eq!(status, StatusCode::OK, "{trust}"); + assert_eq!( + reauth(&app, token, &fresh, "device_trust").await.0, + StatusCode::UNAUTHORIZED, + "upstream proof replay" + ); + let (status,trusted) = request(app.clone(),"POST",&format!("{path}/trust"),token,json!({"device_id":"laptop","device_fingerprint":"fp-1","challenge_id":challenge["challenge_id"],"device_proof":device_proof,"reauthentication_proof":trust["proof"]})).await; + assert_eq!(status, StatusCode::OK, "{trusted}"); + let fresh = evidence(&proofs, &owner, "device_reenroll"); + let (status, reenroll) = reauth(&app, token, &fresh, "device_reenroll").await; + assert_eq!(status, StatusCode::OK, "{reenroll}"); + let body = json!({"device_id":"laptop","device_fingerprint":"fp-2","reauthentication_proof":reenroll["proof"]}); + let (status, result) = request( + app.clone(), + "POST", + &format!("{path}/enroll"), + token, + body.clone(), + ) + .await; + assert!(status.is_success(), "{status}: {result}"); + assert_eq!( + request(app.clone(), "POST", &format!("{path}/enroll"), token, body) + .await + .0, + StatusCode::FORBIDDEN, + "Astra proof replay" + ); + + let (status, attachment) = request( + app.clone(), + "POST", + &format!("/sessions/{session}/attachments"), + token, + json!({"idempotency_key":uuid::Uuid::new_v4().to_string(),"placement":"server"}), + ) + .await; + assert_eq!(status, StatusCode::OK, "{attachment}"); + let fresh = evidence(&proofs, &owner, "session_forced_takeover"); + let (status, takeover) = reauth(&app, token, &fresh, "session_forced_takeover").await; + assert_eq!(status, StatusCode::OK, "{takeover}"); + let attachment_id = attachment["attachment"]["attachment_id"] + .as_str() + .expect("attachment response"); + let mut handoff = json!({"idempotency_key":uuid::Uuid::new_v4().to_string(),"mode":"forced","to_attachment_id":attachment_id,"from_placement":"cli","reason":"verified recovery","reauthentication_proof":takeover["proof"]}); + let (status, result) = request( + app.clone(), + "POST", + &format!("/sessions/{session}/handoffs"), + token, + handoff.clone(), + ) + .await; + assert_eq!(status, StatusCode::OK, "{result}"); + handoff["idempotency_key"] = json!(uuid::Uuid::new_v4().to_string()); + assert_eq!( + request( + app.clone(), + "POST", + &format!("/sessions/{session}/handoffs"), + token, + handoff + ) + .await + .0, + StatusCode::FORBIDDEN + ); + let fresh = evidence(&proofs, &owner, "device_reenroll"); + let (_, pending) = reauth(&app, token, &fresh, "device_reenroll").await; + active.store(false, Ordering::SeqCst); + assert_eq!( + reauth( + &app, + token, + &evidence(&proofs, &owner, "device_trust"), + "device_trust" + ) + .await + .0, + StatusCode::UNAUTHORIZED + ); + assert_eq!(request(app.clone(),"POST",&format!("{path}/enroll"),token,json!({"device_id":"laptop","device_fingerprint":"fp-3","reauthentication_proof":pending["proof"]})).await.0,StatusCode::UNAUTHORIZED,"revocation after proof issuance"); + active.store(true, Ordering::SeqCst); + + // A routine login must not invalidate an uninterrupted connection's proof. + let (status, again) = request( + app.clone(), + "POST", + "/auth/memoria", + "", + json!({"connection_key":"connection-key"}), + ) + .await; + assert_eq!(status, StatusCode::OK, "{again}"); + assert_eq!(again["user_id"], user); + let (status, result) = request(app.clone(),"POST",&format!("{path}/enroll"),token,json!({"device_id":"laptop","device_fingerprint":"fp-3","reauthentication_proof":pending["proof"]})).await; + assert!(status.is_success(), "routine login: {status}: {result}"); + + let (status, pending) = reauth( + &app, + token, + &evidence(&proofs, &owner, "device_reenroll"), + "device_reenroll", + ) + .await; + assert_eq!(status, StatusCode::OK); + let (status, pending_takeover) = reauth( + &app, + token, + &evidence(&proofs, &owner, "session_forced_takeover"), + "session_forced_takeover", + ) + .await; + assert_eq!(status, StatusCode::OK); + + // Pause the upstream response after the first binding read. Reconnecting + // the exact same upstream key must still change Astra's local lifecycle. + pause_attestation.store(true, Ordering::SeqCst); + let racing_app = app.clone(); + let racing_token = token.to_owned(); + let racing_proof = evidence(&proofs, &owner, "device_reenroll"); + let racing = tokio::spawn(async move { + reauth(&racing_app, &racing_token, &racing_proof, "device_reenroll").await + }); + tokio::time::timeout( + std::time::Duration::from_secs(10), + attestation_entered.notified(), + ) + .await + .expect("attester reached"); + assert!( + request(app.clone(), "DELETE", "/auth/memoria", token, json!({})) + .await + .0 + .is_success() + ); + let (status, relogin) = request( + app.clone(), + "POST", + "/auth/memoria", + "", + json!({"connection_key":"connection-key"}), + ) + .await; + assert_eq!(status, StatusCode::OK, "{relogin}"); + assert_eq!(relogin["user_id"], user); + let token = relogin["access_token"].as_str().unwrap(); + attestation_release.notify_one(); + let (status, result) = racing.await.unwrap(); + assert_eq!( + status, + StatusCode::UNAUTHORIZED, + "in-flight proof crossed disconnect: {result}" + ); + let (status, result) = request(app.clone(),"POST",&format!("{path}/enroll"),token,json!({"device_id":"laptop","device_fingerprint":"fp-4","reauthentication_proof":pending["proof"]})).await; + assert_eq!(status, StatusCode::FORBIDDEN, "old proof revived: {result}"); + let (status, result) = request(app.clone(),"POST",&format!("/sessions/{session}/handoffs"),token,json!({"idempotency_key":uuid::Uuid::new_v4().to_string(),"mode":"forced","to_attachment_id":attachment_id,"from_placement":"cli","reason":"recovery","reauthentication_proof":pending_takeover["proof"]})).await; + assert_eq!( + status, + StatusCode::FORBIDDEN, + "old takeover proof revived: {result}" + ); + let (status, fresh) = reauth( + &app, + token, + &evidence(&proofs, &owner, "device_reenroll"), + "device_reenroll", + ) + .await; + assert_eq!(status, StatusCode::OK, "{fresh}"); + let (status, result) = request(app.clone(),"POST",&format!("{path}/enroll"),token,json!({"device_id":"laptop","device_fingerprint":"fp-4","reauthentication_proof":fresh["proof"]})).await; + assert!( + status.is_success(), + "fresh proof after reconnect: {status}: {result}" + ); + upstream_task.abort(); +} diff --git a/crates/runtime/tests/memory_model_eligibility.rs b/crates/runtime/tests/memory_model_eligibility.rs new file mode 100644 index 0000000000..6cf3f83bde --- /dev/null +++ b/crates/runtime/tests/memory_model_eligibility.rs @@ -0,0 +1,194 @@ +//! Background extraction must obey the same owner/model boundary as chat. +use astra_core::{JwtSettings, MatrixOneSettings, SharedPool}; +use astra_runtime::{ + matrix_cloud_runtime::PoolMemoryInferenceResolver, + memory_hooks::MemoryInferenceRequest, + session_memory::{ + BackgroundActivityBroker, ExtractionRequest, MemoryExtractionService, + MemoryInferenceResolver, SpawnDecision, + }, + turn::cloud::memoria_compact::{MemoriaPort, UserScopedMemoriaPort}, +}; +use astra_services::{ + AuthService, DatabaseAuthService, FernetTokenEncryptor, event_ingestion::IngestionSender, +}; +use astra_turn_types::{InferenceInvocationScope, InferencePurpose}; +use axum::{ + Json, Router, + routing::{get, post}, +}; +use serde_json::json; +use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, +}; +use std::time::Duration; + +#[path = "../../services/tests/common/isolated_database.rs"] +mod isolated_database; + +#[tokio::test] +#[ignore = "requires ASTRA_TEST_DB_IT=1 and explicitly isolated ASTRA_TEST_DATABASE"] +async fn readwrite_memoria_extraction_cannot_spend_deployment_credentials() { + let settings = MatrixOneSettings::from_env(); + isolated_database::require_isolated_database(&settings.database); + astra_services::storage::ensure_core_schema(&settings, "mysql") + .await + .unwrap(); + let pool = SharedPool::new(&settings).await.unwrap(); + let encryptor = Arc::new(FernetTokenEncryptor::new("memory-owner-test-encryption").unwrap()); + let subject = uuid::Uuid::new_v4().to_string(); + let whoami = json!({"user_id":subject, "key_id":"fresh-rw-key", "is_active":true, "is_master":false, + "scope":{"type":"personal","id":subject}, "api_version":"1", + "capabilities":["api_key_scopes","memory_filters_v1"], "granted_scopes":["identity:read","memory:read","memory:write"]}); + let hits = Arc::new(AtomicUsize::new(0)); + let requests = hits.clone(); + let router = Router::new() + .route("/auth/whoami", get(move || { let whoami = whoami.clone(); async { Json(whoami) } })) + .route("/v1/chat/completions", post(move || { + requests.fetch_add(1, Ordering::SeqCst); + async { Json(json!({"choices":[{"index":0,"message":{"role":"assistant","content":"NO_CHANGE"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}})) } + })); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let server = tokio::spawn(async move { + axum::serve(listener, router).await.unwrap(); + }); + let auth = DatabaseAuthService::new( + settings.clone(), + JwtSettings { + secret_key: "memory-owner-test-jwt".into(), + algorithm: "HS256".into(), + access_token_expire_minutes: 15, + refresh_token_expire_days: 1, + }, + ) + .with_pool(pool.clone()) + .with_encryptor((*encryptor).clone()) + .with_memoria_base_url(base.clone()); + let login = auth.login_memoria("rw-test-key").await.unwrap(); + let user = login.tokens.user_id; + let memoria = Arc::new(UserScopedMemoriaPort::new( + auth.memoria_credentials().unwrap(), + user.clone(), + )); + assert!( + memoria.admits_operation(true).await.unwrap(), + "test must exercise read-write consent" + ); + let model_id = uuid::Uuid::new_v4().to_string(); + let model_name = format!("selector-{model_id}"); + sqlx::query("INSERT INTO infra_llm_models (model_id,model_name,provider,base_url,is_active,context_window,api_key_encrypted,input_modalities,output_modalities,supported_parameters,pricing,tags,quirks) VALUES (?, ?, 'openai', ?, 1, 128000, ?, '[\"text\"]', '[\"text\"]', '[]', '{}', '[\"selector\"]', '{}')") + .bind(&model_id).bind(&model_name).bind(format!("{base}/v1")).bind(encryptor.encrypt("deployment-only-key").unwrap()) + .execute(pool.get()).await.unwrap(); + let resolver = Arc::new(PoolMemoryInferenceResolver::new(pool.clone(), encryptor)); + assert!(resolver.resolve_candidates(&user).await.is_empty()); + let scope = |owner: &str| InferenceInvocationScope::Session { + session_id: owner.into(), + turn: 1, + round: 0, + operation_id: uuid::Uuid::new_v4().to_string(), + logical_attempt: 0, + }; + let messages = vec![ + json!({"role":"user","content":"Remember that this project requires regression tests before release."}), + json!({"role":"assistant","content":"I will add tests and verify the implementation."}), + ]; + let (ingestion, mut events) = IngestionSender::for_tests(64); + let service = Arc::new(MemoryExtractionService::new( + resolver.clone(), + memoria, + ingestion, + user.clone(), + Arc::new(BackgroundActivityBroker::new()), + )); + let memory_session = uuid::Uuid::new_v4().to_string(); + sqlx::query("INSERT INTO agent_sessions (session_id,user_id) VALUES (?, ?)") + .bind(&memory_session) + .bind(&user) + .execute(pool.get()) + .await + .unwrap(); + assert_eq!( + service.maybe_spawn(ExtractionRequest { + inference_scope: scope(&memory_session), + messages: messages.clone(), + session_facts: Default::default(), + had_error: false, + reanchors_current_objective: false, + }), + SpawnDecision::Spawned + ); + assert_eq!(service.wait_for_pending(Duration::from_secs(10)).await, 0); + assert!( + events.try_recv().is_ok(), + "extraction must emit its deterministic/degraded outcome" + ); + assert_eq!( + hits.load(Ordering::SeqCst), + 0, + "memory consent must not authorize deployment inference" + ); + + let control = format!("local-{}", uuid::Uuid::new_v4()); + sqlx::query("INSERT INTO auth_users (user_id, username, email, password_hash, is_active) VALUES (?, ?, ?, '', 1)") + .bind(&control).bind(&control).bind(format!("{control}@test.invalid")).execute(pool.get()).await.unwrap(); + let candidates = resolver.resolve_candidates(&control).await; + if std::env::var("ASTRA_DEPLOYMENT_MODE").as_deref() == Ok("cloud-byok") { + assert!(candidates.is_empty()); + } else { + let client = candidates + .iter() + .find(|c| c.model_name() == model_name) + .expect("permitted self-hosted selector"); + sqlx::query("INSERT INTO agent_sessions (session_id,user_id) VALUES (?, ?)") + .bind(&control) + .bind(&control) + .execute(pool.get()) + .await + .unwrap(); + let control_scope = scope(&control); + let request = MemoryInferenceRequest { + purpose: InferencePurpose::MemoryExtraction, + invocation_scope: &control_scope, + messages: &messages, + max_output_tokens: 128, + temperature: 0.0, + deadline: Duration::from_secs(5), + }; + assert_eq!(client.complete(request).await.unwrap(), "NO_CHANGE"); + assert_eq!(hits.load(Ordering::SeqCst), 1); + // An already constructed background client must revalidate before I/O. + sqlx::query("INSERT INTO auth_external_identities (provider_id,external_subject,astra_user_id) VALUES ('memoria:test-background-owner', ?, ?)") + .bind(&control).bind(&control).execute(pool.get()).await.unwrap(); + assert!(client.complete(request).await.is_err()); + assert_eq!( + hits.load(Ordering::SeqCst), + 1, + "owner eligibility changed but cached client sent another request" + ); + sqlx::query("DELETE FROM auth_external_identities WHERE astra_user_id=?") + .bind(&control) + .execute(pool.get()) + .await + .unwrap(); + sqlx::query("UPDATE infra_llm_models SET is_active=0 WHERE model_id=?") + .bind(&model_id) + .execute(pool.get()) + .await + .unwrap(); + assert!(client.complete(request).await.is_err()); + assert_eq!( + hits.load(Ordering::SeqCst), + 1, + "revoked model received another request" + ); + } + sqlx::query("DELETE FROM infra_llm_models WHERE model_id=?") + .bind(&model_id) + .execute(pool.get()) + .await + .unwrap(); + auth.disconnect_memoria(&user).await.unwrap(); + server.abort(); +} diff --git a/crates/runtime/tests/user_models_http.rs b/crates/runtime/tests/user_models_http.rs new file mode 100644 index 0000000000..496f51b6ae --- /dev/null +++ b/crates/runtime/tests/user_models_http.rs @@ -0,0 +1,586 @@ +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, +}; + +use astra_core::{ErrorResponse, error_response}; +use astra_runtime::{AppState, HealthChecker, ServiceInfo, build_app}; +use astra_services::{ + ModelCreateRequestData, ModelListItem, ModelRecord, ModelService, ModelUpdateRequestData, + ResolvedModelOffering, UserModelCreateRequestData, UserModelRecord, UserModelUpdateRequestData, + auth::{ + AuthLoginRequestData, AuthRefreshRequestData, AuthRegisterRequestData, AuthService, + AuthTokenRecord, AuthUserRecord, + }, +}; +use async_trait::async_trait; +use axum::{ + Json, + body::{self, Body}, + http::{HeaderMap, Request, StatusCode}, +}; +use tower::ServiceExt; + +type HttpError = (StatusCode, Json); + +#[derive(Clone, Default)] +struct TestModelService { + rows: Arc>>, + preflights: Arc>>, + preflight_network_failure: bool, + allows_deployment: bool, + catalog: Vec, +} + +fn unsupported() -> Result { + Err(error_response( + StatusCode::NOT_IMPLEMENTED, + "unsupported in test", + )) +} + +#[async_trait] +impl ModelService for TestModelService { + async fn allows_deployment_models(&self, _: String) -> Result { + Ok(self.allows_deployment) + } + + async fn validate_user_model_endpoint( + &self, + user_id: String, + base_url: String, + ) -> Result<(), HttpError> { + astra_services::byok_endpoint::parse_endpoint(&base_url) + .map_err(|error| error_response(StatusCode::BAD_REQUEST, error))?; + if self.preflight_network_failure { + return Err(astra_core::error_response_coded( + StatusCode::BAD_GATEWAY, + "Astra Server could not resolve the model endpoint", + "model_endpoint_network", + )); + } + self.preflights.lock().unwrap().push((user_id, base_url)); + Ok(()) + } + + async fn create_user_model( + &self, + user_id: String, + request: UserModelCreateRequestData, + ) -> Result { + let mut rows = self.rows.lock().expect("model rows lock"); + if rows + .iter() + .any(|((owner, _), row)| owner == &user_id && row.name == request.name) + { + return Err(error_response(StatusCode::CONFLICT, "duplicate user model")); + } + let model_id = format!("model-{}-{}", user_id, rows.len() + 1); + let record = UserModelRecord { + model_id: model_id.clone(), + name: request.name, + provider: request.provider, + model: request.model, + base_url: request + .base_url + .unwrap_or_else(|| "https://provider.example/v1".into()), + context_window: request.context_window, + is_default: request.is_default, + is_active: true, + credential_configured: !request.api_key.is_empty(), + created_at: "2026-09-06 00:00:00.000000".into(), + updated_at: "2026-09-06 00:00:00.000000".into(), + }; + rows.insert((user_id, model_id), record.clone()); + Ok(record) + } + + async fn list_user_models(&self, user_id: String) -> Result, HttpError> { + let rows = self.rows.lock().expect("model rows lock"); + Ok(rows + .iter() + .filter(|((owner, _), _)| owner == &user_id) + .map(|(_, row)| row.clone()) + .collect()) + } + + async fn get_user_model( + &self, + user_id: String, + model_id: String, + ) -> Result { + self.rows + .lock() + .expect("model rows lock") + .get(&(user_id, model_id)) + .cloned() + .ok_or_else(|| error_response(StatusCode::NOT_FOUND, "User model not found")) + } + + async fn update_user_model( + &self, + user_id: String, + model_id: String, + request: UserModelUpdateRequestData, + ) -> Result { + let mut rows = self.rows.lock().expect("model rows lock"); + let row = rows + .get_mut(&(user_id, model_id)) + .ok_or_else(|| error_response(StatusCode::NOT_FOUND, "User model not found"))?; + if let Some(context_window) = request.context_window { + row.context_window = context_window; + } + if let Some(is_default) = request.is_default { + row.is_default = is_default; + } + if let Some(is_active) = request.is_active { + row.is_active = is_active; + } + if let Some(api_key) = request.api_key { + row.credential_configured = !api_key.is_empty(); + } + Ok(row.clone()) + } + + async fn delete_user_model(&self, user_id: String, model_id: String) -> Result<(), HttpError> { + self.rows + .lock() + .expect("model rows lock") + .remove(&(user_id, model_id)) + .map(|_| ()) + .ok_or_else(|| error_response(StatusCode::NOT_FOUND, "User model not found")) + } + + async fn check_user_model( + &self, + user_id: String, + model_id: String, + ) -> Result { + self.get_user_model(user_id, model_id).await + } + + async fn create_model( + &self, + _: String, + _: ModelCreateRequestData, + ) -> Result { + unsupported() + } + async fn list_models(&self, _: String, _: bool) -> Result, HttpError> { + Ok(self.catalog.clone()) + } + async fn get_model(&self, _: String) -> Result { + unsupported() + } + async fn resolve_model_offering(&self, _: String) -> Result { + unsupported() + } + async fn update_model( + &self, + _: String, + _: ModelUpdateRequestData, + ) -> Result { + unsupported() + } + async fn delete_model(&self, _: String) -> Result<(), HttpError> { + unsupported() + } + async fn check_model(&self, _: String) -> Result { + unsupported() + } +} + +#[tokio::test] +async fn model_access_matches_run_eligibility_across_catalog_pages() { + use astra_services::{ModelAccessKind, ModelExecutionPlacement}; + let offering = |name: &str, kind: ModelAccessKind| ModelListItem { + offering_id: name.into(), + access_id: if kind == ModelAccessKind::CloudByok { + "cloud-byok" + } else { + "self-hosted" + } + .into(), + access_kind: kind, + access_label: if kind == ModelAccessKind::CloudByok { + "Cloud BYOK" + } else { + "Self-hosted" + } + .into(), + execution_placement: ModelExecutionPlacement::Server, + name: name.into(), + provider: "mock".into(), + description: None, + is_active: true, + context_window: 128000, + max_completion_tokens: None, + architecture: None, + thinking_capability: None, + }; + for (allows_deployment, catalog, expected) in [ + ( + true, + vec![offering("a-deployment", ModelAccessKind::SelfHosted)], + vec!["self-hosted"], + ), + (false, vec![], vec!["cloud-byok"]), + ( + true, + vec![ + offering("a-deployment", ModelAccessKind::SelfHosted), + offering("z-personal", ModelAccessKind::CloudByok), + ], + vec!["cloud-byok", "self-hosted"], + ), + ] { + let service = TestModelService { + allows_deployment, + catalog, + ..Default::default() + }; + let (status, body) = request( + app(service), + "GET", + "/model-access?limit=1", + Some("ordinary-user"), + None, + ) + .await; + assert_eq!(status, StatusCode::OK, "{body}"); + let body: serde_json::Value = serde_json::from_str(&body).unwrap(); + let mut ids = body["accesses"] + .as_array() + .unwrap() + .iter() + .map(|x| x["id"].as_str().unwrap()) + .collect::>(); + ids.sort_unstable(); + assert_eq!(ids, expected); + } +} + +struct HeaderAuthService; + +#[async_trait] +impl AuthService for HeaderAuthService { + async fn register(&self, _: AuthRegisterRequestData) -> Result { + unsupported() + } + async fn login(&self, _: AuthLoginRequestData) -> Result { + unsupported() + } + async fn refresh(&self, _: AuthRefreshRequestData) -> Result { + unsupported() + } + async fn logout(&self, _: AuthRefreshRequestData) -> Result<(), HttpError> { + unsupported() + } + + async fn current_user(&self, headers: &HeaderMap) -> Result { + let token = headers + .get("authorization") + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.strip_prefix("Bearer ")) + .filter(|value| !value.is_empty()) + .ok_or_else(|| error_response(StatusCode::UNAUTHORIZED, "Missing bearer token"))?; + Ok(AuthUserRecord { + user_id: token.to_string(), + username: token.to_string(), + email: format!("{token}@example.com"), + display_name: None, + }) + } +} + +struct Healthy; + +#[async_trait] +impl HealthChecker for Healthy { + async fn database_healthy(&self) -> bool { + true + } +} + +fn app(service: TestModelService) -> axum::Router { + build_app( + AppState::new(ServiceInfo::default(), Arc::new(Healthy)) + .with_auth_service(Arc::new(HeaderAuthService)) + .with_model_service(Arc::new(service)), + ) +} + +async fn request( + app: axum::Router, + method: &str, + path: &str, + user: Option<&str>, + json: Option, +) -> (StatusCode, String) { + let mut builder = Request::builder().method(method).uri(path); + if let Some(user) = user { + builder = builder.header("authorization", format!("Bearer {user}")); + } + let body = if let Some(json) = json { + builder = builder.header("content-type", "application/json"); + Body::from(json.to_string()) + } else { + Body::empty() + }; + let response = app + .oneshot(builder.body(body).expect("request")) + .await + .expect("response"); + let status = response.status(); + let bytes = body::to_bytes(response.into_body(), usize::MAX) + .await + .expect("body"); + ( + status, + String::from_utf8(bytes.to_vec()).expect("utf8 body"), + ) +} + +#[tokio::test] +async fn user_model_crud_is_authenticated_owner_scoped_and_secret_negative() { + let service = TestModelService::default(); + let (status, _) = request(app(service.clone()), "GET", "/me/models", None, None).await; + assert_eq!(status, StatusCode::UNAUTHORIZED); + + let secret = "sk-private-user-a"; + let create = serde_json::json!({ + "name": "deepseek", + "provider": "deepseek", + "model": "deepseek-chat", + "api_key": secret, + "is_default": true + }); + let (status, body) = request( + app(service.clone()), + "POST", + "/me/models", + Some("user-a"), + Some(create.clone()), + ) + .await; + assert_eq!(status, StatusCode::CREATED, "{body}"); + assert!(!body.contains(secret)); + let created: serde_json::Value = serde_json::from_str(&body).expect("create json"); + let model_id = created["model_id"].as_str().expect("model id"); + assert_eq!(created["credential_configured"], true); + + let (status, body) = request( + app(service.clone()), + "GET", + "/me/models", + Some("user-a"), + None, + ) + .await; + assert_eq!(status, StatusCode::OK); + assert!(!body.contains(secret)); + assert_eq!( + serde_json::from_str::(&body).unwrap()["items"] + .as_array() + .unwrap() + .len(), + 1 + ); + + let path = format!("/me/models/{model_id}"); + let (status, _) = request(app(service.clone()), "GET", &path, Some("user-b"), None).await; + assert_eq!(status, StatusCode::NOT_FOUND); + let (status, _) = request( + app(service.clone()), + "PUT", + &path, + Some("user-b"), + Some(serde_json::json!({"is_active": false})), + ) + .await; + assert_eq!(status, StatusCode::NOT_FOUND); + let (status, _) = request(app(service.clone()), "DELETE", &path, Some("user-b"), None).await; + assert_eq!(status, StatusCode::NOT_FOUND); + + let (status, body) = request( + app(service.clone()), + "POST", + "/me/models", + Some("user-b"), + Some(create), + ) + .await; + assert_eq!( + status, + StatusCode::CREATED, + "same alias must be allowed across users: {body}" + ); + + let check_path = format!("{path}/check"); + let (status, body) = request( + app(service.clone()), + "POST", + &check_path, + Some("user-a"), + None, + ) + .await; + assert_eq!(status, StatusCode::OK, "{body}"); + assert!(!body.contains(secret)); + + let (status, body) = request(app(service), "DELETE", &path, Some("user-a"), None).await; + assert_eq!(status, StatusCode::NO_CONTENT, "{body}"); +} + +#[tokio::test] +async fn user_model_payload_rejects_unknown_fields_without_echoing_secret() { + let secret = "sk-unknown-field-secret"; + let (status, body) = request( + app(TestModelService::default()), + "POST", + "/me/models", + Some("user-a"), + Some(serde_json::json!({ + "name": "deepseek", + "provider": "deepseek", + "model": "deepseek-chat", + "api_key": secret, + "owner": "user-b" + })), + ) + .await; + assert_eq!(status, StatusCode::UNPROCESSABLE_ENTITY); + assert!(!body.contains(secret)); +} + +#[tokio::test] +async fn compatible_model_configuration_reaches_owner_scoped_service() { + let (status, body) = request( + app(TestModelService::default()), + "POST", + "/me/models", + Some("user-a"), + Some(serde_json::json!({ + "name": "gateway", "provider": "openai-compatible", "model": "custom-model", + "base_url": "https://gateway.example/v1", "api_key": "sk-test-secret" + })), + ) + .await; + assert_eq!(status, StatusCode::CREATED, "{body}"); + let result: serde_json::Value = serde_json::from_str(&body).unwrap(); + assert_eq!(result["provider"], "openai-compatible"); + assert_eq!(result["base_url"], "https://gateway.example/v1"); + assert!(!body.contains("sk-test-secret")); +} + +#[tokio::test] +async fn endpoint_preflight_is_authenticated_credential_free_and_does_not_create_a_model() { + let service = TestModelService::default(); + let path = "/me/models/validate-endpoint"; + let endpoint = "https://api.moonshot.cn/v1"; + let payload = serde_json::json!({ "base_url": endpoint }); + let (status, _) = request( + app(service.clone()), + "POST", + path, + None, + Some(payload.clone()), + ) + .await; + assert_eq!(status, StatusCode::UNAUTHORIZED); + assert!(service.preflights.lock().unwrap().is_empty()); + let (status, body) = request( + app(service.clone()), + "POST", + path, + Some("user-a"), + Some(payload), + ) + .await; + assert_eq!(status, StatusCode::NO_CONTENT, "{body}"); + assert_eq!( + *service.preflights.lock().unwrap(), + vec![("user-a".into(), endpoint.into())] + ); + assert!(service.rows.lock().unwrap().is_empty()); + let (status, _) = request( + app(service.clone()), + "POST", + path, + Some("user-a"), + Some(serde_json::json!({ "base_url": "https://127.0.0.1/v1" })), + ) + .await; + assert_eq!(status, StatusCode::BAD_REQUEST); + for extra in ["api_key", "user_id"] { + let mut payload = serde_json::json!({ "base_url": endpoint }); + payload[extra] = "do-not-accept-or-echo".into(); + let (status, body) = request( + app(service.clone()), + "POST", + path, + Some("user-a"), + Some(payload), + ) + .await; + assert_eq!(status, StatusCode::UNPROCESSABLE_ENTITY); + assert!(!body.contains("do-not-accept-or-echo")); + } + assert_eq!(service.preflights.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn endpoint_preflight_preserves_server_network_error_code() { + let service = TestModelService { + preflight_network_failure: true, + ..Default::default() + }; + let (status, body) = request( + app(service.clone()), + "POST", + "/me/models/validate-endpoint", + Some("user-a"), + Some(serde_json::json!({ "base_url": "https://provider.example/v1" })), + ) + .await; + assert_eq!(status, StatusCode::BAD_GATEWAY); + let response: serde_json::Value = serde_json::from_str(&body).unwrap(); + assert_eq!(response["error_code"], "model_endpoint_network"); + assert!(service.rows.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn cloud_byok_completion_without_personal_default_does_not_fall_back() { + let service = TestModelService::default(); + let payload = serde_json::json!({ + "operation": "memory_extraction", + "session_id": "test-session", + "turn": 1, + "round": 1, + "logical_attempt": 1, + "messages": [{"role": "user", "content": "test"}] + }); + let (status, _) = request( + app(service.clone()), + "POST", + "/v1/chat/completions", + None, + Some(payload.clone()), + ) + .await; + assert_eq!(status, StatusCode::UNAUTHORIZED); + + // No database or provider is configured. Reaching either fallback would + // produce a different error instead of the required selection response. + let (status, body) = request( + app(service), + "POST", + "/v1/chat/completions", + Some("user-a"), + Some(payload), + ) + .await; + assert_eq!(status, StatusCode::BAD_REQUEST, "{body}"); + let response: serde_json::Value = serde_json::from_str(&body).unwrap(); + assert_eq!(response["error_code"], "missing_model_selection"); +} diff --git a/crates/services/Cargo.toml b/crates/services/Cargo.toml index 11537f7da3..f60913d576 100644 --- a/crates/services/Cargo.toml +++ b/crates/services/Cargo.toml @@ -8,6 +8,20 @@ publish = false [features] openssl-vendored = ["dep:openssl", "openssl/vendored"] e2e-hooks = [] +# External services are not provisioned by the ordinary MatrixOne CI lanes. +external-contract-tests = [] + +[[test]] +name = "memoria_live_contract_it" +required-features = ["external-contract-tests"] + +[[test]] +name = "byok_live_network" +required-features = ["external-contract-tests"] + +[[test]] +name = "user_model_probe_db_it" +required-features = ["external-contract-tests"] [dependencies] astra-config.workspace = true @@ -28,6 +42,11 @@ hmac.workspace = true jsonwebtoken.workspace = true regex.workspace = true reqwest.workspace = true +hickory-resolver = { version = "0.26.1", default-features = false, features = ["system-config", "tokio"] } +rustls.workspace = true +tokio-rustls = { version = "0.26", default-features = false, features = ["ring", "tls12"] } +webpki-roots = "1" +percent-encoding = "2" serde.workspace = true serde_json.workspace = true serde_yaml_ng.workspace = true @@ -48,6 +67,7 @@ libc.workspace = true openssl = { version = "0.10.79", optional = true } [dev-dependencies] +rcgen = "0.13.2" dotenvy.workspace = true astra-core = { workspace = true, features = ["dev-defaults"] } filetime = "0.2" diff --git a/crates/services/src/agent_bindings.rs b/crates/services/src/agent_bindings.rs index 8a089c8d90..87d4cac471 100644 --- a/crates/services/src/agent_bindings.rs +++ b/crates/services/src/agent_bindings.rs @@ -39,6 +39,13 @@ impl AgentBindingOwnerScope { crate::AuthPrincipalOrigin::Internal => { Self::for_internal_user(&principal.user.user_id) } + crate::AuthPrincipalOrigin::VerifiedProvider { + provider_id, + external_subject, + } => Self { + owner_user_id: principal.user.user_id.clone(), + principal_scope_id: format!("verified:{provider_id}:{external_subject}"), + }, crate::AuthPrincipalOrigin::ProviderAuthorizedRequest(context) => { let mut hasher = Sha256::new(); hasher.update(b"astra.agent-binding-principal.v1\0"); diff --git a/crates/services/src/auth/jwt.rs b/crates/services/src/auth/jwt.rs index a9c448896b..b9c3e4b150 100644 --- a/crates/services/src/auth/jwt.rs +++ b/crates/services/src/auth/jwt.rs @@ -44,6 +44,7 @@ pub(super) struct JwtClaims { pub(super) token_type: Option, pub(super) sid: Option, pub(super) origin: Option, + pub(super) iat: Option, } #[derive(Clone, Debug, Serialize)] diff --git a/crates/services/src/auth/memoria.rs b/crates/services/src/auth/memoria.rs new file mode 100644 index 0000000000..7f2a5427af --- /dev/null +++ b/crates/services/src/auth/memoria.rs @@ -0,0 +1,793 @@ +//! Application-scoped Memoria verification and credential lifecycle. +use super::{AuthHttpError, AuthTokenRecord, DatabaseAuthService, sha256_hex}; +use crate::FernetTokenEncryptor; +use astra_core::{MemoriaSettings, SharedPool, error_response, internal_error}; +use axum::http::StatusCode; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +pub const ACCESS_TTL_SECONDS: u32 = 900; + +#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum MemoryAccess { + None, + ReadOnly, + ReadWrite, +} +impl MemoryAccess { + pub fn allows(self, write: bool) -> bool { + self == Self::ReadWrite || (!write && self == Self::ReadOnly) + } + pub fn as_str(self) -> &'static str { + match self { + Self::None => "none", + Self::ReadOnly => "read_only", + Self::ReadWrite => "read_write", + } + } +} +pub fn memory_access_for_scopes(scopes: &[String]) -> Option { + let mut scopes: Vec<&str> = scopes.iter().map(String::as_str).collect(); + scopes.sort_unstable(); + scopes.dedup(); + match scopes.as_slice() { + ["identity:read"] => Some(MemoryAccess::None), + ["identity:read", "memory:read"] => Some(MemoryAccess::ReadOnly), + ["identity:read", "memory:read", "memory:write"] => Some(MemoryAccess::ReadWrite), + _ => None, + } +} + +#[derive(Clone, Debug)] +pub struct MemoriaProvider { + pub base_url: String, + pub issuer: String, + pub provider_id: String, + pub web_url: Option, + legacy_issuer: Option, +} +fn normalize_url(value: &str) -> Result { + let url = reqwest::Url::parse(value).map_err(|_| "Memoria URL must be absolute".to_string())?; + if !matches!(url.scheme(), "http" | "https") + || url.host_str().is_none() + || !url.username().is_empty() + || url.password().is_some() + || url.query().is_some() + || url.fragment().is_some() + { + return Err("Memoria URL must be HTTP(S), without credentials, query or fragment".into()); + } + Ok(url.to_string().trim_end_matches('/').to_string()) +} +impl MemoriaProvider { + pub fn new(settings: &MemoriaSettings) -> Result { + let base_url = normalize_url(&settings.base_url)?; + let issuer = normalize_url(settings.issuer.as_deref().unwrap_or(&base_url))?; + let web_url = settings.web_url.as_deref().map(normalize_url).transpose()?; + if let Some(web) = &web_url { + let url = reqwest::Url::parse(web).map_err(|e| e.to_string())?; + let loopback = url.host_str().is_some_and(|h| { + h == "localhost" + || h.trim_matches(['[', ']']) + .parse::() + .is_ok_and(|ip| ip.is_loopback()) + }); + if url.scheme() != "https" && !loopback { + return Err("MEMORIA_WEB_URL requires HTTPS except on loopback".into()); + } + } + Ok(Self { + provider_id: format!("memoria:{}", sha256_hex(&issuer)), + legacy_issuer: settings + .legacy_issuer + .as_deref() + .map(normalize_url) + .transpose()?, + base_url, + issuer, + web_url, + }) + } + async fn verify(&self, key: &str) -> Result { + if key.trim().is_empty() || key.len() > 4096 { + return Err(error_response( + StatusCode::BAD_REQUEST, + "Invalid Memoria connection key", + )); + } + let response = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .timeout(std::time::Duration::from_secs(10)) + .build() + .map_err(|_| unavailable())? + .get(format!("{}/auth/whoami", self.base_url)) + .bearer_auth(key) + .send() + .await + .map_err(|_| unavailable())?; + if matches!(response.status().as_u16(), 401 | 403) { + return Err(reconnect()); + } + if !response.status().is_success() { + return Err(unavailable()); + } + let value: Value = response.json().await.map_err(|_| unavailable())?; + let user_id = value["user_id"] + .as_str() + .filter(|s| !s.trim().is_empty() && s.len() <= 128) + .ok_or_else(reconnect)?; + let key_id = value["key_id"] + .as_str() + .filter(|s| !s.is_empty() && s.len() <= 128) + .ok_or_else(reconnect)?; + let scopes: Vec = + serde_json::from_value(value["granted_scopes"].clone()).map_err(|_| reconnect())?; + let access = memory_access_for_scopes(&scopes).ok_or_else(reconnect)?; + if value["is_active"] != true + || value["is_master"] != false + || value["scope"]["type"] != "personal" + || value["scope"]["id"] != user_id + || value["api_version"] != "1" + || !value["capabilities"].as_array().is_some_and(|c| { + c.iter().any(|v| v == "api_key_scopes") + && c.iter().any(|v| v == "memory_filters_v1") + }) + { + return Err(reconnect()); + } + Ok(VerifiedMemoriaIdentity { + memoria_user_id: user_id.to_string(), + key_id: key_id.to_string(), + memory_access: access, + granted_scopes: scopes, + issuer: self.issuer.clone(), + connection_generation: None, + }) + } +} +#[derive(Clone, Deserialize, Serialize)] +struct VerifiedMemoriaIdentity { + memoria_user_id: String, + key_id: String, + memory_access: MemoryAccess, + granted_scopes: Vec, + issuer: String, + /// Astra-owned lifecycle nonce, never supplied by the upstream verifier. + #[serde(default, skip_serializing_if = "Option::is_none")] + connection_generation: Option, +} +pub struct MemoriaLogin { + pub tokens: AuthTokenRecord, + pub memory_access: MemoryAccess, + pub granted_scopes: Vec, +} +/// Deliberately not Debug/Serialize: plaintext credentials must not be logged. +pub struct MemoriaCredential { + pub key: String, + pub owner: String, + pub generation: String, + pub access: MemoryAccess, + connection_generation: Option, +} + +#[derive(PartialEq, Eq)] +pub(super) struct ReauthenticationBinding { + provider_id: String, + owner: String, + generation: String, + connection_generation: Option, +} + +pub(super) fn reauthentication_proof_hash( + proof: &str, + binding: Option<&ReauthenticationBinding>, +) -> String { + match binding { + Some(binding) => sha256_hex( + &serde_json::json!([ + proof, + binding.provider_id, + binding.owner, + binding.generation, + binding.connection_generation + ]) + .to_string(), + ), + None => sha256_hex(proof), + } +} +#[derive(Clone)] +pub struct MemoriaCredentialResolver { + pub provider: MemoriaProvider, + pool: SharedPool, + encryptor: FernetTokenEncryptor, +} +impl MemoriaCredentialResolver { + pub fn new( + provider: MemoriaProvider, + pool: SharedPool, + encryptor: FernetTokenEncryptor, + ) -> Self { + Self { + provider, + pool, + encryptor, + } + } + fn token_id(&self, user: &str) -> String { + sha256_hex(&format!( + "memoria-binding\0{}\0{user}", + self.provider.provider_id + )) + } + pub async fn resolve(&self, user: &str) -> Result, String> { + let row: Option<(Option, Option)> = sqlx::query_as( + "SELECT encrypted_value, CAST(metadata AS CHAR) FROM auth_tokens WHERE token_id = ? AND type = 'memoria_connection' AND provider = 'memoria' AND scope_user_id = ? AND is_active = 1 AND EXISTS (SELECT 1 FROM auth_users WHERE user_id = auth_tokens.scope_user_id AND is_active = 1)") + .bind(self.token_id(user)).bind(user).fetch_optional(self.pool.get()).await + .map_err(|_| "Memoria credential lookup failed".to_string())?; + let Some((ciphertext, metadata)) = row else { + return Ok(None); + }; + let identity: VerifiedMemoriaIdentity = + serde_json::from_str(metadata.as_deref().unwrap_or("")) + .map_err(|_| "Invalid Memoria binding metadata".to_string())?; + if identity.issuer != self.provider.issuer { + return Err("Memoria binding issuer mismatch".into()); + } + let key = self + .encryptor + .decrypt(ciphertext.as_deref().ok_or("Missing Memoria credential")?) + .map_err(|_| "Memoria credential decryption failed".to_string())?; + Ok(Some(MemoriaCredential { + key, + owner: identity.memoria_user_id, + generation: identity.key_id, + access: identity.memory_access, + connection_generation: identity.connection_generation, + })) + } +} +impl DatabaseAuthService { + pub(super) async fn reauthentication_binding( + &self, + pool: &sqlx::MySqlPool, + user: &str, + ) -> Result, AuthHttpError> { + let Some(owner) = self.memoria_owner(pool, user).await? else { + return Ok(None); + }; + let resolver = self.credential_resolver().ok_or_else(reconnect)?; + let credential = resolver + .resolve(user) + .await + .map_err(|_| unavailable())? + .ok_or_else(reconnect)?; + let verified = resolver.provider.verify(&credential.key).await?; + if verified.memoria_user_id != owner + || credential.owner != owner + || verified.key_id != credential.generation + { + return Err(reconnect()); + } + Ok(Some(ReauthenticationBinding { + provider_id: resolver.provider.provider_id, + owner, + generation: credential.generation, + connection_generation: credential.connection_generation, + })) + } + + pub(super) async fn verify_memoria_step_up( + &self, + binding: &ReauthenticationBinding, + proof: &str, + purpose: super::ReauthenticationPurpose, + ) -> Result<(), AuthHttpError> { + if !proof.starts_with("msu_") + || proof.len() != 68 + || !proof[4..].bytes().all(|b| b.is_ascii_hexdigit()) + { + return Err(error_response( + StatusCode::UNAUTHORIZED, + "Fresh account verification is required", + )); + } + let provider = self.memoria_provider.as_ref().ok_or_else(reconnect)?; + let web = provider.web_url.as_ref().ok_or_else(|| { + error_response( + StatusCode::SERVICE_UNAVAILABLE, + "Account reauthentication website is not configured", + ) + })?; + // Only the composition-time trusted website may attest fresh authentication. + // A caller cannot supply a verifier URL or a different identity authority. + let mut response = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .timeout(std::time::Duration::from_secs(10)) + .build().map_err(|_| unavailable())? + .post(format!("{web}/api/auth/astra/reauthentication/consume")) + .json(&serde_json::json!({"proof":proof,"subject":binding.owner,"key_id":binding.generation,"purpose":purpose})) + .send().await.map_err(|_| unavailable())?; + if matches!(response.status().as_u16(), 400 | 401 | 403 | 409) { + return Err(error_response( + StatusCode::UNAUTHORIZED, + "Account verification is invalid, expired or already used", + )); + } + if !response.status().is_success() { + return Err(unavailable()); + } + let mut bytes = Vec::new(); + while let Some(chunk) = response.chunk().await.map_err(|_| unavailable())? { + if bytes.len() + chunk.len() > 4096 { + return Err(unavailable()); + } + bytes.extend_from_slice(&chunk); + } + let value: Value = serde_json::from_slice(&bytes).map_err(|_| unavailable())?; + let now = chrono::Utc::now().timestamp(); + let authenticated_at = value["authenticated_at"].as_i64().ok_or_else(reconnect)?; + let expires_at = value["expires_at"].as_i64().ok_or_else(reconnect)?; + if value["subject"] != binding.owner + || value["key_id"] != binding.generation + || value["purpose"] != purpose.as_str() + || authenticated_at > now + 5 + || now - authenticated_at > 120 + || expires_at <= now + || expires_at > authenticated_at + 120 + { + return Err(reconnect()); + } + Ok(()) + } + + pub fn with_memoria_settings(mut self, settings: &MemoriaSettings) -> Result { + self.memoria_provider = Some(MemoriaProvider::new(settings)?); + Ok(self) + } + pub fn with_memoria_base_url(self, base_url: String) -> Self { + self.with_memoria_settings(&MemoriaSettings { + base_url, + master_key: None, + issuer: None, + web_url: None, + legacy_issuer: None, + }) + .expect("valid Memoria URL") + } + pub(super) fn credential_resolver(&self) -> Option { + Some(MemoriaCredentialResolver::new( + self.memoria_provider.clone()?, + self.pool.clone()?, + self.encryptor.as_ref()?.clone(), + )) + } + pub(super) async fn memoria_owner( + &self, + pool: &sqlx::MySqlPool, + user: &str, + ) -> Result, AuthHttpError> { + let identity: Option<(String, String)> = sqlx::query_as( + "SELECT provider_id, external_subject FROM auth_external_identities WHERE astra_user_id = ? AND provider_id LIKE 'memoria:%' LIMIT 1") + .bind(user).fetch_optional(pool).await.map_err(internal_error)?; + if let Some((provider, subject)) = identity { + if self + .memoria_provider + .as_ref() + .is_none_or(|p| p.provider_id != provider) + { + return Err(reconnect()); + } + return Ok(Some(subject)); + } + let legacy: Option = sqlx::query_scalar( + "SELECT memoria_user_id FROM auth_memoria_identities WHERE astra_user_id = ? LIMIT 1", + ) + .bind(user) + .fetch_optional(pool) + .await + .map_err(internal_error)?; + if legacy.is_some() { + return Err(reconnect()); + } + Ok(None) + } + pub(super) async fn revalidate_memoria_connection( + &self, + _pool: &sqlx::MySqlPool, + user: &str, + owner: &str, + ) -> Result<(), AuthHttpError> { + let resolver = self.credential_resolver().ok_or_else(unavailable)?; + let credential = resolver + .resolve(user) + .await + .map_err(|_| unavailable())? + .ok_or_else(reconnect)?; + let identity = resolver.provider.verify(&credential.key).await?; + if identity.memoria_user_id != owner || identity.key_id != credential.generation { + return Err(reconnect()); + } + Ok(()) + } + pub(super) async fn memoria_login(&self, key: &str) -> Result { + let resolver = self.credential_resolver().ok_or_else(unavailable)?; + let identity = resolver.provider.verify(key).await?; + let ciphertext = resolver.encryptor.encrypt(key).map_err(internal_error)?; + let pool = self.get_pool().await.map_err(internal_error)?; + self.ensure_default_roles(&pool) + .await + .map_err(internal_error)?; + // Retry aborted identity insert / serialization transactions, never only + // the credential phase. No durable session exists on a failed attempt. + for attempt in 0..3 { + match self + .persist_memoria_login(&pool, &resolver, &identity, &ciphertext, key) + .await + { + Ok(tokens) => { + return Ok(MemoriaLogin { + tokens, + memory_access: identity.memory_access, + granted_scopes: identity.granted_scopes, + }); + } + Err(error) if attempt < 2 && error.0 == StatusCode::INTERNAL_SERVER_ERROR => { + continue; + } + Err(error) => return Err(error), + } + } + unreachable!() + } + async fn persist_memoria_login( + &self, + pool: &sqlx::MySqlPool, + resolver: &MemoriaCredentialResolver, + identity: &VerifiedMemoriaIdentity, + ciphertext: &str, + key: &str, + ) -> Result { + let mut tx = pool.begin().await.map_err(internal_error)?; + let legacy: Option = sqlx::query_scalar( + "SELECT astra_user_id FROM auth_memoria_identities WHERE memoria_user_id = ? LIMIT 1", + ) + .bind(&identity.memoria_user_id) + .fetch_optional(&mut *tx) + .await + .map_err(internal_error)?; + if legacy.is_some() + && resolver.provider.legacy_issuer.as_deref() != Some(&resolver.provider.issuer) + { + return Err(error_response( + StatusCode::CONFLICT, + "Legacy Memoria identity requires administrator-configured MEMORIA_LEGACY_ISSUER migration", + )); + } + let user = self + .resolve_verified_provider_identity( + &mut tx, + &resolver.provider.provider_id, + &identity.memoria_user_id, + legacy.as_deref(), + ) + .await?; + // Reject credentials revoked while waiting on a concurrent link/unlink. + let fresh = resolver.provider.verify(key).await?; + if fresh.memoria_user_id != identity.memoria_user_id + || fresh.key_id != identity.key_id + || fresh.memory_access != identity.memory_access + { + return Err(reconnect()); + } + // The canonical account lock serializes login with disconnect. Preserve + // the lifecycle only for an uninterrupted binding to the same key. + // Upstream key IDs can be reused after disconnect; this nonce cannot. + let previous: Option = sqlx::query_scalar("SELECT CAST(metadata AS CHAR) FROM auth_tokens WHERE token_id = ? AND type = 'memoria_connection' AND provider = 'memoria' AND scope_user_id = ? AND is_active = 1") + .bind(resolver.token_id(&user.user_id)).bind(&user.user_id) + .fetch_optional(&mut *tx).await.map_err(internal_error)?; + let previous = + previous.and_then(|value| serde_json::from_str::(&value).ok()); + let mut stored_identity = identity.clone(); + stored_identity.connection_generation = Some( + previous + .filter(|old| { + old.issuer == identity.issuer + && old.memoria_user_id == identity.memoria_user_id + && old.key_id == identity.key_id + }) + .and_then(|old| old.connection_generation) + .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()), + ); + sqlx::query("DELETE FROM auth_tokens WHERE type = 'memoria_connection' AND provider = 'memoria' AND scope_user_id = ?") + .bind(&user.user_id).execute(&mut *tx).await.map_err(internal_error)?; + sqlx::query("INSERT INTO auth_tokens (token_id,type,provider,encrypted_value,is_active,scope_user_id,metadata) VALUES (?, 'memoria_connection', 'memoria', ?, 1, ?, ?)") + .bind(resolver.token_id(&user.user_id)).bind(ciphertext).bind(&user.user_id) + .bind(serde_json::to_string(&stored_identity).map_err(internal_error)?).execute(&mut *tx).await.map_err(internal_error)?; + if legacy.is_some() { + sqlx::query("DELETE FROM auth_memoria_identities WHERE astra_user_id = ?") + .bind(&user.user_id) + .execute(&mut *tx) + .await + .map_err(internal_error)?; + sqlx::query("UPDATE auth_refresh_tokens SET is_revoked = 1 WHERE user_id = ?") + .bind(&user.user_id) + .execute(&mut *tx) + .await + .map_err(internal_error)?; + } + let session = uuid::Uuid::new_v4().to_string(); + let origin = format!("verified:{}", resolver.provider.provider_id); + let access_token = self + .create_access_token(&user.user_id, &user.username, &session, &origin) + .map_err(internal_error)?; + let refresh_token = self + .create_refresh_token(&user.user_id, &session, &origin) + .map_err(internal_error)?; + sqlx::query("INSERT INTO auth_refresh_tokens (token_id,user_id,session_id,token_hash,expires_at,is_revoked) VALUES (?, ?, ?, ?, ?, 0)") + .bind(uuid::Uuid::new_v4().to_string()).bind(&user.user_id).bind(&session).bind(sha256_hex(&refresh_token)) + .bind(self.refresh_token_expires_at_string(chrono::Utc::now())).execute(&mut *tx).await.map_err(internal_error)?; + tx.commit().await.map_err(internal_error)?; + Ok(AuthTokenRecord { + user_id: user.user_id, + access_token, + refresh_token, + token_type: "bearer".into(), + expires_in: self + .access_token_expires_in_seconds() + .min(ACCESS_TTL_SECONDS), + }) + } + pub(super) async fn memoria_disconnect(&self, user: &str) -> Result<(), AuthHttpError> { + let pool = self.get_pool().await.map_err(internal_error)?; + let mut tx = pool.begin().await.map_err(internal_error)?; + sqlx::query("SELECT user_id FROM auth_users WHERE user_id = ? FOR UPDATE") + .bind(user) + .fetch_one(&mut *tx) + .await + .map_err(internal_error)?; + sqlx::query("DELETE FROM auth_tokens WHERE type = 'memoria_connection' AND provider = 'memoria' AND scope_user_id = ?") + .bind(user).execute(&mut *tx).await.map_err(internal_error)?; + sqlx::query("UPDATE auth_refresh_tokens SET is_revoked = 1 WHERE user_id = ?") + .bind(user) + .execute(&mut *tx) + .await + .map_err(internal_error)?; + sqlx::query("DELETE FROM auth_reauthentication_proofs WHERE user_id = ?") + .bind(user) + .execute(&mut *tx) + .await + .map_err(internal_error)?; + tx.commit().await.map_err(internal_error)?; + Ok(()) + } +} +fn reconnect() -> AuthHttpError { + error_response( + StatusCode::UNAUTHORIZED, + "Memoria connection expired or revoked; sign in again", + ) +} +fn unavailable() -> AuthHttpError { + error_response( + StatusCode::SERVICE_UNAVAILABLE, + "Memoria identity verification is unavailable; retry later", + ) +} + +#[cfg(test)] +async fn verify_connection( + base_url: &str, + key: &str, + owner: &str, + key_id: &str, +) -> Result<(), AuthHttpError> { + let provider = MemoriaProvider::new(&MemoriaSettings { + base_url: base_url.into(), + master_key: None, + issuer: None, + web_url: None, + legacy_issuer: None, + }) + .unwrap(); + let identity = provider.verify(key).await?; + if identity.memoria_user_id != owner || identity.key_id != key_id { + return Err(reconnect()); + } + Ok(()) +} +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn step_up_proof_hash_binds_issuer_subject_and_generation_without_ambiguous_fields() { + let binding = ReauthenticationBinding { + provider_id: "memoria:one".into(), + owner: "owner".into(), + generation: "key".into(), + connection_generation: Some("lifecycle-1".into()), + }; + let expected = reauthentication_proof_hash("rp_test", Some(&binding)); + for other in [ + ReauthenticationBinding { + provider_id: "memoria:two".into(), + owner: "owner".into(), + generation: "key".into(), + connection_generation: Some("lifecycle-1".into()), + }, + ReauthenticationBinding { + provider_id: "memoria:one".into(), + owner: "other".into(), + generation: "key".into(), + connection_generation: Some("lifecycle-1".into()), + }, + ReauthenticationBinding { + provider_id: "memoria:one".into(), + owner: "owner".into(), + generation: "rotated".into(), + connection_generation: Some("lifecycle-1".into()), + }, + ReauthenticationBinding { + provider_id: "memoria:one".into(), + owner: "owner".into(), + generation: "key".into(), + connection_generation: Some("lifecycle-2".into()), + }, + ReauthenticationBinding { + provider_id: "memoria:one".into(), + owner: "owner".into(), + generation: "key".into(), + connection_generation: None, + }, + ] { + assert_ne!( + expected, + reauthentication_proof_hash("rp_test", Some(&other)) + ); + } + assert_ne!(expected, reauthentication_proof_hash("rp_test", None)); + let left = ReauthenticationBinding { + provider_id: "p".into(), + owner: "x\0y".into(), + generation: "z".into(), + connection_generation: None, + }; + let right = ReauthenticationBinding { + provider_id: "p".into(), + owner: "x".into(), + generation: "y\0z".into(), + connection_generation: None, + }; + assert_ne!( + reauthentication_proof_hash("rp_test", Some(&left)), + reauthentication_proof_hash("rp_test", Some(&right)) + ); + } + use axum::{Json, Router, routing::get}; + use std::sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }; + + #[tokio::test] + async fn refresh_identity_rechecks_revocation_owner_and_outages() { + let revoked = Arc::new(AtomicBool::new(false)); + let flag = revoked.clone(); + let app = Router::new().route("/auth/whoami", get(move || { + let flag = flag.clone(); + async move { + (if flag.load(Ordering::SeqCst) { StatusCode::UNAUTHORIZED } else { StatusCode::OK }, Json(serde_json::json!({ + "user_id":"owner", "key_id":"key", "is_active":true, "is_master":false, + "scope":{"type":"personal","id":"owner"}, "api_version":"1", + "capabilities":["api_key_scopes","memory_filters_v1"], "granted_scopes":["identity:read"] + }))) + } + })); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}", listener.local_addr().unwrap()); + let task = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + assert!( + verify_connection(&url, "secret", "owner", "key") + .await + .is_ok() + ); + assert_eq!( + verify_connection(&url, "secret", "other-owner", "key") + .await + .unwrap_err() + .0, + StatusCode::UNAUTHORIZED + ); + assert_eq!( + verify_connection(&url, "secret", "owner", "other-key") + .await + .unwrap_err() + .0, + StatusCode::UNAUTHORIZED + ); + revoked.store(true, Ordering::SeqCst); + assert_eq!( + verify_connection(&url, "secret", "owner", "key") + .await + .unwrap_err() + .0, + StatusCode::UNAUTHORIZED + ); + task.abort(); + let error = verify_connection(&url, "secret", "owner", "key") + .await + .unwrap_err(); + assert_eq!(error.0, StatusCode::SERVICE_UNAVAILABLE); + assert!(!error.1.detail.contains("secret")); + } +} + +#[cfg(test)] +mod provider_contract_tests { + use super::*; + fn settings(base: &str, web: Option<&str>) -> MemoriaSettings { + MemoriaSettings { + base_url: base.into(), + master_key: None, + issuer: None, + web_url: web.map(str::to_string), + legacy_issuer: None, + } + } + #[test] + fn provider_identity_is_canonical_and_namespaced() { + let a = MemoriaProvider::new(&settings("https://A.example:443/", None)).unwrap(); + let same = MemoriaProvider::new(&settings("https://a.example", None)).unwrap(); + let other = MemoriaProvider::new(&settings("https://b.example", None)).unwrap(); + assert_eq!(a.provider_id, same.provider_id); + assert_ne!(a.provider_id, other.provider_id); + assert!(a.web_url.is_none()); + } + #[test] + fn web_login_rejects_plaintext_remote_and_ambiguous_urls() { + for bad in [ + "http://cloud.example", + "https://user:pass@cloud.example", + "https://cloud.example/?x=1", + "file:///tmp/login", + ] { + assert!( + MemoriaProvider::new(&settings("http://memoria:8100", Some(bad))).is_err(), + "{bad}" + ); + } + for good in [ + "https://thememoria.ai", + "http://localhost", + "http://127.0.0.1:3000", + "http://[::1]:3000", + ] { + assert!( + MemoriaProvider::new(&settings("http://memoria:8100", Some(good))).is_ok(), + "{good}" + ); + } + } + #[test] + fn scoped_access_is_typed_and_fail_closed() { + assert!(!MemoryAccess::None.allows(false)); + assert!(!MemoryAccess::None.allows(true)); + assert!(MemoryAccess::ReadOnly.allows(false)); + assert!(!MemoryAccess::ReadOnly.allows(true)); + assert!(MemoryAccess::ReadWrite.allows(true)); + for scopes in [ + vec![], + vec!["identity:read", "keys:manage"], + vec!["identity:read", "memory:write"], + ] { + assert!( + memory_access_for_scopes( + &scopes.into_iter().map(str::to_string).collect::>() + ) + .is_none() + ); + } + } +} diff --git a/crates/services/src/auth/mod.rs b/crates/services/src/auth/mod.rs index 1ad1a07f68..b15c4571b5 100644 --- a/crates/services/src/auth/mod.rs +++ b/crates/services/src/auth/mod.rs @@ -43,9 +43,11 @@ mod admin; mod encryption; pub mod external; mod jwt; +pub mod memoria; pub mod provider_request; pub mod session; mod validation; +mod verified; pub use admin::{ DatabaseAdminAuditReader, DatabaseAdminAuthorizer, DatabaseAdminFeedbackStatsReader, @@ -399,6 +401,12 @@ pub enum EdgeTokenBinding { #[async_trait] pub trait AuthService: Send + Sync { + async fn reauthentication_options( + &self, + _user_id: &str, + ) -> Result { + Ok(serde_json::json!({"method":"password"})) + } async fn register( &self, request: AuthRegisterRequestData, @@ -409,6 +417,31 @@ pub trait AuthService: Send + Sync { request: AuthLoginRequestData, ) -> Result)>; + /// The authentication service owns verification, identity, session and binding. + async fn login_memoria( + &self, + _connection_key: &str, + ) -> Result)> { + Err(error_response( + StatusCode::NOT_IMPLEMENTED, + "Verified identity login is not configured", + )) + } + + fn memoria_credentials(&self) -> Option { + None + } + + async fn disconnect_memoria( + &self, + _user_id: &str, + ) -> Result<(), (StatusCode, Json)> { + Err(error_response( + StatusCode::NOT_IMPLEMENTED, + "Memoria connection is not configured", + )) + } + async fn refresh( &self, request: AuthRefreshRequestData, @@ -588,6 +621,7 @@ impl ReauthenticationPurpose { #[derive(Clone, PartialEq, Eq)] pub struct ReauthenticationRequestData { pub password: String, + pub memoria_proof: Option, pub purpose: ReauthenticationPurpose, } @@ -643,6 +677,10 @@ impl AuthPrincipal { #[derive(Clone, Debug, PartialEq, Eq)] pub enum AuthPrincipalOrigin { Internal, + VerifiedProvider { + provider_id: String, + external_subject: String, + }, ProviderAuthorizedRequest(AuthProviderAuthorizedRequestContext), } @@ -687,6 +725,7 @@ pub struct DatabaseAuthService { control_pool: Option, jwt: JwtSettings, encryptor: Option, + memoria_provider: Option, ext_providers: Vec, external_client: std::sync::Arc, provider_request_auth: Vec, @@ -762,6 +801,7 @@ impl DatabaseAuthService { pool: None, control_pool: None, encryptor: None, + memoria_provider: None, ext_providers: Vec::new(), external_client: HttpExternalProviderClient::shared(), provider_request_auth: Vec::new(), @@ -889,7 +929,14 @@ impl DatabaseAuthService { iat: 0, jti: String::new(), }, - ChronoDuration::minutes(i64::from(self.jwt.access_token_expire_minutes)), + ChronoDuration::seconds(i64::from( + if origin == "memoria" || origin.starts_with("verified:memoria:") { + self.access_token_expires_in_seconds() + .min(memoria::ACCESS_TTL_SECONDS) + } else { + self.access_token_expires_in_seconds() + }, + )), ) } @@ -1236,7 +1283,8 @@ impl DatabaseAuthService { .clone() .ok_or_else(|| error_response(StatusCode::UNAUTHORIZED, "Invalid token origin"))?; match origin.as_str() { - "internal" => Ok((session_id, origin, None)), + "internal" | "memoria" => Ok((session_id, origin, None)), + value if value.starts_with("verified:memoria:") => Ok((session_id, origin, None)), _ => Err(error_response( StatusCode::UNAUTHORIZED, "Invalid token origin", @@ -1285,69 +1333,19 @@ impl DatabaseAuthService { let now = Utc::now(); let external_expires_at = self.parse_provider_expires_at(&response.expires_at)?; let astra_session_id = Uuid::new_v4().to_string(); - let astra_user_id = Uuid::new_v4().to_string(); - let internal_username = format!("ext_{}", astra_user_id.replace('-', "")); - let internal_email = format!("{internal_username}@external.astra.invalid"); let mut tx = pool .begin() .await .map_err(|e| map_auth_sqlx(e, "external.begin_tx", Some(&pool)))?; - let existing_identity = query( - "SELECT astra_user_id FROM auth_external_identities \ - WHERE provider_id = ? AND external_subject = ? LIMIT 1", - ) - .bind(&provider.id) - .bind(&external_subject) - .fetch_optional(&mut *tx) - .await - .map_err(|e| map_auth_sqlx(e, "external.fetch_identity", Some(&pool)))?; - - let astra_user_id = if let Some(row) = existing_identity { - let existing_user_id: String = row.try_get("astra_user_id").unwrap_or_default(); - query( - "UPDATE auth_external_identities \ - SET username = ?, email = ?, display_name = ?, updated_at = NOW() \ - WHERE provider_id = ? AND external_subject = ?", - ) - .bind(&external_username) - .bind(&external_email) - .bind(&external_display_name) - .bind(&provider.id) - .bind(&external_subject) - .execute(&mut *tx) - .await + let user = self + .resolve_verified_provider_identity(&mut tx, &provider.id, &external_subject, None) + .await?; + let astra_user_id = user.user_id; + query("UPDATE auth_external_identities SET username = ?, email = ?, display_name = ?, updated_at = NOW() WHERE provider_id = ? AND external_subject = ?") + .bind(&external_username).bind(&external_email).bind(&external_display_name) + .bind(&provider.id).bind(&external_subject).execute(&mut *tx).await .map_err(|e| map_auth_sqlx(e, "external.update_identity", Some(&pool)))?; - existing_user_id - } else { - query( - "INSERT INTO auth_users \ - (user_id, username, email, password_hash, display_name, is_active) \ - VALUES (?, ?, ?, '', ?, 1)", - ) - .bind(&astra_user_id) - .bind(&internal_username) - .bind(&internal_email) - .bind(&external_display_name) - .execute(&mut *tx) - .await - .map_err(|e| map_auth_sqlx(e, "external.insert_auth_user", Some(&pool)))?; - query( - "INSERT INTO auth_external_identities \ - (provider_id, external_subject, astra_user_id, username, email, display_name) \ - VALUES (?, ?, ?, ?, ?, ?)", - ) - .bind(&provider.id) - .bind(&external_subject) - .bind(&astra_user_id) - .bind(&external_username) - .bind(&external_email) - .bind(&external_display_name) - .execute(&mut *tx) - .await - .map_err(|e| map_auth_sqlx(e, "external.insert_identity", Some(&pool)))?; - astra_user_id - }; query( "INSERT INTO auth_external_sessions \ @@ -1562,6 +1560,29 @@ fn map_auth_sqlx( #[async_trait] impl AuthService for DatabaseAuthService { + async fn reauthentication_options( + &self, + user_id: &str, + ) -> Result { + let pool = self.get_pool().await.map_err(internal_error)?; + if self.memoria_owner(&pool, user_id).await?.is_some() { + let web = self + .memoria_provider + .as_ref() + .and_then(|p| p.web_url.as_ref()) + .ok_or_else(|| { + error_response( + StatusCode::SERVICE_UNAVAILABLE, + "Account reauthentication website is not configured", + ) + })?; + Ok( + serde_json::json!({"method":"memoria","verification_url":format!("{web}/astra/reauthenticate")}), + ) + } else { + Ok(serde_json::json!({"method":"password"})) + } + } async fn register( &self, request: AuthRegisterRequestData, @@ -1745,6 +1766,18 @@ impl AuthService for DatabaseAuthService { }) } + async fn login_memoria(&self, key: &str) -> Result { + self.memoria_login(key).await + } + + fn memoria_credentials(&self) -> Option { + self.credential_resolver() + } + + async fn disconnect_memoria(&self, user_id: &str) -> Result<(), AuthHttpError> { + self.memoria_disconnect(user_id).await + } + async fn refresh( &self, request: AuthRefreshRequestData, @@ -1792,6 +1825,29 @@ impl AuthService for DatabaseAuthService { .map_err(|e| map_auth_sqlx(e, "refresh.fetch_user_by_id", Some(&pool)))? .ok_or_else(|| error_response(StatusCode::NOT_FOUND, "User not found"))?; + if !user.is_active { + return Err(error_response(StatusCode::UNAUTHORIZED, "User is inactive")); + } + let memoria_owner = self.memoria_owner(&pool, &user_id).await?; + let origin = if let Some(owner) = memoria_owner.as_deref() { + self.revalidate_memoria_connection(&pool, &user_id, owner) + .await?; + format!( + "verified:{}", + self.memoria_provider + .as_ref() + .ok_or_else(|| internal_error("Memoria provider unavailable"))? + .provider_id + ) + } else if origin == "memoria" || origin.starts_with("verified:memoria:") { + return Err(error_response( + StatusCode::UNAUTHORIZED, + "Memoria identity binding is missing", + )); + } else { + origin + }; + let access_token = self .create_access_token(&user.user_id, &user.username, &session_id, &origin) .map_err(internal_error)?; @@ -1801,15 +1857,27 @@ impl AuthService for DatabaseAuthService { let new_refresh_token_hash = sha256_hex(&new_refresh_token); let expires_at = self.refresh_token_expires_at_string(Utc::now()); - let mut tx = pool - .begin() - .await - .map_err(|e| map_auth_sqlx(e, "refresh.begin_tx", Some(&pool)))?; - query("UPDATE auth_refresh_tokens SET is_revoked = 1 WHERE token_hash = ?") - .bind(&refresh_token_hash) - .execute(&mut *tx) + let mut tx = pool.begin().await.map_err(internal_error)?; + query("SELECT user_id FROM auth_users WHERE user_id = ? FOR UPDATE") + .bind(&user_id) + .fetch_one(&mut *tx) .await - .map_err(|e| map_auth_sqlx(e, "refresh.revoke_old_token", Some(&pool)))?; + .map_err(internal_error)?; + // A concurrent disconnect revokes this row under the same user lock; + // the compare-and-revoke prevents refresh from resurrecting it. + let revoked = query( + "UPDATE auth_refresh_tokens SET is_revoked = 1 WHERE token_hash = ? AND is_revoked = 0", + ) + .bind(&refresh_token_hash) + .execute(&mut *tx) + .await + .map_err(|e| map_auth_sqlx(e, "refresh.revoke_old_token", Some(&pool)))?; + if revoked.rows_affected() != 1 { + return Err(error_response( + StatusCode::UNAUTHORIZED, + "Token expired or revoked", + )); + } query( "INSERT INTO auth_refresh_tokens (token_id, user_id, session_id, token_hash, expires_at, is_revoked) \ VALUES (?, ?, ?, ?, ?, 0)", @@ -1831,7 +1899,12 @@ impl AuthService for DatabaseAuthService { access_token, refresh_token: new_refresh_token, token_type: "bearer".to_string(), - expires_in: self.access_token_expires_in_seconds(), + expires_in: if origin == "memoria" || origin.starts_with("verified:memoria:") { + self.access_token_expires_in_seconds() + .min(memoria::ACCESS_TTL_SECONDS) + } else { + self.access_token_expires_in_seconds() + }, }) } @@ -1900,12 +1973,6 @@ impl AuthService for DatabaseAuthService { user_id: &str, request: ReauthenticationRequestData, ) -> Result)> { - if request.password.is_empty() || request.password.len() > AUTH_PASSWORD_MAX_BYTES { - return Err(error_response( - StatusCode::UNAUTHORIZED, - "Reauthentication failed", - )); - } let pool = self .get_pool() .await @@ -1915,7 +1982,41 @@ impl AuthService for DatabaseAuthService { .await .map_err(|e| map_auth_sqlx(e, "reauthenticate.fetch_user", Some(&pool)))? .ok_or_else(|| error_response(StatusCode::UNAUTHORIZED, "Reauthentication failed"))?; - if !user.is_active + if !user.is_active { + return Err(error_response( + StatusCode::UNAUTHORIZED, + "Reauthentication failed", + )); + } + let binding = self.reauthentication_binding(&pool, user_id).await?; + if let Some(binding) = &binding { + if !request.password.is_empty() { + return Err(error_response( + StatusCode::BAD_REQUEST, + "Choose one reauthentication method", + )); + } + self.verify_memoria_step_up( + binding, + request.memoria_proof.as_deref().unwrap_or(""), + request.purpose, + ) + .await?; + // Disconnect/rotation while the provider was verifying must fail closed. + if self + .reauthentication_binding(&pool, user_id) + .await? + .as_ref() + != Some(binding) + { + return Err(error_response( + StatusCode::UNAUTHORIZED, + "Reauthentication binding changed", + )); + } + } else if request.memoria_proof.is_some() + || request.password.is_empty() + || request.password.len() > AUTH_PASSWORD_MAX_BYTES || !bcrypt_verify(request.password.as_str(), &user.password_hash).unwrap_or(false) { return Err(error_response( @@ -1925,7 +2026,7 @@ impl AuthService for DatabaseAuthService { } let proof = format!("rp_{}_{}", Uuid::new_v4().simple(), Uuid::new_v4().simple()); - let proof_hash = sha256_hex(&proof); + let proof_hash = memoria::reauthentication_proof_hash(&proof, binding.as_ref()); let expires_at = Utc::now() + ChronoDuration::seconds(REAUTHENTICATION_PROOF_TTL_SECONDS); query( "INSERT INTO auth_reauthentication_proofs @@ -1964,6 +2065,7 @@ impl AuthService for DatabaseAuthService { .get_pool() .await .map_err(|e| map_auth_sqlx(e, "auth.get_pool", None))?; + let binding = self.reauthentication_binding(&pool, user_id).await?; let result = query( "UPDATE auth_reauthentication_proofs SET consumed_at = NOW(6) @@ -1972,7 +2074,10 @@ impl AuthService for DatabaseAuthService { ) .bind(user_id) .bind(purpose.as_str()) - .bind(sha256_hex(proof)) + .bind(memoria::reauthentication_proof_hash( + proof, + binding.as_ref(), + )) .execute(&pool) .await .map_err(|e| map_auth_sqlx(e, "reauthenticate.consume_proof", Some(&pool)))?; @@ -2031,7 +2136,7 @@ impl AuthService for DatabaseAuthService { .sub .clone() .ok_or_else(|| error_response(StatusCode::UNAUTHORIZED, "Invalid token"))?; - let (session_id, _, _) = self.parse_token_session(&claims)?; + let (session_id, token_origin, _) = self.parse_token_session(&claims)?; let pool = self .get_pool() .await @@ -2047,12 +2152,38 @@ impl AuthService for DatabaseAuthService { )); } + // Bound legacy Memoria tokens too: earlier builds issued them with + // `internal` origin and the self-hosted deployment's longer TTL. + let memoria_owner = self.memoria_owner(&pool, &user_id).await?; + if (token_origin == "memoria" || token_origin.starts_with("verified:memoria:")) + && memoria_owner.is_none() + { + return Err(error_response( + StatusCode::UNAUTHORIZED, + "Memoria identity binding is missing", + )); + } + if memoria_owner.is_some() + && !claims.iat.is_some_and(|iat| { + Utc::now().timestamp() - iat < i64::from(memoria::ACCESS_TTL_SECONDS) + }) + { + return Err(error_response( + StatusCode::UNAUTHORIZED, + "Memoria access token expired; refresh required", + )); + } + let user = self .fetch_user_by_id_or_username(&pool, &user_id, None) .await .map_err(|e| map_auth_sqlx(e, "current_user.fetch_user", Some(&pool)))? .ok_or_else(|| error_response(StatusCode::UNAUTHORIZED, "User not found"))?; + if !user.is_active { + return Err(error_response(StatusCode::UNAUTHORIZED, "User is inactive")); + } + Ok(AuthPrincipal { user: AuthUserRecord { user_id: user.user_id, @@ -2061,7 +2192,28 @@ impl AuthService for DatabaseAuthService { display_name: user.display_name, }, session_id: Some(session_id), - origin: AuthPrincipalOrigin::Internal, + origin: if let Some(subject) = memoria_owner { + let resolver = self + .credential_resolver() + .ok_or_else(|| internal_error("Memoria provider unavailable"))?; + if resolver + .resolve(&user_id) + .await + .map_err(internal_error)? + .is_none() + { + return Err(error_response( + StatusCode::UNAUTHORIZED, + "Memoria connection disconnected", + )); + } + AuthPrincipalOrigin::VerifiedProvider { + provider_id: resolver.provider.provider_id, + external_subject: subject, + } + } else { + AuthPrincipalOrigin::Internal + }, }) } diff --git a/crates/services/src/auth/verified.rs b/crates/services/src/auth/verified.rs new file mode 100644 index 0000000000..bbeb20105b --- /dev/null +++ b/crates/services/src/auth/verified.rs @@ -0,0 +1,51 @@ +//! Canonical identity mapping for already verified provider subjects. +use super::{AuthHttpError, DatabaseAuthService, sha256_hex}; +use astra_core::{error_response, internal_error}; +use axum::http::StatusCode; +use sqlx::Row; + +impl DatabaseAuthService { + /// Canonical provider-scoped mapping, also used by provider-session login. + pub(super) async fn resolve_verified_provider_identity( + &self, + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, + provider: &str, + subject: &str, + legacy_user: Option<&str>, + ) -> Result { + let existing: Option = sqlx::query_scalar("SELECT astra_user_id FROM auth_external_identities WHERE provider_id = ? AND external_subject = ?") + .bind(provider).bind(subject).fetch_optional(&mut **tx).await.map_err(internal_error)?; + let new_account = existing.is_none() && legacy_user.is_none(); + let user = existing.unwrap_or_else(|| { + legacy_user + .map(str::to_string) + .unwrap_or_else(|| format!("ext_{}", sha256_hex(&format!("{provider}\0{subject}")))) + }); + let username = format!("ext_{}", &sha256_hex(&user)[..24]); + if new_account { + sqlx::query("INSERT IGNORE INTO auth_users (user_id,username,email,password_hash,display_name,is_active) VALUES (?, ?, ?, '', 'Astra user', 1)") + .bind(&user).bind(&username).bind(format!("{username}@external.astra.invalid")) + .execute(&mut **tx).await.map_err(internal_error)?; + } + let row = sqlx::query("SELECT user_id,username,email,password_hash,display_name,is_active FROM auth_users WHERE user_id = ? FOR UPDATE") + .bind(&user).fetch_optional(&mut **tx).await.map_err(internal_error)? + .ok_or_else(|| error_response(StatusCode::FORBIDDEN, "Linked account no longer exists"))?; + if row.try_get::("is_active").unwrap_or(0) == 0 { + return Err(error_response(StatusCode::FORBIDDEN, "User is inactive")); + } + sqlx::query("INSERT IGNORE INTO auth_external_identities (provider_id,external_subject,astra_user_id) VALUES (?, ?, ?)") + .bind(provider).bind(subject).bind(&user).execute(&mut **tx).await.map_err(internal_error)?; + if new_account { + sqlx::query("INSERT IGNORE INTO auth_user_roles (user_id,role_id) SELECT ?, role_id FROM auth_roles WHERE role_name = 'astra_user'") + .bind(&user).execute(&mut **tx).await.map_err(internal_error)?; + } + Ok(super::DatabaseUserRecord { + user_id: user, + username: row.try_get("username").map_err(internal_error)?, + email: row.try_get("email").map_err(internal_error)?, + password_hash: String::new(), + display_name: row.try_get("display_name").map_err(internal_error)?, + is_active: true, + }) + } +} diff --git a/crates/services/src/byok_endpoint.rs b/crates/services/src/byok_endpoint.rs new file mode 100644 index 0000000000..f2ae3a338f --- /dev/null +++ b/crates/services/src/byok_endpoint.rs @@ -0,0 +1,330 @@ +//! Network policy for user-configured OpenAI-compatible endpoints. +//! Public HTTPS is the default; optional strict mode uses the administrator registry. + +use std::net::{IpAddr, SocketAddr}; +use std::time::Duration; + +pub const COMPATIBLE_PROVIDER: &str = "openai-compatible"; +pub const ENDPOINT_POLICY_ENV: &str = "ASTRA_BYOK_ENDPOINT_POLICY"; +mod dns; +mod proxy; + +/// Keeps any private CONNECT tunnel alive for the entire response/stream lifetime. +pub struct EndpointClient { + client: reqwest::Client, + tunnel: Option>, +} +impl From for EndpointClient { + fn from(client: reqwest::Client) -> Self { + Self { + client, + tunnel: None, + } + } +} +impl std::ops::Deref for EndpointClient { + type Target = reqwest::Client; + fn deref(&self) -> &Self::Target { + &self.client + } +} +impl Drop for EndpointClient { + fn drop(&mut self) { + if let Some(task) = &self.tunnel { + task.abort(); + } + } +} + +fn configured_value(name: &str) -> Result, String> { + match std::env::var(name) { + Ok(value) => Ok(Some(value)), + Err(std::env::VarError::NotPresent) => Ok(None), + Err(_) => Err(format!("Invalid {name}")), + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum EndpointPolicy { + PublicHttps, + TrustedDomains, +} + +impl EndpointPolicy { + pub fn parse(value: Option<&str>) -> Result { + match value { + None | Some("public-https") => Ok(Self::PublicHttps), + Some("trusted-domains") => Ok(Self::TrustedDomains), + _ => Err(format!( + "Invalid {ENDPOINT_POLICY_ENV}; expected public-https or trusted-domains" + )), + } + } +} + +pub fn parse_endpoint(raw: &str) -> Result { + if raw.len() > 2048 || raw.chars().any(char::is_control) || raw.contains('\\') { + return Err("Invalid model base URL".into()); + } + let url = reqwest::Url::parse(raw).map_err(|_| "Invalid model base URL")?; + if url.scheme() != "https" + || !url.username().is_empty() + || url.password().is_some() + || url.query().is_some() + || url.fragment().is_some() + || url.host_str().is_none() + || url.port_or_known_default() == Some(0) + { + return Err("Model base URL must use HTTPS without credentials, query or fragment".into()); + } + let host = url.host_str().unwrap().trim_end_matches('.'); + if host == "localhost" || host.ends_with(".localhost") || !host.contains('.') { + return Err("Model endpoint must be a public host".into()); + } + if let Ok(ip) = host.trim_matches(['[', ']']).parse::() + && !public_ip(ip) + { + return Err("Model endpoint must use public IP addresses".into()); + } + // Reject encoded path separators/dot segments rather than relying on + // potentially different normalization in upstream gateways. + if url.path().contains('%') { + return Err("Model base URL cannot contain an encoded path".into()); + } + Ok(url) +} + +fn public_ip(ip: IpAddr) -> bool { + match ip { + IpAddr::V4(ip) => { + let [a, b, _, _] = ip.octets(); + !(ip.is_private() + || ip.is_loopback() + || ip.is_link_local() + || ip.is_broadcast() + || ip.is_documentation() + || ip.is_unspecified() + || a == 0 + || a >= 224 + || (a == 100 && (64..=127).contains(&b)) + || (a == 198 && (b == 18 || b == 19)) + || (a == 192 && b == 0)) + && ip.octets() != [168, 63, 129, 16] + } + IpAddr::V6(ip) => { + let s = ip.segments(); + // Global unicast only; exclude special-purpose and transition + // ranges, including IPv4 mapped/NAT64 addresses. + s[0] & 0xe000 == 0x2000 + && !(s[0] == 0x2001 && (s[1] < 0x200 || s[1] == 0xdb8)) + && s[0] != 0x2002 + && !(s[0] == 0x3fff && s[1] < 0x1000) + } + } +} + +pub async fn require_endpoint_policy( + pool: &sqlx::Pool, + raw: &str, +) -> Result<(), String> { + let policy = match std::env::var(ENDPOINT_POLICY_ENV) { + Ok(value) => EndpointPolicy::parse(Some(&value))?, + Err(std::env::VarError::NotPresent) => EndpointPolicy::parse(None)?, + Err(_) => return Err(format!("Invalid {ENDPOINT_POLICY_ENV}")), + }; + require_endpoint_policy_with_mode(pool, raw, policy).await +} + +/// Applies the configured policy without relaxing URL or outbound DNS checks. +pub async fn require_endpoint_policy_with_mode( + pool: &sqlx::Pool, + raw: &str, + policy: EndpointPolicy, +) -> Result<(), String> { + let url = parse_endpoint(raw)?; + if policy == EndpointPolicy::PublicHttps { + return Ok(()); + } + let permitted: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM runtime_llm_trusted_domains WHERE domain_host = ? \ + AND is_enabled = 1 AND (domain_port = ? OR \ + (IFNULL(domain_port, 0) = 0 AND ? = 443))", + ) + .bind(url.host_str().unwrap()) + .bind(i32::from(url.port_or_known_default().unwrap())) + .bind(i32::from(url.port_or_known_default().unwrap())) + .fetch_one(pool) + .await + .map_err(|_| "Unable to check the model endpoint policy")?; + if permitted == 0 { + return Err(format!( + "This deployment requires an approved model endpoint; {}:{} is not enabled. Contact the administrator.", + url.host_str().unwrap(), + url.port_or_known_default().unwrap() + )); + } + Ok(()) +} + +fn transport_builder() -> reqwest::ClientBuilder { + reqwest::Client::builder() + .no_proxy() + .redirect(reqwest::redirect::Policy::none()) + .connect_timeout(Duration::from_secs(15)) + .timeout(Duration::from_secs(360)) +} + +fn validate_addresses(addresses: &[SocketAddr]) -> Result<(), String> { + if addresses.is_empty() || addresses.iter().any(|addr| !public_ip(addr.ip())) { + return Err("Astra Server DNS returned non-public model addresses (possibly proxy Fake-IP). Configure ASTRA_BYOK_DNS_SERVERS to return real public addresses; the model API key is not the cause.".into()); + } + Ok(()) +} + +fn pinned_client(url: &reqwest::Url, addresses: &[SocketAddr]) -> Result { + validate_addresses(addresses)?; + transport_builder() + // User-selected destinations require DNS pinning. An ambient proxy + // would resolve the host again, outside this validation boundary. + .resolve_to_addrs(url.host_str().unwrap(), addresses) + .build() + .map_err(|_| "Unable to create model HTTP client".into()) +} + +/// Resolve and pin a fresh public address set for each outbound attempt. +/// TLS still validates the original hostname. Redirects never forward keys. +pub async fn endpoint_client(raw: &str) -> Result { + let url = parse_endpoint(raw)?; + let proxy = proxy::EgressProxy::parse(configured_value(proxy::PROXY_ENV)?.as_deref())?; + let addresses = dns::resolve(&url, configured_value(dns::DNS_SERVERS_ENV)?.as_deref()).await?; + validate_addresses(&addresses)?; + match proxy { + Some(proxy) => proxy::client(&url, &addresses, proxy, transport_builder()).await, + None => pinned_client(&url, &addresses).map(EndpointClient::from), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn endpoint_policy_defaults_to_public_https_and_invalid_config_fails_closed() { + assert_eq!( + EndpointPolicy::parse(None).unwrap(), + EndpointPolicy::PublicHttps + ); + assert_eq!( + EndpointPolicy::parse(Some("trusted-domains")).unwrap(), + EndpointPolicy::TrustedDomains + ); + for invalid in ["", "public", "allow-all", "trusted-domain"] { + assert!(EndpointPolicy::parse(Some(invalid)).is_err()); + } + } + + #[tokio::test] + async fn public_policy_does_not_require_a_registry_or_database_connection() { + let pool = sqlx::mysql::MySqlPoolOptions::new() + .connect_lazy("mysql://unused:unused@127.0.0.1:1/unused") + .unwrap(); + for raw in [ + "https://api.moonshot.cn/v1", + "https://another-provider.example:8443/v1", + ] { + require_endpoint_policy_with_mode(&pool, raw, EndpointPolicy::PublicHttps) + .await + .unwrap(); + } + for policy in [EndpointPolicy::PublicHttps, EndpointPolicy::TrustedDomains] { + assert!( + require_endpoint_policy_with_mode(&pool, "https://127.0.0.1/v1", policy) + .await + .is_err() + ); + } + } + + #[test] + fn endpoint_rejects_private_and_ambiguous_urls() { + for raw in [ + "http://api.example.com/v1", + "https://127.0.0.1/v1", + "https://169.254.169.254", + "https://10.1.1.1", + "https://[::1]", + "https://localhost", + "https://x.localhost", + "https://user:secret@api.example.com", + "https://api.example.com?key=secret", + "https://api.example.com/#x", + "https://api.example.com/%2fadmin", + "https://100.64.0.1", + "https://2130706433", + "https://0x7f000001", + "https://168.63.129.16", + "https://[::ffff:127.0.0.1]", + ] { + assert!(parse_endpoint(raw).is_err(), "accepted {raw}"); + } + assert!(parse_endpoint("https://api.example.com:8443/compatible/v1").is_ok()); + } + + #[test] + fn dns_pinning_rejects_mixed_and_rebound_answers() { + let url = parse_endpoint("https://api.example.com/v1").unwrap(); + let public = "8.8.8.8:443".parse().unwrap(); + assert!(pinned_client(&url, &[public]).is_ok()); + for blocked in [ + "127.0.0.1:443", + "10.1.1.1:443", + "169.254.169.254:443", + "[::ffff:127.0.0.1]:443", + "[64:ff9b::a00:1]:443", + "[fc00::1]:443", + "168.63.129.16:443", + "100.100.100.200:443", + "[2001:db8::1]:443", + ] { + assert!(pinned_client(&url, &[public, blocked.parse().unwrap()]).is_err()); + } + assert!(pinned_client(&url, &[]).is_err()); + } + + #[tokio::test] + async fn transport_never_follows_provider_redirects() { + use axum::{Router, routing::get}; + use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }; + let hits = Arc::new(AtomicUsize::new(0)); + let observed = hits.clone(); + let app = Router::new() + .route( + "/redirect", + get(|| async { axum::response::Redirect::temporary("/destination") }), + ) + .route( + "/destination", + get(move || async move { + observed.fetch_add(1, Ordering::SeqCst); + "unexpected" + }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + let response = transport_builder() + .build() + .unwrap() + .get(format!("http://{addr}/redirect")) + .bearer_auth("test-secret") + .send() + .await + .unwrap(); + assert_eq!(response.status(), reqwest::StatusCode::TEMPORARY_REDIRECT); + assert_eq!(hits.load(Ordering::SeqCst), 0); + server.abort(); + } +} diff --git a/crates/services/src/byok_endpoint/dns.rs b/crates/services/src/byok_endpoint/dns.rs new file mode 100644 index 0000000000..6c2b4fb3c0 --- /dev/null +++ b/crates/services/src/byok_endpoint/dns.rs @@ -0,0 +1,319 @@ +//! Query configured DNS servers directly; never reuse OS synthetic-address caches. +use hickory_resolver::{ + Resolver, TokioResolver, + config::{ + LookupIpStrategy, NameServerConfig, ResolveHosts, ResolverConfig, ServerOrderingStrategy, + }, + net::runtime::TokioRuntimeProvider, +}; +use std::{ + net::{IpAddr, SocketAddr}, + time::Duration, +}; + +pub(super) const DNS_SERVERS_ENV: &str = "ASTRA_BYOK_DNS_SERVERS"; + +fn resolver(servers: Option<&str>) -> Result { + let config = if let Some(servers) = servers { + let mut config = ResolverConfig::from_parts(None, vec![], vec![]); + if servers.len() > 2048 { + return Err("Invalid ASTRA_BYOK_DNS_SERVERS".into()); + } + for value in servers.split(',') { + let value = value.trim(); + // A TCP-only route must not send a UDP query first: transparent + // proxies may return a syntactically valid synthetic UDP answer. + let (tcp_only, value) = match value.strip_prefix("tcp://") { + Some(address) => (true, address), + None => (false, value), + }; + let address: SocketAddr = value.parse().or_else(|_| value.parse::().map(|ip| SocketAddr::new(ip, 53))) + .map_err(|_| "ASTRA_BYOK_DNS_SERVERS must contain comma-separated DNS IP addresses, optionally with ports and a tcp:// prefix")?; + if address.port() == 0 || address.ip().is_unspecified() || address.ip().is_multicast() { + return Err("Invalid ASTRA_BYOK_DNS_SERVERS address".into()); + } + let mut ns = if tcp_only { + NameServerConfig::tcp(address.ip()) + } else { + NameServerConfig::udp_and_tcp(address.ip()) + }; + for connection in &mut ns.connections { + connection.port = address.port(); + } + config.add_name_server(ns); + } + config + } else { + hickory_resolver::system_conf::read_system_conf() + .map_err( + |_| "Astra Server could not read DNS servers; configure ASTRA_BYOK_DNS_SERVERS", + )? + .0 + }; + if config.name_servers().is_empty() { + return Err("Astra Server has no DNS servers; configure ASTRA_BYOK_DNS_SERVERS".into()); + } + let mut builder = Resolver::builder_with_config(config, TokioRuntimeProvider::default()); + let options = builder.options_mut(); + options.use_hosts_file = ResolveHosts::Never; + // Respect DNS priority. Racing a VPN/system resolver against a secondary + // resolver can otherwise choose a fast synthetic answer nondeterministically. + options.server_ordering_strategy = ServerOrderingStrategy::UserProvidedOrder; + options.num_concurrent_reqs = 1; + options.ip_strategy = LookupIpStrategy::Ipv4AndIpv6; + options.timeout = Duration::from_secs(2); + options.attempts = 1; + builder + .build() + .map_err(|_| "Unable to initialize Astra Server DNS resolver".into()) +} + +pub(super) async fn resolve( + url: &reqwest::Url, + servers: Option<&str>, +) -> Result, String> { + // Validate operator configuration even for literal addresses. + let resolver = resolver(servers)?; + let host = url.host_str().unwrap().trim_matches(['[', ']']); + let port = url.port_or_known_default().unwrap(); + if let Ok(ip) = host.parse::() { + return Ok(vec![SocketAddr::new(ip, port)]); + } + let fqdn = format!("{}.", host.trim_end_matches('.')); + let lookup = tokio::time::timeout(Duration::from_secs(5), resolver.lookup_ip(fqdn)).await + .map_err(|_| "Astra Server DNS lookup timed out; check its DNS/network configuration, not the model API key")? + .map_err(|error| { + // Log resolver evidence server-side only. Never include the URL path, + // proxy credentials or model API key in the public error. + tracing::warn!(host, error = ?error, "BYOK endpoint DNS lookup failed"); + "Astra Server could not resolve the model endpoint; check its DNS configuration and the provider hostname" + })?; + Ok(lookup.iter().map(|ip| SocketAddr::new(ip, port)).collect()) +} + +#[cfg(test)] +mod tests { + use super::*; + use hickory_resolver::proto::{ + op::Message, + rr::{ + RData, Record, RecordType, + rdata::{A, AAAA}, + }, + }; + use std::sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + #[tokio::test] + async fn tcp_only_avoids_fake_udp_and_still_rejects_private_answers() { + let tcp = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = tcp.local_addr().unwrap(); + let udp = tokio::net::UdpSocket::bind(address).await.unwrap(); + let udp_seen = Arc::new(AtomicBool::new(false)); + let seen = udp_seen.clone(); + let udp_task = tokio::spawn(async move { + let mut packet = [0; 4096]; + loop { + let (size, peer) = udp.recv_from(&mut packet).await.unwrap(); + seen.store(true, Ordering::SeqCst); + let mut reply = Message::from_vec(&packet[..size]).unwrap().into_response(); + let query = reply.queries[0].clone(); + if query.query_type() == RecordType::A { + reply.add_answer(Record::from_rdata( + query.name().clone(), + 60, + RData::A(A("198.18.4.197".parse().unwrap())), + )); + } + udp.send_to(&reply.to_vec().unwrap(), peer).await.unwrap(); + } + }); + let private = Arc::new(AtomicBool::new(false)); + let state = private.clone(); + let tcp_task = tokio::spawn(async move { + loop { + let (mut stream, _) = tcp.accept().await.unwrap(); + let state = state.clone(); + tokio::spawn(async move { + while let Ok(size) = stream.read_u16().await { + let mut packet = vec![0; usize::from(size)]; + stream.read_exact(&mut packet).await.unwrap(); + let mut reply = Message::from_vec(&packet).unwrap().into_response(); + let query = reply.queries[0].clone(); + let data = match query.query_type() { + RecordType::A => RData::A(A("8.8.8.8".parse().unwrap())), + RecordType::AAAA => RData::AAAA(AAAA( + if state.load(Ordering::SeqCst) { + "::1" + } else { + "2606:4700:4700::1111" + } + .parse() + .unwrap(), + )), + other => panic!("unexpected query {other}"), + }; + reply.add_answer(Record::from_rdata(query.name().clone(), 60, data)); + let packet = reply.to_vec().unwrap(); + stream.write_u16(packet.len() as u16).await.unwrap(); + stream.write_all(&packet).await.unwrap(); + } + }); + } + }); + let url = reqwest::Url::parse("https://byok.test/v1").unwrap(); + let fake = resolve(&url, Some(&address.to_string())).await.unwrap(); + assert!(super::super::validate_addresses(&fake).is_err()); + assert!(udp_seen.swap(false, Ordering::SeqCst)); + let servers = format!("tcp://{address}"); + let real = resolve(&url, Some(&servers)).await.unwrap(); + assert_eq!(real.len(), 2); + assert!(super::super::validate_addresses(&real).is_ok()); + private.store(true, Ordering::SeqCst); + let rebound = resolve(&url, Some(&servers)).await.unwrap(); + assert!(super::super::validate_addresses(&rebound).is_err()); + assert!(!udp_seen.load(Ordering::SeqCst), "TCP-only mode used UDP"); + tcp_task.abort(); + tcp_task.await.unwrap_err(); + // UDP is still available, but a failed TCP route must fail closed. + assert!(resolve(&url, Some(&servers)).await.is_err()); + assert!( + !udp_seen.load(Ordering::SeqCst), + "TCP failure fell back to UDP" + ); + udp_task.abort(); + } + + #[tokio::test] + async fn configured_dns_returns_all_addresses_and_rechecks_changed_answers() { + let socket = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let servers = socket.local_addr().unwrap().to_string(); + let private = Arc::new(AtomicBool::new(false)); + let state = private.clone(); + let task = tokio::spawn(async move { + let mut packet = [0; 4096]; + loop { + let (size, peer) = socket.recv_from(&mut packet).await.unwrap(); + let mut reply = Message::from_vec(&packet[..size]).unwrap().into_response(); + let query = reply.queries[0].clone(); + assert_eq!(query.name().to_string(), "byok.test."); + let data = match query.query_type() { + RecordType::A => RData::A(A("8.8.8.8".parse().unwrap())), + RecordType::AAAA => RData::AAAA(AAAA( + if state.load(Ordering::SeqCst) { + "::1" + } else { + "2606:4700:4700::1111" + } + .parse() + .unwrap(), + )), + other => panic!("unexpected query {other}"), + }; + reply.add_answer(Record::from_rdata(query.name().clone(), 60, data)); + socket + .send_to(&reply.to_vec().unwrap(), peer) + .await + .unwrap(); + } + }); + let url = reqwest::Url::parse("https://byok.test:8443/v1").unwrap(); + let public = resolve(&url, Some(&servers)).await.unwrap(); + assert_eq!(public.len(), 2); + assert!(public.iter().all(|address| address.port() == 8443)); + assert!(super::super::validate_addresses(&public).is_ok()); + private.store(true, Ordering::SeqCst); + let rebound = resolve(&url, Some(&servers)).await.unwrap(); + assert_eq!(rebound.len(), 2); + assert!( + super::super::validate_addresses(&rebound).is_err(), + "a cached public answer must not hide a new private AAAA record" + ); + task.abort(); + } + + #[tokio::test] + async fn unavailable_dns_fails_without_system_resolver_fallback() { + let listener = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let url = reqwest::Url::parse("https://byok.test/v1").unwrap(); + let error = tokio::time::timeout( + Duration::from_secs(6), + resolve(&url, Some(&listener.local_addr().unwrap().to_string())), + ) + .await + .unwrap() + .unwrap_err(); + assert!(error.contains("DNS")); + } + + #[tokio::test] + async fn dns_priority_is_preserved_instead_of_racing_secondary_servers() { + let primary = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let secondary = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let servers = format!( + "{},{}", + primary.local_addr().unwrap(), + secondary.local_addr().unwrap() + ); + let task = tokio::spawn(async move { + let mut packet = [0; 4096]; + loop { + let (size, peer) = primary.recv_from(&mut packet).await.unwrap(); + let mut reply = Message::from_vec(&packet[..size]).unwrap().into_response(); + let query = reply.queries[0].clone(); + let data = match query.query_type() { + RecordType::A => RData::A(A("8.8.8.8".parse().unwrap())), + RecordType::AAAA => RData::AAAA(AAAA("2606:4700:4700::1111".parse().unwrap())), + other => panic!("unexpected query {other}"), + }; + reply.add_answer(Record::from_rdata(query.name().clone(), 60, data)); + tokio::time::sleep(Duration::from_millis(25)).await; + primary + .send_to(&reply.to_vec().unwrap(), peer) + .await + .unwrap(); + } + }); + let addresses = resolve( + &reqwest::Url::parse("https://byok.test/v1").unwrap(), + Some(&servers), + ) + .await + .unwrap(); + assert_eq!(addresses.len(), 2); + assert!(super::super::validate_addresses(&addresses).is_ok()); + let mut packet = [0; 4096]; + assert!( + tokio::time::timeout(Duration::from_millis(30), secondary.recv_from(&mut packet)) + .await + .is_err(), + "a successful primary DNS must not race secondary resolvers" + ); + task.abort(); + } + + #[tokio::test] + async fn explicit_dns_configuration_is_validated() { + for invalid in [ + "", + "dns.example", + "https://dns.example", + "127.0.0.1:0", + "0.0.0.0", + "1.1.1.1,", + "tcp://", + "tcp://dns.example", + "tcp://127.0.0.1:0", + "tcp://127.0.0.1:53/path", + "udp://127.0.0.1", + "tcp://user:password@127.0.0.1", + ] { + assert!(resolver(Some(invalid)).is_err(), "accepted {invalid}"); + } + assert!(resolver(Some("127.0.0.1:5353,[::1]:5353")).is_ok()); + assert!(resolver(Some("tcp://127.0.0.1:5353,tcp://[::1]:5353")).is_ok()); + } +} diff --git a/crates/services/src/byok_endpoint/proxy.rs b/crates/services/src/byok_endpoint/proxy.rs new file mode 100644 index 0000000000..47450d21c6 --- /dev/null +++ b/crates/services/src/byok_endpoint/proxy.rs @@ -0,0 +1,316 @@ +//! A request-owned CONNECT adapter. Upstream proxies receive validated IPs, +//! never provider hostnames; reqwest retains end-to-end origin TLS/SNI. +use super::EndpointClient; +use base64::{Engine, engine::general_purpose::STANDARD}; +use std::{ + io, + net::{IpAddr, SocketAddr}, + sync::Arc, + time::Duration, +}; +use tokio::{ + io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}, + net::{TcpListener, TcpStream}, + task::JoinSet, +}; + +pub(super) const PROXY_ENV: &str = "ASTRA_BYOK_PROXY_URL"; +const CONNECT_TIMEOUT: Duration = Duration::from_secs(15); +const MAX_HEAD: usize = 8192; + +trait TunnelIo: AsyncRead + AsyncWrite + Unpin + Send {} +impl TunnelIo for T {} +type Stream = Box; + +#[derive(Clone)] +pub(super) struct EgressProxy { + url: reqwest::Url, + username: String, + password: String, + tls: Arc, +} + +fn invalid() -> String { + "Invalid ASTRA_BYOK_PROXY_URL; use an explicit http://, https:// or socks5:// proxy URL".into() +} +fn failed(reason: &'static str) -> io::Error { + io::Error::other(reason) +} + +impl EgressProxy { + pub(super) fn parse(raw: Option<&str>) -> Result, String> { + let Some(raw) = raw.filter(|value| !value.is_empty()) else { + return Ok(None); + }; + if raw.len() > 2048 || raw.chars().any(char::is_control) { + return Err(invalid()); + } + let url = reqwest::Url::parse(raw).map_err(|_| invalid())?; + if !matches!(url.scheme(), "http" | "https" | "socks5" | "socks5h") + || url.host_str().is_none() + || url.port() == Some(0) + || !matches!(url.path(), "" | "/") + || url.query().is_some() + || url.fragment().is_some() + { + return Err(invalid()); + } + let decode = |value: &str| { + percent_encoding::percent_decode_str(value) + .decode_utf8() + .map(|value| value.into_owned()) + .map_err(|_| invalid()) + }; + let username = decode(url.username())?; + let password = decode(url.password().unwrap_or_default())?; + if username.len() > 255 + || password.len() > 255 + || username.contains(':') + || username + .chars() + .chain(password.chars()) + .any(char::is_control) + { + return Err(invalid()); + } + let roots = + rustls::RootCertStore::from_iter(webpki_roots::TLS_SERVER_ROOTS.iter().cloned()); + let tls = rustls::ClientConfig::builder_with_provider(Arc::new( + rustls::crypto::ring::default_provider(), + )) + .with_safe_default_protocol_versions() + .map_err(|_| invalid())? + .with_root_certificates(roots) + .with_no_client_auth(); + Ok(Some(Self { + url, + username, + password, + tls: Arc::new(tls), + })) + } + + async fn connect(&self, target: SocketAddr) -> io::Result { + let host = self.url.host_str().unwrap().trim_matches(['[', ']']); + let port = self.url.port_or_known_default().unwrap_or(1080); + // The proxy is operator-owned configuration, not a user-selected endpoint. + let tcp = TcpStream::connect((host, port)).await?; + tcp.set_nodelay(true)?; + let mut stream: Stream = if self.url.scheme() == "https" { + let name = rustls::pki_types::ServerName::try_from(host.to_owned()) + .map_err(|_| failed("invalid proxy TLS name"))?; + Box::new( + tokio_rustls::TlsConnector::from(self.tls.clone()) + .connect(name, tcp) + .await?, + ) + } else { + Box::new(tcp) + }; + if matches!(self.url.scheme(), "socks5" | "socks5h") { + socks_connect(&mut stream, target, &self.username, &self.password).await?; + } else { + let auth = if !self.username.is_empty() || !self.password.is_empty() { + format!( + "Proxy-Authorization: Basic {}\r\n", + STANDARD.encode(format!("{}:{}", self.username, self.password)) + ) + } else { + String::new() + }; + stream + .write_all( + format!("CONNECT {target} HTTP/1.1\r\nHost: {target}\r\n{auth}\r\n").as_bytes(), + ) + .await?; + let head = read_head(&mut stream).await?; + let first = head.lines().next().unwrap_or_default(); + let fields: Vec<_> = first.split_whitespace().collect(); + if fields.len() < 2 + || !matches!(fields[0], "HTTP/1.0" | "HTTP/1.1") + || fields[1] != "200" + { + return Err(failed("configured proxy refused CONNECT")); + } + } + Ok(stream) + } +} + +async fn read_head(stream: &mut (impl AsyncRead + Unpin + ?Sized)) -> io::Result { + let mut head = Vec::new(); + loop { + if head.len() >= MAX_HEAD { + return Err(failed("proxy header exceeds limit")); + } + head.push(stream.read_u8().await?); + if head.ends_with(b"\r\n\r\n") { + break; + } + } + String::from_utf8(head).map_err(|_| failed("invalid proxy header")) +} + +async fn socks_connect( + stream: &mut Stream, + target: SocketAddr, + username: &str, + password: &str, +) -> io::Result<()> { + let auth = !username.is_empty() || !password.is_empty(); + stream + .write_all(if auth { &[5, 1, 2] } else { &[5, 1, 0] }) + .await?; + let mut selection = [0; 2]; + stream.read_exact(&mut selection).await?; + if selection != [5, if auth { 2 } else { 0 }] { + return Err(failed("SOCKS proxy authentication method rejected")); + } + if auth { + let mut credentials = vec![1, username.len() as u8]; + credentials.extend_from_slice(username.as_bytes()); + credentials.push(password.len() as u8); + credentials.extend_from_slice(password.as_bytes()); + stream.write_all(&credentials).await?; + stream.read_exact(&mut selection).await?; + if selection != [1, 0] { + return Err(failed("SOCKS proxy authentication failed")); + } + } + let mut request = vec![5, 1, 0]; + match target.ip() { + IpAddr::V4(ip) => { + request.push(1); + request.extend_from_slice(&ip.octets()); + } + IpAddr::V6(ip) => { + request.push(4); + request.extend_from_slice(&ip.octets()); + } + } + request.extend_from_slice(&target.port().to_be_bytes()); + stream.write_all(&request).await?; + let mut reply = [0; 4]; + stream.read_exact(&mut reply).await?; + if reply[..3] != [5, 0, 0] { + return Err(failed("SOCKS proxy refused CONNECT")); + } + let length = match reply[3] { + 1 => 4, + 4 => 16, + 3 => usize::from(stream.read_u8().await?), + _ => return Err(failed("invalid SOCKS reply")), + }; + let mut bound = vec![0; length + 2]; + stream.read_exact(&mut bound).await?; + Ok(()) +} + +async fn bridge( + mut incoming: TcpStream, + authority: &str, + authorization: &str, + targets: &[SocketAddr], + proxy: &EgressProxy, +) -> io::Result<()> { + let setup = tokio::time::timeout(CONNECT_TIMEOUT, async { + let head = read_head(&mut incoming).await?; + let mut lines = head.lines(); + if lines.next() != Some(format!("CONNECT {authority} HTTP/1.1").as_str()) { + return Err(failed("unadmitted tunnel destination")); + } + let authenticated = lines + .filter_map(|line| line.split_once(':')) + .any(|(name, value)| { + name.eq_ignore_ascii_case("proxy-authorization") && value.trim() == authorization + }); + if !authenticated { + return Err(failed("unauthorized tunnel")); + } + for &target in targets { + if let Ok(stream) = proxy.connect(target).await { + return Ok(stream); + } + } + Err(failed( + "configured proxy could not connect to a validated public address", + )) + }) + .await; + let mut upstream = match setup { + Ok(Ok(stream)) => stream, + _ => { + incoming + .write_all( + b"HTTP/1.1 502 Bad Gateway\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", + ) + .await?; + return Err(failed("model proxy tunnel unavailable")); + } + }; + incoming + .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n") + .await?; + // Only opaque, end-to-end TLS travels through the adapter. Neither proxy + // authentication header is forwarded to the model endpoint. + tokio::time::timeout( + Duration::from_secs(360), + tokio::io::copy_bidirectional(&mut incoming, &mut upstream), + ) + .await + .map_err(|_| failed("model proxy tunnel deadline"))??; + Ok(()) +} + +pub(super) async fn client( + url: &reqwest::Url, + targets: &[SocketAddr], + proxy: EgressProxy, + builder: reqwest::ClientBuilder, +) -> Result { + super::validate_addresses(targets)?; + let authority = format!( + "{}:{}", + url.host_str().unwrap(), + url.port_or_known_default().unwrap() + ); + let listener = TcpListener::bind("127.0.0.1:0") + .await + .map_err(|_| "Unable to initialize model proxy tunnel")?; + let address = listener + .local_addr() + .map_err(|_| "Unable to initialize model proxy tunnel")?; + let token = uuid::Uuid::new_v4().to_string(); + let authorization = format!("Basic {}", STANDARD.encode(format!("astra:{token}"))); + let local_proxy = reqwest::Proxy::https(format!("http://{address}")) + .map_err(|_| "Unable to initialize model proxy tunnel")? + .basic_auth("astra", &token); + let client = builder + .no_proxy() + .https_only(true) + .proxy(local_proxy) + .build() + .map_err(|_| "Unable to initialize model proxy client")?; + let targets = targets.to_vec(); + let task = tokio::spawn(async move { + let mut connections = JoinSet::new(); + loop { + tokio::select! { + accepted = listener.accept(), if connections.len() < 4 => { + let Ok((incoming, _)) = accepted else { break; }; + let (authority, authorization, targets, proxy) = (authority.clone(), authorization.clone(), targets.clone(), proxy.clone()); + connections.spawn(async move { let _ = bridge(incoming, &authority, &authorization, &targets, &proxy).await; }); + } + _ = connections.join_next(), if !connections.is_empty() => {} + } + } + // Dropping this task drops JoinSet, cancelling all owned connections. + }); + Ok(EndpointClient { + client, + tunnel: Some(task), + }) +} + +#[cfg(test)] +mod tests; diff --git a/crates/services/src/byok_endpoint/proxy/tests.rs b/crates/services/src/byok_endpoint/proxy/tests.rs new file mode 100644 index 0000000000..ab2ebb43fd --- /dev/null +++ b/crates/services/src/byok_endpoint/proxy/tests.rs @@ -0,0 +1,430 @@ +use super::*; +use tokio::sync::{mpsc, oneshot}; + +struct TestTls { + server: Arc, + client: Arc, + certificate: reqwest::Certificate, +} +fn test_tls() -> TestTls { + let rcgen::CertifiedKey { cert, key_pair } = rcgen::generate_simple_self_signed(vec![ + "byok.test".into(), + "localhost".into(), + "127.0.0.1".into(), + ]) + .unwrap(); + let provider = Arc::new(rustls::crypto::ring::default_provider()); + let key = rustls::pki_types::PrivatePkcs8KeyDer::from(key_pair.serialize_der()).into(); + let server = rustls::ServerConfig::builder_with_provider(provider.clone()) + .with_safe_default_protocol_versions() + .unwrap() + .with_no_client_auth() + .with_single_cert(vec![cert.der().clone()], key) + .unwrap(); + let mut roots = rustls::RootCertStore::empty(); + roots.add(cert.der().clone()).unwrap(); + let client = rustls::ClientConfig::builder_with_provider(provider) + .with_safe_default_protocol_versions() + .unwrap() + .with_root_certificates(roots) + .with_no_client_auth(); + TestTls { + server: Arc::new(server), + client: Arc::new(client), + certificate: reqwest::Certificate::from_der(cert.der()).unwrap(), + } +} + +async fn origin( + tls: &TestTls, + reply: &'static [u8], +) -> ( + SocketAddr, + oneshot::Receiver<(String, String)>, + tokio::task::JoinHandle<()>, +) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let tls = tls.server.clone(); + let (tx, rx) = oneshot::channel(); + let task = tokio::spawn(async move { + let (tcp, _) = listener.accept().await.unwrap(); + let Ok(mut stream) = tokio_rustls::TlsAcceptor::from(tls).accept(tcp).await else { + return; + }; + let sni = stream + .get_ref() + .1 + .server_name() + .unwrap_or_default() + .to_owned(); + let Ok(head) = read_head(&mut stream).await else { + return; + }; + let _ = tx.send((sni, head)); + let _ = stream.write_all(reply).await; + let _ = stream.shutdown().await; + }); + (address, rx, task) +} + +// A controllable test proxy maps an explicitly asserted public target to a local +// TLS origin. Production code still performs its real public-address checks. +async fn mock_proxy( + scheme: &str, + tls: &TestTls, + origin: SocketAddr, + refuse: bool, +) -> ( + EgressProxy, + mpsc::Receiver, + tokio::task::JoinHandle<()>, +) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let mut proxy = EgressProxy::parse(Some(&format!( + "{scheme}://test-user:test-password@{address}" + ))) + .unwrap() + .unwrap(); + proxy.tls = tls.client.clone(); + let server_tls = tls.server.clone(); + let scheme = scheme.to_owned(); + let (tx, rx) = mpsc::channel(8); + let task = tokio::spawn(async move { + let mut connections = JoinSet::new(); + loop { + tokio::select! { + accepted = listener.accept() => { + let Ok((tcp, _)) = accepted else { break; }; + let (tls, tx, scheme) = (server_tls.clone(), tx.clone(), scheme.clone()); + connections.spawn(async move { + let mut stream: Stream = if scheme == "https" { + Box::new(tokio_rustls::TlsAcceptor::from(tls).accept(tcp).await.unwrap()) + } else { Box::new(tcp) }; + if scheme.starts_with("socks") { + let mut greeting = [0; 3]; + stream.read_exact(&mut greeting).await.unwrap(); + assert_eq!(greeting, [5, 1, 2]); + stream.write_all(&[5, 2]).await.unwrap(); + assert_eq!(stream.read_u8().await.unwrap(), 1); + let size = stream.read_u8().await.unwrap(); + let mut user = vec![0; usize::from(size)]; + stream.read_exact(&mut user).await.unwrap(); + assert_eq!(user, b"test-user"); + let size = stream.read_u8().await.unwrap(); + let mut password = vec![0; usize::from(size)]; + stream.read_exact(&mut password).await.unwrap(); + assert_eq!(password, b"test-password"); + stream.write_all(&[1, 0]).await.unwrap(); + let mut connect = [0; 10]; + stream.read_exact(&mut connect).await.unwrap(); + assert_eq!(connect, [5, 1, 0, 1, 8, 8, 8, 8, 1, 187], "SOCKS must receive a pinned IP, not a hostname"); + tx.send("8.8.8.8:443".into()).await.unwrap(); + stream.write_all(&[5, if refuse { 2 } else { 0 }, 0, 1, 0, 0, 0, 0, 0, 0]).await.unwrap(); + } else { + let head = read_head(&mut stream).await.unwrap(); + assert!(head.starts_with("CONNECT 8.8.8.8:443 HTTP/1.1\r\n"), "proxy destination must be pinned"); + assert!(head.contains(&STANDARD.encode("test-user:test-password"))); + assert!(!head.contains("provider-key")); + tx.send(head).await.unwrap(); + stream.write_all(if refuse { b"HTTP/1.1 407 Authentication Required\r\n\r\n" } else { b"HTTP/1.1 200 Connection Established\r\n\r\n" }).await.unwrap(); + } + if refuse { return; } + let mut upstream = TcpStream::connect(origin).await.unwrap(); + let _ = tokio::io::copy_bidirectional(&mut stream, &mut upstream).await; + }); + } + result = connections.join_next(), if !connections.is_empty() => { result.unwrap().unwrap(); } + } + } + }); + (proxy, rx, task) +} + +fn test_builder(tls: &TestTls) -> reqwest::ClientBuilder { + super::super::transport_builder() + .add_root_certificate(tls.certificate.clone()) + .timeout(Duration::from_secs(3)) +} +fn endpoint() -> reqwest::Url { + super::super::parse_endpoint("https://byok.test/v1").unwrap() +} +fn target() -> SocketAddr { + "8.8.8.8:443".parse().unwrap() +} + +#[tokio::test] +async fn http_https_and_socks_proxies_preserve_tls_host_auth_and_streaming() { + for scheme in ["http", "https", "socks5", "socks5h"] { + let tls = test_tls(); + let (origin, observed, origin_task) = origin(&tls, b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nTransfer-Encoding: chunked\r\n\r\n9\r\ndata: a\n\n\r\n9\r\ndata: b\n\n\r\n0\r\n\r\n").await; + let (proxy, mut forwarded, proxy_task) = mock_proxy(scheme, &tls, origin, false).await; + let client = client(&endpoint(), &[target()], proxy, test_builder(&tls)) + .await + .unwrap(); + let response = client + .get("https://byok.test/v1") + .bearer_auth("provider-key") + .send() + .await + .unwrap(); + assert_eq!(response.status(), 200, "{scheme}"); + assert_eq!(response.text().await.unwrap(), "data: a\n\ndata: b\n\n"); + let (sni, headers) = observed.await.unwrap(); + assert_eq!(sni, "byok.test"); + assert!(headers.to_lowercase().contains("host: byok.test")); + assert!(headers.contains("provider-key")); + assert!(!headers.to_lowercase().contains("proxy-authorization")); + assert!(!headers.contains("test-password")); + assert!(forwarded.recv().await.is_some()); + drop(client); + origin_task.await.unwrap(); + proxy_task.abort(); + } +} + +#[tokio::test] +async fn proxy_cannot_expand_destination_or_follow_redirects() { + let tls = test_tls(); + let (origin, observed, origin_task) = origin( + &tls, + b"HTTP/1.1 302 Found\r\nLocation: https://private.test/\r\nContent-Length: 0\r\n\r\n", + ) + .await; + let (proxy, mut forwarded, proxy_task) = mock_proxy("http", &tls, origin, false).await; + let client = client(&endpoint(), &[target()], proxy, test_builder(&tls)) + .await + .unwrap(); + assert!(client.get("https://unadmitted.test/").send().await.is_err()); + assert!( + forwarded.try_recv().is_err(), + "wrong destinations must not reach the proxy" + ); + let response = client.get("https://byok.test/v1").send().await.unwrap(); + assert_eq!(response.status(), 302); + observed.await.unwrap(); + forwarded.recv().await.unwrap(); + assert!(forwarded.try_recv().is_err()); + drop(client); + origin_task.await.unwrap(); + proxy_task.abort(); +} + +#[tokio::test] +async fn proxy_failure_does_not_fall_back_to_direct_or_leak_credentials() { + for scheme in ["http", "socks5"] { + let tls = test_tls(); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let (proxy, mut forwarded, task) = + mock_proxy(scheme, &tls, listener.local_addr().unwrap(), true).await; + let client = client(&endpoint(), &[target()], proxy, test_builder(&tls)) + .await + .unwrap(); + let error = client + .get("https://byok.test/v1") + .bearer_auth("provider-key") + .send() + .await + .unwrap_err() + .to_string(); + assert!(!error.contains("provider-key")); + assert!(!error.contains("test-password")); + forwarded.recv().await.unwrap(); + assert!( + tokio::time::timeout(Duration::from_millis(30), listener.accept()) + .await + .is_err() + ); + drop(client); + task.abort(); + } +} + +#[tokio::test] +async fn proxy_never_disables_origin_certificate_validation() { + let tls = test_tls(); + let (origin, observed, origin_task) = + origin(&tls, b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n").await; + let (proxy, mut forwarded, proxy_task) = mock_proxy("http", &tls, origin, false).await; + // Do not add the fixture's certificate to the origin trust store. + let client = client( + &endpoint(), + &[target()], + proxy, + super::super::transport_builder().timeout(Duration::from_secs(3)), + ) + .await + .unwrap(); + assert!( + client + .get("https://byok.test/v1") + .bearer_auth("provider-key") + .send() + .await + .is_err() + ); + assert!( + tokio::time::timeout(Duration::from_secs(3), observed) + .await + .unwrap() + .is_err(), + "untrusted TLS must stop before HTTP/credentials" + ); + forwarded.recv().await.unwrap(); + drop(client); + origin_task.await.unwrap(); + proxy_task.abort(); +} + +#[tokio::test] +async fn https_proxy_certificate_is_checked_before_proxy_credentials_are_sent() { + let tls = test_tls(); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy = EgressProxy::parse(Some(&format!( + "https://user:proxy-secret@{}", + listener.local_addr().unwrap() + ))) + .unwrap() + .unwrap(); + let server_tls = tls.server.clone(); + let proxy_task = tokio::spawn(async move { + let (tcp, _) = listener.accept().await.unwrap(); + assert!( + tokio_rustls::TlsAcceptor::from(server_tls) + .accept(tcp) + .await + .is_err() + ); + }); + // Trust the origin fixture but not the HTTPS proxy certificate. + let guard = client(&endpoint(), &[target()], proxy, test_builder(&tls)) + .await + .unwrap(); + let error = guard + .get("https://byok.test/v1") + .bearer_auth("provider-key") + .send() + .await + .unwrap_err() + .to_string(); + assert!(!error.contains("proxy-secret")); + assert!(!error.contains("provider-key")); + tokio::time::timeout(Duration::from_secs(3), proxy_task) + .await + .unwrap() + .unwrap(); +} + +#[tokio::test] +async fn private_addresses_are_rejected_before_a_proxy_tunnel_is_created() { + let tls = test_tls(); + let proxy = EgressProxy::parse(Some("http://127.0.0.1:1")) + .unwrap() + .unwrap(); + for blocked in [ + "127.0.0.1:443", + "169.254.169.254:443", + "198.18.4.197:443", + "[::ffff:127.0.0.1]:443", + ] { + assert!( + client( + &endpoint(), + &[target(), blocked.parse().unwrap()], + proxy.clone(), + test_builder(&tls) + ) + .await + .is_err() + ); + } +} + +#[test] +fn proxy_configuration_errors_do_not_echo_secrets_or_allow_ambiguous_urls() { + for raw in [ + "proxy:8080", + "ftp://user:secret@proxy", + "http://user:secret@proxy/path", + "http://proxy?key=secret", + "http://user:secret@proxy:0", + "http://u%0d:p@proxy", + ] { + let error = EgressProxy::parse(Some(raw)).err().unwrap(); + assert!(!error.contains("secret")); + } + assert!(EgressProxy::parse(None).unwrap().is_none()); + assert!(EgressProxy::parse(Some("")).unwrap().is_none()); + for raw in [ + "http://localhost:8080", + "https://proxy.example", + "socks5://[::1]:1080", + "socks5h://proxy.example:1080", + ] { + assert!(EgressProxy::parse(Some(raw)).unwrap().is_some()); + } +} + +#[tokio::test] +async fn dropping_client_cancels_a_stalled_proxy_connection() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy = EgressProxy::parse(Some(&format!("http://{}", listener.local_addr().unwrap()))) + .unwrap() + .unwrap(); + let guard = client( + &endpoint(), + &[target()], + proxy, + super::super::transport_builder(), + ) + .await + .unwrap(); + let requester = guard.client.clone(); + let request = tokio::spawn(async move { requester.get("https://byok.test/v1").send().await }); + let (mut stream, _) = tokio::time::timeout(Duration::from_secs(3), listener.accept()) + .await + .unwrap() + .unwrap(); + read_head(&mut stream).await.unwrap(); + drop(guard); + let mut byte = [0]; + assert_eq!( + tokio::time::timeout(Duration::from_secs(3), stream.read(&mut byte)) + .await + .unwrap() + .unwrap(), + 0 + ); + assert!( + tokio::time::timeout(Duration::from_secs(3), request) + .await + .unwrap() + .unwrap() + .is_err() + ); +} + +#[tokio::test] +async fn stalled_proxy_obeys_request_timeout() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy = EgressProxy::parse(Some(&format!("http://{}", listener.local_addr().unwrap()))) + .unwrap() + .unwrap(); + let guard = client( + &endpoint(), + &[target()], + proxy, + super::super::transport_builder().timeout(Duration::from_millis(100)), + ) + .await + .unwrap(); + let error = tokio::time::timeout( + Duration::from_secs(3), + guard.get("https://byok.test/v1").send(), + ) + .await + .unwrap() + .unwrap_err(); + assert!(error.is_timeout()); +} diff --git a/crates/services/src/lib.rs b/crates/services/src/lib.rs index e7128c4aa5..2c7005b11b 100644 --- a/crates/services/src/lib.rs +++ b/crates/services/src/lib.rs @@ -6,6 +6,7 @@ pub mod agents; pub mod artifact_policy; pub mod auth; pub mod branches; +pub mod byok_endpoint; pub mod config_version_cloud; pub mod context; pub mod context_manifest; @@ -261,11 +262,13 @@ pub use models::{ ModelUpdateRequestData, PricingData, PromptCacheCapabilityData, PromptCacheProtocolData, PromptCacheReuseScopeData, PromptCacheVolatileDeliveryData, PromptCacheVolatilePlacementData, QuirksData, ResolvedActiveLlmModel, ResolvedModelOffering, UnconfiguredModelService, + UserModelCreateRequestData, UserModelRecord, UserModelUpdateRequestData, model_catalog_revision, project_model_access, project_model_access_page, project_model_access_page_with_default_catalog, project_model_access_with_default, prompt_cache_capability_from_models_yaml, resolve_active_llm_model, resolve_active_llm_offering, resolve_memory_offerings, resolve_reasoning_offering, - revalidate_active_llm_offering, validate_model_offering_id, + revalidate_active_llm_offering, revalidate_admitted_model_execution, + validate_model_offering_id, }; pub use multi_agent::{ DatabaseEdgeDispatchService, DatabaseEdgeRegistryService, EdgeAgentRecord, diff --git a/crates/services/src/models.rs b/crates/services/src/models.rs index e14de33e23..db4cd505db 100644 --- a/crates/services/src/models.rs +++ b/crates/services/src/models.rs @@ -3,7 +3,7 @@ use axum::{Json, http::StatusCode}; use serde::{Deserialize, Deserializer, Serialize, Serializer, ser::SerializeStruct}; use serde_json::{Map, Value}; use sha2::{Digest, Sha256}; -use sqlx::{Row, query}; +use sqlx::{Row, query, query_scalar}; use std::{ collections::{BTreeMap, BTreeSet, HashMap}, path::{Path, PathBuf}, @@ -636,6 +636,7 @@ pub struct ResolvedModelOffering { #[serde(rename_all = "snake_case")] pub enum ModelAccessKind { AstraCloud, + CloudByok, Workspace, ThisDevice, SelfHosted, @@ -646,6 +647,7 @@ impl ModelAccessKind { pub fn as_str(self) -> &'static str { match self { Self::AstraCloud => "astra_cloud", + Self::CloudByok => "cloud_byok", Self::Workspace => "workspace", Self::ThisDevice => "this_device", Self::SelfHosted => "self_hosted", @@ -1569,6 +1571,153 @@ pub async fn revalidate_active_llm_offering( } } +/// Revalidate execution material for an authenticated user's effective +/// catalog. Personal Cloud BYOK Offerings are owner-scoped; deployment +/// Offerings retain the existing global catalog behavior. +pub async fn revalidate_admitted_model_execution( + matrixone: &MatrixOneSettings, + encryptor: &FernetTokenEncryptor, + user_id: &str, + offering_id: &str, + pool: Option<&sqlx::Pool>, +) -> Result { + let offering_id = validate_model_offering_id(offering_id)?; + let pool = require_pool(pool, matrixone) + .await + .map_err(ModelOfferingResolutionError::Backend)?; + let row = query( + "SELECT model_alias, model_name, provider, api_key_encrypted, base_url, \ + context_window, is_active FROM user_llm_models \ + WHERE user_id = ? AND model_id = ? LIMIT 1", + ) + .bind(user_id) + .bind(offering_id) + .fetch_optional(&pool) + .await + .map_err(|error| ModelOfferingResolutionError::Backend(format!("DB query: {error}")))?; + + if let Some(row) = row { + let alias: String = row.try_get("model_alias").map_err(|error| { + ModelOfferingResolutionError::Backend(format!( + "invalid user_llm_models.model_alias: {error}" + )) + })?; + let is_active: i16 = row.try_get("is_active").map_err(|error| { + ModelOfferingResolutionError::Backend(format!( + "invalid user_llm_models.is_active: {error}" + )) + })?; + if is_active == 0 { + return Err(ModelOfferingResolutionError::Inactive { + offering_id: offering_id.to_string(), + model_name: alias, + }); + } + let encrypted: String = row.try_get("api_key_encrypted").map_err(|error| { + ModelOfferingResolutionError::Backend(format!( + "invalid user_llm_models.api_key_encrypted: {error}" + )) + })?; + let api_key = encryptor + .decrypt(&encrypted) + .map_err(ModelOfferingResolutionError::Backend)?; + let context_window: i32 = row.try_get("context_window").map_err(|error| { + ModelOfferingResolutionError::Backend(format!( + "invalid user_llm_models.context_window: {error}" + )) + })?; + let context_window = u32::try_from(context_window).map_err(|_| { + ModelOfferingResolutionError::Backend( + "invalid user_llm_models.context_window: must be positive".to_string(), + ) + })?; + let provider: String = row + .try_get("provider") + .map_err(|error| ModelOfferingResolutionError::Backend(error.to_string()))?; + let base_url: String = row + .try_get("base_url") + .map_err(|error| ModelOfferingResolutionError::Backend(error.to_string()))?; + if provider == crate::byok_endpoint::COMPATIBLE_PROVIDER { + crate::byok_endpoint::require_endpoint_policy(&pool, &base_url) + .await + .map_err(ModelOfferingResolutionError::Backend)?; + } + return Ok(AdmittedModelExecution { + offering_id: offering_id.to_string(), + access_kind: ModelAccessKind::CloudByok, + execution_placement: ModelExecutionPlacement::Server, + model_name: alias, + wire_model_name: Some(row.try_get("model_name").map_err(|error| { + ModelOfferingResolutionError::Backend(format!( + "invalid user_llm_models.model_name: {error}" + )) + })?), + api_key, + base_url: row.try_get("base_url").map_err(|error| { + ModelOfferingResolutionError::Backend(format!( + "invalid user_llm_models.base_url: {error}" + )) + })?, + provider: row.try_get("provider").map_err(|error| { + ModelOfferingResolutionError::Backend(format!( + "invalid user_llm_models.provider: {error}" + )) + })?, + cache_capability: None, + thinking_capability: None, + request_body_overrides: None, + context_window: Some(context_window), + max_completion_tokens: None, + header_overrides: HashMap::new(), + completions_url_override: None, + request_timeout_ms: None, + }); + } + + if !deployment_models_allowed(&pool, user_id) + .await + .map_err(ModelOfferingResolutionError::Backend)? + { + return Err(ModelOfferingResolutionError::NotFound { + offering_id: offering_id.to_string(), + }); + } + let offering = + revalidate_active_llm_offering(matrixone, encryptor, offering_id, Some(&pool)).await?; + AdmittedModelExecution::from_offering(offering).map_err(ModelOfferingResolutionError::Backend) +} + +/// Shared eligibility gate for catalog and execution, including resumed runs. +/// Memoria identities are personal BYOK even on a self-hosted deployment. +async fn deployment_models_allowed(pool: &sqlx::MySqlPool, user_id: &str) -> Result { + let mode = std::env::var("ASTRA_DEPLOYMENT_MODE") + .map(Some) + .or_else(|error| match error { + std::env::VarError::NotPresent => Ok(None), + _ => Err("Invalid ASTRA_DEPLOYMENT_MODE".to_string()), + })?; + if !deployment_mode_allows_shared_models(mode.as_deref())? { + return Ok(false); + } + let mapped: Option = sqlx::query_scalar( + "SELECT external_subject FROM auth_external_identities WHERE astra_user_id = ? AND provider_id LIKE 'memoria:%' UNION ALL SELECT memoria_user_id FROM auth_memoria_identities WHERE astra_user_id = ? LIMIT 1", + ) + .bind(user_id) + .bind(user_id) + .fetch_optional(pool) + .await + .map_err(|error| format!("Model ownership lookup failed: {error}"))?; + Ok(mapped.is_none()) +} + +fn deployment_mode_allows_shared_models(mode: Option<&str>) -> Result { + match mode { + None | Some("self-hosted") => Ok(true), + Some("cloud-byok") => Ok(false), + _ => Err("Invalid ASTRA_DEPLOYMENT_MODE; expected cloud-byok or self-hosted".into()), + } +} + async fn resolve_active_llm_offering_uncached( encryptor: &FernetTokenEncryptor, offering_id: &str, @@ -1843,10 +1992,18 @@ fn rank_memory_model_candidate_indices( pub async fn resolve_memory_offerings( matrixone: &MatrixOneSettings, encryptor: &FernetTokenEncryptor, + user_id: &str, pool: Option<&sqlx::Pool>, ) -> Result, String> { let pool = require_pool(pool, matrixone).await?; + // A memory-write grant is not authorization to spend deployment credentials. + // Personal BYOK has no implicit background selector binding: use the + // existing deterministic extraction path until one is explicitly admitted. + if !deployment_models_allowed(&pool, user_id).await? { + return Ok(Vec::new()); + } + let rows = sqlx::query(&format!( "SELECT model_id, {RESOLVE_COLS} FROM infra_llm_models WHERE is_active = 1" )) @@ -2054,10 +2211,157 @@ fn model_list_page_from_items( }) } +/// User-owned Server-side BYOK model configuration. +/// +/// The credential is deliberately absent. It is accepted only on create or +/// rotation and is never projected back through the API. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct UserModelRecord { + pub model_id: String, + pub name: String, + pub provider: String, + pub model: String, + pub base_url: String, + pub context_window: i32, + pub is_default: bool, + pub is_active: bool, + pub credential_configured: bool, + pub created_at: String, + pub updated_at: String, +} + +#[derive(Clone, PartialEq)] +pub struct UserModelCreateRequestData { + pub name: String, + pub provider: String, + pub model: String, + pub base_url: Option, + pub api_key: String, + pub context_window: i32, + pub is_default: bool, +} + +impl std::fmt::Debug for UserModelCreateRequestData { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("UserModelCreateRequestData") + .field("name", &self.name) + .field("provider", &self.provider) + .field("model", &self.model) + .field("api_key", &"") + .field("context_window", &self.context_window) + .field("is_default", &self.is_default) + .finish() + } +} + +#[derive(Clone, Default, PartialEq)] +pub struct UserModelUpdateRequestData { + pub api_key: Option, + pub context_window: Option, + pub is_default: Option, + pub is_active: Option, +} + +impl std::fmt::Debug for UserModelUpdateRequestData { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("UserModelUpdateRequestData") + .field("api_key", &self.api_key.as_ref().map(|_| "")) + .field("context_window", &self.context_window) + .field("is_default", &self.is_default) + .field("is_active", &self.is_active) + .finish() + } +} + // ── Trait ───────────────────────────────────────────────────────────────────── #[async_trait] pub trait ModelService: Send + Sync { + /// Credential-free preflight. This is not authorization for later requests. + async fn validate_user_model_endpoint( + &self, + _user_id: String, + _base_url: String, + ) -> Result<(), (StatusCode, Json)> { + Err(internal_error("user model service not configured")) + } + + async fn create_user_model( + &self, + _user_id: String, + _request: UserModelCreateRequestData, + ) -> Result)> { + Err(internal_error("user model service not configured")) + } + + async fn list_user_models( + &self, + _user_id: String, + ) -> Result, (StatusCode, Json)> { + Ok(Vec::new()) + } + + async fn get_user_model( + &self, + _user_id: String, + _model_id: String, + ) -> Result)> { + Err(internal_error("user model service not configured")) + } + + async fn update_user_model( + &self, + _user_id: String, + _model_id: String, + _request: UserModelUpdateRequestData, + ) -> Result)> { + Err(internal_error("user model service not configured")) + } + + async fn delete_user_model( + &self, + _user_id: String, + _model_id: String, + ) -> Result<(), (StatusCode, Json)> { + Err(internal_error("user model service not configured")) + } + + async fn check_user_model( + &self, + _user_id: String, + _model_id: String, + ) -> Result)> { + Err(internal_error("user model service not configured")) + } + + async fn default_user_model_offering_id( + &self, + _user_id: String, + ) -> Result, (StatusCode, Json)> { + Ok(None) + } + + /// Revalidate and materialize an Offering for one authenticated user. + async fn allows_deployment_models( + &self, + _user_id: String, + ) -> Result)> { + Ok(true) + } + + /// Revalidate and materialize an Offering for one authenticated user. + /// The default implementation preserves the deployment catalog behavior; + /// database-backed services additionally resolve user-owned BYOK rows. + async fn admit_model_offering( + &self, + _user_id: String, + offering_id: String, + ) -> Result)> { + let offering = self.revalidate_model_offering(offering_id).await?; + AdmittedModelExecution::from_offering(offering).map_err(internal_error) + } + async fn create_model( &self, user_id: String, @@ -2299,29 +2603,416 @@ impl DatabaseModelService { thinking_capability, }) } -} -pub const MODEL_SELECT_COLS: &str = "\ - model_id, model_name, provider, base_url, description, is_active, \ - context_window, max_completion_tokens, architecture, \ - CAST(input_modalities AS CHAR) AS input_modalities_json, \ - CAST(output_modalities AS CHAR) AS output_modalities_json, \ - CAST(supported_parameters AS CHAR) AS supported_parameters_json, \ - CAST(pricing AS CHAR) AS pricing_json, \ - CAST(tags AS CHAR) AS tags_json, \ - CAST(quirks AS CHAR) AS quirks_json, \ - thinking_capability, thinking_probe_error"; -const MODEL_LIST_SELECT_COLS: &str = "\ - model_id, model_name, provider, description, is_active, \ - context_window, max_completion_tokens, architecture, \ - thinking_capability"; -const MODEL_LIST_CURSOR_SQL: &str = " AND (provider > ? \ - OR (provider = ? AND model_name > ?) \ - OR (provider = ? AND model_name = ? AND model_id > ?))"; -const MODEL_LIST_ORDER_SQL: &str = " ORDER BY provider ASC, model_name ASC, model_id ASC LIMIT ?"; + fn user_model_record_from_row( + row: &sqlx::mysql::MySqlRow, + ) -> Result)> { + let context_window: i32 = row.try_get("context_window").map_err(internal_error)?; + validate_context_window_value(context_window, "user model context_window") + .map_err(internal_error)?; + let is_default: i16 = row.try_get("is_default").map_err(internal_error)?; + let is_active: i16 = row.try_get("is_active").map_err(internal_error)?; + Ok(UserModelRecord { + model_id: row.try_get("model_id").map_err(internal_error)?, + name: row.try_get("model_alias").map_err(internal_error)?, + provider: row.try_get("provider").map_err(internal_error)?, + model: row.try_get("model_name").map_err(internal_error)?, + base_url: row.try_get("base_url").map_err(internal_error)?, + context_window, + is_default: is_default != 0, + is_active: is_active != 0, + credential_configured: true, + created_at: row.try_get("created_at_text").map_err(internal_error)?, + updated_at: row.try_get("updated_at_text").map_err(internal_error)?, + }) + } + + async fn user_model_row( + &self, + user_id: &str, + model_id: &str, + ) -> Result, (StatusCode, Json)> { + let pool = self.get_pool().await.map_err(internal_error)?; + query(&format!( + "SELECT {USER_MODEL_SELECT_COLS} FROM user_llm_models WHERE user_id = ? AND model_id = ?" + )) + .bind(user_id) + .bind(model_id) + .fetch_optional(&pool) + .await + .map_err(internal_error) + } +} + +pub const MODEL_SELECT_COLS: &str = "\ + model_id, model_name, provider, base_url, description, is_active, \ + context_window, max_completion_tokens, architecture, \ + CAST(input_modalities AS CHAR) AS input_modalities_json, \ + CAST(output_modalities AS CHAR) AS output_modalities_json, \ + CAST(supported_parameters AS CHAR) AS supported_parameters_json, \ + CAST(pricing AS CHAR) AS pricing_json, \ + CAST(tags AS CHAR) AS tags_json, \ + CAST(quirks AS CHAR) AS quirks_json, \ + thinking_capability, thinking_probe_error"; +const MODEL_LIST_SELECT_COLS: &str = "\ + model_id, model_name, provider, description, is_active, \ + context_window, max_completion_tokens, architecture, \ + thinking_capability"; +const MODEL_LIST_CURSOR_SQL: &str = " AND (provider > ? \ + OR (provider = ? AND model_name > ?) \ + OR (provider = ? AND model_name = ? AND model_id > ?))"; +const MODEL_LIST_ORDER_SQL: &str = " ORDER BY provider ASC, model_name ASC, model_id ASC LIMIT ?"; +const USER_MODEL_SELECT_COLS: &str = "model_id, model_alias, model_name, provider, base_url, \ + context_window, is_default, is_active, \ + CAST(created_at AS CHAR) AS created_at_text, CAST(updated_at AS CHAR) AS updated_at_text"; + +#[async_trait] +impl ModelService for DatabaseModelService { + async fn allows_deployment_models( + &self, + user_id: String, + ) -> Result)> { + let pool = self.get_pool().await.map_err(internal_error)?; + deployment_models_allowed(&pool, &user_id) + .await + .map_err(internal_error) + } + async fn validate_user_model_endpoint( + &self, + _user_id: String, + base_url: String, + ) -> Result<(), (StatusCode, Json)> { + let pool = self.get_pool().await.map_err(internal_error)?; + crate::byok_endpoint::require_endpoint_policy(&pool, &base_url) + .await + .map_err(|error| error_response(StatusCode::BAD_REQUEST, error))?; + // Validate DNS without sending an HTTP request or accepting any secret. + crate::byok_endpoint::endpoint_client(&base_url) + .await + .map_err(|error| { + astra_core::error_response_coded( + StatusCode::BAD_GATEWAY, + error, + "model_endpoint_network", + ) + })?; + Ok(()) + } + + async fn create_user_model( + &self, + user_id: String, + request: UserModelCreateRequestData, + ) -> Result)> { + let name = validate_user_model_identifier("name", &request.name)?; + let model = validate_user_model_identifier("model", &request.model)?; + let provider = request.provider.trim().to_ascii_lowercase(); + let base_url = user_byok_base_url(&provider, request.base_url.as_deref())?; + validate_context_window_value(request.context_window, "context_window") + .map_err(|error| error_response(StatusCode::BAD_REQUEST, error))?; + validate_user_model_api_key(&request.api_key)?; + + let pool = self.get_pool().await.map_err(internal_error)?; + let duplicate = + query("SELECT 1 FROM user_llm_models WHERE user_id = ? AND model_alias = ? LIMIT 1") + .bind(&user_id) + .bind(&name) + .fetch_optional(&pool) + .await + .map_err(internal_error)?; + if duplicate.is_some() { + return Err(error_response( + StatusCode::CONFLICT, + format!("User model '{name}' already exists"), + )); + } + + if provider == crate::byok_endpoint::COMPATIBLE_PROVIDER { + crate::byok_endpoint::require_endpoint_policy(&pool, &base_url) + .await + .map_err(|error| error_response(StatusCode::BAD_REQUEST, error))?; + } + + let encrypted_key = self + .encryptor + .encrypt(&request.api_key) + .map_err(internal_error)?; + let connectivity = validate_connectivity( + &provider, + &model, + &request.api_key, + Some(&base_url), + None, + None, + ) + .await; + if let Some(reason) = connectivity { + return Err(error_response( + StatusCode::BAD_REQUEST, + format!("Model credential or endpoint check failed: {reason}"), + )); + } + + let model_id = Uuid::new_v4().to_string(); + let mut tx = pool.begin().await.map_err(internal_error)?; + if request.is_default { + query( + "UPDATE user_llm_models SET is_default = 0, updated_at = NOW(6) WHERE user_id = ?", + ) + .bind(&user_id) + .execute(&mut *tx) + .await + .map_err(internal_error)?; + } + query( + "INSERT INTO user_llm_models \ + (model_id, user_id, model_alias, model_name, provider, api_key_encrypted, base_url, \ + context_window, is_default, is_active, created_at, updated_at) \ + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, 1, NOW(6), NOW(6))", + ) + .bind(&model_id) + .bind(&user_id) + .bind(&name) + .bind(&model) + .bind(&provider) + .bind(&encrypted_key) + .bind(&base_url) + .bind(request.context_window) + .bind(if request.is_default { 1_i16 } else { 0_i16 }) + .execute(&mut *tx) + .await + .map_err(internal_error)?; + tx.commit().await.map_err(internal_error)?; + + self.get_user_model(user_id, model_id).await + } + + async fn list_user_models( + &self, + user_id: String, + ) -> Result, (StatusCode, Json)> { + let pool = self.get_pool().await.map_err(internal_error)?; + let rows = query(&format!( + "SELECT {USER_MODEL_SELECT_COLS} FROM user_llm_models \ + WHERE user_id = ? ORDER BY is_default DESC, model_alias ASC, model_id ASC" + )) + .bind(user_id) + .fetch_all(&pool) + .await + .map_err(internal_error)?; + rows.iter().map(Self::user_model_record_from_row).collect() + } + + async fn get_user_model( + &self, + user_id: String, + model_id: String, + ) -> Result)> { + let row = self + .user_model_row(&user_id, &model_id) + .await? + .ok_or_else(user_model_not_found)?; + Self::user_model_record_from_row(&row) + } + + async fn update_user_model( + &self, + user_id: String, + model_id: String, + request: UserModelUpdateRequestData, + ) -> Result)> { + if request.api_key.is_none() + && request.context_window.is_none() + && request.is_default.is_none() + && request.is_active.is_none() + { + return Err(error_response( + StatusCode::BAD_REQUEST, + "At least one user model field must be supplied", + )); + } + if let Some(context_window) = request.context_window { + validate_context_window_value(context_window, "context_window") + .map_err(|error| error_response(StatusCode::BAD_REQUEST, error))?; + } + if request.is_default == Some(true) && request.is_active == Some(false) { + return Err(error_response( + StatusCode::BAD_REQUEST, + "A disabled user model cannot be the default", + )); + } + + let existing = self + .user_model_row(&user_id, &model_id) + .await? + .ok_or_else(user_model_not_found)?; + let provider: String = existing.try_get("provider").map_err(internal_error)?; + let model: String = existing.try_get("model_name").map_err(internal_error)?; + let base_url: String = existing.try_get("base_url").map_err(internal_error)?; + let pool = self.get_pool().await.map_err(internal_error)?; + if provider == crate::byok_endpoint::COMPATIBLE_PROVIDER + && (request.api_key.is_some() + || request.is_active == Some(true) + || request.is_default == Some(true)) + { + crate::byok_endpoint::require_endpoint_policy(&pool, &base_url) + .await + .map_err(|error| error_response(StatusCode::BAD_REQUEST, error))?; + } + let encrypted_key = if let Some(api_key) = request.api_key.as_deref() { + validate_user_model_api_key(api_key)?; + if let Some(reason) = + validate_connectivity(&provider, &model, api_key, Some(&base_url), None, None).await + { + return Err(error_response( + StatusCode::BAD_REQUEST, + format!("Model credential or endpoint check failed: {reason}"), + )); + } + Some(self.encryptor.encrypt(api_key).map_err(internal_error)?) + } else { + None + }; + + let mut tx = pool.begin().await.map_err(internal_error)?; + if request.is_default == Some(true) { + query( + "UPDATE user_llm_models SET is_default = 0, updated_at = NOW(6) WHERE user_id = ?", + ) + .bind(&user_id) + .execute(&mut *tx) + .await + .map_err(internal_error)?; + } + if let Some(encrypted_key) = encrypted_key { + query("UPDATE user_llm_models SET api_key_encrypted = ?, updated_at = NOW(6) WHERE user_id = ? AND model_id = ?") + .bind(encrypted_key) + .bind(&user_id) + .bind(&model_id) + .execute(&mut *tx) + .await + .map_err(internal_error)?; + } + if let Some(context_window) = request.context_window { + query("UPDATE user_llm_models SET context_window = ?, updated_at = NOW(6) WHERE user_id = ? AND model_id = ?") + .bind(context_window) + .bind(&user_id) + .bind(&model_id) + .execute(&mut *tx) + .await + .map_err(internal_error)?; + } + if let Some(is_active) = request.is_active { + query("UPDATE user_llm_models SET is_active = ?, is_default = CASE WHEN ? = 0 THEN 0 ELSE is_default END, updated_at = NOW(6) WHERE user_id = ? AND model_id = ?") + .bind(if is_active { 1_i16 } else { 0_i16 }) + .bind(if is_active { 1_i16 } else { 0_i16 }) + .bind(&user_id) + .bind(&model_id) + .execute(&mut *tx) + .await + .map_err(internal_error)?; + } + if let Some(is_default) = request.is_default { + query("UPDATE user_llm_models SET is_default = ?, is_active = CASE WHEN ? = 1 THEN 1 ELSE is_active END, updated_at = NOW(6) WHERE user_id = ? AND model_id = ?") + .bind(if is_default { 1_i16 } else { 0_i16 }) + .bind(if is_default { 1_i16 } else { 0_i16 }) + .bind(&user_id) + .bind(&model_id) + .execute(&mut *tx) + .await + .map_err(internal_error)?; + } + tx.commit().await.map_err(internal_error)?; + self.get_user_model(user_id, model_id).await + } + + async fn delete_user_model( + &self, + user_id: String, + model_id: String, + ) -> Result<(), (StatusCode, Json)> { + let pool = self.get_pool().await.map_err(internal_error)?; + let result = query("DELETE FROM user_llm_models WHERE user_id = ? AND model_id = ?") + .bind(user_id) + .bind(model_id) + .execute(&pool) + .await + .map_err(internal_error)?; + if result.rows_affected() == 0 { + return Err(user_model_not_found()); + } + Ok(()) + } + + async fn check_user_model( + &self, + user_id: String, + model_id: String, + ) -> Result)> { + let pool = self.get_pool().await.map_err(internal_error)?; + let row = query( + "SELECT model_name, provider, api_key_encrypted, base_url \ + FROM user_llm_models WHERE user_id = ? AND model_id = ?", + ) + .bind(&user_id) + .bind(&model_id) + .fetch_optional(&pool) + .await + .map_err(internal_error)? + .ok_or_else(user_model_not_found)?; + let model: String = row.try_get("model_name").map_err(internal_error)?; + let provider: String = row.try_get("provider").map_err(internal_error)?; + let encrypted: String = row.try_get("api_key_encrypted").map_err(internal_error)?; + let base_url: String = row.try_get("base_url").map_err(internal_error)?; + if provider == crate::byok_endpoint::COMPATIBLE_PROVIDER { + crate::byok_endpoint::require_endpoint_policy(&pool, &base_url) + .await + .map_err(|error| error_response(StatusCode::BAD_REQUEST, error))?; + } + let api_key = self.encryptor.decrypt(&encrypted).map_err(internal_error)?; + if let Some(reason) = + validate_connectivity(&provider, &model, &api_key, Some(&base_url), None, None).await + { + return Err(error_response( + StatusCode::BAD_GATEWAY, + format!("Model credential or endpoint check failed: {reason}"), + )); + } + self.get_user_model(user_id, model_id).await + } + + async fn default_user_model_offering_id( + &self, + user_id: String, + ) -> Result, (StatusCode, Json)> { + let pool = self.get_pool().await.map_err(internal_error)?; + query_scalar( + "SELECT model_id FROM user_llm_models \ + WHERE user_id = ? AND is_default = 1 AND is_active = 1 \ + ORDER BY updated_at DESC, model_id ASC LIMIT 1", + ) + .bind(user_id) + .fetch_optional(&pool) + .await + .map_err(internal_error) + } + + async fn admit_model_offering( + &self, + user_id: String, + offering_id: String, + ) -> Result)> { + revalidate_admitted_model_execution( + &self.matrixone, + self.encryptor.as_ref(), + &user_id, + &offering_id, + self.pool.as_ref().map(SharedPool::get), + ) + .await + .map_err(model_offering_resolution_error_response) + } -#[async_trait] -impl ModelService for DatabaseModelService { async fn create_model( &self, user_id: String, @@ -2450,7 +3141,7 @@ impl ModelService for DatabaseModelService { async fn list_models( &self, - _user_id: String, + user_id: String, is_admin: bool, ) -> Result, (StatusCode, Json)> { let pool = self.get_pool().await.map_err(internal_error)?; @@ -2466,7 +3157,11 @@ impl ModelService for DatabaseModelService { MODEL_LIST_SELECT_COLS ) }; - let rows = query(&sql).fetch_all(&pool).await.map_err(internal_error)?; + let rows = if is_admin || self.allows_deployment_models(user_id.clone()).await? { + query(&sql).fetch_all(&pool).await.map_err(internal_error)? + } else { + Vec::new() + }; let mut models = Vec::with_capacity(rows.len()); for row in rows { @@ -2498,16 +3193,50 @@ impl ModelService for DatabaseModelService { }, }); } + if !is_admin && !user_id.is_empty() { + let user_rows = query( + "SELECT model_id, model_alias, provider, context_window, is_active \ + FROM user_llm_models WHERE user_id = ? AND is_active = 1 \ + ORDER BY provider, model_alias, model_id", + ) + .bind(&user_id) + .fetch_all(&pool) + .await + .map_err(internal_error)?; + for row in user_rows { + let is_active: i16 = row.try_get("is_active").map_err(internal_error)?; + models.push(ModelListItem { + offering_id: row.try_get("model_id").map_err(internal_error)?, + access_id: "cloud-byok".to_string(), + access_kind: ModelAccessKind::CloudByok, + access_label: "Cloud BYOK".to_string(), + execution_placement: ModelExecutionPlacement::Server, + name: row.try_get("model_alias").map_err(internal_error)?, + provider: row.try_get("provider").map_err(internal_error)?, + description: Some("Personal BYOK model".to_string()), + is_active: is_active != 0, + context_window: row.try_get("context_window").map_err(internal_error)?, + max_completion_tokens: None, + architecture: None, + thinking_capability: None, + }); + } + } + sort_model_list_items(&mut models); Ok(models) } async fn list_models_page( &self, - _user_id: String, + user_id: String, is_admin: bool, limit: u32, cursor: Option, ) -> Result)> { + if !is_admin && !user_id.is_empty() { + let items = self.list_models(user_id, false).await?; + return model_list_page_from_items(items, limit, cursor); + } let limit = validate_model_list_limit(limit); let cursor = cursor .as_ref() @@ -2571,9 +3300,14 @@ impl ModelService for DatabaseModelService { async fn model_catalog_revision( &self, - _user_id: String, + user_id: String, is_admin: bool, ) -> Result)> { + if !is_admin && !user_id.is_empty() { + return Ok(model_catalog_revision( + &self.list_models(user_id, false).await?, + )); + } let pool = self.get_pool().await.map_err(internal_error)?; let visibility = if is_admin { "1 = 1" } else { "is_active = 1" }; let fingerprint_sql = format!( @@ -2969,6 +3703,111 @@ pub fn resolve_provider_base_url(provider: &str) -> Option { } } +fn user_byok_base_url( + provider: &str, + requested: Option<&str>, +) -> Result)> { + match (provider, requested) { + (crate::byok_endpoint::COMPATIBLE_PROVIDER, Some(raw)) => { + Ok(crate::byok_endpoint::parse_endpoint(raw) + .map_err(|error| error_response(StatusCode::BAD_REQUEST, error))? + .as_str() + .trim_end_matches('/') + .to_string()) + } + (crate::byok_endpoint::COMPATIBLE_PROVIDER, None) => Err(error_response( + StatusCode::BAD_REQUEST, + "base_url is required for OpenAI-compatible", + )), + (_, Some(_)) => Err(error_response( + StatusCode::BAD_REQUEST, + "Use provider openai-compatible to configure a custom base_url", + )), + (_, None) => user_byok_provider_base_url(provider), + } +} + +fn user_byok_provider_base_url( + provider: &str, +) -> Result)> { + if std::env::var("ASTRA_ALLOW_INSECURE_DEFAULTS").as_deref() == Ok("1") + && provider == "deepseek" + && let Ok(base_url) = std::env::var("ASTRA_BYOK_DEEPSEEK_BASE_URL") + { + let parsed = reqwest::Url::parse(&base_url).map_err(|_| { + error_response( + StatusCode::INTERNAL_SERVER_ERROR, + "ASTRA_BYOK_DEEPSEEK_BASE_URL is invalid", + ) + })?; + let loopback = parsed.host_str().is_some_and(|host| { + host == "localhost" + || host + .parse::() + .is_ok_and(|ip| ip.is_loopback()) + }); + if parsed.scheme() != "http" + || !loopback + || !parsed.username().is_empty() + || parsed.password().is_some() + || parsed.query().is_some() + || parsed.fragment().is_some() + { + return Err(error_response( + StatusCode::INTERNAL_SERVER_ERROR, + "ASTRA_BYOK_DEEPSEEK_BASE_URL must be an unauthenticated loopback HTTP origin", + )); + } + return Ok(base_url.trim_end_matches('/').to_string()); + } + let base_url = match provider { + "openai" => "https://api.openai.com/v1", + "anthropic" => "https://api.anthropic.com", + "deepseek" => "https://api.deepseek.com", + _ => { + return Err(error_response( + StatusCode::BAD_REQUEST, + format!( + "Unsupported BYOK provider '{provider}'. Supported providers: openai, anthropic, deepseek, openai-compatible" + ), + )); + } + }; + Ok(base_url.to_string()) +} + +fn validate_user_model_identifier( + field: &str, + value: &str, +) -> Result)> { + let value = value.trim(); + if value.is_empty() || value.len() > 100 || value.chars().any(char::is_control) { + return Err(error_response( + StatusCode::BAD_REQUEST, + format!("{field} must contain 1 to 100 non-control characters"), + )); + } + Ok(value.to_string()) +} + +fn validate_user_model_api_key(api_key: &str) -> Result<(), (StatusCode, Json)> { + if api_key.is_empty() + || api_key.len() > 8_192 + || api_key.trim() != api_key + || api_key.chars().any(char::is_control) + { + return Err(error_response( + StatusCode::BAD_REQUEST, + "api_key must be a non-empty credential without surrounding whitespace or control characters", + )); + } + Ok(()) +} + +fn user_model_not_found() -> (StatusCode, Json) { + error_response(StatusCode::NOT_FOUND, "User model not found") +} + /// Full URL for a minimal Anthropic Messages API probe (`POST`, JSON body). /// /// - Official Anthropic: `https://api.anthropic.com` → `.../v1/messages`. @@ -3015,13 +3854,21 @@ pub async fn validate_connectivity( // Connectivity probes reach external provider endpoints (Anthropic, Bedrock, // OpenAI-compatible base_urls) — same class of traffic as the LLM client. - // They are NOT "internal connections" in the sense of 3e3d6fa8, so they - // share the LLM client's proxy policy via the single authoritative - // implementation in `astra_core::net::apply_env_proxy`. + // Native providers share the LLM client's ambient proxy policy. User-owned + // OpenAI-compatible endpoints instead use the pinned BYOK egress transport + // below, matching their runtime path without delegating destination DNS. let probe_builder = astra_core::net::apply_env_proxy( reqwest::Client::builder().timeout(std::time::Duration::from_secs(15)), ); - let client = match probe_builder.build() { + let client_result = if provider == crate::byok_endpoint::COMPATIBLE_PROVIDER { + crate::byok_endpoint::endpoint_client(base_url.unwrap_or_default()).await + } else { + probe_builder + .build() + .map(crate::byok_endpoint::EndpointClient::from) + .map_err(|e| e.to_string()) + }; + let client = match client_result { Ok(c) => c, Err(e) => return Some(format!("Client error: {}", e)), }; @@ -3040,14 +3887,12 @@ pub async fn validate_connectivity( } } } - let send_result = req - .json(&serde_json::json!({ - "model": model_name, - "max_tokens": 1, - "messages": [{"role": "user", "content": "hi"}] - })) - .send() - .await; + let mut body = serde_json::json!({ + "model": model_name, + "messages": [{"role": "user", "content": "hi"}] + }); + astra_core::model_wire::apply_chat_output_token_limit(&mut body, provider, 32); + let send_result = req.json(&body).send().await; (send_result, probe) } else if provider == "bedrock" { let Some(base_url_value) = base_url.map(str::trim).filter(|url| !url.is_empty()) else { @@ -3115,21 +3960,26 @@ pub async fn validate_connectivity( } } } - let send_result = req - .json(&serde_json::json!({ - "model": model_name, - "max_tokens": 1, - "messages": [{"role": "user", "content": "hi"}] - })) - .send() - .await; + let mut body = serde_json::json!({ + "model": model_name, + "messages": [{"role": "user", "content": "hi"}] + }); + astra_core::model_wire::apply_chat_output_token_limit(&mut body, provider, 32); + let send_result = req.json(&body).send().await; (send_result, probe) }; match result { - (Ok(resp), _) if resp.status().as_u16() < 400 => None, + (Ok(resp), _) if resp.status().is_success() => None, (Ok(resp), _) => { let status = resp.status().as_u16(); + if provider == crate::byok_endpoint::COMPATIBLE_PROVIDER { + // Arbitrary compatible endpoints may reflect credentials in + // their error body. Return status without forwarding that body. + return Some(format!( + "HTTP {status}: check the model ID, API key and endpoint" + )); + } let text = resp.text().await.unwrap_or_default(); let detail = serde_json::from_str::(&text) .ok() @@ -3632,6 +4482,50 @@ impl ModelService for UnconfiguredModelService { // ── HTTP types ─────────────────────────────────────────────────────────────── +fn default_user_model_context_window() -> i32 { + 128_000 +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +pub struct UserModelCreateRequest { + pub name: String, + pub provider: String, + pub model: String, + pub base_url: Option, + pub api_key: String, + #[serde(default = "default_user_model_context_window")] + pub context_window: i32, + #[serde(default)] + pub is_default: bool, +} + +impl std::fmt::Debug for UserModelCreateRequest { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("UserModelCreateRequest") + .field("name", &self.name) + .field("provider", &self.provider) + .field("model", &self.model) + .field("api_key", &"") + .finish_non_exhaustive() + } +} + +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct UserModelUpdateRequest { + pub api_key: Option, + pub context_window: Option, + pub is_default: Option, + pub is_active: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct UserModelListResponse { + pub items: Vec, +} + #[derive(Deserialize)] pub struct ModelCreateRequest { pub name: String, @@ -3834,6 +4728,7 @@ pub enum ModelAccessAction { ContactAdministrator, ReconnectDevice, ConfigureDeviceModels, + ConfigureCloudByokModels, Reauthenticate, ManageBilling, Retry, @@ -3977,6 +4872,9 @@ fn recovery_actions( (_, Some(ModelAccessReason::NoEligibleOfferings), ModelAccessKind::ThisDevice) => { vec![ModelAccessAction::ConfigureDeviceModels] } + (_, Some(ModelAccessReason::NoEligibleOfferings), ModelAccessKind::CloudByok) => { + vec![ModelAccessAction::ConfigureCloudByokModels] + } ( _, Some(ModelAccessReason::NoEligibleOfferings | ModelAccessReason::PolicyDisabled), @@ -4177,7 +5075,6 @@ pub fn project_model_access_page_with_default_catalog( sort_model_list_item_responses(&mut offerings); let mut offering_ids = BTreeSet::new(); - let mut counts = BTreeMap::::new(); for offering in &offerings { validate_model_offering_id(&offering.offering_id).map_err(|_| { crate::service_error::ServiceError::invalid(format!( @@ -4212,15 +5109,45 @@ pub fn project_model_access_page_with_default_catalog( offering.offering_id, offering.access_id ))); } + } + + // Access counts describe the complete effective catalog, not the current + // seek-paginated page. Counting the page (or assigning the aggregate + // total to every access) makes a mixed self-hosted/Cloud-BYOK catalog + // internally inconsistent and causes strict clients to reject it. + let mut counts = BTreeMap::::new(); + for offering in default_catalog { + let Some(access) = accesses.get(&offering.access_id) else { + return Err(crate::service_error::ServiceError::invalid(format!( + "Offering '{}' references undeclared Model Access '{}'", + offering.offering_id, offering.access_id + ))); + }; + if access.kind != offering.access_kind + || access.label != offering.access_label + || access.execution_placement != offering.execution_placement + { + return Err(crate::service_error::ServiceError::conflict(format!( + "Offering '{}' conflicts with Model Access '{}'", + offering.offering_id, offering.access_id + ))); + } let count = counts.entry(offering.access_id.clone()).or_default(); *count = count.saturating_add(1); } + if let Some(total) = total_offerings + && usize::try_from(total).ok() != Some(default_catalog.len()) + { + return Err(crate::service_error::ServiceError::conflict(format!( + "effective catalog advertised {total} Offerings but supplied {} for projection", + default_catalog.len() + ))); + } let accesses: Vec = accesses .into_values() .map(|access| { - let available_model_count = total_offerings - .unwrap_or_else(|| counts.get(&access.id).copied().unwrap_or_default()); + let available_model_count = counts.get(&access.id).copied().unwrap_or_default(); let availability = project_access_availability(&access, available_model_count)?; Ok(ModelAccessViewResponse { id: access.id, @@ -4339,6 +5266,22 @@ fn resolve_model_default( mod tests { use super::*; + #[test] + fn deployment_mode_gate_is_explicit_and_fails_closed() { + assert_eq!(deployment_mode_allows_shared_models(None), Ok(true)); + assert_eq!( + deployment_mode_allows_shared_models(Some("self-hosted")), + Ok(true) + ); + assert_eq!( + deployment_mode_allows_shared_models(Some("cloud-byok")), + Ok(false) + ); + for invalid in ["", "cloud", "CLOUD-BYOK", " cloud-byok "] { + assert!(deployment_mode_allows_shared_models(Some(invalid)).is_err()); + } + } + fn test_model_list_item(provider: &str, name: &str, offering_id: &str) -> ModelListItem { ModelListItem { offering_id: offering_id.to_string(), @@ -5829,6 +6772,57 @@ mod tests { ); } + #[test] + fn paged_model_access_reports_counts_per_access_from_complete_catalog() { + let self_hosted = DeclaredModelAccess { + id: "self-hosted".into(), + kind: ModelAccessKind::SelfHosted, + label: "Self-hosted".into(), + execution_placement: ModelExecutionPlacement::Server, + availability: ModelAccessAvailability::Ready, + }; + let cloud_byok = DeclaredModelAccess { + id: "cloud-byok".into(), + kind: ModelAccessKind::CloudByok, + label: "Cloud BYOK".into(), + execution_placement: ModelExecutionPlacement::Server, + availability: ModelAccessAvailability::Ready, + }; + let byok_offering = ModelListItemResponse { + offering_id: "byok-1".into(), + access_id: cloud_byok.id.clone(), + access_kind: cloud_byok.kind, + access_label: cloud_byok.label.clone(), + execution_placement: cloud_byok.execution_placement, + name: "deepseek-real".into(), + provider: "deepseek".into(), + description: None, + is_active: true, + context_window: 128_000, + max_completion_tokens: None, + architecture: None, + thinking_capability: None, + }; + + let projection = project_model_access_page_with_default_catalog( + vec![self_hosted, cloud_byok], + vec![byok_offering.clone()], + Some(1), + &[byok_offering], + None, + "2026-09-07T00:00:00Z".into(), + ) + .expect("mixed access projection"); + + let counts = projection + .accesses + .iter() + .map(|access| (access.id.as_str(), access.available_model_count)) + .collect::>(); + assert_eq!(counts.get("self-hosted"), Some(&0)); + assert_eq!(counts.get("cloud-byok"), Some(&1)); + } + #[test] fn model_access_projection_honors_provider_default_offering() { let declared = DeclaredModelAccess { @@ -6490,6 +7484,98 @@ mod tests { } // ── Provider-aware probe regression tests ───────────────────────── + + #[tokio::test] + async fn connectivity_probe_uses_native_token_budget_and_preserves_provider_errors() { + use axum::{Json, Router, http::HeaderMap, routing::post}; + for (provider, model, path, token_field, wrong_field) in [ + ( + "openai", + "o3", + "/chat/completions", + "max_completion_tokens", + "max_tokens", + ), + ( + "openai", + "gpt-4o", + "/chat/completions", + "max_completion_tokens", + "max_tokens", + ), + ( + "deepseek", + "deepseek-chat", + "/chat/completions", + "max_tokens", + "max_completion_tokens", + ), + ( + "deepseek", + "deepseek-v4-flash", + "/chat/completions", + "max_tokens", + "max_completion_tokens", + ), + ( + "anthropic", + "claude-sonnet-4-5", + "/v1/messages", + "max_tokens", + "max_completion_tokens", + ), + ] { + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let captured = calls.clone(); + let app = Router::new().route( + path, + post(move |headers: HeaderMap, Json(body): Json| { + captured.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + async move { + let auth = if provider == "anthropic" { + assert_eq!(headers["anthropic-version"], "2023-06-01"); + headers.get("x-api-key").and_then(|v| v.to_str().ok()) + == Some("valid-key") + } else { + headers.get("authorization").and_then(|v| v.to_str().ok()) + == Some("Bearer valid-key") + }; + let status = if !auth { + StatusCode::UNAUTHORIZED + } else if body["model"] != model { + StatusCode::NOT_FOUND + } else if body[token_field] != 32 || body.get(wrong_field).is_some() { + StatusCode::BAD_REQUEST + } else { + StatusCode::OK + }; + ( + status, + Json(serde_json::json!({"error":{"message":"fixture rejection"}})), + ) + } + }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let task = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + let result = + validate_connectivity(provider, model, "valid-key", Some(&base), None, None).await; + assert!(result.is_none(), "{provider}/{model}: {result:?}"); + for (key, model_id, expected) in [ + ("invalid-key", model, "HTTP 401"), + ("valid-key", "missing-model", "HTTP 404"), + ] { + let result = + validate_connectivity(provider, model_id, key, Some(&base), None, None) + .await + .unwrap(); + assert!(result.contains(expected), "{provider}: {result}"); + } + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 3); + task.abort(); + } + } // // Based on real API recordings from 2026-05-04. // Each mock simulates the provider's actual response pattern. @@ -6689,4 +7775,84 @@ mod tests { "mock provider should skip cleanly, not error" ); } + + #[test] + fn user_model_requests_redact_credentials_in_debug_output() { + let create = UserModelCreateRequestData { + name: "deepseek".into(), + provider: "deepseek".into(), + model: "deepseek-chat".into(), + base_url: None, + api_key: "sk-user-secret".into(), + context_window: 128_000, + is_default: true, + }; + let update = UserModelUpdateRequestData { + api_key: Some("sk-rotated-secret".into()), + ..Default::default() + }; + assert!(!format!("{create:?}").contains("sk-user-secret")); + assert!(!format!("{update:?}").contains("sk-rotated-secret")); + } + + #[test] + fn cloud_byok_accepts_only_fixed_provider_endpoints() { + assert_eq!( + user_byok_provider_base_url("deepseek").unwrap(), + "https://api.deepseek.com" + ); + let error = user_byok_provider_base_url("custom").unwrap_err(); + assert_eq!(error.0, StatusCode::BAD_REQUEST); + assert!(error.1.detail.contains("Unsupported BYOK provider")); + } + + #[test] + fn empty_cloud_byok_access_projects_a_user_configuration_action() { + let projection = project_model_access( + vec![DeclaredModelAccess { + id: "cloud-byok".into(), + kind: ModelAccessKind::CloudByok, + label: "Cloud BYOK".into(), + execution_placement: ModelExecutionPlacement::Server, + availability: ModelAccessAvailability::Ready, + }], + Vec::new(), + "2026-09-06T00:00:00Z".into(), + ) + .expect("Cloud BYOK setup projection"); + assert_eq!(projection.accesses.len(), 1); + assert_eq!( + projection.accesses[0].actions, + vec![ModelAccessAction::ConfigureCloudByokModels] + ); + } + + #[test] + fn user_model_api_keys_reject_ambiguous_whitespace_and_controls() { + assert!(validate_user_model_api_key("sk-valid").is_ok()); + for invalid in ["", " sk-key", "sk-key ", "sk\nkey"] { + assert!(validate_user_model_api_key(invalid).is_err(), "{invalid:?}"); + } + } + + #[test] + fn user_byok_endpoint_selection_requires_explicit_compatible_mode() { + assert_eq!( + user_byok_base_url("openai", None).unwrap(), + "https://api.openai.com/v1" + ); + assert_eq!( + user_byok_base_url("anthropic", None).unwrap(), + "https://api.anthropic.com" + ); + for provider in ["openai", "anthropic", "deepseek", "unknown"] { + assert!(user_byok_base_url(provider, Some("https://gateway.example/v1")).is_err()); + } + assert!(user_byok_base_url("openai-compatible", None).is_err()); + assert!(user_byok_base_url("openai-compatible", Some("http://127.0.0.1:1234/v1")).is_err()); + assert_eq!( + user_byok_base_url("openai-compatible", Some("https://gateway.example/v1/")).unwrap(), + "https://gateway.example/v1" + ); + } } diff --git a/crates/services/src/runs.rs b/crates/services/src/runs.rs index 9b2aa81518..b4c9a30740 100644 --- a/crates/services/src/runs.rs +++ b/crates/services/src/runs.rs @@ -10366,9 +10366,9 @@ impl DatabaseRunStateStore { OR CAST(? AS SIGNED) > total_completion_tokens OR CAST(? AS SIGNED) > total_tool_calls, NOW(6), updated_at), - total_prompt_tokens = GREATEST(total_prompt_tokens, CAST(? AS SIGNED)), - total_completion_tokens = GREATEST(total_completion_tokens, CAST(? AS SIGNED)), - total_tool_calls = GREATEST(total_tool_calls, CAST(? AS SIGNED)) + total_prompt_tokens = IF(CAST(? AS SIGNED) > total_prompt_tokens, CAST(? AS SIGNED), total_prompt_tokens), + total_completion_tokens = IF(CAST(? AS SIGNED) > total_completion_tokens, CAST(? AS SIGNED), total_completion_tokens), + total_tool_calls = IF(CAST(? AS SIGNED) > total_tool_calls, CAST(? AS SIGNED), total_tool_calls) WHERE user_id = ? AND session_id = ? AND run_id = ?", ) .bind(prompt_tokens as i64) @@ -10379,8 +10379,11 @@ impl DatabaseRunStateStore { .bind(completion_tokens as i64) .bind(tool_calls as i64) .bind(prompt_tokens as i64) + .bind(prompt_tokens as i64) + .bind(completion_tokens as i64) .bind(completion_tokens as i64) .bind(tool_calls as i64) + .bind(tool_calls as i64) .bind(user_id) .bind(expected_session_id) .bind(run_id) @@ -10510,12 +10513,12 @@ impl DatabaseRunStateStore { waiting_for = IF(VALUES(projection_event_idx) > projection_event_idx, VALUES(waiting_for), waiting_for), error_message = IF(VALUES(projection_event_idx) > projection_event_idx, VALUES(error_message), error_message), latest_event_type = IF(VALUES(projection_event_idx) > projection_event_idx, VALUES(latest_event_type), latest_event_type), - total_prompt_tokens = GREATEST(total_prompt_tokens, VALUES(total_prompt_tokens)), - total_completion_tokens = GREATEST(total_completion_tokens, VALUES(total_completion_tokens)), - total_tool_calls = GREATEST(total_tool_calls, VALUES(total_tool_calls)), + total_prompt_tokens = IF(VALUES(total_prompt_tokens) > total_prompt_tokens, VALUES(total_prompt_tokens), total_prompt_tokens), + total_completion_tokens = IF(VALUES(total_completion_tokens) > total_completion_tokens, VALUES(total_completion_tokens), total_completion_tokens), + total_tool_calls = IF(VALUES(total_tool_calls) > total_tool_calls, VALUES(total_tool_calls), total_tool_calls), projection_hash = IF(VALUES(projection_event_idx) > projection_event_idx, VALUES(projection_hash), projection_hash), updated_at = IF(VALUES(projection_event_idx) > projection_event_idx, NOW(6), updated_at), - projection_event_idx = GREATEST(projection_event_idx, VALUES(projection_event_idx))", + projection_event_idx = IF(VALUES(projection_event_idx) > projection_event_idx, VALUES(projection_event_idx), projection_event_idx)", [run.waiting_for.is_some(), run.error_message.is_some()], ); sqlx::query(&upsert_sql) @@ -10787,12 +10790,12 @@ impl DatabaseRunStateStore { latest_checkpoint_id = IF(VALUES(projection_event_idx) >= projection_event_idx, VALUES(latest_checkpoint_id), latest_checkpoint_id), latest_checkpoint_kind = IF(VALUES(projection_event_idx) >= projection_event_idx, VALUES(latest_checkpoint_kind), latest_checkpoint_kind), latest_checkpoint_version = IF(VALUES(projection_event_idx) >= projection_event_idx, VALUES(latest_checkpoint_version), latest_checkpoint_version), - total_prompt_tokens = GREATEST(total_prompt_tokens, VALUES(total_prompt_tokens)), - total_completion_tokens = GREATEST(total_completion_tokens, VALUES(total_completion_tokens)), - total_tool_calls = GREATEST(total_tool_calls, VALUES(total_tool_calls)), + total_prompt_tokens = IF(VALUES(total_prompt_tokens) > total_prompt_tokens, VALUES(total_prompt_tokens), total_prompt_tokens), + total_completion_tokens = IF(VALUES(total_completion_tokens) > total_completion_tokens, VALUES(total_completion_tokens), total_completion_tokens), + total_tool_calls = IF(VALUES(total_tool_calls) > total_tool_calls, VALUES(total_tool_calls), total_tool_calls), projection_hash = IF(VALUES(projection_event_idx) >= projection_event_idx, VALUES(projection_hash), projection_hash), updated_at = IF(VALUES(projection_event_idx) >= projection_event_idx, NOW(6), updated_at), - projection_event_idx = GREATEST(projection_event_idx, VALUES(projection_event_idx))", + projection_event_idx = IF(VALUES(projection_event_idx) > projection_event_idx, VALUES(projection_event_idx), projection_event_idx)", [ projection.waiting_for.is_some(), projection.error_message.is_some(), @@ -15940,17 +15943,20 @@ impl RunStateStore for DatabaseRunStateStore { OR CAST(? AS SIGNED) > total_completion_tokens OR CAST(? AS SIGNED) > total_tool_calls, NOW(6), updated_at), - total_prompt_tokens = GREATEST(total_prompt_tokens, CAST(? AS SIGNED)), - total_completion_tokens = GREATEST(total_completion_tokens, CAST(? AS SIGNED)), - total_tool_calls = GREATEST(total_tool_calls, CAST(? AS SIGNED)) + total_prompt_tokens = IF(CAST(? AS SIGNED) > total_prompt_tokens, CAST(? AS SIGNED), total_prompt_tokens), + total_completion_tokens = IF(CAST(? AS SIGNED) > total_completion_tokens, CAST(? AS SIGNED), total_completion_tokens), + total_tool_calls = IF(CAST(? AS SIGNED) > total_tool_calls, CAST(? AS SIGNED), total_tool_calls) WHERE user_id = ? AND session_id = ? AND run_id = ?", ) .bind(prompt_tokens as i64) .bind(completion_tokens as i64) .bind(tool_calls as i64) .bind(prompt_tokens as i64) + .bind(prompt_tokens as i64) + .bind(completion_tokens as i64) .bind(completion_tokens as i64) .bind(tool_calls as i64) + .bind(tool_calls as i64) .bind(user_id) .bind(expected_session_id) .bind(run_id) @@ -16341,16 +16347,19 @@ impl RunStateStore for DatabaseRunStateStore { OR CAST(? AS SIGNED) > total_completion_tokens OR CAST(? AS SIGNED) > total_tool_calls, NOW(6), updated_at), - total_prompt_tokens = GREATEST(total_prompt_tokens, CAST(? AS SIGNED)), - total_completion_tokens = GREATEST(total_completion_tokens, CAST(? AS SIGNED)), - total_tool_calls = GREATEST(total_tool_calls, CAST(? AS SIGNED)) + total_prompt_tokens = IF(CAST(? AS SIGNED) > total_prompt_tokens, CAST(? AS SIGNED), total_prompt_tokens), + total_completion_tokens = IF(CAST(? AS SIGNED) > total_completion_tokens, CAST(? AS SIGNED), total_completion_tokens), + total_tool_calls = IF(CAST(? AS SIGNED) > total_tool_calls, CAST(? AS SIGNED), total_tool_calls) WHERE user_id = ? AND session_id = ? AND run_id = ? AND run_generation = ?", ) .bind(prompt_tokens as i64) .bind(completion_tokens as i64) .bind(tool_calls as i64) .bind(prompt_tokens as i64) + .bind(prompt_tokens as i64) .bind(completion_tokens as i64) + .bind(completion_tokens as i64) + .bind(tool_calls as i64) .bind(tool_calls as i64) .bind(user_id) .bind(expected_session_id) diff --git a/crates/services/src/storage.rs b/crates/services/src/storage.rs index 4d5a3e42eb..ad97f600ad 100644 --- a/crates/services/src/storage.rs +++ b/crates/services/src/storage.rs @@ -2723,6 +2723,15 @@ fn inference_canonical_transition_wal_definition_mismatches(create_sql: &str) -> } } +fn database_reports_check_constraints_in_show_create(version: &str) -> bool { + // MatrixOne currently accepts CHECK syntax but omits CHECK clauses from + // SHOW CREATE TABLE. Treating the omitted text as proof of an obsolete + // table makes every freshly bootstrapped MatrixOne database fail startup. + // MySQL-compatible engines that do report the definition retain the + // strict recovery-mode contract check below. + !version.to_ascii_lowercase().contains("matrixone") +} + fn inference_invocation_schema_mismatches( columns: &BTreeMap, ) -> Vec { @@ -3055,9 +3064,12 @@ async fn verify_inference_canonical_transition_wal_schema_contract( let show_create = format!("SHOW CREATE TABLE `{database}`.`{table}`"); let create_row = query(&show_create).fetch_one(pool).await?; let create_sql: String = create_row.try_get(1_usize)?; - reasons.extend(inference_canonical_transition_wal_definition_mismatches( - &create_sql, - )); + let database_version: String = query_scalar("SELECT VERSION()").fetch_one(pool).await?; + if database_reports_check_constraints_in_show_create(&database_version) { + reasons.extend(inference_canonical_transition_wal_definition_mismatches( + &create_sql, + )); + } if reasons.is_empty() { return Ok(()); } @@ -3657,6 +3669,39 @@ async fn ensure_core_schema_while_leased( .execute(&pool) .await?; + core_schema_create!( + pool, + "auth_external_identities", + "CREATE TABLE IF NOT EXISTS auth_external_identities ( + provider_id VARCHAR(128) NOT NULL, + external_subject VARCHAR(255) NOT NULL, + astra_user_id VARCHAR(128) NOT NULL, + username VARCHAR(255) NULL, + email VARCHAR(255) NULL, + display_name VARCHAR(255) NULL, + created_at DATETIME(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), + updated_at DATETIME(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), + PRIMARY KEY (provider_id, external_subject), + INDEX idx_external_identity_user (astra_user_id) + )", + ) + .execute(&pool) + .await?; + + // Read-only migration source. Never create new unnamespaced identities. + core_schema_create!( + pool, + "auth_memoria_identities", + "CREATE TABLE IF NOT EXISTS auth_memoria_identities ( + memoria_user_id VARCHAR(128) PRIMARY KEY, + astra_user_id VARCHAR(128) NOT NULL UNIQUE, + created_at DATETIME(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), + INDEX idx_auth_memoria_astra_user (astra_user_id) + )", + ) + .execute(&pool) + .await?; + core_schema_create!(pool, "auth_provider_request_replay", "CREATE TABLE IF NOT EXISTS auth_provider_request_replay ( provider VARCHAR(64) NOT NULL, @@ -6741,6 +6786,32 @@ async fn ensure_core_schema_while_leased( .execute(&pool) .await?; + core_schema_create!( + pool, + "user_llm_models", + "CREATE TABLE IF NOT EXISTS user_llm_models ( + model_id VARCHAR(64) NOT NULL, + user_id VARCHAR(128) NOT NULL, + model_alias VARCHAR(100) NOT NULL, + model_name VARCHAR(255) NOT NULL, + provider VARCHAR(50) NOT NULL, + api_key_encrypted TEXT NOT NULL, + base_url VARCHAR(500) NOT NULL, + context_window INT NOT NULL, + is_default SMALLINT NOT NULL DEFAULT 0, + is_active SMALLINT NOT NULL DEFAULT 1, + created_at DATETIME(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), + updated_at DATETIME(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), + PRIMARY KEY (user_id, model_id), + UNIQUE KEY uq_user_llm_models_alias (user_id, model_alias), + CONSTRAINT chk_user_llm_models_context_window CHECK (context_window > 0), + INDEX idx_user_llm_models_catalog (user_id, is_active, provider, model_alias, model_id), + INDEX idx_user_llm_models_default (user_id, is_default, is_active) + )", + ) + .execute(&pool) + .await?; + // Canonical inference execution ledger. Admission writes the immutable route // and logical invocation together; each physical attempt is then committed // before its provider I/O. Route rows contain no credential or endpoint @@ -8995,6 +9066,13 @@ impl Default for RetentionPolicy { } } +/// Remove one bounded batch of credentials owned by inactive or deleted accounts. +pub async fn cleanup_inactive_memoria_credentials(pool: &sqlx::MySqlPool) -> Result { + sqlx::query("DELETE FROM auth_tokens WHERE type = 'memoria_connection' AND provider = 'memoria' AND NOT EXISTS (SELECT 1 FROM auth_users WHERE user_id = auth_tokens.scope_user_id AND is_active = 1) ORDER BY created_at ASC, token_id ASC LIMIT 100") + .execute(pool).await.map(|result| result.rows_affected()) + .map_err(|error| format!("cleanup inactive Memoria credentials: {error}")) +} + /// Purge expired authentication and operational records with TTL/expiry semantics. /// /// Returns a list of per-table cleanup results showing how many rows were deleted. @@ -9063,7 +9141,9 @@ pub async fn cleanup_expired_data( rows_deleted: deleted, }); - // 2. Inactive auth tokens + // 2. Inactive auth tokens, plus credentials whose account was removed or + // deactivated. Scoped resolvers already deny these before cleanup runs. + let orphaned_memoria = cleanup_inactive_memoria_credentials(pool).await?; let deleted = sqlx::query( "DELETE FROM auth_tokens \ WHERE is_active = 0 \ @@ -9075,7 +9155,7 @@ pub async fn cleanup_expired_data( .bind(AUTH_TOKEN_BATCH_LIMIT) .execute(pool) .await - .map(|r| r.rows_affected()) + .map(|r| r.rows_affected() + orphaned_memoria) .map_err(|e| format!("cleanup auth_tokens: {e}"))?; results.push(CleanupResult { table: "auth_tokens", @@ -9783,6 +9863,12 @@ mod tests { #[test] fn canonical_transition_wal_schema_contract_is_exact_and_fail_closed() { + assert!(!database_reports_check_constraints_in_show_create( + "8.0.30-MatrixOne-v2.1.0" + )); + assert!(database_reports_check_constraints_in_show_create( + "8.0.36 MySQL Community Server" + )); let exact_columns = canonical_transition_wal_columns(); let exact_indexes = canonical_transition_wal_indexes(); assert!( diff --git a/crates/services/tests/auth_reauthentication_db_it.rs b/crates/services/tests/auth_reauthentication_db_it.rs index fd94c17556..4ba71538e7 100644 --- a/crates/services/tests/auth_reauthentication_db_it.rs +++ b/crates/services/tests/auth_reauthentication_db_it.rs @@ -42,6 +42,7 @@ async fn reauthentication_proofs_are_owner_purpose_expiry_and_one_time_bound() { .reauthenticate( &user.user_id, ReauthenticationRequestData { + memoria_proof: None, password: "wrong-password".to_string(), purpose: ReauthenticationPurpose::DeviceTrust, }, @@ -57,6 +58,7 @@ async fn reauthentication_proofs_are_owner_purpose_expiry_and_one_time_bound() { .reauthenticate( &user.user_id, ReauthenticationRequestData { + memoria_proof: None, password: password.to_string(), purpose: ReauthenticationPurpose::DeviceTrust, }, @@ -97,6 +99,7 @@ async fn reauthentication_proofs_are_owner_purpose_expiry_and_one_time_bound() { .reauthenticate( &user.user_id, ReauthenticationRequestData { + memoria_proof: None, password: password.to_string(), purpose: ReauthenticationPurpose::DeviceReenroll, }, @@ -125,6 +128,7 @@ async fn reauthentication_proofs_are_owner_purpose_expiry_and_one_time_bound() { .reauthenticate( &user.user_id, ReauthenticationRequestData { + memoria_proof: None, password: password.to_string(), purpose: ReauthenticationPurpose::DeviceTrust, }, diff --git a/crates/services/tests/byok_live_network.rs b/crates/services/tests/byok_live_network.rs new file mode 100644 index 0000000000..134ea06e3d --- /dev/null +++ b/crates/services/tests/byok_live_network.rs @@ -0,0 +1,24 @@ +//! Explicit opt-in network smoke: no database, model key or inference request. +//! Set ASTRA_TEST_BYOK_MODELS_URL to a provider's public /models endpoint that +//! requires authentication, plus the Server's ASTRA_BYOK_DNS_SERVERS/proxy settings. + +#[tokio::test] +#[ignore = "requires explicit public provider URL and working outbound DNS/HTTPS"] +async fn byok_live_network_reaches_provider_without_credentials() { + let url = std::env::var("ASTRA_TEST_BYOK_MODELS_URL") + .expect("set ASTRA_TEST_BYOK_MODELS_URL to the provider's HTTPS /models endpoint"); + let client = astra_services::byok_endpoint::endpoint_client(&url) + .await + .expect("real BYOK DNS resolution and public-address validation"); + let response = client + .get(&url) + .timeout(std::time::Duration::from_secs(15)) + .send() + .await + .expect("pinned HTTPS connection with origin certificate validation"); + assert_eq!( + response.status(), + reqwest::StatusCode::UNAUTHORIZED, + "provider should reject the intentionally unauthenticated request" + ); +} diff --git a/crates/services/tests/common/isolated_database.rs b/crates/services/tests/common/isolated_database.rs new file mode 100644 index 0000000000..92fb37ef9b --- /dev/null +++ b/crates/services/tests/common/isolated_database.rs @@ -0,0 +1,49 @@ +//! Shared isolation contract for tests that exercise identity/credential writes. +//! The online runner explicitly designates its database via ASTRA_TEST_DATABASE. + +pub fn require_isolated_database(database: &str) { + assert_eq!(std::env::var("ASTRA_TEST_DB_IT").as_deref(), Ok("1")); + let designated = std::env::var("ASTRA_TEST_DATABASE") + .expect("explicitly designate an isolated database with ASTRA_TEST_DATABASE"); + assert!( + is_designated_database(database, &designated), + "effective ASTRA_DATABASE must match the explicitly designated ASTRA_TEST_DATABASE" + ); +} + +fn is_designated_database(database: &str, designated: &str) -> bool { + !database.is_empty() + && database == designated + && database.len() <= 64 + && database + .bytes() + .all(|b| b.is_ascii_alphanumeric() || b == b'_') + && !matches!( + database, + "mysql" | "information_schema" | "mo_catalog" | "system" | "astra_runtime" + ) +} + +#[test] +fn isolated_database_contract_accepts_runner_names_and_rejects_implicit_targets() { + for name in [ + "astra_runtime_test_integration", + "astra_runtime_test_runtime_ignored", + "review_local", + ] { + assert!(is_designated_database(name, name)); + } + for name in [ + "", + "mysql", + "mo_catalog", + "astra_runtime", + "test;DROP DATABASE mysql", + ] { + assert!(!is_designated_database(name, name)); + } + assert!(!is_designated_database( + "production", + "astra_runtime_test_integration" + )); +} diff --git a/crates/services/tests/memoria_auth_db_it.rs b/crates/services/tests/memoria_auth_db_it.rs new file mode 100644 index 0000000000..52583dec12 --- /dev/null +++ b/crates/services/tests/memoria_auth_db_it.rs @@ -0,0 +1,531 @@ +mod common; +#[path = "common/isolated_database.rs"] +mod isolated_database; +use astra_core::JwtSettings; +use astra_services::{ + DatabaseModelService, FernetTokenEncryptor, ModelService, + auth::{AuthRefreshRequestData, AuthService, DatabaseAuthService}, +}; +use axum::{ + Json, Router, + http::{HeaderMap, HeaderValue, StatusCode}, + routing::get, +}; +use serde_json::json; +use sha2::Digest; +use std::sync::{ + Arc, + atomic::{AtomicBool, Ordering}, +}; +use uuid::Uuid; + +#[tokio::test] +#[ignore = "requires ASTRA_TEST_DB_IT=1 and explicitly isolated ASTRA_TEST_DATABASE"] +async fn memoria_refresh_revocation_and_deployment_model_isolation() { + isolated_database::require_isolated_database(&common::require_db_it_env().database); + let (shared, settings) = common::setup_pool_and_settings().await; + let pool = shared.get(); + let owner = Uuid::new_v4().to_string(); + let key_id = Uuid::new_v4().to_string(); + let revoked = Arc::new(AtomicBool::new(false)); + let flag = revoked.clone(); + let whoami = json!({"user_id":owner, "key_id":key_id, "is_active":true, "is_master":false, + "scope":{"type":"personal","id":owner},"api_version":"1", + "capabilities":["api_key_scopes","memory_filters_v1"],"granted_scopes":["identity:read"]}); + let app = Router::new().route( + "/auth/whoami", + get(move || { + let flag = flag.clone(); + let body = whoami.clone(); + async move { + ( + if flag.load(Ordering::SeqCst) { + StatusCode::UNAUTHORIZED + } else { + StatusCode::OK + }, + Json(body), + ) + } + }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}", listener.local_addr().unwrap()); + let server = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + let encryptor = FernetTokenEncryptor::new("test-only-key").unwrap(); + let jwt = JwtSettings { + secret_key: "test-only-jwt-secret".into(), + algorithm: "HS256".into(), + access_token_expire_minutes: 90, + refresh_token_expire_days: 7, + }; + let auth = DatabaseAuthService::new(settings.clone(), jwt.clone()) + .with_pool(shared.clone()) + .with_encryptor(encryptor.clone()) + .with_memoria_base_url(url); + let tokens = auth.login_memoria("test-key").await.unwrap().tokens; + assert_eq!(tokens.expires_in, 900); + + // Actual JWT -> stored refresh session -> mapping -> stored credential -> HTTP. + let refreshed = auth + .refresh(AuthRefreshRequestData { + refresh_token: tokens.refresh_token, + }) + .await + .unwrap(); + assert_eq!(refreshed.expires_in, 900); + let mut headers = HeaderMap::new(); + headers.insert( + "authorization", + HeaderValue::from_str(&format!("Bearer {}", refreshed.access_token)).unwrap(), + ); + auth.current_user(&headers).await.unwrap(); + revoked.store(true, Ordering::SeqCst); + assert_eq!( + auth.refresh(AuthRefreshRequestData { + refresh_token: refreshed.refresh_token.clone() + }) + .await + .unwrap_err() + .0, + StatusCode::UNAUTHORIZED + ); + + // Legacy `internal` tokens cannot evade refresh validation or the 15-min + // access window when the self-hosted JWT configuration used a longer TTL. + let decoded: jsonwebtoken::TokenData = jsonwebtoken::decode( + &refreshed.access_token, + &jsonwebtoken::DecodingKey::from_secret(jwt.secret_key.as_bytes()), + &jsonwebtoken::Validation::default(), + ) + .unwrap(); + for kind in ["refresh", "access"] { + let mut claims = decoded.claims.clone(); + claims["origin"] = json!("internal"); + claims["type"] = json!(kind); + claims["iat"] = json!(chrono::Utc::now().timestamp() - 901); + claims["exp"] = json!(chrono::Utc::now().timestamp() + 3600); + let legacy = jsonwebtoken::encode( + &jsonwebtoken::Header::default(), + &claims, + &jsonwebtoken::EncodingKey::from_secret(jwt.secret_key.as_bytes()), + ) + .unwrap(); + if kind == "refresh" { + use sha2::{Digest, Sha256}; + sqlx::query("INSERT INTO auth_refresh_tokens (token_id,user_id,session_id,token_hash,expires_at,is_revoked) VALUES (?, ?, ?, ?, DATE_ADD(NOW(), INTERVAL 1 DAY), 0)") + .bind(Uuid::new_v4().to_string()).bind(&tokens.user_id).bind(claims["sid"].as_str().unwrap()) + .bind(format!("{:x}", Sha256::digest(legacy.as_bytes()))).execute(pool).await.unwrap(); + assert_eq!( + auth.refresh(AuthRefreshRequestData { + refresh_token: legacy + }) + .await + .unwrap_err() + .0, + StatusCode::UNAUTHORIZED + ); + } else { + headers.insert( + "authorization", + HeaderValue::from_str(&format!("Bearer {legacy}")).unwrap(), + ); + assert_eq!( + auth.current_user(&headers).await.unwrap_err().0, + StatusCode::UNAUTHORIZED + ); + } + } + + let deployment = Uuid::new_v4().to_string(); + sqlx::query("INSERT INTO infra_llm_models (model_id,model_name,provider,base_url,is_active,context_window,api_key_encrypted,input_modalities,output_modalities,supported_parameters,pricing,tags,quirks) VALUES (?, ?, 'mock', 'http://127.0.0.1:1', 1, 128000, ?, ?, ?, '[]', '{}', '[]', '{}')") + .bind(&deployment).bind(format!("deployment-{deployment}")) + .bind(encryptor.encrypt("deployment-test-secret").unwrap()).bind(r#"["text"]"#).bind(r#"["text"]"#) + .execute(pool).await.unwrap(); + let models = DatabaseModelService::new(settings, Arc::new(encryptor)).with_pool(shared.clone()); + assert!( + !models + .allows_deployment_models(tokens.user_id.clone()) + .await + .unwrap() + ); + assert!( + models + .list_models(tokens.user_id.clone(), false) + .await + .unwrap() + .is_empty() + ); + assert_eq!( + models + .admit_model_offering(tokens.user_id.clone(), deployment.clone()) + .await + .unwrap_err() + .0, + StatusCode::NOT_FOUND + ); + // The test is run in both deployment modes: non-Memoria accounts keep the + // original self-hosted behavior, but cannot bypass explicit cloud-byok. + let self_hosted = std::env::var("ASTRA_DEPLOYMENT_MODE").as_deref() != Ok("cloud-byok"); + assert_eq!( + models + .allows_deployment_models("local-test-user".into()) + .await + .unwrap(), + self_hosted + ); + assert_eq!( + models + .admit_model_offering("local-test-user".into(), deployment.clone()) + .await + .is_ok(), + self_hosted + ); + sqlx::query("DELETE FROM infra_llm_models WHERE model_id = ?") + .bind(deployment) + .execute(pool) + .await + .unwrap(); + for table in [ + "auth_tokens", + "auth_refresh_tokens", + "auth_user_roles", + "auth_memoria_identities", + "auth_external_identities", + "auth_users", + ] { + let column = match table { + "auth_tokens" => "scope_user_id", + "auth_memoria_identities" | "auth_external_identities" => "astra_user_id", + _ => "user_id", + }; + sqlx::query(&format!("DELETE FROM {table} WHERE {column} = ?")) + .bind(&tokens.user_id) + .execute(pool) + .await + .unwrap(); + } + server.abort(); +} + +#[tokio::test] +#[ignore = "requires ASTRA_TEST_DB_IT=1 and explicitly isolated ASTRA_TEST_DATABASE"] +async fn memoria_issuer_atomicity_concurrent_binding_and_disconnect() { + isolated_database::require_isolated_database(&common::require_db_it_env().database); + let (shared, db) = common::setup_pool_and_settings().await; + let subject = format!("review-{}", Uuid::new_v4()); + let subject_for_http = subject.clone(); + let revoked = Arc::new(AtomicBool::new(false)); + let flag = revoked.clone(); + let app = Router::new().route("/auth/whoami", get(move |headers: HeaderMap| { + let subject = subject_for_http.clone(); + let revoked = flag.load(Ordering::SeqCst); + async move { + let key = headers.get("authorization").and_then(|h| h.to_str().ok()) + .unwrap_or("").trim_start_matches("Bearer "); + let owner = match key { + "atomic-failure" => format!("{subject}-failure"), + "legacy-key" => format!("{subject}-legacy"), + _ => subject, + }; + (if revoked { StatusCode::UNAUTHORIZED } else { StatusCode::OK }, + Json(json!({"user_id": owner, "key_id": key, "is_active": true, "is_master": false, + "scope": {"type": "personal", "id": owner}, "api_version": "1", + "capabilities": ["api_key_scopes", "memory_filters_v1"], + "granted_scopes": ["identity:read", "memory:read"]}))) + } + })); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let server = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + let settings = astra_core::MemoriaSettings { + base_url: base.clone(), + master_key: None, + issuer: None, + web_url: Some("http://localhost".into()), + legacy_issuer: None, + }; + let jwt = JwtSettings { + secret_key: "review-704-jwt".into(), + algorithm: "HS256".into(), + access_token_expire_minutes: 90, + refresh_token_expire_days: 7, + }; + let encryptor = FernetTokenEncryptor::new("review-704-encryption").unwrap(); + let auth = DatabaseAuthService::new(db.clone(), jwt.clone()) + .with_pool(shared.clone()) + .with_encryptor(encryptor.clone()) + .with_memoria_settings(&settings) + .unwrap(); + let mut logins = Vec::new(); + for i in 0..6 { + let auth = auth.clone(); + logins.push(tokio::spawn(async move { + auth.login_memoria(&format!("key-{i}")).await + })); + } + let mut tokens = Vec::new(); + for login in logins { + tokens.push(login.await.unwrap().unwrap().tokens); + } + let user = tokens[0].user_id.clone(); + assert!(tokens.iter().all(|t| t.user_id == user)); + let count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM auth_tokens WHERE type = 'memoria_connection' AND scope_user_id = ? AND is_active = 1") + .bind(&user).fetch_one(shared.get()).await.unwrap(); + assert_eq!( + count, 1, + "one durable binding, even for concurrent first login/relink" + ); + let identities: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM auth_external_identities WHERE external_subject = ?", + ) + .bind(&subject) + .fetch_one(shared.get()) + .await + .unwrap(); + assert_eq!(identities, 1); + let resolver = auth.memoria_credentials().unwrap(); + let grant = resolver.resolve(&user).await.unwrap().unwrap(); + assert_eq!( + grant.access, + astra_services::auth::memoria::MemoryAccess::ReadOnly + ); + assert!(grant.generation.starts_with("key-")); + assert_eq!(grant.owner, subject); + let encrypted: String = sqlx::query_scalar("SELECT encrypted_value FROM auth_tokens WHERE type = 'memoria_connection' AND scope_user_id = ?") + .bind(&user).fetch_one(shared.get()).await.unwrap(); + assert_ne!(encrypted, grant.key); + let mut headers = HeaderMap::new(); + headers.insert( + "authorization", + HeaderValue::from_str(&format!("Bearer {}", tokens[0].access_token)).unwrap(), + ); + assert!(matches!( + auth.current_principal(&headers).await.unwrap().origin, + astra_services::AuthPrincipalOrigin::VerifiedProvider { .. } + )); + + // Losing a provider mapping must not downgrade a verified JWT to an + // internal principal while its session row is still active. + sqlx::query( + "DELETE FROM auth_external_identities WHERE provider_id = ? AND external_subject = ?", + ) + .bind(&resolver.provider.provider_id) + .bind(&subject) + .execute(shared.get()) + .await + .unwrap(); + assert_eq!( + auth.current_user(&headers).await.unwrap_err().0, + StatusCode::UNAUTHORIZED + ); + sqlx::query("INSERT INTO auth_external_identities (provider_id, external_subject, astra_user_id) VALUES (?, ?, ?)") + .bind(&resolver.provider.provider_id).bind(&subject).bind(&user) + .execute(shared.get()).await.unwrap(); + + // Same subject at a different issuer is a different account, and cannot + // refresh / access a session from the old issuer after configuration changes. + let other = DatabaseAuthService::new(db.clone(), jwt.clone()) + .with_pool(shared.clone()) + .with_encryptor(encryptor.clone()) + .with_memoria_settings(&astra_core::MemoriaSettings { + issuer: Some("https://another-issuer.example".into()), + ..settings.clone() + }) + .unwrap(); + let other_login = other.login_memoria("key-other").await.unwrap(); + assert_ne!(other_login.tokens.user_id, user); + assert_eq!( + other.current_user(&headers).await.err().unwrap().0, + StatusCode::UNAUTHORIZED + ); + assert_eq!( + other + .refresh(AuthRefreshRequestData { + refresh_token: tokens[0].refresh_token.clone() + }) + .await + .err() + .unwrap() + .0, + StatusCode::UNAUTHORIZED + ); + assert!( + other + .memoria_credentials() + .unwrap() + .resolve(&user) + .await + .unwrap() + .is_none() + ); + + // Token creation fails AFTER identity and credential SQL, but rolls them + // both back. This is not a fixture that manually seeds an already valid grant. + let bad = DatabaseAuthService::new( + db.clone(), + JwtSettings { + algorithm: "unsupported".into(), + ..jwt + }, + ) + .with_pool(shared.clone()) + .with_encryptor(encryptor) + .with_memoria_settings(&settings) + .unwrap(); + assert!(bad.login_memoria("atomic-failure").await.is_err()); + let orphan_count: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM auth_external_identities WHERE external_subject = ?", + ) + .bind(format!("{subject}-failure")) + .fetch_one(shared.get()) + .await + .unwrap(); + assert_eq!( + orphan_count, 0, + "failed login cannot leave an identity or its account" + ); + let failed_user = format!( + "ext_{:x}", + sha2::Sha256::digest( + format!("{}\0{subject}-failure", resolver.provider.provider_id).as_bytes() + ) + ); + for (table, column) in [ + ("auth_users", "user_id"), + ("auth_tokens", "scope_user_id"), + ("auth_refresh_tokens", "user_id"), + ] { + let rows: i64 = + sqlx::query_scalar(&format!("SELECT COUNT(*) FROM {table} WHERE {column} = ?")) + .bind(&failed_user) + .fetch_one(shared.get()) + .await + .unwrap(); + assert_eq!(rows, 0, "failed login left rows in {table}"); + } + + // Legacy subjects have no issuer. They are never silently assigned to + // whichever instance happens to answer the next login. + let legacy_user = format!("legacy-{}", Uuid::new_v4()); + let legacy_subject = format!("{subject}-legacy"); + sqlx::query("INSERT INTO auth_users (user_id,username,email,password_hash,is_active) VALUES (?, ?, ?, '', 1)") + .bind(&legacy_user).bind(&legacy_user).bind(format!("{legacy_user}@test.invalid")) + .execute(shared.get()).await.unwrap(); + sqlx::query( + "INSERT INTO auth_memoria_identities (memoria_user_id,astra_user_id) VALUES (?, ?)", + ) + .bind(&legacy_subject) + .bind(&legacy_user) + .execute(shared.get()) + .await + .unwrap(); + assert_eq!( + auth.login_memoria("legacy-key").await.err().unwrap().0, + StatusCode::CONFLICT + ); + let migrator = auth + .clone() + .with_memoria_settings(&astra_core::MemoriaSettings { + legacy_issuer: Some(base), + ..settings.clone() + }) + .unwrap(); + assert_eq!( + migrator + .login_memoria("legacy-key") + .await + .unwrap() + .tokens + .user_id, + legacy_user + ); + assert_eq!( + migrator + .login_memoria("legacy-key") + .await + .unwrap() + .tokens + .user_id, + legacy_user + ); + let remaining: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM auth_memoria_identities WHERE astra_user_id = ?") + .bind(&legacy_user) + .fetch_one(shared.get()) + .await + .unwrap(); + assert_eq!(remaining, 0); + migrator.disconnect_memoria(&legacy_user).await.unwrap(); + + auth.disconnect_memoria(&user).await.unwrap(); + auth.disconnect_memoria(&user).await.unwrap(); // idempotent service operation + assert!(resolver.resolve(&user).await.unwrap().is_none()); + assert_eq!( + auth.current_user(&headers).await.err().unwrap().0, + StatusCode::UNAUTHORIZED + ); + for token in &tokens { + assert_eq!( + auth.refresh(AuthRefreshRequestData { + refresh_token: token.refresh_token.clone() + }) + .await + .err() + .unwrap() + .0, + StatusCode::UNAUTHORIZED + ); + } + let relink = auth.login_memoria("new-key").await.unwrap(); + assert_eq!( + relink.tokens.user_id, user, + "unlink must not destroy account continuity" + ); + revoked.store(true, Ordering::SeqCst); + assert_eq!( + auth.refresh(AuthRefreshRequestData { + refresh_token: relink.tokens.refresh_token + }) + .await + .err() + .unwrap() + .0, + StatusCode::UNAUTHORIZED + ); + auth.disconnect_memoria(&user).await.unwrap(); + other + .disconnect_memoria(&other_login.tokens.user_id) + .await + .unwrap(); + // Account deactivation blocks existing runtime bindings immediately, and + // maintenance removes the encrypted secret without deleting Work/history. + revoked.store(false, Ordering::SeqCst); + auth.login_memoria("retention-key").await.unwrap(); + sqlx::query("UPDATE auth_users SET is_active = 0 WHERE user_id = ?") + .bind(&user) + .execute(shared.get()) + .await + .unwrap(); + assert!(resolver.resolve(&user).await.unwrap().is_none()); + assert_eq!( + auth.login_memoria("retention-key").await.err().unwrap().0, + StatusCode::FORBIDDEN + ); + astra_services::storage::cleanup_inactive_memoria_credentials(shared.get()) + .await + .unwrap(); + let secrets: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM auth_tokens WHERE scope_user_id = ? AND type = 'memoria_connection'", + ) + .bind(&user) + .fetch_one(shared.get()) + .await + .unwrap(); + assert_eq!(secrets, 0); + server.abort(); +} diff --git a/crates/services/tests/memoria_live_contract_it.rs b/crates/services/tests/memoria_live_contract_it.rs new file mode 100644 index 0000000000..6a183b7815 --- /dev/null +++ b/crates/services/tests/memoria_live_contract_it.rs @@ -0,0 +1,121 @@ +//! Opt-in contract against a running Memoria API, not a whoami fixture. +mod common; +#[path = "common/isolated_database.rs"] +mod isolated_database; + +use astra_core::JwtSettings; +use astra_services::{ + AuthService, DatabaseAuthService, FernetTokenEncryptor, + auth::{AuthRefreshRequestData, memoria::MemoryAccess}, +}; +use axum::http::StatusCode; +use serde_json::{Value, json}; + +#[tokio::test] +#[ignore = "requires isolated ASTRA_TEST_DATABASE and ASTRA_TEST_MEMORIA_URL / ASTRA_TEST_MEMORIA_MASTER_KEY"] +async fn scoped_key_api_v1_preserves_identity_modes_and_revocation() { + isolated_database::require_isolated_database(&common::require_db_it_env().database); + let base = std::env::var("ASTRA_TEST_MEMORIA_URL").unwrap(); + let master = std::env::var("ASTRA_TEST_MEMORIA_MASTER_KEY").unwrap(); + let (pool, settings) = common::setup_pool_and_settings().await; + let auth = DatabaseAuthService::new( + settings, + JwtSettings { + secret_key: "isolated-live-contract-jwt".into(), + algorithm: "HS256".into(), + access_token_expire_minutes: 90, + refresh_token_expire_days: 7, + }, + ) + .with_pool(pool) + .with_encryptor(FernetTokenEncryptor::new("isolated-live-contract-encryption").unwrap()) + .with_memoria_base_url(base.clone()); + let client = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .timeout(std::time::Duration::from_secs(10)) + .build() + .unwrap(); + let owner = format!("review-704-{}", uuid::Uuid::new_v4()); + let mut account = None; + + // Broad/master keys cannot become ordinary Astra account credentials. + assert_eq!( + auth.login_memoria(&master).await.err().unwrap().0, + StatusCode::UNAUTHORIZED + ); + for (access, scopes) in [ + (MemoryAccess::None, vec!["identity:read"]), + (MemoryAccess::ReadOnly, vec!["identity:read", "memory:read"]), + ( + MemoryAccess::ReadWrite, + vec!["identity:read", "memory:read", "memory:write"], + ), + ] { + let issued: Value = client + .post(format!("{base}/auth/keys")) + .bearer_auth(&master) + .json(&json!({"user_id": owner, "name": "isolated-astra-contract", "scopes": scopes})) + .send() + .await + .unwrap() + .error_for_status() + .unwrap() + .json() + .await + .unwrap(); + let key = issued["raw_key"].as_str().unwrap(); + let key_id = issued["key_id"].as_str().unwrap(); + let login = auth.login_memoria(key).await.unwrap(); + assert_eq!(login.memory_access, access); + if let Some(previous) = account.as_ref() { + assert_eq!(&login.tokens.user_id, previous); + } else { + account = Some(login.tokens.user_id.clone()); + } + let credential = auth + .memoria_credentials() + .unwrap() + .resolve(&login.tokens.user_id) + .await + .unwrap() + .unwrap(); + assert_eq!(credential.owner, owner); + assert_eq!(credential.generation, key_id); + assert_eq!(credential.access, access); + let refreshed = auth + .refresh(AuthRefreshRequestData { + refresh_token: login.tokens.refresh_token, + }) + .await + .unwrap(); + assert_eq!(refreshed.expires_in, 900); + + client + .delete(format!("{base}/auth/keys/{key_id}")) + .bearer_auth(&master) + .send() + .await + .unwrap() + .error_for_status() + .unwrap(); + assert_eq!( + auth.refresh(AuthRefreshRequestData { + refresh_token: refreshed.refresh_token, + }) + .await + .unwrap_err() + .0, + StatusCode::UNAUTHORIZED + ); + } + let account = account.unwrap(); + auth.disconnect_memoria(&account).await.unwrap(); + assert!( + auth.memoria_credentials() + .unwrap() + .resolve(&account) + .await + .unwrap() + .is_none() + ); +} diff --git a/crates/services/tests/models_db_it.rs b/crates/services/tests/models_db_it.rs index b98219f5b4..43975fe3dd 100644 --- a/crates/services/tests/models_db_it.rs +++ b/crates/services/tests/models_db_it.rs @@ -3,8 +3,8 @@ mod common; use std::sync::Arc; use astra_services::{ - DatabaseModelService, FernetTokenEncryptor, ModelOfferingResolutionError, ModelService, - resolve_active_llm_offering, revalidate_active_llm_offering, + DatabaseModelService, FernetTokenEncryptor, ModelAccessKind, ModelOfferingResolutionError, + ModelService, resolve_active_llm_offering, revalidate_active_llm_offering, }; use axum::http::StatusCode; use serial_test::serial; @@ -177,3 +177,209 @@ async fn effective_offering_resolution_is_exact_active_and_secret_safe() { .expect("clean Offering fixture"); astra_services::models::invalidate_active_llm_model_resolution_cache(); } + +#[tokio::test] +#[ignore = "requires live DB: run with ASTRA_TEST_DB_IT=1"] +#[serial] +async fn user_byok_models_are_owner_scoped_encrypted_and_admitted_for_owner_only() { + let (shared_pool, settings) = common::setup_pool_and_settings().await; + let pool = shared_pool.get().clone(); + let encryptor = FernetTokenEncryptor::new("user-models-db-it-key").expect("test encryptor"); + let service = + DatabaseModelService::new(settings, Arc::new(encryptor.clone())).with_pool(shared_pool); + let alias = format!("byok_{}", Uuid::new_v4().simple()); + let user_a = format!("user_a_{}", Uuid::new_v4().simple()); + let user_b = format!("user_b_{}", Uuid::new_v4().simple()); + let model_a = Uuid::new_v4().to_string(); + let model_b = Uuid::new_v4().to_string(); + let secret_a = "user-a-provider-secret"; + let secret_b = "user-b-provider-secret"; + + for (user_id, model_id, secret) in + [(&user_a, &model_a, secret_a), (&user_b, &model_b, secret_b)] + { + sqlx::query( + "INSERT INTO user_llm_models \ + (model_id, user_id, model_alias, model_name, provider, api_key_encrypted, base_url, \ + context_window, is_default, is_active) \ + VALUES (?, ?, ?, 'deepseek-chat', 'deepseek', ?, 'https://api.deepseek.com', \ + 128000, 1, 1)", + ) + .bind(model_id) + .bind(user_id) + .bind(&alias) + .bind(encryptor.encrypt(secret).expect("encrypt secret")) + .execute(&pool) + .await + .expect("seed user model"); + } + + let stored: String = sqlx::query_scalar( + "SELECT api_key_encrypted FROM user_llm_models WHERE user_id = ? AND model_id = ?", + ) + .bind(&user_a) + .bind(&model_a) + .fetch_one(&pool) + .await + .expect("read encrypted secret"); + assert_ne!(stored, secret_a); + assert!(!stored.contains(secret_a)); + + let listed_a = service + .list_user_models(user_a.clone()) + .await + .expect("list user A"); + let listed_b = service + .list_user_models(user_b.clone()) + .await + .expect("list user B"); + assert_eq!(listed_a.len(), 1); + assert_eq!(listed_b.len(), 1); + assert_eq!(listed_a[0].name, alias); + assert_eq!( + listed_b[0].name, alias, + "aliases are unique per owner, not globally" + ); + + let admitted = service + .admit_model_offering(user_a.clone(), model_a.clone()) + .await + .expect("owner admits model"); + assert_eq!(admitted.access_kind, ModelAccessKind::CloudByok); + assert_eq!(admitted.api_key, secret_a); + assert_eq!(admitted.wire_model_name.as_deref(), Some("deepseek-chat")); + assert!(!format!("{admitted:?}").contains(secret_a)); + + let error = service + .admit_model_offering(user_b.clone(), model_a.clone()) + .await + .expect_err("another user cannot admit the Offering"); + assert_eq!(error.0, StatusCode::NOT_FOUND); + + for (user_id, model_id) in [(&user_a, &model_a), (&user_b, &model_b)] { + sqlx::query("DELETE FROM user_llm_models WHERE user_id = ? AND model_id = ?") + .bind(user_id) + .bind(model_id) + .execute(&pool) + .await + .expect("clean user model fixture"); + } +} + +#[tokio::test] +#[ignore = "requires live DB: run with ASTRA_TEST_DB_IT=1"] +#[serial] +async fn compatible_byok_admission_rechecks_trust_and_owner() { + let (shared_pool, settings) = common::setup_pool_and_settings().await; + // Run this same production entrypoint in both deployment modes. + let strict = std::env::var("ASTRA_BYOK_ENDPOINT_POLICY").as_deref() == Ok("trusted-domains"); + let pool = shared_pool.get().clone(); + let encryptor = FernetTokenEncryptor::new("compatible-db-it-key").unwrap(); + let service = + DatabaseModelService::new(settings, Arc::new(encryptor.clone())).with_pool(shared_pool); + let owner = format!("compatible_{}", Uuid::new_v4().simple()); + let model_id = Uuid::new_v4().to_string(); + let host = format!("{}.example.com", Uuid::new_v4().simple()); + let domain_id = Uuid::new_v4().to_string(); + sqlx::query( + "INSERT INTO user_llm_models (model_id, user_id, model_alias, model_name, provider, \ + api_key_encrypted, base_url, context_window, is_default, is_active) \ + VALUES (?, ?, 'gateway', 'upstream-model', 'openai-compatible', ?, ?, 128000, 1, 1)", + ) + .bind(&model_id) + .bind(&owner) + .bind(encryptor.encrypt("test-secret").unwrap()) + .bind(format!("https://{host}/v1")) + .execute(&pool) + .await + .unwrap(); + let initial = service + .admit_model_offering(owner.clone(), model_id.clone()) + .await; + assert_eq!( + initial.is_err(), + strict, + "default public policy must not require registry setup" + ); + sqlx::query("INSERT INTO runtime_llm_trusted_domains (domain_id, domain_host, domain_port, is_enabled) VALUES (?, ?, 443, 1)") + .bind(&domain_id).bind(&host).execute(&pool).await.unwrap(); + let admitted = service + .admit_model_offering(owner.clone(), model_id.clone()) + .await + .unwrap(); + assert_eq!(admitted.wire_model_name.as_deref(), Some("upstream-model")); + assert_eq!(admitted.provider, "openai-compatible"); + sqlx::query("UPDATE user_llm_models SET base_url = ? WHERE model_id = ?") + .bind(format!("https://{host}:8443/v1")) + .bind(&model_id) + .execute(&pool) + .await + .unwrap(); + assert_eq!( + service + .admit_model_offering(owner.clone(), model_id.clone()) + .await + .is_err(), + strict, + "strict mode must not allow unapproved ports" + ); + sqlx::query("UPDATE user_llm_models SET base_url = ? WHERE model_id = ?") + .bind(format!("https://{host}/v1")) + .bind(&model_id) + .execute(&pool) + .await + .unwrap(); + assert!( + service + .admit_model_offering("another-user".into(), model_id.clone()) + .await + .is_err() + ); + sqlx::query("UPDATE runtime_llm_trusted_domains SET is_enabled = 0 WHERE domain_id = ?") + .bind(&domain_id) + .execute(&pool) + .await + .unwrap(); + let revoked = service + .admit_model_offering(owner.clone(), model_id.clone()) + .await; + assert_eq!( + revoked.is_err(), + strict, + "strict policy must recheck revocation at admission" + ); + sqlx::query("UPDATE user_llm_models SET base_url = 'https://127.0.0.1/v1' WHERE model_id = ?") + .bind(&model_id) + .execute(&pool) + .await + .unwrap(); + assert!( + service + .admit_model_offering(owner.clone(), model_id.clone()) + .await + .is_err(), + "unsafe persisted endpoints must be denied in both modes" + ); + let error = service + .check_user_model(owner.clone(), model_id.clone()) + .await + .unwrap_err(); + assert_eq!(error.0, StatusCode::BAD_REQUEST); + assert!(!error.1.detail.contains("test-secret")); + let error = service + .validate_user_model_endpoint(owner.clone(), "https://169.254.169.254/v1".into()) + .await + .unwrap_err(); + assert_eq!(error.0, StatusCode::BAD_REQUEST); + sqlx::query("DELETE FROM user_llm_models WHERE user_id = ? AND model_id = ?") + .bind(&owner) + .bind(&model_id) + .execute(&pool) + .await + .unwrap(); + sqlx::query("DELETE FROM runtime_llm_trusted_domains WHERE domain_id = ?") + .bind(&domain_id) + .execute(&pool) + .await + .unwrap(); +} diff --git a/crates/services/tests/user_model_probe_db_it.rs b/crates/services/tests/user_model_probe_db_it.rs new file mode 100644 index 0000000000..7ac32e7857 --- /dev/null +++ b/crates/services/tests/user_model_probe_db_it.rs @@ -0,0 +1,206 @@ +//! Real model service + isolated DB + strict, non-billable provider fixture. +//! Run with ASTRA_TEST_DB_IT=1, ASTRA_TEST_DATABASE=$ASTRA_DATABASE, +//! ASTRA_ALLOW_INSECURE_DEFAULTS=1 and ASTRA_BYOK_DEEPSEEK_BASE_URL set to +//! an unused loopback HTTP origin (e.g. http://127.0.0.1:18994). +//! Explicitly selected with --features external-contract-tests; ordinary +//! online lanes do not provide this operator-only endpoint override. +mod common; +#[path = "common/isolated_database.rs"] +mod isolated_database; + +use astra_services::{ + DatabaseModelService, FernetTokenEncryptor, ModelService, + models::{UserModelCreateRequestData, UserModelUpdateRequestData}, +}; +use axum::{Json, Router, http::StatusCode, routing::post}; +use serde_json::{Value, json}; +use std::sync::Arc; + +#[tokio::test] +#[ignore = "requires isolated MatrixOne DB and loopback DeepSeek fixture override"] +async fn user_model_create_rotate_and_probe_enforce_provider_wire_contract() { + let settings = astra_core::MatrixOneSettings::from_env(); + isolated_database::require_isolated_database(&settings.database); + assert_eq!( + std::env::var("ASTRA_ALLOW_INSECURE_DEFAULTS").as_deref(), + Ok("1") + ); + let base = std::env::var("ASTRA_BYOK_DEEPSEEK_BASE_URL").expect("loopback fixture origin"); + let url = reqwest::Url::parse(&base).unwrap(); + assert_eq!(url.scheme(), "http"); + assert_eq!(url.host_str(), Some("127.0.0.1")); + let app = Router::new() + .route( + "/chat/completions", + post( + |headers: axum::http::HeaderMap, Json(body): Json| async move { + // Each simulated upstream enforces its own documented + // field; do not reuse the serializer under test here. + let (limit, forbidden) = match body["model"].as_str() { + Some("deepseek-chat") => ("max_tokens", "max_completion_tokens"), + _ => ("max_completion_tokens", "max_tokens"), + }; + let status = if headers + .get("authorization") + .is_none_or(|v| v != "Bearer valid-key" && v != "Bearer rotated-key") + { + StatusCode::UNAUTHORIZED + } else if body["model"] != "deepseek-chat" && body["model"] != "o3" { + StatusCode::NOT_FOUND + } else if body[limit] != 32 || body.get(forbidden).is_some() { + StatusCode::BAD_REQUEST + } else { + StatusCode::OK + }; + ( + status, + Json(json!({"error":{"message":"strict fixture rejection"}})), + ) + }, + ), + ) + .route( + "/v1/messages", + post( + |headers: axum::http::HeaderMap, Json(body): Json| async move { + let status = if headers + .get("x-api-key") + .is_none_or(|v| v != "valid-key" && v != "rotated-key") + { + StatusCode::UNAUTHORIZED + } else if body["model"] != "claude-sonnet-4-5" { + StatusCode::NOT_FOUND + } else if headers + .get("anthropic-version") + .is_none_or(|v| v != "2023-06-01") + || body["max_tokens"] != 32 + || body.get("max_completion_tokens").is_some() + { + StatusCode::BAD_REQUEST + } else { + StatusCode::OK + }; + ( + status, + Json(json!({"error":{"message":"strict fixture rejection"}})), + ) + }, + ), + ); + let listener = tokio::net::TcpListener::bind(("127.0.0.1", url.port().unwrap())) + .await + .unwrap(); + let fixture = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + let (pool, settings) = common::setup_pool_and_settings().await; + let encryptor = FernetTokenEncryptor::new("provider-wire-test-only-key").unwrap(); + let service = + DatabaseModelService::new(settings, Arc::new(encryptor.clone())).with_pool(pool.clone()); + let owner = uuid::Uuid::new_v4().to_string(); + let make_request = |name: &str, model: &str, key: &str| UserModelCreateRequestData { + name: name.into(), + provider: "deepseek".into(), + model: model.into(), + base_url: None, + api_key: key.into(), + context_window: 128000, + is_default: true, + }; + for (model, key) in [ + ("deepseek-chat", "invalid-key"), + ("missing-model", "valid-key"), + ] { + let error = service + .create_user_model(owner.clone(), make_request("rejected", model, key)) + .await + .unwrap_err(); + assert_eq!(error.0, StatusCode::BAD_REQUEST); + assert!( + service + .list_user_models(owner.clone()) + .await + .unwrap() + .is_empty(), + "failed create persisted data" + ); + } + let created = service + .create_user_model( + owner.clone(), + make_request("fixture", "deepseek-chat", "valid-key"), + ) + .await + .unwrap(); + // Fixed official endpoints are not user-overridable. Seed their *test row* + // with this loopback fixture to exercise the real rotate/probe service paths + // without adding a production transport bypass or spending a live API key. + for (provider, model) in [ + ("deepseek", "deepseek-chat"), + ("openai", "o3"), + ("anthropic", "claude-sonnet-4-5"), + ] { + sqlx::query("UPDATE user_llm_models SET provider = ?, model_name = ?, api_key_encrypted = ? WHERE user_id = ? AND model_id = ?") + .bind(provider).bind(model).bind(encryptor.encrypt("valid-key").unwrap()).bind(&owner).bind(&created.model_id).execute(pool.get()).await.unwrap(); + service + .check_user_model(owner.clone(), created.model_id.clone()) + .await + .unwrap(); + let before: String = sqlx::query_scalar( + "SELECT api_key_encrypted FROM user_llm_models WHERE user_id = ? AND model_id = ?", + ) + .bind(&owner) + .bind(&created.model_id) + .fetch_one(pool.get()) + .await + .unwrap(); + let error = service + .update_user_model( + owner.clone(), + created.model_id.clone(), + UserModelUpdateRequestData { + api_key: Some("invalid-key".into()), + ..Default::default() + }, + ) + .await + .unwrap_err(); + assert_eq!(error.0, StatusCode::BAD_REQUEST, "{provider}: {error:?}"); + let after: String = sqlx::query_scalar( + "SELECT api_key_encrypted FROM user_llm_models WHERE user_id = ? AND model_id = ?", + ) + .bind(&owner) + .bind(&created.model_id) + .fetch_one(pool.get()) + .await + .unwrap(); + assert_eq!(before, after, "failed rotation overwrote credential"); + service + .update_user_model( + owner.clone(), + created.model_id.clone(), + UserModelUpdateRequestData { + api_key: Some("rotated-key".into()), + ..Default::default() + }, + ) + .await + .unwrap(); + let after: String = sqlx::query_scalar( + "SELECT api_key_encrypted FROM user_llm_models WHERE user_id = ? AND model_id = ?", + ) + .bind(&owner) + .bind(&created.model_id) + .fetch_one(pool.get()) + .await + .unwrap(); + assert_eq!(encryptor.decrypt(&after).unwrap(), "rotated-key"); + service + .check_user_model(owner.clone(), created.model_id.clone()) + .await + .unwrap(); + } + service + .delete_user_model(owner, created.model_id) + .await + .unwrap(); + fixture.abort(); +} diff --git a/deployment/all-in-one/README.md b/deployment/all-in-one/README.md index e8dfd664b9..1dc30af2f5 100644 --- a/deployment/all-in-one/README.md +++ b/deployment/all-in-one/README.md @@ -132,6 +132,8 @@ astra login ASTRA_EDGE_WORKSPACE_DIR=/path/to/repo make stack-up-server-edge ``` +With the default self-hosted configuration, `astra login` prompts for the local Astra username and password. A hosted browser login is enabled only when that Server is deliberately configured with a matching Memoria issuer/API and `MEMORIA_WEB_URL`; a hosted website key cannot authenticate against this stack's unrelated local Memoria instance. + `stack-up-server-edge` starts the same compose stack and then launches a local host `astra-edge` process connected to `/edge/ws`. The edge process reads the selected Astra CLI profile token by default; set `ASTRA_TOKEN` if you need to diff --git a/docs/design/README.md b/docs/design/README.md index e43db56735..26e468ec4a 100644 --- a/docs/design/README.md +++ b/docs/design/README.md @@ -12,6 +12,7 @@ These documents describe target contracts. They should not be read as proof that | Domain | Canonical document | Owns | | --- | --- | --- | | System architecture | [ARCHITECTURE.md](ARCHITECTURE.md) | Runtime shape, state layers, system-wide invariants. | +| Authentication | [authentication.md](authentication.md) | Issuer-scoped identity, scoped credentials, session revocation and login discovery. | | Documentation architecture | [documentation-architecture.md](documentation-architecture.md) | Documentation class rules, domain ownership, and migration policy. | | Agent/provider model | [agent-backbone-capacity-provider.md](agent-backbone-capacity-provider.md) | Shared backbone semantics and capacity provider contract. | | Runtime lifecycle | [runtime-lifecycle.md](runtime-lifecycle.md) | Session, run, turn, task, plan, cancel, resume, recovery. | diff --git a/docs/design/authentication.md b/docs/design/authentication.md new file mode 100644 index 0000000000..ac984ced62 --- /dev/null +++ b/docs/design/authentication.md @@ -0,0 +1,106 @@ +# Authentication + +Authentication owns identity, session issuance and credential binding. Request-authorized external providers keep their existing authorization protocol; verified scoped-key providers reuse the canonical identity mapping and Astra refresh-session lifecycle. + +## Identity and provenance + +External principals are identified by issuer/provider and subject, never subject or email alone. Memoria's provider ID is `memoria:` followed by the SHA-256 of its normalized issuer URL. The issuer defaults to the configured API base URL; an explicitly stable `MEMORIA_ISSUER` lets administrators move the transport without changing the identity authority. + +`auth_external_identities` has primary key `(provider_id, external_subject)`. The mapping resolves an Astra account; it does not replace its roles or lifecycle. JWT origin and the runtime principal retain the verified provider identity. Provider-request authorization is a different principal variant and is not granted by scoped-key login. + +## One composition boundary + +The auth service captures validated Memoria settings during application composition and supplies one credential resolver to login, refresh, proxy routes and memory runtime consumers. A pool setter never selects a memory transport. Explicit fixture/admin transport overrides cannot redirect the user-scoped credential. + +The provider verifies the scoped-key API contract through `/auth/whoami`: active non-master personal key, exact owner, nonempty key ID, API version 1, scopes capability and memory-filter capability. HTTP redirects are not followed. Identity, model-provider and Memoria secrets must not be logged. + +## Atomic binding and sessions + +Login verifies the key, takes the canonical account lock, rechecks the key after waiting, and commits identity, encrypted binding and refresh session together. Concurrent first logins resolve one account. A deterministic credential token primary key identifies one binding per provider/account; replacements remove superseded ciphertext in the same transaction. The verified key ID identifies the upstream key generation. Astra also persists a local connection lifecycle nonce in credential metadata: ordinary login with the same active binding preserves it; key replacement or reconnect after disconnect generates a new nonce, even when the upstream key ID is unchanged. Failed issuance rolls back every login-owned row. + +Refresh validates the current Memoria credential online. Compare-and-revoke of the old refresh token prevents a concurrent disconnect from resurrecting its session. Access TTL is bounded to 15 minutes. Runtime consumers independently resolve current consent and deny inactive/missing accounts. + +## Disconnect and retention + +- Logout revokes one Astra session. It does not revoke other devices or turn off memory sharing. +- Disconnect removes Astra's stored Memoria credential and revokes the account's Astra refresh sessions atomically. It retains identity mapping, models and Work so explicit relinking recovers the same account. +- Disabling memory sharing is different: an identity-only key remains usable for sign-in. +- Memoria owns upstream key revocation and deletion of its accounts/memories. Astra disconnect must not claim to perform those actions. +- Deactivated/deleted Astra accounts cannot resolve their scoped credential. Bounded maintenance removes their stored ciphertext. Existing external mappings do not automatically recreate a deleted local account. +- Account-wide Work/history erasure remains governed by the account retention contract; sign-out, disconnect and temporary verification outages never erase it. + +## Legacy migration + +The old `auth_memoria_identities` table is a read-only migration source. Its missing issuer cannot safely be guessed. Administrators must set `MEMORIA_LEGACY_ISSUER` to the known original issuer, matching the configured issuer. A fresh verified login then moves that mapping to the canonical table, preserves the Astra user ID and its models/Work, replaces the credential, and revokes old sessions in one transaction. Without that assertion, login fails with 409 and legacy refresh sessions fail closed. Do not set this option when pointing an old Astra database at a different Memoria instance. + +## Fresh reauthentication + +Normal sign-in is separate from authorizing device trust, device re-enrollment, +or forced session takeover. `GET /auth/reauthenticate` is authenticated discovery: +local accounts receive `method: password`; Memoria accounts receive +`method: memoria` and a verification page under the configured `MEMORIA_WEB_URL`. +This URL is a trusted identity-verification endpoint, not a caller-selected URL. +It requires HTTPS except for explicit loopback development. + +The website requires a fresh email code sent to the current account's email. +This also works for GitHub/Google-created accounts with an accessible email; +an existing OAuth login session by itself is not fresh verification. Private +API-key login deployments keep their existing local-password reauthentication. +The email names the sensitive action. Codes are purpose/account/key-generation +bound, expire in five minutes, allow five attempts, and have a 60-second resend +cooldown. Code hashes use an application-keyed HMAC. Normal login codes cannot +be used here. Verification neither rotates the connection key nor changes +memory consent. + +Successful verification returns an opaque `msu_` proof, valid for two minutes. +The user copies it back to the originating action. It stays out of URLs, +browser storage, and logs. The SDK accepts +`reauthenticate({ memoriaProof }, purpose)`; the HTTP request is +`POST /auth/reauthenticate` with `{ "memoria_proof": "msu_…", "purpose": "device_trust" }`. +The existing password request remains supported for local accounts. Mixing +methods is rejected. + +Astra validates the active issuer/subject/connection generation online, then +consumes the website proof over the fixed HTTPS backchannel +`/api/auth/astra/reauthentication/consume`, with redirects disabled and a bounded +response. It checks the returned subject, key generation, purpose and timestamps +and revalidates its binding before issuing the existing five-minute `rp_` proof. +That proof is single-use, purpose-bound and additionally bound to the Memoria +identity, upstream key generation and Astra connection lifecycle. Disconnect +deletes pending proofs in the same account-locked transaction as the binding. +The lifecycle binding also prevents an in-flight proof issuance from becoming +usable after reconnect with the same upstream key. Legacy credential metadata +without a lifecycle nonce remains readable; the next login assigns one. Proofs +issued before this binding-format upgrade must be obtained again (their maximum +lifetime is five minutes). Device trust still requires the separate +device-possession challenge. Both proof exchanges fail closed on upstream +revocation, disconnect, identity mismatch or replay; failed exchanges require +fresh evidence rather than bypassing proof checks. + +The website must be deployed with the reauthentication API before Astra clients +can use this path. An unavailable verifier blocks only the sensitive operation, +not ordinary sign-in or chat. Astra does not need a Memoria master key or a new +shared signing secret. + +## Client/deployment contract + +`GET /auth/methods` advertises the Server's website and issuer. An unset website retains interactive password login, including all-in-one deployments. The CLI falls back to the older password journey only on discovery 404, not on outages or malformed configuration. Explicit username/password and manual scoped-key login remain available. + +Browser login URLs require HTTPS except for explicit loopback development addresses. The callback remains bound to 127.0.0.1 with exact Origin, nonce, method, content-type and bounded request validation. Windows passes the URL as child-process environment data, not shell source. + +## Verification + +Focused coverage lives in `memoria_auth_db_it`, `memoria_auth_http`, CLI auth-flow tests and runtime consent-admission tests. It includes issuer separation, concurrent login/relink, post-write failure rollback, legacy migration, disconnect, inactive-account retention, source-configuration mismatch, read-only extraction and login discovery. Actual Windows browser launch and live OAuth callbacks require platform/deployment testing in addition to deterministic contracts. + +`memoria_live_contract_it` additionally runs against a real Memoria API implementing scoped-key API version 1. It creates disposable identity-only, read-only and read-write keys, verifies stable Astra account mapping, and revokes each key through Memoria before checking refresh rejection. Run only against an isolated Memoria test service and an explicitly designated test database: + +```bash +# Set the isolated MatrixOne connection variables as in the DB testing guide. +export ASTRA_TEST_DB_IT=1 +export ASTRA_TEST_DATABASE=review_memoria_contract +export ASTRA_TEST_MEMORIA_URL=http://127.0.0.1:18104 +# Set ASTRA_TEST_MEMORIA_MASTER_KEY to the isolated Memoria service's test key. +make test-memoria-auth-online-contract +``` + +The test issues and revokes upstream keys; never point it at a production service. diff --git a/docs/design/memory.md b/docs/design/memory.md index c1fedfdba6..0fd7e9adf6 100644 --- a/docs/design/memory.md +++ b/docs/design/memory.md @@ -38,6 +38,16 @@ Memory loading is intent-driven: The design does not require one physical backend. Vector, fulltext, graph, tabular, or MCP-backed memory can coexist as long as they satisfy the same provenance, confidence, deletion, and trace contract. +## Scoped credential admission + +The Server uses the application-scoped credential resolver owned by authentication. Each operation resolves the current owner binding and generation; no master-key fallback is inferred by a generic pool builder. + +- Missing binding or `none`: normal disabled state. Prompt recall reports `NotAttempted` and does not contact Memoria. +- `read_only`: recall is allowed; write-oriented extraction, reflection and session-end cleanup are not admitted. +- `read_write`: read and write operations are allowed. Transport checks remain in place to catch revocation or changes after admission. + +The background coordinator may launch a lightweight admission task, but it checks consent before loading snapshots, resolving an LLM, generating memory, or scheduling persistence. See [authentication](authentication.md) for issuer, credential replacement and retention. + ## Learning boundary Memory is not training data by default. Learning artifacts require consent, redaction, quality gate, lineage, and deletion propagation. diff --git a/docs/design/model-access-and-inference.md b/docs/design/model-access-and-inference.md index c500c2db21..d8917b217d 100644 --- a/docs/design/model-access-and-inference.md +++ b/docs/design/model-access-and-inference.md @@ -1,7 +1,7 @@ # Model access and inference > Status: target design contract. -> Last updated: 2026-07-20. +> Last updated: 2026-09-06. Model access and inference defines how Astra presents model capability as a product, binds cloud accounts, resolves an eligible model to a trusted execution path, and records inference usage consistently across Web, CLI, Server, and Edge. @@ -30,7 +30,7 @@ It does not own: - A new Astra Cloud user can use an administrator-approved model without understanding provider wiring. - Astra can link a user or organization to TaaS without making TaaS a special agent runtime. -- A user can use a non-TaaS provider or local model without uploading its secret to Astra Server. +- A user can explicitly choose Cloud BYOK, where Astra Server encrypts the provider credential and executes inference, or This device, where the secret is never uploaded. - Web, CLI, Server Only, and Edge + Server expose the same model and run semantics. - Every inference has an immutable execution placement, credential owner, billing owner, policy decision, and durable identity. - Account, entitlement, credential, endpoint, and billing failures are visible and recoverable without corrupting run state. @@ -39,7 +39,8 @@ It does not own: ## Non-goals - A generic provider plugin marketplace or routing DSL. -- Arbitrary user-supplied inference URLs on Astra Cloud Server. +- Ungoverned or request-scoped inference URLs on Astra Cloud Server. Personal + BYOK configuration may register public HTTPS endpoints under the policy below. - Separate agent loops for TaaS, direct providers, or Edge models. - Silent fallback across billing or data boundaries. - Treating cached balance, health, or entitlement data as authoritative without freshness. @@ -78,6 +79,7 @@ must preserve the first-page server decision. | Product source | Credential owner | Billing owner | Execution | Availability | | --- | --- | --- | --- | --- | | Astra Cloud | User-linked TaaS account | User TaaS account | Server | All clients | +| Cloud BYOK | User | User external account | Server | All clients | | Workspace | Organization | Organization account | Server | Authorized workspace clients | | This device | User/device | User external account or local compute | Edge | While the bound Edge is available | | Self-hosted | Deployment administrator | Deployment administrator | Server | Self-hosted deployment | @@ -103,6 +105,7 @@ struct ModelAccessView { enum ModelAccessKind { AstraCloud, + CloudByok, Workspace, ThisDevice, SelfHosted, @@ -123,6 +126,84 @@ Astra Cloud is the default personal model-access product. It is backed by a TaaS The model picker must not show both `Astra Cloud` and `My TaaS` for the same personal account. TaaS appears where the account or billing relationship matters, not as a duplicate model source. +### Cloud BYOK + +Cloud BYOK is an explicit personal access source for users who want Astra +Server to remain available without a connected Edge while billing inference to +their own provider account. + +- `POST /me/models` accepts a provider credential once over an authenticated + TLS connection. Owner identity is derived from the Astra access token; the + request cannot select another owner. +- Astra Server encrypts the credential with the configured secret backend and + stores only ciphertext in `user_llm_models`. +- `GET /me/models`, `GET /me/models/{model_id}`, catalog, trace, error, and + audit projections never return credential material or a reversible prefix. +- `PUT /me/models/{model_id}` rotates the credential or changes activation and + default state. `DELETE /me/models/{model_id}` removes only the authenticated + user's row. Cross-owner access returns `404` so resource existence is not + disclosed. +- User configuration presents OpenAI, Anthropic, DeepSeek, and + OpenAI-compatible. The first three use fixed official endpoints. The + `openai-compatible` provider requires an HTTPS `base_url` and a model ID; + public HTTPS endpoints do not require administrator registration by default. + `ASTRA_BYOK_ENDPOINT_POLICY=trusted-domains` opts a deployment into the + existing administrator registry; a host-only entry admits port 443 only. + The default is `public-https`; invalid configuration fails closed. +- Custom endpoint policy is rechecked at create, credential rotation, probe, + activation, and inference admission. Each outbound attempt pins a freshly + validated public DNS address set and forbids redirects. DNS queries go directly + to system-configured nameservers instead of OS synthetic-address caches; + `ASTRA_BYOK_DNS_SERVERS` optionally selects comma-separated DNS IPs (ports + optional); prefix an entry with `tcp://` for TCP-only DNS without a UDP + attempt or UDP fallback. This is an explicit operator route, not automatic + retry of rejected private/Fake-IP answers. Nameserver priority is preserved + instead of racing answers from different DNS servers. There is no hardcoded + public DNS or OS resolver fallback. + Direct egress is the default. Operators may set `ASTRA_BYOK_PROXY_URL` to an + HTTP/HTTPS CONNECT or SOCKS5/SOCKS5h proxy. A request-owned authenticated + loopback adapter sends only validated public IP targets to that proxy, + retaining the origin hostname for end-to-end TLS, SNI and Host. SOCKS5h does + not delegate target DNS. Dropping the client cancels its tunnel tasks. + Ambient proxy variables cannot override this route; proxy failure never + falls back to direct access. Private/Fake-IP DNS answers remain rejected. +- Authenticated `POST /me/models/validate-endpoint` accepts only `base_url` and + returns `204` after URL, policy, and DNS checks, without provider HTTP traffic, + credentials, or stored model state. CLI runs it immediately after custom URL + input, before asking for the key. This preflight is advisory: every later + operation still enforces policy and every outbound attempt revalidates DNS. + Server DNS/egress configuration failures return `502` with code + `model_endpoint_network`; malformed or forbidden URLs remain `400`. +- `astra model add` prompts for provider, model ID, alias, context window, + default selection and a hidden key. Explicit flags and `--api-key-stdin` + support non-interactive configuration. Custom endpoints reuse the existing + OpenAI transport; adding user ownership does not create another agent loop. +- Inference admission revalidates `(user_id, offering_id, is_active)` and + decrypts the current credential immediately before provider execution. +- Create, credential rotation and explicit probe validate connectivity with a + small output budget using the same provider-specific wire-field rule as + inference: OpenAI (including o-series) and the generic OpenAI-compatible + adapter use `max_completion_tokens`; native DeepSeek Chat Completions and + Anthropic Messages use `max_tokens`. Sharing the Chat Completions message + format does not imply identical optional parameters. Provider credential/model + errors remain failures; failed checks do not persist a new or rotated credential. +- Memoria-authenticated identities can only use their own BYOK Offerings; + deployment Offerings are excluded from both their catalog and execution. + `ASTRA_DEPLOYMENT_MODE=cloud-byok` applies the same restriction to all + Server-authenticated users. `self-hosted` (the default) retains deployment + models for non-Memoria users. Admin registry management is separate from + end-user inference eligibility. Unknown mode values do not grant access. +- Background memory selectors obey the same owner eligibility and fresh + admission checks. A memory read/write grant never authorizes spending a + deployment model credential. Without an explicitly eligible background model + route, personal Cloud BYOK uses the existing deterministic/degraded path; + it does not silently select the deployment registry or a personal default. +- Missing a personal default in Cloud BYOK requires model configuration or + selection; it never falls back to a deployment reasoning model. The same + owner gate applies when resuming runs and executing child runs. +- Cloud BYOK never silently imports an existing This device credential. Moving + a credential between those access sources requires an explicit user action. + ### Workspace Workspace access is owned and paid for by an organization. Personal Cloud and Workspace remain distinguishable even when they expose the same upstream model. @@ -133,6 +214,9 @@ Routing between them is allowed only when policy explicitly permits crossing the Non-TaaS provider credentials, private endpoints, Ollama, LM Studio, and other user-local models belong to This device. +This device remains distinct from Cloud BYOK. Selecting it is the explicit +choice that keeps a personal provider credential outside Astra Server. + - Secrets remain in the device vault. - Edge advertises a typed, leased capability and non-secret model metadata. - Server remains authoritative for the canonical run, transcript, task, route summary, and usage projection. @@ -140,7 +224,7 @@ Non-TaaS provider credentials, private endpoints, Ollama, LM Studio, and other u ### Self-hosted -A deployment administrator may register Server-local or organization-trusted inference endpoints. This does not authorize ordinary Astra Cloud users to upload arbitrary Server-side provider URLs or keys. +A deployment administrator may register Server-local or organization-trusted inference endpoints. Ordinary users remain subject to the public-network Cloud BYOK policy; administrator access does not implicitly authorize their private-network endpoints. ## Product surfaces @@ -838,6 +922,9 @@ Retry is coordinated at one layer. Server, gateway, and provider adapters cannot - TaaS instance registration validates origin, redirects, DNS results, and network policy. - Normal users cannot turn a binding request into arbitrary Server egress. - Secret material never appears in profile responses, route records, transcript, journal, SSE, traces, errors, or snapshots. +- Cloud BYOK ciphertext is owner-scoped at every query and mutation boundary; + plaintext exists only in bounded process memory while checking or executing + the selected provider request. - Effective Offering IDs are principal-bound and revalidated. - Tenant/organization/user/device ownership is enforced at query and mutation boundaries. - Cache namespaces include trust, connection, credential generation, and tenant scope where content may be sensitive. @@ -954,7 +1041,9 @@ Cover Web, CLI + Server, Server Only, and Edge + Server: - A normal user can understand available models without provider configuration knowledge. - The UI always states execution placement and billing owner when it affects a decision. - TaaS account handling never becomes a special branch in the agent loop. -- Non-TaaS personal credentials remain on Edge. +- Non-TaaS personal credentials remain on Edge when the user selects This + device; Cloud BYOK credentials are explicitly uploaded and encrypted on + Astra Server. - All inference purposes use one resolver and invocation contract. - Every upstream request is attributable to a durable provider attempt. - No client can select an endpoint, credential, or placement directly. diff --git a/docs/guides/testing.md b/docs/guides/testing.md index 16cc15de47..8b5c56ca2a 100644 --- a/docs/guides/testing.md +++ b/docs/guides/testing.md @@ -42,6 +42,52 @@ cargo check --manifest-path Cargo.toml ## Live MatrixOne system E2E +Memoria identity/credential fixtures require `ASTRA_TEST_DB_IT=1` and an +explicit `ASTRA_TEST_DATABASE` matching the effective database name. The normal +online runner supplies its isolated lane database; local callers must designate +their disposable database rather than relying on a developer-specific prefix. + +Contracts requiring an external scoped-key Memoria API or a real provider are +separately gated by the services `external-contract-tests` feature. They are not +selected by ordinary MatrixOne online lanes. To run the scoped Memoria contract, +provision a test Memoria API with scoped keys, set `ASTRA_TEST_MEMORIA_URL`, +`ASTRA_TEST_MEMORIA_MASTER_KEY`, `ASTRA_TEST_DATABASE` and the test MatrixOne +connection variables, then run `make test-memoria-auth-online-contract`. +Missing configuration fails the explicitly selected contract; it is not reported +as a passing test. No production credentials are needed in fork CI. + +The key-free BYOK network smoke is also explicit: + +```bash +ASTRA_BYOK_DNS_SERVERS=tcp://223.5.5.5 \ +ASTRA_TEST_BYOK_MODELS_URL=https://api.moonshot.cn/v1/models \ +CARGO_INCREMENTAL=0 cargo test --locked -p astra-services \ + --features external-contract-tests --test byok_live_network -- --ignored +``` + +It expects HTTP 401 with no provider key and proves DNS/TLS/HTTP reachability, +not successful model inference. + +The provider-wire regression uses a strict loopback HTTP fixture and a disposable +MatrixOne database, without real provider keys. Create the designated database +first and supply its `MATRIXONE_*` connection settings. Choose an unused loopback +port for the fixture: + +```bash +ASTRA_TEST_DB_IT=1 ASTRA_DATABASE_PREFIX= \ +ASTRA_DATABASE=astra_probe_test ASTRA_TEST_DATABASE=astra_probe_test \ +ASTRA_ALLOW_INSECURE_DEFAULTS=1 \ +ASTRA_BYOK_DEEPSEEK_BASE_URL=http://127.0.0.1:18994 \ +CARGO_INCREMENTAL=0 cargo test --locked -p astra-services \ + --features external-contract-tests --test user_model_probe_db_it -- --ignored +``` + +This covers create, credential rotation, explicit probe and failed-write +preservation. Official OpenAI/Anthropic probe and rotation tests seed only their +fixture rows with loopback endpoints; production official endpoints remain fixed. +`memoria_reauthentication_http` separately covers same-key reconnect, pending +proof invalidation and an in-flight verification crossing disconnect/reconnect. + ```bash ASTRA_TEST_DB_IT=1 \ ASTRA_TEST_E2E_SECRET=system-matrix-e2e-secret \ diff --git a/docs/reference/api-reference.md b/docs/reference/api-reference.md index be062d519f..c443aebf9d 100644 --- a/docs/reference/api-reference.md +++ b/docs/reference/api-reference.md @@ -7,7 +7,7 @@ Interactive docs: `http://localhost:17001/docs` (Swagger UI) | `http://localhost ## Authentication All protected endpoints require a JWT. The public authentication exceptions -are `/live`, `/ready`, `/health`, `/auth/register`, `/auth/login`, `/auth/refresh`, and +are `/live`, `/ready`, `/health`, `/auth/register`, `/auth/login`, `GET /auth/methods`, `POST /auth/memoria`, `/auth/refresh`, and `/auth/logout` (refresh/logout authenticate the supplied refresh token): ``` @@ -48,6 +48,28 @@ Authorization: Bearer Returns current user info. +### GET /auth/methods + +Returns Server-owned login discovery. With no configured browser integration: + +```json +{"password":true,"memoria":null} +``` + +When enabled, `memoria` contains `issuer` and `authorization_url` (the website base URL). CLI login uses this discovery; it does not assume a hosted website for a self-hosted Server. + +### POST /auth/memoria + +Accepts `{"connection_key":""}`. The auth service verifies the key online, then commits the provider-scoped identity, encrypted credential and refresh session in one transaction. Response: `user_id`, `access_token`, `refresh_token`, `token_type`, `expires_in` (at most 900 seconds), `memory_access`, and `granted_scopes`. No Memoria secret is returned. + +Invalid keys fail with 401; unavailable verification fails with 503; legacy identity mappings without explicitly configured provenance fail with 409. An identity-only key can log in but cannot access memory. + +### DELETE /auth/memoria + +Requires an Astra access token. Removes the account's stored Memoria credential and revokes all its Astra refresh sessions in one transaction; returns 204. Existing access tokens then fail session validation. The service operation is idempotent; retrying with an already revoked access token returns 401. + +Account identity, Work/history and Memoria memories are retained. This disconnects Astra; it does not delete the Memoria account or revoke keys at their issuer. Upstream key revocation remains owned by Memoria/its integration settings. Ordinary `/auth/logout` signs out only the submitted session and does not disconnect other devices. + --- ## Agents diff --git a/docs/reference/cli-commands.md b/docs/reference/cli-commands.md index 95571da124..512325e2c2 100644 --- a/docs/reference/cli-commands.md +++ b/docs/reference/cli-commands.md @@ -2,6 +2,55 @@ Current Rust CLI reference for the single `astra` CLI, including `astra admin`. +## Personal BYOK model setup + +Run `astra model add` after logging in to configure a personal model interactively. + +- **Model ID from your provider** is the exact identifier documented by the + provider. Astra sends it unchanged to the provider API. +- **Configuration alias in Astra** is the account-local name used to select that + configuration. Press Enter to use the model ID, or enter a unique alias to + distinguish different credentials/endpoints for the same model. + +For example, an alias `work-model` for provider model ID `deepseek-v4-flash` +is selected with `astra chat --model work-model`. Commands such as +`astra model show work-model` and `astra model probe work-model` use the alias. +The alias is not an account username or a provider model ID. + +Scripts can still supply both explicitly: + +```bash +astra model add work-model --provider deepseek --model deepseek-v4-flash --api-key-stdin +``` + +With `--api-key-stdin`, the alias and required provider/model arguments remain +mandatory; only the interactive wizard offers a default alias. Keys should be +entered at the hidden prompt or supplied via stdin, never as command arguments. + +For OpenAI-compatible services, enter the provider's public HTTPS base URL +(including `/v1` if required). The CLI checks URL policy and DNS with the Server +before requesting the API key. This does not send a request to the provider or +prove that the key/model is valid; adding the model performs that check next. +The key prompt hides typed input; press Enter to submit it. + +Public HTTPS endpoints do not require individual administrator registration by +default. Operators can opt into strict host/port approval by setting +`ASTRA_BYOK_ENDPOINT_POLICY=trusted-domains` on every Server replica and using +`PUT /admin/llm/trusted-domains`. Both modes block private/metadata addresses, +validate and pin DNS results, and disable redirects. Custom BYOK uses direct +queries to the Server's configured DNS servers, not OS Fake-IP caches. Operators +may override DNS IPs with `ASTRA_BYOK_DNS_SERVERS` and select an HTTP(S) CONNECT +or SOCKS5 proxy with `ASTRA_BYOK_PROXY_URL`. Proxies connect to validated public +IPs, retaining origin TLS verification; ambient proxy variables are not used. +Direct UDP DNS can still be intercepted by a VPN/proxy. Use an explicit +TCP-only entry such as `ASTRA_BYOK_DNS_SERVERS=tcp://10.0.0.53:53` when a +reachable DNS server supports TCP. TCP-only entries never fall back to UDP; +private/Fake-IP answers remain rejected. DNS failures retain resolver diagnostics +in Server logs without exposing them in public API errors. +These are Server settings, not options ordinary CLI users need to fill in. +The default policy is `public-https`; invalid values fail closed. Upgrade the +Server together with the CLI to provide `/me/models/validate-endpoint`. + ## Installation For day-to-day development builds: @@ -42,6 +91,7 @@ Commands: ```bash # Auth +astra login # Server discovery: configured browser sign-in, otherwise password prompts astra register --username alice --email alice@example.com --password '***' astra login --username alice --password '***' astra interactive @@ -74,6 +124,8 @@ astra skill show [--version 1.0.0] astra skill status [--per-group 50] ``` +`astra login --username alice` explicitly selects password login. `astra login --manual` accepts a scoped connection key when browser handoff is unavailable. Older Servers returning 404 for `/auth/methods` retain the password journey; network errors do not silently select another provider. Browser addresses come from the target Server's `MEMORIA_WEB_URL`, not the CLI environment. + ## astra admin Global options: diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index 2ecb12c925..22daa0a30b 100644 --- a/docs/reference/configuration.md +++ b/docs/reference/configuration.md @@ -124,8 +124,13 @@ If `reasoning_offering_id` is not set, the server applies its governed default a ### Memoria - `MEMORIA_BASE_URL`, `MEMORIA_MASTER_KEY` +- `MEMORIA_ISSUER` — stable identity issuer URL; defaults to normalized `MEMORIA_BASE_URL`. Changing the issuer creates a different identity namespace. Keep it stable when changing only the service transport address. +- `MEMORIA_WEB_URL` — Server-owned browser sign-in website, advertised through `GET /auth/methods`. Unset preserves password login. Requires HTTPS except for explicit loopback development URLs. The CLI does not read this environment variable. +- `MEMORIA_LEGACY_ISSUER` — explicit administrator assertion of the issuer that owned pre-issuer Memoria identities. Migration is allowed only when it equals the configured issuer, after fresh key verification. Leave unset unless the provenance of the old database is known. - `MEMORIA_EMBEDDING_PROVIDER`, `MEMORIA_EMBEDDING_MODEL`, `MEMORIA_EMBEDDING_DIM`, `MEMORIA_EMBEDDING_API_KEY`, `MEMORIA_EMBEDDING_BASE_URL` +Scoped login, refresh, memory proxy, recall, extraction and session-end governance share the authentication service's provider configuration and credential resolver. Runtime builders do not independently select a transport from environment variables or fall back to a master key. See [authentication](../design/authentication.md). + ### Runtime tuning (optional) - `ASTRA_MAX_TURNS`, `ASTRA_PLAN_SUBTASK_MAX_TURNS`, `ASTRA_TURN_TIMEOUT_S` diff --git a/packages/sdk/src/__tests__/work-contract.test.ts b/packages/sdk/src/__tests__/work-contract.test.ts index 08c0afa6c7..227bd4afcf 100644 --- a/packages/sdk/src/__tests__/work-contract.test.ts +++ b/packages/sdk/src/__tests__/work-contract.test.ts @@ -678,6 +678,28 @@ test("reauthenticate returns one bounded purpose-bound proof", async () => { }); }); +test("Memoria reauthentication sends fresh evidence without a password", async () => { + const proof = { proof: "rp_opaque", purpose: "device_trust", expires_in: 300 }; + const fetchMock = vi.fn().mockResolvedValue(response(200, proof)); + globalThis.fetch = fetchMock; + const client = new AstraClient({ baseUrl: "https://astra.example" }); + const fresh = `msu_${"a".repeat(64)}`; + await expect(client.reauthenticate({ memoriaProof: fresh }, "device_trust")).resolves.toEqual(proof); + expect(JSON.parse(String(fetchMock.mock.calls[0][1].body))).toEqual({ memoria_proof: fresh, purpose: "device_trust" }); + await expect(client.reauthenticate({ memoriaProof: "stored-connection-key" }, "device_trust")).rejects.toThrow("fresh Memoria"); + expect(fetchMock).toHaveBeenCalledTimes(1); +}); + +test("reauthentication discovery rejects unsafe verification URLs", async () => { + const client = new AstraClient({ baseUrl: "https://astra.example" }); + for (const verification_url of ["javascript:alert(1)", "http://remote.example/verify", "https://user:password@example.com/verify"]) { + globalThis.fetch = vi.fn().mockResolvedValue(response(200, { method: "memoria", verification_url })); + await expect(client.getReauthenticationOptions()).rejects.toThrow(); + } + globalThis.fetch = vi.fn().mockResolvedValue(response(200, { method: "memoria", verification_url: "https://thememoria.ai/astra/reauthenticate" })); + await expect(client.getReauthenticationOptions()).resolves.toMatchObject({ method: "memoria" }); +}); + test("reauthenticate rejects a proof not bound to the requested purpose", async () => { globalThis.fetch = vi.fn().mockResolvedValue( response(200, { diff --git a/packages/sdk/src/client.ts b/packages/sdk/src/client.ts index d9415bdba8..7e5b966066 100644 --- a/packages/sdk/src/client.ts +++ b/packages/sdk/src/client.ts @@ -550,14 +550,36 @@ export class AstraClient { return this.fetch(PATH_AUTH_ME); } + async getReauthenticationOptions(): Promise<{ method: "password" } | { method: "memoria"; verification_url: string }> { + const value = await this.fetch<{ method: string; verification_url?: unknown }>(PATH_AUTH_REAUTHENTICATE); + if (value?.method === "password") return { method: "password" }; + if (value?.method === "memoria" && typeof value.verification_url === "string") { + const url = new URL(value.verification_url); + if (!url.username && !url.password && !url.hash && !url.search && + (url.protocol === "https:" || (url.protocol === "http:" && ["localhost", "127.0.0.1", "[::1]"].includes(url.hostname)))) { + return { method: "memoria", verification_url: url.toString() }; + } + } + throw new TypeError("reauthentication options are invalid"); + } + async reauthenticate( - password: string, + credential: string | { memoriaProof: string }, purpose: ReauthenticationPurpose, ): Promise { - if (password.length === 0 || password.length > 4096) { - throw new TypeError("password must be non-empty and bounded"); + let evidence: { password: string } | { memoria_proof: string }; + if (typeof credential === "string") { + if (credential.length === 0 || credential.length > 4096) { + throw new TypeError("password must be non-empty and bounded"); + } + evidence = { password: credential }; + } else { + if (!credential || Object.keys(credential).join() !== "memoriaProof" || !/^msu_[a-f0-9]{64}$/.test(credential.memoriaProof)) { + throw new TypeError("a fresh Memoria verification proof is required"); + } + evidence = { memoria_proof: credential.memoriaProof }; } - const raw = await this.post(PATH_AUTH_REAUTHENTICATE, { password, purpose }); + const raw = await this.post(PATH_AUTH_REAUTHENTICATE, { ...evidence, purpose }); if (!raw || typeof raw !== "object" || Array.isArray(raw)) { throw new TypeError("reauthentication response must be an object"); } diff --git a/scripts/dev/start-api.sh b/scripts/dev/start-api.sh index ef695b47c5..4cd19ed3ee 100755 --- a/scripts/dev/start-api.sh +++ b/scripts/dev/start-api.sh @@ -46,9 +46,26 @@ fi # Clean up old process record rm -f "$PID_FILE" -# Load .env early so DB host/port are available for the readiness check -if [ -f .env ]; then - set -a; source .env; set +a +# Load the selected env file early so DB host/port are available for the +# readiness check. ASTRA_ENV_FILE lets cross-repository local harnesses use an +# isolated configuration without rewriting a developer's normal .env. +ENV_FILE="${ASTRA_ENV_FILE:-.env}" +if [ -f "$ENV_FILE" ]; then + set -a; source "$ENV_FILE"; set +a +fi + +# A caller-selected env file is an explicit configuration boundary. Prevent +# the server's own config loader from filling missing values from the repo +# .env or user/system config after this script has deliberately omitted them. +if [ -n "${ASTRA_ENV_FILE:-}" ]; then + export ASTRA_CONFIG_SOURCE="${ASTRA_CONFIG_SOURCE:-explicit-env}" +fi + +# Cloud BYOK must resolve a per-user Memoria credential. This explicit switch +# guarantees an inherited shell variable cannot silently re-enable master-key +# fallback in the local cross-repository test topology. +if [ "${ASTRA_DISABLE_MEMORIA_MASTER_KEY:-}" = "1" ]; then + unset MEMORIA_MASTER_KEY fi API_PORT="${ASTRA_API_PORT:-17001}" diff --git a/scripts/schema/schema_inventory.py b/scripts/schema/schema_inventory.py index 23a6ddb92a..7f329ca015 100755 --- a/scripts/schema/schema_inventory.py +++ b/scripts/schema/schema_inventory.py @@ -940,6 +940,36 @@ class AutoIncrementMetadata: migration_owner="astra_services::storage / auth", product_owner="admin secret/token management and scoped provider credentials", ), + "user_llm_models": TableMetadata( + semantic_owner="astra_services::models", + state_class="durable owner-scoped BYOK model configuration and encrypted credential", + primary_query="owner-scoped catalog, alias lookup, default-model selection and inference admission", + retention_policy="retain while configured; explicit model removal deletes its credential; account erasure follows account retention", + rebuildability="not rebuildable without user-provided model settings and API key", + merge_guidance="personal offerings use the canonical model admission path; never expose them as deployment-wide models", + migration_owner="astra_services::storage / models", + product_owner="Cloud BYOK model configuration", + ), + "auth_external_identities": TableMetadata( + semantic_owner="astra_services::auth::verified", + state_class="durable issuer-scoped external identity fact", + primary_query="resolve Astra account by provider_id and external_subject; query provenance by astra_user_id", + retention_policy="retain through credential disconnect for account continuity; account erasure follows the account retention workflow", + rebuildability="not safely rebuildable without the trusted issuer and original account binding", + merge_guidance="canonical mapping shared by verified providers; do not add provider-specific identity state machines", + migration_owner="astra_services::storage / auth", + product_owner="authentication and account continuity", + ), + "auth_memoria_identities": TableMetadata( + semantic_owner="astra_services::auth / Memoria integration", + state_class="durable external-to-Astra identity mapping fact", + primary_query="read-only legacy migration lookup; never insert new identities", + retention_policy="remove the legacy row atomically after explicitly issuer-authorized migration to auth_external_identities", + rebuildability="cannot infer missing issuer; require administrator-confirmed original provenance", + merge_guidance="migration source only; canonical provider identity belongs to auth_external_identities", + migration_owner="astra_services::storage / auth", + product_owner="Memoria sign-in and Astra account continuity", + ), "auth_audit_logs": TableMetadata( semantic_owner="astra_services::auth::admin / auth session audit", state_class="durable auth audit event", diff --git a/scripts/schema/test_schema_inventory.py b/scripts/schema/test_schema_inventory.py index 276a914174..27bc2e7b7e 100755 --- a/scripts/schema/test_schema_inventory.py +++ b/scripts/schema/test_schema_inventory.py @@ -528,6 +528,7 @@ def test_auth_admin_model_config_tables_have_semantic_metadata(self) -> None: "auth_roles", "auth_refresh_tokens", "auth_tokens", + "auth_memoria_identities", "auth_audit_logs", "infra_llm_models", "runtime_llm_trusted_domains", diff --git a/web/__tests__/components/app/work-turn-composer.test.tsx b/web/__tests__/components/app/work-turn-composer.test.tsx index 9895c936e9..73673b17c4 100644 --- a/web/__tests__/components/app/work-turn-composer.test.tsx +++ b/web/__tests__/components/app/work-turn-composer.test.tsx @@ -31,10 +31,12 @@ vi.mock("next/navigation", () => ({ const forceTakeover = vi.hoisted(() => vi.fn()); const observeControl = vi.hoisted(() => vi.fn()); const abortControl = vi.hoisted(() => vi.fn()); +const reauthOptions = vi.hoisted(() => vi.fn()); vi.mock("@/app/(workspace)/works/[workId]/actions", () => ({ forceTakeoverWorkBranchAction: forceTakeover, observeWorkBranchControlAction: observeControl, abortWorkBranchControlAction: abortControl, + getWorkReauthenticationOptionsAction: reauthOptions, })); import { StrictMode } from "react"; @@ -43,6 +45,7 @@ import { WorkTurnComposer } from "@/components/app/work-turn-composer"; beforeEach(() => { vi.clearAllMocks(); + reauthOptions.mockResolvedValue({ method: "password" }); streamHarness.instances.length = 0; Object.defineProperty(globalThis.crypto, "randomUUID", { configurable: true, @@ -59,6 +62,29 @@ test("keeps continuation unavailable without a durable attachment", () => { expect(streamHarness.instances).toHaveLength(0); }); +test("Memoria takeover requests fresh verification and submits no password", async () => { + reauthOptions.mockResolvedValue({ method: "memoria", verification_url: "https://thememoria.ai/astra/reauthenticate" }); + forceTakeover.mockResolvedValue({ ok:false, status:403, code:"reauthentication_required", retryable:false }); + render(); + fireEvent.change(screen.getByRole("textbox", { name:"Guide this Work" }), {target:{value:"Continue"}}); + fireEvent.click(screen.getByRole("button", {name:"Send guidance"})); + act(() => { + streamHarness.instances[0]!.options.onEvent({ type:"error", code:"writer_conflict", message:"Active elsewhere", retryable:false, http_status:409, action_hints:["refresh_work"] }); + streamHarness.instances[0]!.options.onStateChange("disconnected"); + }); + fireEvent.click(screen.getByRole("button", {name:"Continue here"})); + const input = await screen.findByLabelText("One-time verification proof"); + expect(screen.queryByLabelText("Password")).not.toBeInTheDocument(); + expect(screen.getByRole("link", {name:/Verify your identity/})).toHaveAttribute("href", "https://thememoria.ai/astra/reauthenticate?purpose=session_forced_takeover"); + const fresh = `msu_${"a".repeat(64)}`; + fireEvent.change(input, {target:{value:fresh}}); + fireEvent.click(screen.getByRole("button", {name:"Confirm"})); + await waitFor(() => expect(forceTakeover).toHaveBeenCalled()); + expect(forceTakeover.mock.calls[0][0]).toMatchObject({ memoriaProof:fresh }); + expect(forceTakeover.mock.calls[0][0]).not.toHaveProperty("password"); + expect(await screen.findByLabelText("One-time verification proof")).toHaveValue(""); +}); + test("submits one typed Work turn and applies only decoded visible text", async () => { render( { expect(screen.queryByRole("button", { name: "Reconnect" })).not.toBeInTheDocument(); fireEvent.click(screen.getByRole("button", { name: "Continue here" })); expect(forceTakeover).not.toHaveBeenCalled(); - fireEvent.change(screen.getByLabelText("Password"), { + fireEvent.change(await screen.findByLabelText("Password"), { target: { value: "correct horse battery staple" }, }); fireEvent.click(screen.getByRole("button", { name: "Confirm" })); @@ -350,7 +376,7 @@ test("stops a durable takeover only while the server marks it abortable", async }), ); fireEvent.click(screen.getByRole("button", { name: "Continue here" })); - fireEvent.change(screen.getByLabelText("Password"), { target: { value: "password" } }); + fireEvent.change(await screen.findByLabelText("Password"), { target: { value: "password" } }); fireEvent.click(screen.getByRole("button", { name: "Confirm" })); fireEvent.click(await screen.findByRole("button", { name: "Stop moving" })); diff --git a/web/__tests__/lib/work-criteria-actions.test.ts b/web/__tests__/lib/work-criteria-actions.test.ts index 21bc8dd197..c724826adf 100644 --- a/web/__tests__/lib/work-criteria-actions.test.ts +++ b/web/__tests__/lib/work-criteria-actions.test.ts @@ -513,7 +513,10 @@ test("acquires branch control with exact attachment and causal basis", async () }); }); -test("reauthenticates before sending one sealed forced takeover", async () => { +test.each([ + { password: "correct horse battery staple" }, + { memoriaProof: `msu_${"a".repeat(64)}` }, +])("reauthenticates before sending one sealed forced takeover: %j", async credential => { const authorization = { proof: "opaque-step-up-proof", purpose: "session_forced_takeover", @@ -540,11 +543,11 @@ test("reauthenticates before sending one sealed forced takeover", async () => { writer_epoch: 4, canonical_root_hash: "a".repeat(64), }, - password: "correct horse battery staple", + ...credential, }), ).resolves.toEqual({ ok: true, operation }); expect(reauthenticate).toHaveBeenCalledWith( - "correct horse battery staple", + "memoriaProof" in credential ? credential : credential.password, "session_forced_takeover", ); expect(controlWorkBranch).toHaveBeenCalledWith("work-1", "branch-1", { @@ -565,6 +568,10 @@ test("reauthenticates before sending one sealed forced takeover", async () => { ); }); +test.each([null, "invalid", 42, [], {}])("rejects malformed reauthentication input before contacting Astra: %j", async input => { + await expect(forceTakeoverWorkBranchAction(input as never)).resolves.toMatchObject({ ok:false, status:400 }); +}); + test("observes and aborts the same durable control operation", async () => { const operation = { schema_version: 2, diff --git a/web/app/(workspace)/works/[workId]/actions.ts b/web/app/(workspace)/works/[workId]/actions.ts index 8fd530382d..84991217d7 100644 --- a/web/app/(workspace)/works/[workId]/actions.ts +++ b/web/app/(workspace)/works/[workId]/actions.ts @@ -737,7 +737,8 @@ type AcquireWorkBranchControlInput = { }; type ForceTakeoverWorkBranchInput = AcquireWorkBranchControlInput & { - password: string; + password?: string; + memoriaProof?: string; }; type WorkBranchControlOperationInput = { @@ -975,6 +976,11 @@ export async function acquireWorkBranchControlAction( } } +export async function getWorkReauthenticationOptionsAction() { + const runtime = await requireRuntimeClient({ auth: "required", operation: "verify identity before moving Work" }); + return runtime.sdk.getReauthenticationOptions(); +} + export async function forceTakeoverWorkBranchAction( input: ForceTakeoverWorkBranchInput, ): Promise { @@ -984,13 +990,13 @@ export async function forceTakeoverWorkBranchAction( "branchId", "expectedBranchRevision", "expectedControlBasis", - "password", + input && typeof input === "object" && "memoriaProof" in input ? "memoriaProof" : "password", "requestId", "workId", ]) || - typeof input.password !== "string" || - input.password.length < 1 || - input.password.length > 4096 + ("memoriaProof" in input + ? typeof input.memoriaProof !== "string" || !/^msu_[a-f0-9]{64}$/.test(input.memoriaProof) + : typeof input.password !== "string" || input.password.length < 1 || input.password.length > 4096) ) { return { ok: false, @@ -1005,7 +1011,7 @@ export async function forceTakeoverWorkBranchAction( operation: "continue Work on this device", }); const authorization = await runtime.sdk.reauthenticate( - input.password, + "memoriaProof" in input ? { memoriaProof: input.memoriaProof! } : input.password!, "session_forced_takeover", ); return { diff --git a/web/components/app/work-turn-composer.tsx b/web/components/app/work-turn-composer.tsx index f2629c3cdb..87cb19ea42 100644 --- a/web/components/app/work-turn-composer.tsx +++ b/web/components/app/work-turn-composer.tsx @@ -16,6 +16,7 @@ import { Button } from "@/components/ui/button"; import { abortWorkBranchControlAction, forceTakeoverWorkBranchAction, + getWorkReauthenticationOptionsAction, observeWorkBranchControlAction, } from "@/app/(workspace)/works/[workId]/actions"; @@ -96,6 +97,7 @@ export function WorkTurnComposer({ const [takingControl, setTakingControl] = useState(false); const [confirmingTakeover, setConfirmingTakeover] = useState(false); const [takeoverPassword, setTakeoverPassword] = useState(""); + const [takeoverMethod, setTakeoverMethod] = useState<{ method: "password" } | { method: "memoria"; verification_url: string }>({ method: "password" }); const [controlOperation, setControlOperation] = useState(null); const [abortingControl, setAbortingControl] = useState(false); @@ -397,6 +399,9 @@ export function WorkTurnComposer({ controlRequestId.current ?? `web-work-control:${crypto.randomUUID()}`; controlRequestId.current = requestId; try { + // A failed exchange may already have consumed the one-time evidence. + // Do not leave it available for an accidental retry. + if (takeoverMethod.method === "memoria") setTakeoverPassword(""); const result = await forceTakeoverWorkBranchAction({ workId, branchId, @@ -404,7 +409,7 @@ export function WorkTurnComposer({ requestId, expectedBranchRevision: branchRevision, expectedControlBasis: currentControlBasis, - password: takeoverPassword, + ...(takeoverMethod.method === "memoria" ? { memoriaProof: takeoverPassword } : { password: takeoverPassword }), }); if (!mounted.current) return; if (!result.ok) { @@ -417,9 +422,9 @@ export function WorkTurnComposer({ } setError( result.status === 401 - ? "Your password was not accepted or your sign-in expired. Nothing was moved." + ? "Identity verification was not accepted or your sign-in expired. Nothing was moved." : result.code === "reauthentication_required" || result.status === 403 - ? "Your password was not accepted. Nothing was moved." + ? "Identity verification was not accepted. Nothing was moved." : result.retryable ? "This Work could not move here yet. You can safely try again." : "This Work could not continue on this device.", @@ -508,6 +513,11 @@ export function WorkTurnComposer({ effects are kept for review and are not repeated automatically.

+ {takeoverMethod.method === "memoria" ? ( + + Verify your identity by email, then paste the one-time proof below + + ) : null} -