|
| 1 | +use anyhow::Context; |
| 2 | +use rmcp::{ |
| 3 | + handler::server::wrapper::Parameters, |
| 4 | + model::CallToolResult, |
| 5 | + schemars::{self, JsonSchema}, |
| 6 | + tool, tool_router, |
| 7 | +}; |
| 8 | +use serde::Deserialize; |
| 9 | +use sift_rs::{ |
| 10 | + assets::v1::{ListAssetsRequest, ListAssetsResponse, asset_service_client::AssetServiceClient}, |
| 11 | + runs::v2::{ListRunsRequest, ListRunsResponse, run_service_client::RunServiceClient}, |
| 12 | +}; |
| 13 | + |
| 14 | +use crate::{error, server::SiftMcpServer}; |
| 15 | + |
| 16 | +#[cfg(test)] |
| 17 | +mod test; |
| 18 | + |
| 19 | +const PAGE_SIZE: u32 = 1000; |
| 20 | + |
| 21 | +#[derive(Debug, Deserialize, JsonSchema)] |
| 22 | +pub struct GetParams { |
| 23 | + filter: String, |
| 24 | + limit: Option<u32>, |
| 25 | +} |
| 26 | + |
| 27 | +#[tool_router(router = resource_router, vis = "pub(crate)")] |
| 28 | +impl SiftMcpServer { |
| 29 | + #[tool( |
| 30 | + name = "get_asset", |
| 31 | + description = " |
| 32 | + Retrieve and filter assets in Sift. The `filter` parameter is a Common Expression Language (CEL). |
| 33 | + Available fields to filter by are `asset_id`, `created_by_user_id`, `modified_by_user_id`, |
| 34 | + `created_date`, `modified_date`, `name`, 'name_lower', `tag_id`, `tag_name`, 'archived_date', `is_archived`, and `metadata`. |
| 35 | + Metadata can be used in filters by using `metadata.{metadata_key_name}` as the field name. |
| 36 | + ", |
| 37 | + annotations(title = "Resource/get_asset", read_only_hint = true) |
| 38 | + )] |
| 39 | + pub async fn get_asset(&self, params: Parameters<GetParams>) -> error::McpResult { |
| 40 | + let Parameters(GetParams { filter, limit }) = params; |
| 41 | + let (page_size, record_limit) = paging(limit); |
| 42 | + |
| 43 | + let mut client = AssetServiceClient::new(self.channel.clone()); |
| 44 | + let mut page_token = String::new(); |
| 45 | + let mut results = Vec::new(); |
| 46 | + |
| 47 | + loop { |
| 48 | + let resp = client |
| 49 | + .list_assets(ListAssetsRequest { |
| 50 | + filter: filter.clone(), |
| 51 | + page_size, |
| 52 | + page_token, |
| 53 | + ..Default::default() |
| 54 | + }) |
| 55 | + .await |
| 56 | + .map_err(error::from_grpc_status)?; |
| 57 | + |
| 58 | + let ListAssetsResponse { |
| 59 | + assets, |
| 60 | + next_page_token, |
| 61 | + } = resp.into_inner(); |
| 62 | + if assets.is_empty() { |
| 63 | + break; |
| 64 | + } |
| 65 | + results.extend(assets); |
| 66 | + |
| 67 | + if results.len() >= record_limit || next_page_token.is_empty() { |
| 68 | + break; |
| 69 | + } |
| 70 | + page_token = next_page_token; |
| 71 | + } |
| 72 | + |
| 73 | + results.truncate(record_limit); |
| 74 | + let out = serde_json::to_value(&results) |
| 75 | + .context("failed to serialize assets") |
| 76 | + .map_err(error::from_anyhow)?; |
| 77 | + |
| 78 | + Ok(CallToolResult::structured(out)) |
| 79 | + } |
| 80 | + |
| 81 | + #[tool( |
| 82 | + name = "get_run", |
| 83 | + description = " |
| 84 | + Retrieve and filter runs in Sift. The `filter` parameter is a Common Expression Language (CEL). |
| 85 | + Available fields to filter by are `run_id` `organization_id`, `asset_id`, `asset_name`, `client_key`, `name`, |
| 86 | + `description`, `created_by_user_id`, `modified_by_user_id`, `created_date`, `modified_date`, `start_time`, |
| 87 | + `stop_time`, `tag_id`, `asset_tag_id`, `duration`, 'duration_string', `annotation_comments_count`, `annotation_state`, |
| 88 | + `archived_date`, `is_archived`, and `metadata`. Metadata can be used in filters by using |
| 89 | + `metadata.{metadata_key_name}` as the field name. `duration` is in the format of elapsed seconds and `duration_string` |
| 90 | + allows for `h`, `m`, `s`, `ms` suffixes (example: `duration_string > duration('10h')). |
| 91 | + ", |
| 92 | + annotations(title = "Resource/get_run", read_only_hint = true) |
| 93 | + )] |
| 94 | + pub async fn get_run(&self, params: Parameters<GetParams>) -> error::McpResult { |
| 95 | + let Parameters(GetParams { filter, limit }) = params; |
| 96 | + let (page_size, record_limit) = paging(limit); |
| 97 | + |
| 98 | + let mut client = RunServiceClient::new(self.channel.clone()); |
| 99 | + let mut page_token = String::new(); |
| 100 | + let mut results = Vec::new(); |
| 101 | + |
| 102 | + loop { |
| 103 | + let resp = client |
| 104 | + .list_runs(ListRunsRequest { |
| 105 | + filter: filter.clone(), |
| 106 | + page_size, |
| 107 | + page_token, |
| 108 | + ..Default::default() |
| 109 | + }) |
| 110 | + .await |
| 111 | + .map_err(error::from_grpc_status)?; |
| 112 | + |
| 113 | + let ListRunsResponse { |
| 114 | + runs, |
| 115 | + next_page_token, |
| 116 | + } = resp.into_inner(); |
| 117 | + if runs.is_empty() { |
| 118 | + break; |
| 119 | + } |
| 120 | + results.extend(runs); |
| 121 | + |
| 122 | + if results.len() >= record_limit || next_page_token.is_empty() { |
| 123 | + break; |
| 124 | + } |
| 125 | + page_token = next_page_token; |
| 126 | + } |
| 127 | + |
| 128 | + results.truncate(record_limit); |
| 129 | + let out = serde_json::to_value(&results) |
| 130 | + .context("failed to serialize runs") |
| 131 | + .map_err(error::from_anyhow)?; |
| 132 | + |
| 133 | + Ok(CallToolResult::structured(out)) |
| 134 | + } |
| 135 | +} |
| 136 | + |
| 137 | +fn paging(limit: Option<u32>) -> (u32, usize) { |
| 138 | + match limit { |
| 139 | + Some(lim) if lim <= PAGE_SIZE => (lim, lim as usize), |
| 140 | + _ => (PAGE_SIZE, usize::MAX), |
| 141 | + } |
| 142 | +} |
0 commit comments