From e8db4e1a9b4b328e871b37edeb0ea117b185c9fa Mon Sep 17 00:00:00 2001 From: Release Repro Date: Thu, 23 Jul 2026 13:38:18 -0400 Subject: [PATCH] refactor(llm): simplify session trace header plumbing --- .../fabro-auth/src/extra_headers_source.rs | 99 ++++++++++--------- .../fabro-workflow/src/operations/start.rs | 1 - .../src/pipeline/execute/tests.rs | 3 - .../fabro-workflow/src/pipeline/initialize.rs | 29 +++--- .../fabro-workflow/src/pipeline/types.rs | 1 - 5 files changed, 66 insertions(+), 67 deletions(-) diff --git a/lib/crates/fabro-auth/src/extra_headers_source.rs b/lib/crates/fabro-auth/src/extra_headers_source.rs index 2e57ee6ea..29ad806cd 100644 --- a/lib/crates/fabro-auth/src/extra_headers_source.rs +++ b/lib/crates/fabro-auth/src/extra_headers_source.rs @@ -29,10 +29,14 @@ impl CredentialSource for ExtraHeadersCredentialSource { let mut resolved = self.inner.resolve(catalog).await?; for credential in &mut resolved.credentials { for (name, value) in &self.headers { - credential + if credential .extra_headers - .entry(name.clone()) - .or_insert_with(|| value.clone()); + .keys() + .any(|existing| existing.eq_ignore_ascii_case(name)) + { + continue; + } + credential.extra_headers.insert(name.clone(), value.clone()); } } Ok(resolved) @@ -49,8 +53,9 @@ mod tests { use crate::{ApiCredential, ResolveError}; struct StubSource { - credentials: Vec, - auth_issues: Vec<(ProviderId, ResolveError)>, + credentials: Vec, + auth_issue_provider: Option, + configured_providers: Vec, } #[async_trait] @@ -59,12 +64,12 @@ mod tests { Ok(ResolvedCredentials { credentials: self.credentials.clone(), auth_issues: self - .auth_issues + .auth_issue_provider .iter() - .map(|(provider, _)| { + .map(|provider| { ( provider.clone(), - ResolveError::NotConfigured(provider.clone()), + ResolveError::RefreshTokenMissing(provider.clone()), ) }) .collect(), @@ -72,16 +77,13 @@ mod tests { } async fn configured_providers(&self, _catalog: &Catalog) -> Vec { - self.credentials - .iter() - .map(|c| c.provider.clone()) - .collect() + self.configured_providers.clone() } } - fn credential(provider: &str, extra_headers: HashMap) -> ApiCredential { + fn credential(provider: ProviderId, extra_headers: HashMap) -> ApiCredential { ApiCredential { - provider: ProviderId::new(provider), + provider, auth_header: None, extra_headers, base_url: None, @@ -91,74 +93,83 @@ mod tests { } } - fn catalog() -> Catalog { - Catalog::from_builtin().unwrap() - } - #[tokio::test] async fn appends_headers_to_every_resolved_credential() { let source = ExtraHeadersCredentialSource::new( Arc::new(StubSource { - credentials: vec![ - credential("anthropic", HashMap::new()), - credential("openai", HashMap::new()), + credentials: vec![ + credential(ProviderId::anthropic(), HashMap::new()), + credential(ProviderId::openai(), HashMap::new()), ], - auth_issues: Vec::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()).await.unwrap(); + 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"), - Some(&"run-123".to_string()) + credential + .extra_headers + .get("x-session-id") + .map(String::as_str), + Some("run-123") ); } } #[tokio::test] - async fn preserves_headers_already_set_on_a_credential() { + async fn preserves_case_insensitive_headers_already_set_on_a_credential() { let source = ExtraHeadersCredentialSource::new( Arc::new(StubSource { - credentials: vec![credential( - "openrouter", - HashMap::from([("x-session-id".to_string(), "configured".to_string())]), + credentials: vec![credential( + ProviderId::new("openrouter"), + HashMap::from([("X-Session-Id".to_string(), "configured".to_string())]), )], - auth_issues: Vec::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()).await.unwrap(); + let resolved = source.resolve(Catalog::builtin()).await.unwrap(); assert_eq!( - resolved.credentials[0].extra_headers.get("x-session-id"), - Some(&"configured".to_string()) + 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 provider = ProviderId::new("anthropic"); + let auth_issue_provider = ProviderId::anthropic(); + let configured_provider = ProviderId::gemini(); let source = ExtraHeadersCredentialSource::new( Arc::new(StubSource { - credentials: vec![credential("openai", HashMap::new())], - auth_issues: vec![( - provider.clone(), - ResolveError::NotConfigured(provider.clone()), - )], + 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()).await.unwrap(); - assert_eq!(resolved.auth_issues.len(), 1); - assert_eq!(resolved.auth_issues[0].0, provider); + 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()).await; - assert_eq!(providers, vec![ProviderId::new("openai")]); + let providers = source.configured_providers(Catalog::builtin()).await; + assert_eq!(providers, vec![configured_provider]); } } diff --git a/lib/crates/fabro-workflow/src/operations/start.rs b/lib/crates/fabro-workflow/src/operations/start.rs index 7fa47e938..8d0be3d5f 100644 --- a/lib/crates/fabro-workflow/src/operations/start.rs +++ b/lib/crates/fabro-workflow/src/operations/start.rs @@ -831,7 +831,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 ffe53e180..e20b2fb61 100644 --- a/lib/crates/fabro-workflow/src/pipeline/initialize.rs +++ b/lib/crates/fabro-workflow/src/pipeline/initialize.rs @@ -360,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 @@ -825,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(); @@ -907,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(); @@ -1043,26 +1041,24 @@ mod tests { #[tokio::test] async fn build_llm_source_appends_run_session_trace_header() { - let dir = tempfile::tempdir().unwrap(); - let mut vault = Vault::load(dir.path().join("secrets.json")).unwrap(); - vault - .set( - "ANTHROPIC_API_KEY", - "anthropic-key", - SecretType::Token, - None, - ) + 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), test_run_id()); + 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), - Some(&test_run_id().to_string()) + credential + .extra_headers + .get(SESSION_ID_HEADER) + .map(String::as_str), + Some(expected_session_id.as_str()) ); } } @@ -1135,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(), @@ -1231,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(), @@ -1370,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 8449453a8..69a03c9fe 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,