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