Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
02c7552725 |
@@ -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(())
|
||||
|
||||
@@ -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(),
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user