1use std::collections::{BTreeSet, HashMap};
2use std::path::{Path, PathBuf};
3use std::sync::atomic::{AtomicI64, Ordering};
4use std::sync::Arc;
5
6use anyhow::{anyhow, Context, Result};
7use axum::{
8 body::{self, Body},
9 extract::{Extension, Json, Query, State},
10 http::{header, Request, StatusCode},
11 middleware::{self, Next},
12 response::sse::{Event as SseEvent, KeepAlive, Sse},
13 response::{IntoResponse, Response},
14 routing::get,
15 Router,
16};
17use futures::Stream;
18use md5::{Digest, Md5};
19use rmcp::{
20 handler::server::ServerHandler,
21 model::{
22 CallToolRequest, CallToolRequestParams, CallToolResult, ClientJsonRpcMessage,
23 ClientRequest, JsonRpcRequest, NumberOrString, ServerJsonRpcMessage, ServerResult,
24 },
25 service::{serve_directly, RequestContext, RoleServer},
26 transport::{OneshotTransport, StreamableHttpService},
27};
28use serde::{Deserialize, Serialize};
29use serde_json::{json, Map, Value};
30use sha2::Sha256;
31use tokio::io::AsyncWriteExt;
32use tokio::sync::broadcast;
33use tokio::time::Duration;
34
35use crate::tools::LeanCtxServer;
36
37pub mod roles;
38pub use roles::TeamRole;
39
40#[cfg(test)]
41mod tests;
42
43const WORKSPACE_ARG_KEY: &str = "workspaceId";
44const CHANNEL_ARG_KEY: &str = "channelId";
45const WORKSPACE_HEADER: &str = "x-leanctx-workspace";
46
47#[derive(Clone, Debug, Serialize, Deserialize)]
48#[serde(rename_all = "camelCase")]
49pub struct TeamServerConfig {
50 pub host: String,
51 pub port: u16,
52 pub default_workspace_id: String,
53 pub workspaces: Vec<TeamWorkspaceConfig>,
54 #[serde(default)]
55 pub tokens: Vec<TeamTokenConfig>,
56 pub audit_log_path: PathBuf,
57 #[serde(default)]
58 pub disable_host_check: bool,
59 #[serde(default)]
60 pub allowed_hosts: Vec<String>,
61 #[serde(default = "default_max_body_bytes")]
62 pub max_body_bytes: usize,
63 #[serde(default = "default_max_concurrency")]
64 pub max_concurrency: usize,
65 #[serde(default = "default_max_rps")]
66 pub max_rps: u32,
67 #[serde(default = "default_rate_burst")]
68 pub rate_burst: u32,
69 #[serde(default = "default_request_timeout_ms")]
70 pub request_timeout_ms: u64,
71 #[serde(default)]
72 pub stateful_mode: bool,
73 #[serde(default = "default_true")]
74 pub json_response: bool,
75 #[serde(default, skip_serializing_if = "Option::is_none")]
80 pub storage_quota_bytes: Option<u64>,
81 #[serde(default, skip_serializing_if = "Option::is_none")]
85 pub roi_webhook_url: Option<String>,
86}
87
88fn default_true() -> bool {
89 true
90}
91fn default_max_body_bytes() -> usize {
92 2 * 1024 * 1024
93}
94fn default_max_concurrency() -> usize {
95 32
96}
97fn default_max_rps() -> u32 {
98 50
99}
100fn default_rate_burst() -> u32 {
101 100
102}
103fn default_request_timeout_ms() -> u64 {
104 30_000
105}
106
107#[derive(Clone, Debug, Serialize, Deserialize)]
108#[serde(rename_all = "camelCase")]
109pub struct TeamWorkspaceConfig {
110 pub id: String,
111 pub label: Option<String>,
112 pub root: PathBuf,
113}
114
115#[derive(Clone, Debug, Serialize, Deserialize)]
116#[serde(rename_all = "camelCase")]
117pub struct TeamTokenConfig {
118 pub id: String,
119 pub sha256_hex: String,
121 #[serde(default)]
123 pub scopes: Vec<TeamScope>,
124 #[serde(default, skip_serializing_if = "Option::is_none")]
128 pub role: Option<roles::TeamRole>,
129}
130
131impl TeamTokenConfig {
132 #[must_use]
135 pub fn effective_scopes(&self) -> BTreeSet<TeamScope> {
136 let mut s: BTreeSet<TeamScope> = self.scopes.iter().copied().collect();
137 if let Some(role) = self.role {
138 s.extend(role.scopes());
139 }
140 s
141 }
142}
143
144#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
145#[serde(rename_all = "lowercase")]
146pub enum TeamScope {
147 Search,
148 Graph,
149 Artifacts,
150 Index,
151 Events,
152 SessionMutations,
153 Knowledge,
154 Audit,
155}
156
157impl TeamScope {
158 #[must_use]
160 pub fn all() -> &'static [TeamScope] {
161 &[
162 TeamScope::Search,
163 TeamScope::Graph,
164 TeamScope::Artifacts,
165 TeamScope::Index,
166 TeamScope::Events,
167 TeamScope::SessionMutations,
168 TeamScope::Knowledge,
169 TeamScope::Audit,
170 ]
171 }
172}
173
174impl TeamServerConfig {
175 pub fn load(path: &Path) -> Result<Self> {
176 let s =
177 std::fs::read_to_string(path).with_context(|| format!("read {}", path.display()))?;
178 let cfg: Self =
179 serde_json::from_str(&s).with_context(|| format!("parse {}", path.display()))?;
180 cfg.validate()?;
181 Ok(cfg)
182 }
183
184 pub fn save(&self, path: &Path) -> Result<()> {
185 let s = serde_json::to_string_pretty(self).context("serialize TeamServerConfig")?;
186 std::fs::write(path, format!("{s}\n")).with_context(|| format!("write {}", path.display()))
187 }
188
189 pub fn validate(&self) -> Result<()> {
190 if self.workspaces.is_empty() {
191 return Err(anyhow!("team server requires at least 1 workspace"));
192 }
193 let mut ws_ids = BTreeSet::new();
194 for ws in &self.workspaces {
195 let id = ws.id.trim();
196 if id.is_empty() {
197 return Err(anyhow!("workspace id must be non-empty"));
198 }
199 if !ws_ids.insert(id.to_string()) {
200 return Err(anyhow!("duplicate workspace id: {id}"));
201 }
202 if !ws.root.exists() {
203 return Err(anyhow!(
204 "workspace root does not exist: {}",
205 ws.root.display()
206 ));
207 }
208 }
209 if !ws_ids.contains(self.default_workspace_id.trim()) {
210 return Err(anyhow!(
211 "defaultWorkspaceId '{}' not found in workspaces",
212 self.default_workspace_id
213 ));
214 }
215
216 let mut token_ids = BTreeSet::new();
217 for t in &self.tokens {
218 let id = t.id.trim();
219 if id.is_empty() {
220 return Err(anyhow!("token id must be non-empty"));
221 }
222 if !token_ids.insert(id.to_string()) {
223 return Err(anyhow!("duplicate token id: {id}"));
224 }
225 if t.effective_scopes().is_empty() {
228 return Err(anyhow!("token '{id}' must have at least 1 scope or a role"));
229 }
230 parse_sha256_hex(&t.sha256_hex)
231 .with_context(|| format!("token '{id}' invalid sha256Hex"))?;
232 }
233
234 if let Some(parent) = self.audit_log_path.parent() {
235 if !parent.as_os_str().is_empty() && !parent.exists() {
236 return Err(anyhow!(
237 "auditLogPath parent does not exist: {}",
238 parent.display()
239 ));
240 }
241 }
242 Ok(())
243 }
244
245 pub fn validate_for_serve(&self) -> Result<()> {
246 self.validate()?;
247 if self.tokens.is_empty() {
248 return Err(anyhow!("team server requires at least 1 token"));
249 }
250 Ok(())
251 }
252}
253
254#[derive(Clone)]
255struct TeamAuthContext {
256 token_id: String,
257 scopes: BTreeSet<TeamScope>,
258}
259
260#[derive(Clone)]
261pub struct TeamRequestContext {
262 pub workspace_id: String,
263}
264
265#[derive(Clone)]
266pub struct TeamState {
267 auth: Arc<Vec<TeamTokenConfig>>,
268 engine: Arc<TeamContextEngine>,
269 audit: Arc<tokio::sync::Mutex<tokio::fs::File>>,
270 pub savings_store_dir: Arc<tokio::sync::Mutex<std::path::PathBuf>>,
271 pub storage_roots: super::team_billing::StorageRoots,
273 pub storage_cache: Arc<tokio::sync::Mutex<super::team_billing::StorageCache>>,
275}
276
277#[derive(Clone)]
278pub struct TeamAppState {
279 concurrency: Arc<tokio::sync::Semaphore>,
280 rate: Arc<super::RateLimiter>,
281 timeout: Duration,
282 pub team: Arc<TeamState>,
283 max_body_bytes: usize,
284}
285
286#[derive(Debug, Deserialize)]
287#[serde(rename_all = "camelCase")]
288struct ToolCallBody {
289 name: String,
290 #[serde(default)]
291 arguments: Option<Value>,
292 #[serde(default)]
293 workspace_id: Option<String>,
294 #[serde(default)]
295 channel_id: Option<String>,
296}
297
298#[derive(Debug, Deserialize)]
299#[serde(rename_all = "camelCase")]
300struct ToolsQuery {
301 #[serde(default)]
302 offset: Option<usize>,
303 #[serde(default)]
304 limit: Option<usize>,
305}
306
307#[derive(Debug, Deserialize)]
308#[serde(rename_all = "camelCase")]
309struct EventsQuery {
310 #[serde(default)]
311 since: Option<i64>,
312 #[serde(default)]
313 limit: Option<usize>,
314 #[serde(default)]
315 channel_id: Option<String>,
316}
317
318#[derive(Clone)]
319struct TeamCtxServer {
320 default_workspace_id: String,
321 roots: Arc<HashMap<String, String>>,
322}
323
324impl TeamCtxServer {
325 fn default_root(&self) -> &str {
326 self.roots
327 .get(&self.default_workspace_id)
328 .expect("default workspace root")
329 }
330
331 fn rewrite_dot_paths(args: &mut Map<String, Value>, root: &str) {
332 for k in ["path", "target_directory", "targetDirectory"] {
333 let Some(Value::String(s)) = args.get(k) else {
334 continue;
335 };
336 let t = s.trim();
337 if t.is_empty() || t == "." {
338 args.insert(k.to_string(), Value::String(root.to_string()));
339 }
340 }
341 }
342
343 fn pick_workspace(
344 &self,
345 args: &mut Map<String, Value>,
346 ) -> std::result::Result<(String, String), rmcp::ErrorData> {
347 let ws = args
348 .get(WORKSPACE_ARG_KEY)
349 .and_then(|v| v.as_str())
350 .unwrap_or(self.default_workspace_id.as_str())
351 .to_string();
352 args.remove(WORKSPACE_ARG_KEY);
353
354 let root = self
355 .roots
356 .get(&ws)
357 .cloned()
358 .ok_or_else(|| rmcp::ErrorData::invalid_params("unknown workspaceId", None))?;
359 Self::rewrite_dot_paths(args, &root);
360 Ok((ws, root))
361 }
362
363 fn make_server(&self, workspace_id: &str, channel_id: &str) -> LeanCtxServer {
364 let root = self
365 .roots
366 .get(workspace_id)
367 .cloned()
368 .unwrap_or_else(|| self.default_root().to_string());
369 LeanCtxServer::new_shared_with_context(&root, workspace_id, channel_id)
370 }
371}
372
373impl ServerHandler for TeamCtxServer {
374 fn get_info(&self) -> rmcp::model::ServerInfo {
375 let s = self.make_server(&self.default_workspace_id, "default");
376 <LeanCtxServer as ServerHandler>::get_info(&s)
377 }
378
379 async fn initialize(
380 &self,
381 request: rmcp::model::InitializeRequestParams,
382 context: RequestContext<RoleServer>,
383 ) -> std::result::Result<rmcp::model::InitializeResult, rmcp::ErrorData> {
384 let s = self.make_server(&self.default_workspace_id, "default");
385 <LeanCtxServer as ServerHandler>::initialize(&s, request, context).await
386 }
387
388 async fn list_tools(
389 &self,
390 request: Option<rmcp::model::PaginatedRequestParams>,
391 context: RequestContext<RoleServer>,
392 ) -> std::result::Result<rmcp::model::ListToolsResult, rmcp::ErrorData> {
393 let s = self.make_server(&self.default_workspace_id, "default");
394 <LeanCtxServer as ServerHandler>::list_tools(&s, request, context).await
395 }
396
397 async fn call_tool(
398 &self,
399 mut request: CallToolRequestParams,
400 context: RequestContext<RoleServer>,
401 ) -> std::result::Result<CallToolResult, rmcp::ErrorData> {
402 let mut args = request.arguments.take().unwrap_or_default();
403 let (ws, root) = self.pick_workspace(&mut args)?;
404 let channel = args
405 .get(CHANNEL_ARG_KEY)
406 .and_then(|v| v.as_str())
407 .unwrap_or("default")
408 .to_string();
409 args.remove(CHANNEL_ARG_KEY);
410 Self::rewrite_dot_paths(&mut args, &root);
412 request.arguments = Some(args);
413 let s = LeanCtxServer::new_shared_with_context(&root, &ws, &channel);
414 <LeanCtxServer as ServerHandler>::call_tool(&s, request, context).await
415 }
416}
417
418struct TeamContextEngine {
419 server: TeamCtxServer,
420 next_id: AtomicI64,
421}
422
423impl TeamContextEngine {
424 fn new(server: TeamCtxServer) -> Self {
425 Self {
426 server,
427 next_id: AtomicI64::new(1),
428 }
429 }
430
431 fn manifest_value() -> Value {
432 crate::core::mcp_manifest::manifest_value()
433 }
434
435 async fn call_tool_value(&self, name: &str, arguments: Option<Value>) -> Result<Value> {
436 let result = self.call_tool_result(name, arguments).await?;
437 serde_json::to_value(result).map_err(|e| anyhow!("serialize CallToolResult: {e}"))
438 }
439
440 async fn call_tool_result(
441 &self,
442 name: &str,
443 arguments: Option<Value>,
444 ) -> Result<CallToolResult> {
445 let id = self.next_id.fetch_add(1, Ordering::Relaxed);
446 let req_id = NumberOrString::Number(id);
447
448 let args_obj: Map<String, Value> = match arguments {
449 None => Map::new(),
450 Some(Value::Object(m)) => m,
451 Some(other) => {
452 return Err(anyhow!(
453 "tool arguments must be a JSON object (got {other})"
454 ))
455 }
456 };
457
458 let params = CallToolRequestParams::new(name.to_string()).with_arguments(args_obj);
459 let call: CallToolRequest = CallToolRequest::new(params);
460 let client_req = ClientRequest::CallToolRequest(call);
461 let msg = ClientJsonRpcMessage::Request(JsonRpcRequest::new(req_id, client_req));
462
463 let (transport, mut rx) = OneshotTransport::<RoleServer>::new(msg);
464 let service = serve_directly(self.server.clone(), transport, None);
465 tokio::spawn(async move {
466 let _ = service.waiting().await;
467 });
468
469 let Some(server_msg) = rx.recv().await else {
470 return Err(anyhow!("no response from tool call"));
471 };
472
473 match server_msg {
474 ServerJsonRpcMessage::Response(r) => match r.result {
475 ServerResult::CallToolResult(result) => Ok(result),
476 other => Err(anyhow!("unexpected server result: {other:?}")),
477 },
478 ServerJsonRpcMessage::Error(e) => Err(anyhow!("{e:?}")).context("tool call error"),
479 ServerJsonRpcMessage::Notification(_) => Err(anyhow!("unexpected notification")),
480 ServerJsonRpcMessage::Request(_) => Err(anyhow!("unexpected request")),
481 }
482 }
483}
484
485fn sha256_hex(bytes: &[u8]) -> String {
486 let mut h = Sha256::new();
487 h.update(bytes);
488 let digest = h.finalize();
489 hex_lower(&digest)
490}
491
492fn hex_lower(bytes: &[u8]) -> String {
493 const HEX: &[u8; 16] = b"0123456789abcdef";
494 let mut out = Vec::with_capacity(bytes.len() * 2);
495 for &b in bytes {
496 out.push(HEX[(b >> 4) as usize]);
497 out.push(HEX[(b & 0x0f) as usize]);
498 }
499 String::from_utf8(out).unwrap_or_default()
500}
501
502fn parse_sha256_hex(s: &str) -> Result<Vec<u8>> {
503 let s = s.trim();
504 if s.len() != 64 {
505 return Err(anyhow!("sha256 hex must be 64 chars"));
506 }
507 let mut out = Vec::with_capacity(32);
508 let bytes = s.as_bytes();
509 let to_nibble = |c: u8| -> Option<u8> {
510 match c {
511 b'0'..=b'9' => Some(c - b'0'),
512 b'a'..=b'f' => Some(c - b'a' + 10),
513 b'A'..=b'F' => Some(c - b'A' + 10),
514 _ => None,
515 }
516 };
517 for i in (0..64).step_by(2) {
518 let hi = to_nibble(bytes[i]).ok_or_else(|| anyhow!("invalid hex"))?;
519 let lo = to_nibble(bytes[i + 1]).ok_or_else(|| anyhow!("invalid hex"))?;
520 out.push((hi << 4) | lo);
521 }
522 Ok(out)
523}
524
525fn required_scopes(tool_name: &str, args: Option<&Value>) -> Option<BTreeSet<TeamScope>> {
526 if matches!(tool_name, "ctx_shell" | "ctx_execute" | "ctx_edit") {
527 return None;
528 }
529
530 if tool_name == "ctx" {
531 let Value::Object(m) = args? else {
532 return None;
533 };
534 let sub = m.get("tool")?.as_str()?.trim();
535 if sub.is_empty() {
536 return None;
537 }
538 let canonical = if sub.starts_with("ctx_") {
539 sub.to_string()
540 } else {
541 format!("ctx_{sub}")
542 };
543 let mut m2 = m.clone();
544 m2.remove("tool");
545 return required_scopes(&canonical, Some(&Value::Object(m2)));
546 }
547
548 let mut s = BTreeSet::new();
549 match tool_name {
550 "ctx_read" | "ctx_multi_read" | "ctx_smart_read" | "ctx_search" | "ctx_tree"
552 | "ctx_outline" | "ctx_expand" | "ctx_delta" | "ctx_dedup" | "ctx_prefetch"
553 | "ctx_preload" | "ctx_review" | "ctx_response" | "ctx_task" | "ctx_overview"
554 | "ctx_architecture" | "ctx_benchmark" | "ctx_cost" | "ctx_intent" | "ctx_heatmap"
555 | "ctx_gain" | "ctx_analyze" | "ctx_discover_tools" | "ctx_discover" | "ctx_symbol"
556 | "ctx_index" | "ctx_metrics" | "ctx_cache" | "ctx_agent" => {
557 s.insert(TeamScope::Search);
558 Some(s)
559 }
560 "ctx_pack" => {
562 s.insert(TeamScope::Search);
563 s.insert(TeamScope::Graph);
564 Some(s)
565 }
566 "ctx_graph" | "ctx_impact" | "ctx_callgraph" | "ctx_routes" => {
568 s.insert(TeamScope::Graph);
569
570 if tool_name == "ctx_graph" {
571 let action = args
572 .and_then(|v| v.get("action"))
573 .and_then(|v| v.as_str())
574 .unwrap_or("");
575 if matches!(
576 action,
577 "index-build"
578 | "index-build-full"
579 | "index-build-background"
580 | "index-build-full-background"
581 ) {
582 s.insert(TeamScope::Index);
583 }
584 }
585
586 Some(s)
587 }
588 "ctx_semantic_search" => {
589 s.insert(TeamScope::Search);
590 if args
591 .and_then(|v| v.get("artifacts"))
592 .and_then(Value::as_bool)
593 .unwrap_or(false)
594 {
595 s.insert(TeamScope::Artifacts);
596 }
597 if args
598 .and_then(|v| v.get("action"))
599 .and_then(|v| v.as_str())
600 .is_some_and(|v| v.eq_ignore_ascii_case("reindex"))
601 {
602 s.insert(TeamScope::Index);
603 }
604 Some(s)
605 }
606 "ctx_session" | "ctx_handoff" | "ctx_workflow" | "ctx_compress" | "ctx_share" => {
608 s.insert(TeamScope::SessionMutations);
609 Some(s)
610 }
611 "ctx_knowledge" | "ctx_knowledge_relations" => {
613 s.insert(TeamScope::Knowledge);
614 Some(s)
615 }
616 "ctx_artifacts" | "ctx_proof" | "ctx_verify" => {
618 s.insert(TeamScope::Artifacts);
619 Some(s)
620 }
621 _ => None,
622 }
623}
624
625async fn team_slo_middleware(req: Request<Body>, next: Next) -> Response {
632 let measured = {
633 let p = req.uri().path();
634 p.starts_with("/v1/") || p.starts_with("/api/v1/")
635 };
636 let start = std::time::Instant::now();
637 let res = next.run(req).await;
638 if measured {
639 crate::core::team_slo::global().record_request(
640 u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX),
641 !res.status().is_server_error(),
642 );
643 }
644 res
645}
646
647async fn team_rate_limit_middleware(
648 State(state): State<TeamAppState>,
649 req: Request<Body>,
650 next: Next,
651) -> Response {
652 if req.uri().path() == "/health" {
653 return next.run(req).await;
654 }
655 if !state.rate.allow().await {
656 return StatusCode::TOO_MANY_REQUESTS.into_response();
657 }
658 next.run(req).await
659}
660
661async fn team_concurrency_middleware(
662 State(state): State<TeamAppState>,
663 req: Request<Body>,
664 next: Next,
665) -> Response {
666 if req.uri().path() == "/health" {
667 return next.run(req).await;
668 }
669 let Ok(permit) = state.concurrency.clone().try_acquire_owned() else {
670 return StatusCode::TOO_MANY_REQUESTS.into_response();
671 };
672 let resp = next.run(req).await;
673 drop(permit);
674 resp
675}
676
677async fn team_auth_middleware(
678 State(state): State<TeamAppState>,
679 mut req: Request<Body>,
680 next: Next,
681) -> Response {
682 if req.uri().path() == "/health" {
683 return next.run(req).await;
684 }
685
686 let Some(h) = req.headers().get(header::AUTHORIZATION) else {
687 return super::json_error(
688 StatusCode::UNAUTHORIZED,
689 "unauthorized",
690 "missing Authorization header",
691 );
692 };
693 let Ok(s) = h.to_str() else {
694 return super::json_error(
695 StatusCode::UNAUTHORIZED,
696 "unauthorized",
697 "malformed Authorization header",
698 );
699 };
700 let Some(token) = s
701 .strip_prefix("Bearer ")
702 .or_else(|| s.strip_prefix("bearer "))
703 else {
704 return super::json_error(
705 StatusCode::UNAUTHORIZED,
706 "unauthorized",
707 "Authorization must use the Bearer scheme",
708 );
709 };
710
711 let token_hash = sha256_hex(token.as_bytes());
712 let mut matched: Option<TeamTokenConfig> = None;
713 for t in state.team.auth.iter() {
714 if super::constant_time_eq(token_hash.as_bytes(), t.sha256_hex.as_bytes()) {
715 matched = Some(t.clone());
716 break;
717 }
718 }
719 let Some(tok) = matched else {
720 return super::json_error(
721 StatusCode::UNAUTHORIZED,
722 "unauthorized",
723 "invalid bearer token",
724 );
725 };
726 let tok_scopes: BTreeSet<TeamScope> = tok.effective_scopes();
727
728 let workspace_id = req
729 .headers()
730 .get(WORKSPACE_HEADER)
731 .and_then(|h| h.to_str().ok())
732 .map(|s| s.trim().to_string())
733 .filter(|s| !s.is_empty())
734 .unwrap_or_else(|| state.team.engine.server.default_workspace_id.clone());
735 if !state.team.engine.server.roots.contains_key(&workspace_id) {
736 return super::json_error(
737 StatusCode::BAD_REQUEST,
738 "unknown_workspace",
739 "unknown workspace",
740 );
741 }
742 let workspace_id_for_audit = workspace_id.clone();
743
744 req.extensions_mut().insert(TeamAuthContext {
745 token_id: tok.id.clone(),
746 scopes: tok_scopes.clone(),
747 });
748 req.extensions_mut()
749 .insert(TeamRequestContext { workspace_id });
750
751 let path0 = req.uri().path();
753 if path0 == "/v1/events" {
754 let allow = tok_scopes.contains(&TeamScope::Events);
755 let _ = audit_write(
756 &state.team.audit,
757 &tok.id,
758 &workspace_id_for_audit,
759 None,
760 Some("events"),
761 allow,
762 if allow { None } else { Some("scope_denied") },
763 None,
764 )
765 .await;
766 if !allow {
767 return super::json_error(
768 StatusCode::FORBIDDEN,
769 "scope_denied",
770 "token lacks required scope: events",
771 );
772 }
773 }
774
775 if path0 == "/v1/metrics" {
776 let allow = tok_scopes.contains(&TeamScope::Audit);
777 let _ = audit_write(
778 &state.team.audit,
779 &tok.id,
780 &workspace_id_for_audit,
781 None,
782 Some("metrics"),
783 allow,
784 if allow { None } else { Some("scope_denied") },
785 None,
786 )
787 .await;
788 if !allow {
789 return super::json_error(
790 StatusCode::FORBIDDEN,
791 "scope_denied",
792 "token lacks required scope: audit",
793 );
794 }
795 }
796
797 let audit_gated = match path0 {
800 "/v1/savings/summary" => Some("savings_summary"),
801 "/v1/storage" => Some("storage"),
802 "/v1/usage" => Some("usage"),
803 p if p.starts_with("/v1/savings/member/") => Some("savings_member"),
804 _ => None,
805 };
806 if let Some(action) = audit_gated {
807 let allow = tok_scopes.contains(&TeamScope::Audit);
808 let _ = audit_write(
809 &state.team.audit,
810 &tok.id,
811 &workspace_id_for_audit,
812 None,
813 Some(action),
814 allow,
815 if allow { None } else { Some("scope_denied") },
816 None,
817 )
818 .await;
819 if !allow {
820 return super::json_error(
821 StatusCode::FORBIDDEN,
822 "scope_denied",
823 "token lacks required scope: audit",
824 );
825 }
826 }
827
828 let path = req.uri().path().to_string();
830 if path != "/v1/tools/call"
831 && path != "/v1/tools"
832 && path != "/v1/manifest"
833 && path != "/health"
834 {
835 if req.method() != axum::http::Method::POST {
836 return next.run(req).await;
837 }
838
839 let (parts, body0) = req.into_parts();
840 let Ok(bytes) = body::to_bytes(body0, state.max_body_bytes).await else {
841 return super::json_error(
842 StatusCode::BAD_REQUEST,
843 "invalid_request",
844 "could not read request body",
845 );
846 };
847
848 let mut allow = false;
849 let mut denied_reason: Option<String> = None;
850 if let Ok(v) = serde_json::from_slice::<Value>(&bytes) {
851 if v.is_array() {
852 denied_reason = Some("batch_requests_not_supported".to_string());
853 let _ = audit_write(
854 &state.team.audit,
855 &tok.id,
856 &workspace_id_for_audit,
857 None,
858 None,
859 false,
860 denied_reason.as_deref(),
861 None,
862 )
863 .await;
864 } else {
865 let method = v.get("method").and_then(|m| m.as_str()).unwrap_or("");
866 if method.eq_ignore_ascii_case("tools/call") {
867 let tool = v
868 .pointer("/params/name")
869 .and_then(|x| x.as_str())
870 .unwrap_or("");
871 let args = v.pointer("/params/arguments");
872 let req_scopes = required_scopes(tool, args);
873 allow = match req_scopes {
874 None => false,
875 Some(reqs) => reqs.is_subset(&tok_scopes),
876 };
877 if !allow {
878 denied_reason = Some("scope_denied".to_string());
879 }
880 let _ = audit_write(
881 &state.team.audit,
882 &tok.id,
883 &workspace_id_for_audit,
884 Some(tool),
885 Some(method),
886 allow,
887 denied_reason.as_deref(),
888 args,
889 )
890 .await;
891 } else {
892 allow = true;
893 }
894 }
895 }
896
897 if !allow {
898 return super::json_error(
899 StatusCode::FORBIDDEN,
900 "scope_denied",
901 "token lacks required scope for this tool",
902 );
903 }
904
905 req = Request::from_parts(parts, Body::from(bytes));
906 }
907
908 next.run(req).await
909}
910
911async fn audit_write(
912 file: &tokio::sync::Mutex<tokio::fs::File>,
913 token_id: &str,
914 workspace_id: &str,
915 tool: Option<&str>,
916 method: Option<&str>,
917 allowed: bool,
918 denied_reason: Option<&str>,
919 args: Option<&Value>,
920) -> Result<()> {
921 let args_hash = args
922 .map(|a| {
923 let s = a.to_string();
924 let mut hasher = Md5::new();
925 hasher.update(s.as_bytes());
926 format!("{:x}", hasher.finalize())
927 })
928 .unwrap_or_default();
929
930 let ts = chrono::Local::now().to_rfc3339();
931 let rec = json!({
932 "ts": ts,
933 "tokenId": token_id,
934 "workspaceId": workspace_id,
935 "tool": tool,
936 "method": method,
937 "allowed": allowed,
938 "deniedReason": denied_reason,
939 "argumentsMd5": args_hash,
940 });
941
942 let mut guard = file.lock().await;
943 guard.write_all(rec.to_string().as_bytes()).await?;
944 guard.write_all(b"\n").await?;
945 guard.flush().await?;
946 Ok(())
947}
948
949async fn audit_event(
951 file: &tokio::sync::Mutex<tokio::fs::File>,
952 token_id: &str,
953 workspace_id: &str,
954 channel_id: &str,
955 event_kind: &str,
956 actor: Option<&str>,
957 event_id: i64,
958) -> Result<()> {
959 let ts = chrono::Local::now().to_rfc3339();
960 let rec = json!({
961 "ts": ts,
962 "type": "context_event",
963 "tokenId": token_id,
964 "workspaceId": workspace_id,
965 "channelId": channel_id,
966 "eventKind": event_kind,
967 "actor": actor,
968 "eventId": event_id,
969 });
970
971 let mut guard = file.lock().await;
972 guard.write_all(rec.to_string().as_bytes()).await?;
973 guard.write_all(b"\n").await?;
974 guard.flush().await?;
975 Ok(())
976}
977
978async fn v1_manifest(State(_state): State<TeamAppState>) -> impl IntoResponse {
979 let v = TeamContextEngine::manifest_value();
980 (StatusCode::OK, Json(v))
981}
982
983async fn v1_tools(
984 State(_state): State<TeamAppState>,
985 Query(q): Query<ToolsQuery>,
986) -> impl IntoResponse {
987 let v = TeamContextEngine::manifest_value();
988 let tools = v
989 .get("tools")
990 .and_then(|t| t.get("granular"))
991 .cloned()
992 .unwrap_or(Value::Array(vec![]));
993
994 let all = tools.as_array().cloned().unwrap_or_default();
995 let total = all.len();
996 let offset = q.offset.unwrap_or(0).min(total);
997 let limit = q.limit.unwrap_or(200).min(500);
998 let page = all.into_iter().skip(offset).take(limit).collect::<Vec<_>>();
999
1000 (
1001 StatusCode::OK,
1002 Json(json!({
1003 "tools": page,
1004 "total": total,
1005 "offset": offset,
1006 "limit": limit,
1007 })),
1008 )
1009}
1010
1011async fn v1_tool_call(
1012 State(state): State<TeamAppState>,
1013 Extension(auth): Extension<TeamAuthContext>,
1014 Extension(ctx): Extension<TeamRequestContext>,
1015 Json(body): Json<ToolCallBody>,
1016) -> impl IntoResponse {
1017 let workspace_id = body
1018 .workspace_id
1019 .clone()
1020 .unwrap_or_else(|| ctx.workspace_id.clone());
1021 if !state.team.engine.server.roots.contains_key(&workspace_id) {
1022 let _ = audit_write(
1023 &state.team.audit,
1024 &auth.token_id,
1025 &workspace_id,
1026 Some(&body.name),
1027 Some("/v1/tools/call"),
1028 false,
1029 Some("unknown_workspace"),
1030 body.arguments.as_ref(),
1031 )
1032 .await;
1033 return super::json_error(
1034 StatusCode::BAD_REQUEST,
1035 "unknown_workspace",
1036 "unknown workspace",
1037 );
1038 }
1039
1040 let mut args = match body.arguments {
1041 None => Value::Object(Map::new()),
1042 Some(Value::Object(m)) => Value::Object(m),
1043 Some(other) => {
1044 let _ = audit_write(
1045 &state.team.audit,
1046 &auth.token_id,
1047 &workspace_id,
1048 Some(&body.name),
1049 Some("/v1/tools/call"),
1050 false,
1051 Some("invalid_arguments"),
1052 Some(&other),
1053 )
1054 .await;
1055 return super::json_error(
1056 StatusCode::BAD_REQUEST,
1057 "invalid_arguments",
1058 &format!("tool arguments must be a JSON object (got {other})"),
1059 );
1060 }
1061 };
1062
1063 if let Value::Object(ref mut m) = args {
1064 m.insert(
1065 WORKSPACE_ARG_KEY.to_string(),
1066 Value::String(workspace_id.clone()),
1067 );
1068 if let Some(ch) = body.channel_id.as_deref() {
1069 if !ch.trim().is_empty() {
1070 m.insert(
1071 CHANNEL_ARG_KEY.to_string(),
1072 Value::String(ch.trim().to_string()),
1073 );
1074 }
1075 }
1076 }
1077
1078 let required = required_scopes(&body.name, Some(&args));
1079 let mutates_index = required
1082 .as_ref()
1083 .is_some_and(|reqs| reqs.contains(&TeamScope::Index));
1084 let allowed = match required {
1085 None => false,
1086 Some(reqs) => reqs.is_subset(&auth.scopes),
1087 };
1088 if !allowed {
1089 let _ = audit_write(
1090 &state.team.audit,
1091 &auth.token_id,
1092 &workspace_id,
1093 Some(&body.name),
1094 Some("/v1/tools/call"),
1095 false,
1096 Some("scope_denied"),
1097 Some(&args),
1098 )
1099 .await;
1100 return super::json_error(
1101 StatusCode::FORBIDDEN,
1102 "scope_denied",
1103 "token lacks required scope for this tool",
1104 );
1105 }
1106
1107 let tool_name = body.name.clone();
1108 let call = tokio::time::timeout(
1109 state.timeout,
1110 state
1111 .team
1112 .engine
1113 .call_tool_value(&tool_name, Some(args.clone())),
1114 )
1115 .await;
1116
1117 match call {
1118 Ok(Ok(v)) => {
1119 if mutates_index {
1120 crate::core::team_slo::global().record_index_write();
1121 }
1122 let _ = audit_write(
1123 &state.team.audit,
1124 &auth.token_id,
1125 &workspace_id,
1126 Some(&tool_name),
1127 Some("/v1/tools/call"),
1128 true,
1129 None,
1130 Some(&args),
1131 )
1132 .await;
1133 (StatusCode::OK, Json(json!({ "result": v }))).into_response()
1134 }
1135 Ok(Err(e)) => {
1136 let _ = audit_write(
1137 &state.team.audit,
1138 &auth.token_id,
1139 &workspace_id,
1140 Some(&tool_name),
1141 Some("/v1/tools/call"),
1142 true,
1143 Some("tool_error"),
1144 Some(&args),
1145 )
1146 .await;
1147 {
1148 tracing::warn!("team tool call error: {e}");
1149 super::json_error(
1150 StatusCode::BAD_REQUEST,
1151 "tool_error",
1152 "tool execution failed",
1153 )
1154 }
1155 }
1156 Err(_) => {
1157 let _ = audit_write(
1158 &state.team.audit,
1159 &auth.token_id,
1160 &workspace_id,
1161 Some(&tool_name),
1162 Some("/v1/tools/call"),
1163 true,
1164 Some("request_timeout"),
1165 Some(&args),
1166 )
1167 .await;
1168 super::json_error(
1169 StatusCode::GATEWAY_TIMEOUT,
1170 "request_timeout",
1171 "tool call timed out",
1172 )
1173 }
1174 }
1175}
1176
1177async fn v1_events(
1178 State(state): State<TeamAppState>,
1179 Extension(auth): Extension<TeamAuthContext>,
1180 Extension(ctx): Extension<TeamRequestContext>,
1181 Query(q): Query<EventsQuery>,
1182) -> Sse<impl Stream<Item = Result<SseEvent, std::convert::Infallible>>> {
1183 let ws = ctx.workspace_id;
1184 let ch = q.channel_id.unwrap_or_else(|| "default".to_string());
1185 let since = q.since.unwrap_or(0);
1186 let limit = q.limit.unwrap_or(200).min(1000);
1187
1188 let _ = audit_event(
1189 &state.team.audit,
1190 &auth.token_id,
1191 &ws,
1192 &ch,
1193 "sse_subscribe",
1194 None,
1195 since,
1196 )
1197 .await;
1198
1199 let rt = crate::core::context_os::runtime();
1200 let replay = rt.bus.read(&ws, &ch, since, limit);
1201 let rx = if let Some(rx) = rt.bus.subscribe(&ws, &ch) {
1202 rx
1203 } else {
1204 tracing::warn!("SSE subscriber limit reached for {ws}/{ch}");
1205 let (_, rx) = tokio::sync::broadcast::channel::<crate::core::context_os::ContextEventV1>(1);
1206 rx
1207 };
1208 rt.metrics.record_sse_connect();
1209 rt.metrics.record_events_replayed(replay.len() as u64);
1210 rt.metrics.record_workspace_active(&ws);
1211
1212 let bus = rt.bus.clone();
1213 let metrics = rt.metrics.clone();
1214 let pending: std::collections::VecDeque<crate::core::context_os::ContextEventV1> =
1215 replay.into();
1216
1217 use crate::core::context_os::{redact_event_payload, RedactionLevel};
1218 let redaction = RedactionLevel::RefsOnly;
1219
1220 let stream = futures::stream::unfold(
1221 (
1222 pending,
1223 rx,
1224 ws.clone(),
1225 ch.clone(),
1226 since,
1227 redaction,
1228 bus,
1229 metrics,
1230 ),
1231 |(mut pending, mut rx, ws, ch, mut last_id, redaction, bus, metrics)| async move {
1232 if let Some(mut ev) = pending.pop_front() {
1233 last_id = ev.id;
1234 redact_event_payload(&mut ev, redaction);
1235 let data = serde_json::to_string(&ev).unwrap_or_else(|_| "{}".to_string());
1236 let evt = SseEvent::default()
1237 .id(ev.id.to_string())
1238 .event(ev.kind)
1239 .data(data);
1240 return Some((
1241 Ok(evt),
1242 (pending, rx, ws, ch, last_id, redaction, bus, metrics),
1243 ));
1244 }
1245
1246 loop {
1247 match rx.recv().await {
1248 Ok(mut ev) if ev.id > last_id => {
1249 last_id = ev.id;
1250 redact_event_payload(&mut ev, redaction);
1251 let data = serde_json::to_string(&ev).unwrap_or_else(|_| "{}".to_string());
1252 let evt = SseEvent::default()
1253 .id(ev.id.to_string())
1254 .event(ev.kind)
1255 .data(data);
1256 return Some((
1257 Ok(evt),
1258 (pending, rx, ws, ch, last_id, redaction, bus, metrics),
1259 ));
1260 }
1261 Ok(_) => {}
1262 Err(broadcast::error::RecvError::Closed) => return None,
1263 Err(broadcast::error::RecvError::Lagged(skipped)) => {
1264 let missed = bus.read(&ws, &ch, last_id, skipped as usize);
1265 metrics.record_events_replayed(missed.len() as u64);
1266 for ev in missed {
1267 last_id = last_id.max(ev.id);
1268 pending.push_back(ev);
1269 }
1270 }
1271 }
1272 }
1273 },
1274 );
1275
1276 let metrics_ref = rt.metrics.clone();
1277 let guarded = super::SseDisconnectGuard {
1278 inner: Box::pin(stream),
1279 metrics: metrics_ref,
1280 };
1281
1282 Sse::new(guarded).keep_alive(KeepAlive::new().interval(Duration::from_secs(15)))
1283}
1284
1285#[derive(Debug, Deserialize)]
1286struct MetricsQuery {
1287 #[serde(default)]
1290 format: Option<String>,
1291}
1292
1293async fn v1_team_metrics(
1294 State(_state): State<TeamAppState>,
1295 Query(q): Query<MetricsQuery>,
1296) -> Response {
1297 let slo = crate::core::team_slo::global().snapshot();
1298
1299 if q.format.as_deref() == Some("prometheus") {
1300 return (
1301 StatusCode::OK,
1302 [(
1303 axum::http::header::CONTENT_TYPE,
1304 "text/plain; version=0.0.4",
1305 )],
1306 slo.to_prometheus(),
1307 )
1308 .into_response();
1309 }
1310
1311 let rt = crate::core::context_os::runtime();
1312 let snap = rt.metrics.snapshot();
1313 let mut v = serde_json::to_value(snap).unwrap_or_default();
1314 if let Value::Object(ref mut m) = v {
1315 m.insert(
1316 "slo".to_string(),
1317 serde_json::to_value(&slo).unwrap_or_default(),
1318 );
1319 }
1320 (StatusCode::OK, Json(v)).into_response()
1321}
1322
1323fn streamable_http_config(cfg: &TeamServerConfig) -> rmcp::transport::StreamableHttpServerConfig {
1324 let mut out = rmcp::transport::StreamableHttpServerConfig::default()
1325 .with_stateful_mode(cfg.stateful_mode)
1326 .with_json_response(cfg.json_response);
1327
1328 if cfg.disable_host_check {
1329 out = out.disable_allowed_hosts();
1330 return out;
1331 }
1332 if !cfg.allowed_hosts.is_empty() {
1333 out = out.with_allowed_hosts(cfg.allowed_hosts.clone());
1334 return out;
1335 }
1336 let host = cfg.host.trim();
1337 if host == "127.0.0.1" || host == "localhost" || host == "::1" {
1338 out.allowed_hosts.push(host.to_string());
1339 }
1340 out
1341}
1342
1343pub async fn serve_team(cfg: TeamServerConfig) -> Result<()> {
1344 cfg.validate_for_serve()?;
1345
1346 let addr: std::net::SocketAddr = format!("{}:{}", cfg.host, cfg.port)
1347 .parse()
1348 .context("invalid host/port")?;
1349
1350 let team_server = TeamCtxServer {
1351 default_workspace_id: cfg.default_workspace_id.clone(),
1352 roots: Arc::new(
1353 cfg.workspaces
1354 .iter()
1355 .map(|w| (w.id.clone(), w.root.to_string_lossy().to_string()))
1356 .collect(),
1357 ),
1358 };
1359 let engine = Arc::new(TeamContextEngine::new(team_server.clone()));
1360
1361 let audit_file = tokio::fs::OpenOptions::new()
1362 .create(true)
1363 .append(true)
1364 .open(&cfg.audit_log_path)
1365 .await
1366 .with_context(|| format!("open audit log {}", cfg.audit_log_path.display()))?;
1367
1368 let savings_dir = cfg
1369 .audit_log_path
1370 .parent()
1371 .unwrap_or(std::path::Path::new("."))
1372 .join("savings");
1373 let workspace_roots: Vec<(String, std::path::PathBuf)> = cfg
1374 .workspaces
1375 .iter()
1376 .map(|w| (w.id.clone(), w.root.clone()))
1377 .collect();
1378 let team = Arc::new(TeamState {
1379 auth: Arc::new(cfg.tokens.clone()),
1380 engine,
1381 audit: Arc::new(tokio::sync::Mutex::new(audit_file)),
1382 savings_store_dir: Arc::new(tokio::sync::Mutex::new(savings_dir)),
1383 storage_roots: super::team_billing::storage_roots_from_config(
1384 &cfg.audit_log_path,
1385 &workspace_roots,
1386 cfg.storage_quota_bytes,
1387 ),
1388 storage_cache: Arc::new(tokio::sync::Mutex::new(
1389 super::team_billing::StorageCache::default(),
1390 )),
1391 });
1392
1393 let state = TeamAppState {
1394 concurrency: Arc::new(tokio::sync::Semaphore::new(cfg.max_concurrency.max(1))),
1395 rate: Arc::new(super::RateLimiter::new(cfg.max_rps, cfg.rate_burst)),
1396 timeout: Duration::from_millis(cfg.request_timeout_ms.max(1)),
1397 team,
1398 max_body_bytes: cfg.max_body_bytes,
1399 };
1400
1401 let service_factory =
1402 move || -> std::result::Result<TeamCtxServer, std::io::Error> { Ok(team_server.clone()) };
1403 let mcp_http = StreamableHttpService::new(
1404 service_factory,
1405 Arc::new(
1406 rmcp::transport::streamable_http_server::session::local::LocalSessionManager::default(),
1407 ),
1408 streamable_http_config(&cfg),
1409 );
1410
1411 if let Some(url) = &cfg.roi_webhook_url {
1414 super::roi_webhook::validate_webhook_url(url)
1415 .map_err(|e| anyhow!("invalid roiWebhookUrl in team config: {e}"))?;
1416 super::roi_webhook::spawn_weekly_roi_webhook(state.clone(), url.clone());
1417 tracing::info!("team ROI webhook enabled (weekly)");
1418 }
1419
1420 let app = Router::new()
1421 .route("/health", get(super::health))
1422 .route("/v1/manifest", get(v1_manifest))
1423 .route("/v1/tools", get(v1_tools))
1424 .route("/v1/tools/call", axum::routing::post(v1_tool_call))
1425 .route("/v1/events", get(v1_events))
1426 .route(
1427 "/v1/context/summary",
1428 get(super::context_views::v1_context_summary),
1429 )
1430 .route(
1431 "/v1/events/search",
1432 get(super::context_views::v1_events_search),
1433 )
1434 .route(
1435 "/v1/events/lineage",
1436 get(super::context_views::v1_event_lineage),
1437 )
1438 .route("/v1/metrics", get(v1_team_metrics))
1439 .route(
1440 "/v1/savings/summary",
1441 get(super::savings_summary::v1_savings_summary),
1442 )
1443 .route(
1444 "/v1/savings/member/{signer}",
1445 get(super::savings_summary::v1_savings_member),
1446 )
1447 .route("/v1/storage", get(super::team_billing::v1_storage))
1448 .route("/v1/usage", get(super::team_billing::v1_usage))
1449 .route(
1450 "/api/v1/savings/ingest",
1451 axum::routing::post(super::savings_ingest::v1_savings_ingest),
1452 )
1453 .fallback_service(mcp_http)
1454 .layer(axum::extract::DefaultBodyLimit::max(cfg.max_body_bytes))
1455 .layer(middleware::from_fn_with_state(
1456 state.clone(),
1457 team_rate_limit_middleware,
1458 ))
1459 .layer(middleware::from_fn_with_state(
1460 state.clone(),
1461 team_concurrency_middleware,
1462 ))
1463 .layer(middleware::from_fn_with_state(
1464 state.clone(),
1465 team_auth_middleware,
1466 ))
1467 .layer(middleware::from_fn(team_slo_middleware))
1469 .with_state(state);
1470
1471 crate::core::team_slo::global().mark_started();
1472
1473 let listener = tokio::net::TcpListener::bind(addr)
1474 .await
1475 .with_context(|| format!("bind {addr}"))?;
1476
1477 tracing::info!(
1478 "lean-ctx TEAM server listening on http://{addr} (workspaces={}, audit={})",
1479 cfg.workspaces.len(),
1480 cfg.audit_log_path.display()
1481 );
1482
1483 axum::serve(listener, app)
1484 .with_graceful_shutdown(async move {
1485 let _ = tokio::signal::ctrl_c().await;
1486 })
1487 .await
1488 .context("team http server")?;
1489 Ok(())
1490}
1491
1492pub fn create_token() -> Result<(String, String)> {
1493 let mut bytes = [0u8; 32];
1494 getrandom::fill(&mut bytes).map_err(|e| anyhow!("getrandom: {e}"))?;
1495 let token = hex_lower(&bytes);
1496 let hash = sha256_hex(token.as_bytes());
1497 Ok((token, hash))
1498}