diff --git a/lib/crates/fabro-auth/src/extra_headers_source.rs b/lib/crates/fabro-auth/src/extra_headers_source.rs new file mode 100644 index 000000000..29ad806cd --- /dev/null +++ b/lib/crates/fabro-auth/src/extra_headers_source.rs @@ -0,0 +1,175 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use async_trait::async_trait; +use fabro_model::{Catalog, ProviderId}; + +use crate::credential_source::{CredentialSource, ResolvedCredentials}; + +/// Decorates another [`CredentialSource`] by appending fixed extra headers to +/// every credential it resolves. +/// +/// Headers already present on a credential (for example from explicit +/// provider configuration) are left untouched. +pub struct ExtraHeadersCredentialSource { + inner: Arc, + headers: HashMap, +} + +impl ExtraHeadersCredentialSource { + #[must_use] + pub fn new(inner: Arc, headers: HashMap) -> Self { + Self { inner, headers } + } +} + +#[async_trait] +impl CredentialSource for ExtraHeadersCredentialSource { + async fn resolve(&self, catalog: &Catalog) -> anyhow::Result { + let mut resolved = self.inner.resolve(catalog).await?; + for credential in &mut resolved.credentials { + for (name, value) in &self.headers { + if credential + .extra_headers + .keys() + .any(|existing| existing.eq_ignore_ascii_case(name)) + { + continue; + } + credential.extra_headers.insert(name.clone(), value.clone()); + } + } + Ok(resolved) + } + + async fn configured_providers(&self, catalog: &Catalog) -> Vec { + self.inner.configured_providers(catalog).await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ApiCredential, ResolveError}; + + struct StubSource { + credentials: Vec, + auth_issue_provider: Option, + configured_providers: Vec, + } + + #[async_trait] + impl CredentialSource for StubSource { + async fn resolve(&self, _catalog: &Catalog) -> anyhow::Result { + Ok(ResolvedCredentials { + credentials: self.credentials.clone(), + auth_issues: self + .auth_issue_provider + .iter() + .map(|provider| { + ( + provider.clone(), + ResolveError::RefreshTokenMissing(provider.clone()), + ) + }) + .collect(), + }) + } + + async fn configured_providers(&self, _catalog: &Catalog) -> Vec { + self.configured_providers.clone() + } + } + + fn credential(provider: ProviderId, extra_headers: HashMap) -> ApiCredential { + ApiCredential { + provider, + auth_header: None, + extra_headers, + base_url: None, + codex_mode: false, + org_id: None, + project_id: None, + } + } + + #[tokio::test] + async fn appends_headers_to_every_resolved_credential() { + let source = ExtraHeadersCredentialSource::new( + Arc::new(StubSource { + credentials: vec![ + credential(ProviderId::anthropic(), HashMap::new()), + credential(ProviderId::openai(), HashMap::new()), + ], + auth_issue_provider: None, + configured_providers: Vec::new(), + }), + HashMap::from([("x-session-id".to_string(), "run-123".to_string())]), + ); + + let resolved = source.resolve(Catalog::builtin()).await.unwrap(); + + assert_eq!(resolved.credentials.len(), 2); + for credential in &resolved.credentials { + assert_eq!( + credential + .extra_headers + .get("x-session-id") + .map(String::as_str), + Some("run-123") + ); + } + } + + #[tokio::test] + async fn preserves_case_insensitive_headers_already_set_on_a_credential() { + let source = ExtraHeadersCredentialSource::new( + Arc::new(StubSource { + credentials: vec![credential( + ProviderId::new("openrouter"), + HashMap::from([("X-Session-Id".to_string(), "configured".to_string())]), + )], + auth_issue_provider: None, + configured_providers: Vec::new(), + }), + HashMap::from([("x-session-id".to_string(), "run-123".to_string())]), + ); + + let resolved = source.resolve(Catalog::builtin()).await.unwrap(); + + assert_eq!( + resolved.credentials[0] + .extra_headers + .get("X-Session-Id") + .map(String::as_str), + Some("configured") + ); + assert_eq!(resolved.credentials[0].extra_headers.len(), 1); + } + + #[tokio::test] + async fn passes_through_auth_issues_and_configured_providers() { + let auth_issue_provider = ProviderId::anthropic(); + let configured_provider = ProviderId::gemini(); + let source = ExtraHeadersCredentialSource::new( + Arc::new(StubSource { + credentials: vec![credential(ProviderId::openai(), HashMap::new())], + auth_issue_provider: Some(auth_issue_provider.clone()), + configured_providers: vec![configured_provider.clone()], + }), + HashMap::from([("x-session-id".to_string(), "run-123".to_string())]), + ); + + let resolved = source.resolve(Catalog::builtin()).await.unwrap(); + let [(reported_provider, ResolveError::RefreshTokenMissing(error_provider))] = + resolved.auth_issues.as_slice() + else { + panic!("expected the inner source's refresh-token issue"); + }; + assert_eq!(reported_provider, &auth_issue_provider); + assert_eq!(error_provider, &auth_issue_provider); + + let providers = source.configured_providers(Catalog::builtin()).await; + assert_eq!(providers, vec![configured_provider]); + } +} diff --git a/lib/crates/fabro-auth/src/lib.rs b/lib/crates/fabro-auth/src/lib.rs index 1bc135149..77c217317 100644 --- a/lib/crates/fabro-auth/src/lib.rs +++ b/lib/crates/fabro-auth/src/lib.rs @@ -2,6 +2,7 @@ mod context; mod credential; mod credential_source; mod env_source; +mod extra_headers_source; mod refresh; mod resolve; mod sql_vault_source; @@ -15,6 +16,7 @@ pub use context::{AuthContextRequest, AuthContextResponse}; pub use credential::{ApiKeyHeader, OAuthConfig, OAuthCredential, OAuthTokens}; pub use credential_source::{CredentialSource, ResolvedCredentials}; pub use env_source::EnvCredentialSource; +pub use extra_headers_source::ExtraHeadersCredentialSource; pub use refresh::refresh_oauth_credential; pub use resolve::{ ApiCredential, CredentialResolver, CredentialUsage, EnvLookup, ResolveError, diff --git a/lib/crates/fabro-workflow/src/operations/start.rs b/lib/crates/fabro-workflow/src/operations/start.rs index 25f7a8014..9c2e3ed8d 100644 --- a/lib/crates/fabro-workflow/src/operations/start.rs +++ b/lib/crates/fabro-workflow/src/operations/start.rs @@ -858,7 +858,6 @@ impl RunSession { store_progress_logger.register(self.emitter.as_ref()); let init_options = InitOptions { - run_id: record.run_id, run_store: self.run_store.clone(), dry_run: run_options.dry_run_enabled(), emitter: self.emitter, diff --git a/lib/crates/fabro-workflow/src/pipeline/execute/tests.rs b/lib/crates/fabro-workflow/src/pipeline/execute/tests.rs index e2b96ed37..400050896 100644 --- a/lib/crates/fabro-workflow/src/pipeline/execute/tests.rs +++ b/lib/crates/fabro-workflow/src/pipeline/execute/tests.rs @@ -256,7 +256,6 @@ async fn execute_test_run_with_options( let initialized = initialize( persisted_workflow(graph, String::new(), &run_options.run_dir, run_id_value), InitOptions { - run_id: run_id_value, run_store: run_store.into(), dry_run: false, emitter: emitter.clone(), @@ -317,7 +316,6 @@ async fn execute_runs_start_to_exit_and_returns_final_context() { let initialized = initialize( persisted_workflow(graph, source, &run_dir, test_run_id("run-test")), InitOptions { - run_id: test_run_id("run-test"), run_store: run_store.into(), dry_run: false, emitter: test_emitter_arc("run-test"), @@ -393,7 +391,6 @@ async fn run_with_lifecycle( let initialized = initialize( persisted_workflow(graph.clone(), String::new(), &run_dir, run_id), InitOptions { - run_id, run_store: run_store.into(), dry_run: false, emitter: emitter.clone(), diff --git a/lib/crates/fabro-workflow/src/pipeline/initialize.rs b/lib/crates/fabro-workflow/src/pipeline/initialize.rs index 54a6f1f45..e20b2fb61 100644 --- a/lib/crates/fabro-workflow/src/pipeline/initialize.rs +++ b/lib/crates/fabro-workflow/src/pipeline/initialize.rs @@ -5,7 +5,8 @@ use std::time::Instant; use fabro_agent::{Sandbox, ToolSecrets}; use fabro_auth::{ - CredentialSource, EnvCredentialSource, VaultCredentialSource, auth_issue_message, + CredentialSource, EnvCredentialSource, ExtraHeadersCredentialSource, VaultCredentialSource, + auth_issue_message, }; use fabro_graphviz::graph; use fabro_hooks::{HookContext, HookDecision, HookEvent, HookExecutionContext, HookRunner}; @@ -260,11 +261,23 @@ fn graph_needs_api_backend(graph: &graph::Graph) -> bool { graph.nodes.values().any(routing::node_needs_api_backend) } -fn build_llm_source(vault: Option>>) -> Arc { - match vault { +/// Trace header attached to every LLM request in a run so gateways that +/// understand it (e.g. OpenRouter broadcast) can group the run's requests +/// into one session. Explicit `extra_headers` provider configuration wins. +const SESSION_ID_HEADER: &str = "x-session-id"; + +fn build_llm_source( + vault: Option>>, + run_id: fabro_types::RunId, +) -> Arc { + let inner: Arc = match vault { Some(vault) => Arc::new(VaultCredentialSource::new(vault)), None => Arc::new(EnvCredentialSource::new()), - } + }; + Arc::new(ExtraHeadersCredentialSource::new( + inner, + HashMap::from([(SESSION_ID_HEADER.to_string(), run_id.to_string())]), + )) } /// INITIALIZE phase: prepare the sandbox, env, and handlers for execution. @@ -277,7 +290,7 @@ pub async fn initialize( options.run_options.run_dir = run_dir.clone(); options.run_options.git = options.git.clone(); - let llm_source = build_llm_source(options.vault.clone()); + let llm_source = build_llm_source(options.vault.clone(), options.run_options.run_id); let tool_secrets = tool_secrets_from_configured_sources(options.vault.as_ref()).await; let catalog = Arc::clone(&options.catalog); let sandbox_git = Arc::new(SandboxGitRuntime::new()); @@ -347,7 +360,7 @@ pub async fn initialize( let sandbox = reconnect_for_run_with_callback( instance, daytona_api_key, - Some(options.run_id), + Some(options.run_options.run_id), Some(Arc::clone(&sandbox_event_callback)), ) .await @@ -812,7 +825,6 @@ mod tests { }); let result = initialize(persisted, InitOptions { - run_id: test_run_id(), run_store: { let store = memory_store(); let inner = store.create_run(&test_run_id()).await.unwrap(); @@ -894,7 +906,6 @@ mod tests { let emitter = Arc::new(crate::event::Emitter::new(test_run_id())); let initialized = initialize(persisted, InitOptions { - run_id: test_run_id(), run_store: { let store = memory_store(); let inner = store.create_run(&test_run_id()).await.unwrap(); @@ -1028,6 +1039,30 @@ mod tests { assert!(!effective_dry_run); } + #[tokio::test] + async fn build_llm_source_appends_run_session_trace_header() { + let mut vault = Vault::from_entries(HashMap::new()); + fabro_auth::vault_set_token(&mut vault, EnvVars::ANTHROPIC_API_KEY, "anthropic-key") + .unwrap(); + let vault = Arc::new(AsyncRwLock::new(vault)); + let run_id = test_run_id(); + let expected_session_id = run_id.to_string(); + + let source = build_llm_source(Some(vault), run_id); + let resolved = source.resolve(test_catalog().as_ref()).await.unwrap(); + + assert!(!resolved.credentials.is_empty()); + for credential in &resolved.credentials { + assert_eq!( + credential + .extra_headers + .get(SESSION_ID_HEADER) + .map(String::as_str), + Some(expected_session_id.as_str()) + ); + } + } + #[tokio::test] async fn initialize_executes_acp_backend_node_from_registry() { let temp = tempfile::tempdir().unwrap(); @@ -1096,7 +1131,6 @@ mod tests { let store = memory_store(); let run_store = store.create_run(&test_run_id()).await.unwrap(); let initialized = initialize(test_persisted(graph, source, &run_dir), InitOptions { - run_id: test_run_id(), run_store: run_store.into(), dry_run: false, emitter: emitter.clone(), @@ -1192,7 +1226,6 @@ mod tests { store_logger.register(&emitter); let initialized = initialize(persisted, InitOptions { - run_id: test_run_id(), run_store: run_store.into(), dry_run: false, emitter: emitter.clone(), @@ -1331,7 +1364,6 @@ mod tests { let emitter = Arc::new(crate::event::Emitter::new(test_run_id())); let result = initialize(persisted, InitOptions { - run_id: test_run_id(), run_store: { let store = memory_store(); let inner = store.create_run(&test_run_id()).await.unwrap(); diff --git a/lib/crates/fabro-workflow/src/pipeline/types.rs b/lib/crates/fabro-workflow/src/pipeline/types.rs index dc5b25c9e..2b881a24b 100644 --- a/lib/crates/fabro-workflow/src/pipeline/types.rs +++ b/lib/crates/fabro-workflow/src/pipeline/types.rs @@ -249,7 +249,6 @@ pub struct SandboxEnvSpec { } pub struct InitOptions { - pub run_id: RunId, pub run_store: RunStoreHandle, pub dry_run: bool, pub emitter: Arc,