From 02c7552725252b6c2494d7b9ae60accfd34db448 Mon Sep 17 00:00:00 2001
From: Sharang Parnerkar <30073382+mighty840@users.noreply.github.com>
Date: Wed, 22 Jul 2026 11:23:25 +0200
Subject: [PATCH] =?UTF-8?q?fix(mcp):=20bind=20tenant=20to=20session=20?=
 =?UTF-8?q?=E2=80=94=20bearer=20context=20was=20lost=20over=20HTTP?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Root cause (found while wiring the breakpilot MCP loop, A): every MCP tool call
over Streamable HTTP failed with '-32603 no tenant context — bearer middleware
not in chain'. The bearer middleware sets the tenant in a tokio task_local scoped
around the request, but rmcp's StreamableHttpService spawns the session's
tool-serving task (tower.rs: 'spawn a task to serve the session'), and task_locals
do NOT cross tokio::spawn. So initialize/tools-list worked but every real tool call
(list_findings, oscal_assessment, ...) failed — meaning the loop never worked over
HTTP and breakpilot's client fell back to demo data.

Fix: bind the tenant to the per-session server instance instead of a per-request
task_local. rmcp calls the service factory in the request task (before the spawn)
and still inside the middleware's scope, so the factory reads the bearer-set
tenant once and bakes it into ComplianceMcpServer; tool_db() then reads
self.tenant_id. stdio passes its synthetic tenant the same way. No tool-handler
changes; the middleware still validates (and can revoke) the token per request.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 compliance-mcp/src/main.rs   | 25 ++++++++++++++++---------
 compliance-mcp/src/server.rs | 22 +++++++++-------------
 2 files changed, 25 insertions(+), 22 deletions(-)

diff --git a/compliance-mcp/src/main.rs b/compliance-mcp/src/main.rs
index b35547f..a98e680 100644
--- a/compliance-mcp/src/main.rs
+++ b/compliance-mcp/src/main.rs
@@ -42,7 +42,19 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
         let pool_for_factory = pool.clone();
         let service = StreamableHttpService::new(
-            move || Ok(ComplianceMcpServer::new(pool_for_factory.clone())),
+            move || {
+                // The factory runs in the request task, still inside the bearer
+                // middleware's `TENANT_ID` scope, and BEFORE rmcp spawns the
+                // session task (which would lose the task_local). So bind the
+                // tenant into the session's server instance here, once.
+                let tenant_id = auth::current_tenant_id().ok_or_else(|| {
+                    std::io::Error::other("no tenant context when creating MCP session")
+                })?;
+                Ok(ComplianceMcpServer::new(
+                    pool_for_factory.clone(),
+                    tenant_id,
+                ))
+            },
             Arc::new(LocalSessionManager::default()),
             StreamableHttpServerConfig::default(),
         );
@@ -69,16 +81,11 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
             tenant_id = %synth_tenant,
             "stdio transport — using synthetic tenant id; DO NOT use in production"
         );
-        let server = ComplianceMcpServer::new(pool);
+        let server = ComplianceMcpServer::new(pool, synth_tenant);
         let transport = rmcp::transport::stdio();
         use rmcp::ServiceExt;
-        auth::TENANT_ID
-            .scope(synth_tenant, async {
-                let handle = server.serve(transport).await?;
-                handle.waiting().await?;
-                Ok::<_, Box<dyn std::error::Error>>(())
-            })
-            .await?;
+        let handle = server.serve(transport).await?;
+        handle.waiting().await?;
     }
 
     Ok(())
diff --git a/compliance-mcp/src/server.rs b/compliance-mcp/src/server.rs
index 6e29505..2306877 100644
--- a/compliance-mcp/src/server.rs
+++ b/compliance-mcp/src/server.rs
@@ -2,37 +2,33 @@ use rmcp::{
     handler::server::wrapper::Parameters, model::*, tool, tool_handler, tool_router, ServerHandler,
 };
 
-use crate::auth::current_tenant_id;
 use crate::database::{Database, DatabasePool};
 use crate::tools::{dast, findings, oscal, pentest, sbom};
 
 pub struct ComplianceMcpServer {
     pool: DatabasePool,
+    /// Tenant this session serves. Bound once at session creation (the HTTP
+    /// factory reads the bearer-set tenant while still in the request scope;
+    /// stdio passes a synthetic id) — NOT a per-request `task_local`, which is
+    /// lost across the `tokio::spawn` that runs the Streamable-HTTP session.
+    tenant_id: String,
     #[allow(dead_code)]
     tool_router: rmcp::handler::server::router::tool::ToolRouter<Self>,
 }
 
 impl ComplianceMcpServer {
-    /// Resolve the per-tenant `Database` from the bearer-set
-    /// `task_local`. Every tool handler calls this; missing context
-    /// surfaces as `internal_error` because it means the auth
-    /// middleware was misconfigured (handler ran without scope).
+    /// The per-tenant `Database` for this session.
     fn tenant_db(&self) -> Result<Database, rmcp::ErrorData> {
-        let tenant_id = current_tenant_id().ok_or_else(|| {
-            rmcp::ErrorData::internal_error(
-                "no tenant context — bearer middleware not in chain".to_string(),
-                None,
-            )
-        })?;
-        Ok(self.pool.for_tenant_id(&tenant_id))
+        Ok(self.pool.for_tenant_id(&self.tenant_id))
     }
 }
 
 #[tool_router]
 impl ComplianceMcpServer {
-    pub fn new(pool: DatabasePool) -> Self {
+    pub fn new(pool: DatabasePool, tenant_id: String) -> Self {
         Self {
             pool,
+            tenant_id,
             tool_router: Self::tool_router(),
         }
     }