Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 11 additions & 4 deletions orch8-api/src/mcp_server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,7 @@ type ToolResult = Result<Value, String>;
async fn handle_mcp(
State(state): State<AppState>,
tenant_ctx: OptionalTenant,
admin_ctx: crate::auth::OptionalAdmin,
body: Bytes,
) -> Response {
let msg: Value = match serde_json::from_slice(&body) {
Expand Down Expand Up @@ -132,7 +133,7 @@ async fn handle_mcp(
"initialize" => Ok(initialize_result(&params)),
"ping" => Ok(json!({})),
"tools/list" => Ok(json!({ "tools": tool_catalog() })),
"tools/call" => tools_call(state, tenant_ctx, &params).await,
"tools/call" => tools_call(state, tenant_ctx, admin_ctx, &params).await,
other => Err((METHOD_NOT_FOUND, format!("method not found: {other}"))),
};
match outcome {
Expand Down Expand Up @@ -162,6 +163,7 @@ fn initialize_result(params: &Value) -> Value {
async fn tools_call(
state: AppState,
tenant_ctx: OptionalTenant,
admin_ctx: crate::auth::OptionalAdmin,
params: &Value,
) -> Result<Value, (i64, String)> {
let Some(name) = params.get("name").and_then(Value::as_str) else {
Expand All @@ -180,7 +182,7 @@ async fn tools_call(
"send_signal" => tool_send_signal(state, tenant_ctx, &args).await,
"retry_instance" => tool_retry_instance(state, tenant_ctx, &args).await,
"list_dlq" => tool_list_dlq(state, tenant_ctx, &args).await,
"get_usage" => tool_get_usage(state, tenant_ctx, &args).await,
"get_usage" => tool_get_usage(state, tenant_ctx, admin_ctx, &args).await,
// Unknown tool name β†’ -32602 (documented choice, see module docs):
// the catalog is static, so a bad name is a protocol-level caller
// bug rather than a domain outcome.
Expand Down Expand Up @@ -361,15 +363,20 @@ async fn tool_list_dlq(state: AppState, tenant_ctx: OptionalTenant, args: &Value
}

/// `get_usage`: tenant-scoped LLM token/cost aggregation over a time window.
async fn tool_get_usage(state: AppState, tenant_ctx: OptionalTenant, args: &Value) -> ToolResult {
async fn tool_get_usage(
state: AppState,
tenant_ctx: OptionalTenant,
admin_ctx: crate::auth::OptionalAdmin,
args: &Value,
) -> ToolResult {
let mut query = serde_json::Map::new();
for key in ["tenant", "start", "end"] {
if let Some(v) = args.get(key) {
query.insert(key.into(), v.clone());
}
}
let q: crate::usage::UsageQuery = parse_args(Value::Object(query))?;
rest_json(crate::usage::get_usage(State(state), tenant_ctx, Query(q)).await).await
rest_json(crate::usage::get_usage(State(state), tenant_ctx, admin_ctx, Query(q)).await).await
}

// ---- Plumbing ---------------------------------------------------------------
Expand Down
17 changes: 11 additions & 6 deletions orch8-api/src/usage.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,8 @@ use chrono::{DateTime, Duration, Utc};
use serde::Deserialize;

use crate::AppState;
use crate::auth::TenantContext;
use crate::api_keys::require_admin;
use crate::auth::{OptionalAdmin, TenantContext};
use crate::error::ApiError;
use crate::model_pricing;

Expand Down Expand Up @@ -49,18 +50,22 @@ pub struct UsageQuery {
pub async fn get_usage(
State(state): State<AppState>,
tenant_ctx: Option<axum::Extension<TenantContext>>,
admin_ctx: OptionalAdmin,
Query(q): Query<UsageQuery>,
) -> Result<impl IntoResponse, ApiError> {
// A header-scoped caller is locked to its own tenant (the `?tenant=` param
// is ignored, so it can't read another tenant's usage). Only an unscoped
// caller may select a tenant via the query param.
let tenant = match &tenant_ctx {
Some(axum::Extension(ctx)) => ctx.tenant_id.as_str().to_string(),
None => q.tenant.clone().ok_or_else(|| {
ApiError::InvalidArgument(
"usage requires a tenant (X-Tenant-Id header or ?tenant=)".into(),
)
})?,
None => {
require_admin(&admin_ctx)?;
q.tenant.clone().ok_or_else(|| {
ApiError::InvalidArgument(
"usage requires a tenant (X-Tenant-Id header or ?tenant=)".into(),
)
})?
}
};

let end = q.end.unwrap_or_else(Utc::now);
Expand Down
Loading