scriptoriumadmin/fabro
1
0
Fork 0
mirror of https://github.com/fabro-sh/fabro.git synced 2026-09-11 22:53:00 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

refactor(llm): simplify session trace header plumbing

Browse source
This commit is contained in:
Release Repro 2026-07-23 13:38:18 -04:00
parent 1448d996e2
commit e8db4e1a9b
No known key found for this signature in database
5 changed files with 66 additions and 67 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

99
lib/crates/fabro-auth/src/extra_headers_source.rs
View file

@ -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<ApiCredential>,
auth_issues: Vec<(ProviderId, ResolveError)>,
credentials: Vec<ApiCredential>,
auth_issue_provider: Option<ProviderId>,
configured_providers: Vec<ProviderId>,
}
#[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<ProviderId> {
self.credentials
.iter()
.map(|c| c.provider.clone())
.collect()
self.configured_providers.clone()
}
}
fn credential(provider: &str, extra_headers: HashMap<String, String>) -> ApiCredential {
fn credential(provider: ProviderId, extra_headers: HashMap<String, String>) -> 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]);
}
}

1
lib/crates/fabro-workflow/src/operations/start.rs
View file

@ -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,

3
lib/crates/fabro-workflow/src/pipeline/execute/tests.rs
View file

@ -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(),

29
lib/crates/fabro-workflow/src/pipeline/initialize.rs
View file

@ -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();

1
lib/crates/fabro-workflow/src/pipeline/types.rs
View file

@ -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<Emitter>,