Merge pull request #1034 from fabro-sh/codex/run-client-headers

Restore run session headers using scoped model clients

Cargo.lock

@@ -4289,7 +4289,7 @@ checksum = "6373607a59f0be73a39b6fe456b8192fcc3585f602af20751600e974dd455e77"
 [[package]]
 name = "lithos-llm"
 version = "0.1.0"
-source = "git+https://github.com/lithoscomputer/lithos-llm?branch=main#fd42e6b2805be8c5cf6764c6039cea88703df6bc"
+source = "git+https://github.com/lithoscomputer/lithos-llm?branch=main#b47b0065df903cb759a0911100d151bd04fc5546"
 dependencies = [
  "async-trait",
  "aws-config",

docs/public/core-concepts/models.mdx

@@ -137,7 +137,9 @@ Model roles are separate: the provider's `default_model` controls normal model s
 
 Provider auth has two parts. The lithos `auth` scheme says how a credential is sent: `{ type = "bearer" }`, `{ type = "header", name = "x-api-key" }`, `{ type = "headers" }` for providers that take several secret headers, `{ type = "none" }`, or `{ type = "aws" }`. lithos also says which secret names a provider reads: `OPENAI_API_KEY` for `openai`, `GEMINI_API_KEY` then `GOOGLE_API_KEY` for `gemini`, `MODAL_TOKEN_ID` and `MODAL_TOKEN_SECRET` for `modal`, and `<PROVIDER>_API_KEY` for a provider you define. Fabro looks each name up in the process environment first and the server vault second. Custom headers for any provider go in `default_headers` as literal text or `{{ secrets.NAME }}` tokens; put credentials in secrets and reference them with `{{ secrets.NAME }}` instead of a bare literal.
 
-Workflow runs also add `x-session-id: <run-id>` to every LLM request so compatible gateways can group requests from the same run. An explicitly configured `x-session-id` in provider `default_headers` takes precedence.
+Workflow runs add `x-session-id: <run-id>` to native HTTP model requests, including model-based hooks, run titles, Ask Fabro and pull request descriptions. The ID stays the same across retries, provider fallback and resume; separate runs get separate IDs. Explicit provider `default_headers` take precedence regardless of header-name casing. Calls outside a run do not receive an automatic session header.
+
+Bedrock applies the header before AWS signing. External ACP agents manage their own model traffic and do not receive it automatically.
 
 Provider `metadata.agent.profile` defaults from `adapter` and controls profile-specific behavior such as which tools the agent registers, project-memory filenames, CLI/ACP command selection, and native session routing. Valid values are `anthropic`, `claude-5`, `openai`, `gemini`, `kimi`, `gpt56`, and `gpt6`; model-level values override provider-level values.
 

lib/apps/fabro-cli/src/commands/run/publish.rs

@@ -245,6 +245,7 @@ impl GitHubPublisher {
         let run_state = self.client.get_run_state(&self.run_id).await.ok();
         let origin_url = self.repository.https_url();
         let created = pull_request::open_pull_request(OpenPullRequestRequest {
+            run_id:            self.run_id,
             github:            context,
             origin_url:        &origin_url,
             base_branch:       &self.base_branch,

lib/apps/fabro-server/src/run_title_generation.rs

@@ -2,7 +2,7 @@ use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::Duration;
 
-use fabro_llm::{Client, Request};
+use fabro_llm::{Client, Request, client as llm_client};
 use fabro_template::{TemplateContext, TemplateError};
 use fabro_types::{MAX_RUN_TITLE_CHARS, RunGraph, RunId};
 use fabro_util::error;
@@ -57,8 +57,8 @@ pub(crate) async fn generate_title_or_current(input: GenerateTitleInput<'_>) ->
         }
     };
 
-    let completion = match input
-        .client
+    let client = llm_client::for_run(&input.client, *input.prompt.run_id);
+    let completion = match client
         .complete_object(request, "run_title", title_response_schema())
         .await
     {
@@ -343,6 +343,7 @@ mod tests {
     }
 
     struct CapturedCall {
+        session_id:        String,
         provider:          String,
         model:             String,
         max_output_tokens: Option<u32>,
@@ -377,6 +378,7 @@ mod tests {
             },
         })
         .await;
+        assert_eq!(captured.lock().unwrap()[0].session_id, run_id.to_string());
         (title, captured)
     }
 
@@ -394,6 +396,10 @@ mod tests {
 
         async fn complete(&self, call: &ResolvedCall) -> Result<Response, LlmError> {
             self.captured.lock().unwrap().push(CapturedCall {
+                session_id:        call.context().default_headers()["x-session-id"]
+                    .to_str()
+                    .unwrap()
+                    .to_string(),
                 provider:          call.route().provider().id().to_string(),
                 model:             call.route().model().id().to_string(),
                 max_output_tokens: call.request().max_output_tokens(),

lib/apps/fabro-server/src/server/handler/sessions.rs

@@ -14,7 +14,7 @@ use fabro_api::types::{
     CreateRunSessionRequest, PaginatedSessionEventList, PaginationMeta, SubmitTurnRequest,
 };
 use fabro_llm::lithos_catalog::Catalog;
-use fabro_llm::{FabroClient, ModelSelectionError, selection};
+use fabro_llm::{FabroClient, ModelSelectionError, client as llm_client, selection};
 use fabro_redact::SecretRedactor;
 use fabro_store::{ProjectedRunSession, project_run_session, project_run_sessions};
 use fabro_tool::fabro_client::ClientBackend;
@@ -697,9 +697,10 @@ async fn build_agent(
     session: &ProjectedRunSession,
 ) -> Result<CodingAgent, AskFabroBuildError> {
     let catalog = state.catalog();
-    let llm_result = state.resolve_llm_client().await.map_err(|err| {
+    let mut llm_result = state.resolve_llm_client().await.map_err(|err| {
         AskFabroBuildError::LlmUnconfigured(format!("LLM credentials are not configured: {err}"))
     })?;
+    llm_result.client = llm_client::for_run(&llm_result.client, run_id);
     for (provider, issue) in &llm_result.auth_issues {
         warn!(provider = %provider, error = %issue, "LLM provider unavailable due to auth issue");
     }

lib/apps/fabro-server/src/server/pull_request_supervisor.rs

@@ -175,6 +175,7 @@ async fn attempt_pull_request_creation(
     };
     let catalog = state.catalog();
     let request = pull_request::OpenPullRequestRequest {
+        run_id: *run_id,
         github,
         origin_url: &inputs.normalized_origin,
         base_branch: inputs.base_branch,

lib/components/fabro-llm/src/client.rs

@@ -3,6 +3,7 @@
 use std::sync::Arc;
 use std::time::Duration;
 
+use fabro_types::RunId;
 use lithos_llm::adapter::ProviderAdapter;
 use lithos_llm::catalog::{Catalog, ProviderId};
 use lithos_llm::client::{Client, ClientBuildError, ClientBuilder, ProviderBuildIssue};
@@ -29,6 +30,14 @@ pub fn default_retry_policy() -> RetryPolicy {
         .jitter(true)
 }
 
+/// Scope ancillary model calls to the same session as the run's native steps.
+#[must_use]
+pub fn for_run(client: &Client, run_id: RunId) -> Client {
+    client
+        .with_default_headers([("x-session-id", run_id.to_string())])
+        .expect("a run UUID is a valid HTTP header value")
+}
+
 /// One retry the lithos retry middleware decided on.
 #[derive(Clone, Debug)]
 pub struct RetryNotice {

lib/components/fabro-petri/src/engine.rs

@@ -63,6 +63,7 @@ use fabro_types::settings::run::{
 use fabro_types::settings::size::Size;
 use fabro_types::{FailureReason, RunId, SandboxProviderKind};
 use fabro_util::sync;
+use lithos_llm::types::Error as LlmError;
 use petri_execution::host::{self, HostError, HostRun};
 use petri_execution::inspect::{self, InspectError, RunInspection};
 use petri_execution::watchdog::{StallWatchdog, WatchdogTask};
@@ -167,6 +168,8 @@ pub struct RunOutcome {
 /// Why the run could not be executed or its outcome read.
 #[derive(Debug, thiserror::Error)]
 pub enum RunError {
+    #[error("the run ID could not be applied to model request headers")]
+    ModelHeaders(#[source] LlmError),
     #[error("the run's sandbox provider `{provider}` is not one Petri serves")]
     UnsupportedProvider { provider: SandboxProviderKind },
     #[error("the Daytona CPU allocation must be a non-negative integer: {cpu}")]
@@ -215,7 +218,7 @@ fn network_policy(settings: &EnvironmentNetworkSettings) -> NetworkPolicy {
 }
 
 /// Execute the run to its end and report what the record says.
-pub async fn run(request: RunRequest) -> Result<RunOutcome, RunError> {
+pub async fn run(mut request: RunRequest) -> Result<RunOutcome, RunError> {
     let backend = if request.runtime.dry_run {
         // The dry-run admission pass puts every scope on the host. Keep the
         // real local workspace for Fabro's checkpoint and artifact hooks.
@@ -225,6 +228,15 @@ pub async fn run(request: RunRequest) -> Result<RunOutcome, RunError> {
             provider: request.provider.clone(),
         })?
     };
+    // Scope once before either start or resume; native steps, hooks and
+    // subagents inherit this client through the runtime capability.
+    request.runtime.model_client = request
+        .runtime
+        .model_client
+        .as_ref()
+        .map(|client| client.with_default_headers([("x-session-id", request.run_id.as_str())]))
+        .transpose()
+        .map_err(RunError::ModelHeaders)?;
     let key = RunKey::new(request.run_id.as_str());
     let stall_timeout = match &request.execution {
         Execution::Start(graphs) => graphs.graph.policy.stall_timeout,

lib/components/fabro-petri/tests/model.rs

@@ -18,7 +18,7 @@ use fabro_petri::check::Launch;
 use fabro_petri::engine::{self, RunStatus};
 use fabro_petri::runtime::{self, RuntimeSpec};
 use fabro_test::{TwinScenario, TwinScenarios, twin_openai};
-use fabro_types::SecretType;
+use fabro_types::{RunId, SecretType};
 use fabro_vault::Vault;
 use lithos_llm::catalog::ProviderId;
 use petri_store::MemoryRunStore;
@@ -106,3 +106,116 @@ async fn a_model_call_authenticates_through_the_vault_and_skills_read_the_home()
         "the configured home is searched: {resolved}"
     );
 }
+
+/// The engine binds the durable identity for both native step kinds and both
+/// entry modes, without changing a shared client used outside a run.
+#[tokio::test]
+async fn model_session_headers_follow_run_identity_on_start_and_resume() {
+    use std::sync::Mutex;
+
+    use fabro_llm::test_support::{self, ScriptedAdapter};
+    use lithos_llm::Client;
+    use lithos_llm::middleware::{Call, Observer, ObserverMiddleware};
+    use petri_store::{Access, OwnerId, RunKey, RunStore as _};
+
+    struct Sessions(Arc<Mutex<Vec<Option<String>>>>);
+    impl Observer for Sessions {
+        fn on_start(&self, call: &Call) {
+            self.0.lock().unwrap().push(
+                call.context()
+                    .default_headers()
+                    .get("x-session-id")
+                    .map(|value| value.to_str().unwrap().to_owned()),
+            );
+        }
+    }
+    let seen = Arc::new(Mutex::new(Vec::new()));
+    let client = Client::builder()
+        .catalog(test_support::test_catalog())
+        .adapter(
+            ProviderId::new("openai"),
+            ScriptedAdapter::new(vec![test_support::text_response(
+                "openai",
+                OPENAI_MODEL,
+                "done",
+            )]),
+        )
+        .middleware(ObserverMiddleware::new(Sessions(Arc::clone(&seen))))
+        .build()
+        .unwrap()
+        .client;
+    let runtime = RuntimeSpec {
+        model_client: Some(client.clone()),
+        ..RuntimeSpec::default()
+    };
+    let workflow = r#"digraph Headers {
+        graph [backend="api", default_max_retries=0]
+        start [shape=Mdiamond]
+        exit [shape=Msquare]
+        agent [shape=box, prompt="Say done"]
+        prompt [shape=tab, prompt="Say done"]
+        start -> agent -> prompt -> exit
+    }"#;
+    for resume in [false, true] {
+        let run_id = RunId::new().to_string();
+        let root = tempfile::tempdir().unwrap();
+        let store = Arc::new(MemoryRunStore::new());
+        let graphs = support::admit(
+            &[
+                ("workflow.fabro", workflow),
+                ("workflow.toml", support::SETTINGS),
+            ],
+            Launch {
+                default_model: Some(format!("openai/{OPENAI_MODEL}")),
+                ..Launch::default()
+            },
+            &runtime,
+        );
+        let mut request = run_request(
+            &run_id,
+            root.path(),
+            graphs,
+            store.clone(),
+            runtime.clone(),
+            no_questions(Arc::new(Silent)),
+        );
+        if resume {
+            // Recover the same durable key after creation was interrupted.
+            drop(
+                store
+                    .open(&RunKey::new(&run_id), Access::Create {
+                        owner: OwnerId::new("interrupted"),
+                    })
+                    .await
+                    .unwrap(),
+            );
+            request.execution = match request.execution {
+                engine::Execution::Start(graphs) | engine::Execution::Resume(graphs) => {
+                    engine::Execution::Resume(graphs)
+                }
+            };
+        }
+        seen.lock().unwrap().clear();
+        let outcome = engine::run(request).await.unwrap();
+        assert_eq!(outcome.status, RunStatus::Success, "{outcome:?}");
+        let calls = seen.lock().unwrap();
+        assert!(calls.len() >= 2, "both native stages call the model");
+        assert!(
+            calls
+                .iter()
+                .all(|header| header.as_deref() == Some(&run_id))
+        );
+    }
+    seen.lock().unwrap().clear();
+    client
+        .complete(
+            fabro_llm::Request::builder()
+                .model(format!("openai/{OPENAI_MODEL}"))
+                .user("outside a run")
+                .build()
+                .unwrap(),
+        )
+        .await
+        .unwrap();
+    assert_eq!(*seen.lock().unwrap(), vec![None]);
+}

lib/components/fabro-workflow/src/pull_request.rs

@@ -5,10 +5,12 @@ use std::time::Duration;
 use fabro_github::{self as github_app, ssh_url_to_https};
 use fabro_llm::credentials::CredentialProvider;
 use fabro_llm::lithos_catalog::Catalog;
-use fabro_llm::{Client, ClientOptions, Request, selection};
+use fabro_llm::{Client, ClientOptions, Request, client as llm_client, selection};
 use fabro_store::RunProjection;
 use fabro_types::settings::run::MergeStrategy;
-use fabro_types::{Conclusion, PullRequestLink, RunSpec, format_cost as outcome_format_cost};
+use fabro_types::{
+    Conclusion, PullRequestLink, RunId, RunSpec, format_cost as outcome_format_cost,
+};
 use fabro_util::text::strip_goal_decoration;
 use lithos_llm::catalog::ProviderId;
 use lithos_llm::types::{Cost, Message, Role};
@@ -297,6 +299,7 @@ fn assemble_pr_body(
 /// Build complete PR content by combining LLM-generated narrative with
 /// deterministic fallbacks and programmatic sections.
 pub async fn build_pr_content(
+    run_id: RunId,
     diff: &str,
     goal: &str,
     model: &str,
@@ -321,7 +324,7 @@ pub async fn build_pr_content(
         catalog.as_ref(),
         conclusion,
         run_state,
-        Arc::new(client),
+        Arc::new(llm_client::for_run(&client, run_id)),
     )
     .await
 }
@@ -402,6 +405,7 @@ pub struct AutoMergeOptions {
 
 /// Inputs for [`open_pull_request`].
 pub struct OpenPullRequestRequest<'a> {
+    pub run_id:            RunId,
     pub github:            github_app::GitHubContext<'a>,
     pub origin_url:        &'a str,
     pub base_branch:       &'a str,
@@ -564,6 +568,7 @@ pub async fn open_pull_request(
     }
 
     let content = build_pr_content(
+        req.run_id,
         req.diff,
         req.goal,
         req.model,
@@ -1076,12 +1081,14 @@ capabilities = { text = true, tools = true, response_format = { json_object = tr
 
     #[tokio::test]
     async fn build_pr_content_uses_vault_only_openai_codex_source() {
+        let run_id = RunId::new();
         let server = MockServer::start_async().await;
         let response_mock = server
             .mock_async(|when, then| {
                 when.method(POST)
                     .path("/v1/responses")
-                    .header("authorization", "Bearer vault-openai-key");
+                    .header("authorization", "Bearer vault-openai-key")
+                    .header("x-session-id", run_id.to_string());
                 then.status(200)
                     .header("content-type", "application/json")
                     .json_body(openai_responses_payload(&pr_content_json(
@@ -1101,13 +1108,14 @@ capabilities = { text = true, tools = true, response_format = { json_object = tr
                 None,
             )
             .unwrap();
-        let llm_source: Arc<dyn CredentialProvider> = Arc::new(VaultCredentialSource::new(
+        let llm_source: Arc<dyn CredentialProvider> = Arc::new(VaultCredentialSource::vault_only(
             Arc::new(AsyncRwLock::new(vault)),
         ));
         // Use catalog settings to override base_url instead of env var
         let catalog = test_catalog_with_provider_base_url("openai", &server.url("/v1"));
 
         let PrContent { title, body } = build_pr_content(
+            run_id,
             "diff --git a/src/lib.rs b/src/lib.rs\n+fn new_feature() {}\n",
             "Implement feature",
             "gpt-5.4",
@@ -1265,6 +1273,7 @@ capabilities = { text = true, tools = true, response_format = { json_object = tr
         let harness = setup_fallback_test_harness_with_branch_sha(&payload, "stale-sha").await;
         let github_base_url = harness.github_server.url("");
         let error = open_pull_request(OpenPullRequestRequest {
+            run_id:            RunId::new(),
             github:            fabro_github::GitHubContext::new(&harness.creds, &github_base_url),
             origin_url:        "https://github.com/owner/repo.git",
             base_branch:       "main",
@@ -1547,6 +1556,7 @@ capabilities = { text = true, tools = true, response_format = { json_object = tr
         let github = github_app::GitHubContext::new(&harness.creds, &github_base_url);
 
         let result = open_pull_request(OpenPullRequestRequest {
+            run_id: RunId::new(),
             github,
             origin_url: "https://github.com/owner/repo.git",
             base_branch: "main",
@@ -1594,6 +1604,7 @@ capabilities = { text = true, tools = true, response_format = { json_object = tr
         let github = github_app::GitHubContext::new(&harness.creds, &github_base_url);
 
         let result = open_pull_request(OpenPullRequestRequest {
+            run_id: RunId::new(),
             github,
             origin_url: "https://github.com/owner/repo.git",
             base_branch: "main",
@@ -1630,6 +1641,7 @@ capabilities = { text = true, tools = true, response_format = { json_object = tr
         let goal = "x".repeat(200);
 
         let result = open_pull_request(OpenPullRequestRequest {
+            run_id: RunId::new(),
             github,
             origin_url: "https://github.com/owner/repo.git",
             base_branch: "main",

lib/foundation/fabro-auth/src/extra_headers_source.rs

@@ -1,154 +0,0 @@
-use std::collections::HashMap;
-use std::sync::Arc;
-
-use async_trait::async_trait;
-use lithos_llm::catalog::CatalogProvider;
-use lithos_llm::credentials::{
-    CredentialError, CredentialHeader, CredentialProvider, Credentials, SecretValue,
-};
-
-/// Decorates another [`CredentialProvider`] by appending fixed extra headers
-/// to every HTTP credential it resolves.
-///
-/// Headers already present on a credential (for example from explicit
-/// provider configuration) are left untouched. AWS-signed credentials carry
-/// no header list and pass through unchanged.
-pub struct ExtraHeadersCredentialSource {
-    inner:   Arc<dyn CredentialProvider>,
-    headers: HashMap<String, String>,
-}
-
-impl ExtraHeadersCredentialSource {
-    #[must_use]
-    pub fn new(inner: Arc<dyn CredentialProvider>, headers: HashMap<String, String>) -> Self {
-        Self { inner, headers }
-    }
-}
-
-#[async_trait]
-impl CredentialProvider for ExtraHeadersCredentialSource {
-    async fn credentials(
-        &self,
-        provider: &CatalogProvider,
-    ) -> Result<Credentials, CredentialError> {
-        let mut credentials = self.inner.credentials(provider).await?;
-        if let Credentials::Http(http) = &mut credentials {
-            for (name, value) in &self.headers {
-                if http
-                    .extra_headers
-                    .iter()
-                    .any(|existing| existing.name.eq_ignore_ascii_case(name))
-                {
-                    continue;
-                }
-                http.extra_headers.push(CredentialHeader::new(
-                    name.clone(),
-                    SecretValue::new(value.clone()),
-                ));
-            }
-        }
-        Ok(credentials)
-    }
-
-    async fn is_configured(&self, provider: &CatalogProvider) -> bool {
-        self.inner.is_configured(provider).await
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use lithos_llm::credentials::HttpAuthentication;
-
-    use super::*;
-    use crate::test_support::test_catalog;
-
-    struct StubSource {
-        configured:      bool,
-        existing_header: Option<(String, String)>,
-    }
-
-    #[async_trait]
-    impl CredentialProvider for StubSource {
-        async fn credentials(
-            &self,
-            provider: &CatalogProvider,
-        ) -> Result<Credentials, CredentialError> {
-            if provider.id().as_str() == "bedrock" {
-                return Ok(Credentials::AwsDefaultChain { region: None });
-            }
-            let mut credentials = Credentials::bearer(SecretValue::new("key"));
-            if let (Credentials::Http(http), Some((name, value))) =
-                (&mut credentials, &self.existing_header)
-            {
-                http.extra_headers.push(CredentialHeader::new(
-                    name.clone(),
-                    SecretValue::new(value.clone()),
-                ));
-            }
-            Ok(credentials)
-        }
-
-        async fn is_configured(&self, _provider: &CatalogProvider) -> bool {
-            self.configured
-        }
-    }
-
-    fn header<'a>(credentials: &'a Credentials, name: &str) -> Option<&'a str> {
-        match credentials {
-            Credentials::Http(http) => http
-                .extra_headers
-                .iter()
-                .find(|header| header.name.eq_ignore_ascii_case(name))
-                .map(|header| header.value.expose_secret()),
-            _ => None,
-        }
-    }
-
-    #[tokio::test]
-    async fn appends_headers_to_http_credentials_only() {
-        let catalog = test_catalog();
-        let source = ExtraHeadersCredentialSource::new(
-            Arc::new(StubSource {
-                configured:      false,
-                existing_header: None,
-            }),
-            HashMap::from([("x-session-id".to_string(), "run-123".to_string())]),
-        );
-        let openai = source
-            .credentials(catalog.provider("openai").unwrap())
-            .await
-            .unwrap();
-        assert!(matches!(
-            &openai,
-            Credentials::Http(http) if matches!(http.auth, HttpAuthentication::Bearer(_))
-        ));
-        assert_eq!(header(&openai, "x-session-id"), Some("run-123"));
-        let bedrock = source
-            .credentials(catalog.provider("bedrock").unwrap())
-            .await
-            .unwrap();
-        assert!(matches!(bedrock, Credentials::AwsDefaultChain { .. }));
-    }
-
-    #[tokio::test]
-    async fn preserves_case_insensitive_headers_already_set() {
-        let catalog = test_catalog();
-        let source = ExtraHeadersCredentialSource::new(
-            Arc::new(StubSource {
-                configured:      true,
-                existing_header: Some(("X-Session-Id".to_string(), "configured".to_string())),
-            }),
-            HashMap::from([("x-session-id".to_string(), "run-123".to_string())]),
-        );
-        let credentials = source
-            .credentials(catalog.provider("openai").unwrap())
-            .await
-            .unwrap();
-        assert_eq!(header(&credentials, "x-session-id"), Some("configured"));
-        assert!(
-            source
-                .is_configured(catalog.provider("openai").unwrap())
-                .await
-        );
-    }
-}

lib/foundation/fabro-auth/src/lib.rs

@@ -11,7 +11,6 @@
 mod api_key_source;
 mod context;
 mod credential;
-mod extra_headers_source;
 mod refresh;
 mod secrets;
 mod sql_vault_source;
@@ -26,7 +25,6 @@ pub mod strategies;
 pub use api_key_source::ApiKeyCredentialSource;
 pub use context::{AuthContextRequest, AuthContextResponse};
 pub use credential::{OAuthConfig, OAuthCredential, OAuthTokens};
-pub use extra_headers_source::ExtraHeadersCredentialSource;
 pub use refresh::refresh_oauth_credential;
 pub use secrets::{accepts_api_key, expected_secret_name, secret_names};
 pub use sql_vault_source::SqlVaultCredentialSource;