321 lines
11 KiB
Rust
321 lines
11 KiB
Rust
//! Desktop-owned MCP managed-temp tool provider.
|
|
//!
|
|
//! This module maps public `idea_managed_temp_*` tools to the application
|
|
//! managed-temp use cases. Tool arguments never include filesystem paths.
|
|
|
|
use std::sync::Arc;
|
|
|
|
use application::{
|
|
AllocateManagedTemp, AllocateManagedTempInput, GcManagedTemp, ReleaseManagedTemp,
|
|
ReleaseManagedTempInput, ResolveManagedTemp, ResolveManagedTempInput,
|
|
};
|
|
use async_trait::async_trait;
|
|
use domain::{ManagedTempKind, Project};
|
|
use infrastructure::{ManagedTempToolError, ManagedTempToolProvider};
|
|
use serde_json::{json, Value};
|
|
|
|
use crate::dto::{ManagedTempGcResponseDto, ManagedTempLeaseDto, ManagedTempLeaseResponseDto};
|
|
|
|
/// Concrete provider for the public MCP managed-temp tools.
|
|
#[derive(Clone)]
|
|
pub struct AppManagedTempToolProvider {
|
|
/// Allocation use case.
|
|
pub allocate: Arc<AllocateManagedTemp>,
|
|
/// Resolve use case.
|
|
pub resolve: Arc<ResolveManagedTemp>,
|
|
/// Release use case.
|
|
pub release: Arc<ReleaseManagedTemp>,
|
|
/// Garbage-collection use case.
|
|
pub gc: Arc<GcManagedTemp>,
|
|
}
|
|
|
|
#[async_trait]
|
|
impl ManagedTempToolProvider for AppManagedTempToolProvider {
|
|
async fn handle_managed_temp_tool(
|
|
&self,
|
|
project: &Project,
|
|
_requester: &str,
|
|
name: &str,
|
|
arguments: Value,
|
|
) -> Result<Value, ManagedTempToolError> {
|
|
let result = match name {
|
|
"idea_managed_temp_allocate" => {
|
|
let output = self
|
|
.allocate
|
|
.execute(AllocateManagedTempInput {
|
|
owner_project_id: Some(project.id),
|
|
kind: managed_temp_kind(required_str(&arguments, "kind")?)?,
|
|
purpose: required_str(&arguments, "purpose")?.to_owned(),
|
|
ttl_ms: optional_u64(&arguments, "ttlMs")?,
|
|
})
|
|
.await
|
|
.map_err(managed_temp_error)?;
|
|
json!(ManagedTempLeaseResponseDto {
|
|
lease: ManagedTempLeaseDto::from(output.lease)
|
|
})
|
|
}
|
|
"idea_managed_temp_resolve" => {
|
|
let output = self
|
|
.resolve
|
|
.execute(ResolveManagedTempInput {
|
|
handle: required_str(&arguments, "handle")?.to_owned(),
|
|
})
|
|
.await
|
|
.map_err(managed_temp_error)?;
|
|
json!(ManagedTempLeaseResponseDto {
|
|
lease: ManagedTempLeaseDto::from(output.lease)
|
|
})
|
|
}
|
|
"idea_managed_temp_release" => {
|
|
let output = self
|
|
.release
|
|
.execute(ReleaseManagedTempInput {
|
|
handle: required_str(&arguments, "handle")?.to_owned(),
|
|
})
|
|
.await
|
|
.map_err(managed_temp_error)?;
|
|
json!(ManagedTempLeaseResponseDto {
|
|
lease: ManagedTempLeaseDto::from(output.lease)
|
|
})
|
|
}
|
|
"idea_managed_temp_gc" => {
|
|
let output = self.gc.execute().await.map_err(managed_temp_error)?;
|
|
json!(ManagedTempGcResponseDto {
|
|
deleted: output.report.deleted
|
|
})
|
|
}
|
|
other => {
|
|
return Err(ManagedTempToolError::new(
|
|
"notFound",
|
|
format!("unknown managed temp tool: {other}"),
|
|
))
|
|
}
|
|
};
|
|
Ok(result)
|
|
}
|
|
}
|
|
|
|
fn managed_temp_error(err: application::AppError) -> ManagedTempToolError {
|
|
let code = match err.code() {
|
|
"NOT_FOUND" => "notFound",
|
|
"INVALID" => "invalid",
|
|
"STORE" => "store",
|
|
"FILESYSTEM" => "filesystem",
|
|
_ => "internal",
|
|
};
|
|
ManagedTempToolError::new(code, err.to_string())
|
|
}
|
|
|
|
fn required_str<'a>(arguments: &'a Value, key: &str) -> Result<&'a str, ManagedTempToolError> {
|
|
arguments
|
|
.get(key)
|
|
.and_then(Value::as_str)
|
|
.ok_or_else(|| ManagedTempToolError::new("invalid", format!("missing {key}")))
|
|
}
|
|
|
|
fn optional_u64(arguments: &Value, key: &str) -> Result<Option<u64>, ManagedTempToolError> {
|
|
match arguments.get(key) {
|
|
None | Some(Value::Null) => Ok(None),
|
|
Some(value) => value
|
|
.as_u64()
|
|
.map(Some)
|
|
.ok_or_else(|| ManagedTempToolError::new("invalid", format!("{key} must be a u64"))),
|
|
}
|
|
}
|
|
|
|
fn managed_temp_kind(raw: &str) -> Result<ManagedTempKind, ManagedTempToolError> {
|
|
match raw {
|
|
"directory" => Ok(ManagedTempKind::Directory),
|
|
"file" => Ok(ManagedTempKind::File),
|
|
_ => Err(ManagedTempToolError::new(
|
|
"invalid",
|
|
"kind must be `directory` or `file`",
|
|
)),
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use std::collections::HashMap;
|
|
use std::sync::{Arc, Mutex};
|
|
|
|
use application::{
|
|
AllocateManagedTemp, GcManagedTemp, ReleaseManagedTemp, ResolveManagedTemp,
|
|
MANAGED_TEMP_DEFAULT_TTL_MS,
|
|
};
|
|
use domain::ports::{Clock, IdGenerator, LocalPath, ManagedTempStore};
|
|
use domain::{
|
|
ManagedTempAllocation, ManagedTempError, ManagedTempGcReport, ManagedTempHandle,
|
|
ManagedTempLease, ManagedTempState, ProjectId, ProjectPath, RemoteRef,
|
|
};
|
|
use serde_json::json;
|
|
use uuid::Uuid;
|
|
|
|
use super::*;
|
|
|
|
struct FixedClock;
|
|
|
|
impl Clock for FixedClock {
|
|
fn now_millis(&self) -> i64 {
|
|
50_000
|
|
}
|
|
}
|
|
|
|
struct FixedIds;
|
|
|
|
impl IdGenerator for FixedIds {
|
|
fn new_uuid(&self) -> Uuid {
|
|
Uuid::from_u128(0x189)
|
|
}
|
|
}
|
|
|
|
#[derive(Default)]
|
|
struct FakeStore {
|
|
allocations: Mutex<Vec<ManagedTempAllocation>>,
|
|
leases: Mutex<HashMap<ManagedTempHandle, ManagedTempLease>>,
|
|
}
|
|
|
|
#[async_trait]
|
|
impl ManagedTempStore for FakeStore {
|
|
async fn allocate(
|
|
&self,
|
|
allocation: ManagedTempAllocation,
|
|
) -> Result<ManagedTempLease, ManagedTempError> {
|
|
let lease = ManagedTempLease {
|
|
handle: allocation.handle.clone(),
|
|
kind: allocation.kind,
|
|
owner_project_id: allocation.owner_project_id,
|
|
purpose: allocation.purpose.clone(),
|
|
path: LocalPath::new(format!("/managed/{}", allocation.handle.as_str())),
|
|
state: ManagedTempState::Active,
|
|
created_at_ms: allocation.created_at_ms,
|
|
last_access_at_ms: allocation.created_at_ms,
|
|
expires_at_ms: allocation.expires_at_ms,
|
|
released_at_ms: None,
|
|
};
|
|
self.allocations.lock().unwrap().push(allocation);
|
|
self.leases
|
|
.lock()
|
|
.unwrap()
|
|
.insert(lease.handle.clone(), lease.clone());
|
|
Ok(lease)
|
|
}
|
|
|
|
async fn resolve(
|
|
&self,
|
|
handle: &ManagedTempHandle,
|
|
now_ms: u64,
|
|
) -> Result<ManagedTempLease, ManagedTempError> {
|
|
let mut lease = self
|
|
.leases
|
|
.lock()
|
|
.unwrap()
|
|
.get(handle)
|
|
.cloned()
|
|
.ok_or(ManagedTempError::NotFound)?;
|
|
lease.last_access_at_ms = now_ms;
|
|
Ok(lease)
|
|
}
|
|
|
|
async fn release(
|
|
&self,
|
|
handle: &ManagedTempHandle,
|
|
now_ms: u64,
|
|
) -> Result<ManagedTempLease, ManagedTempError> {
|
|
let mut leases = self.leases.lock().unwrap();
|
|
let lease = leases.get_mut(handle).ok_or(ManagedTempError::NotFound)?;
|
|
lease.state = ManagedTempState::Released;
|
|
lease.released_at_ms = Some(now_ms);
|
|
Ok(lease.clone())
|
|
}
|
|
|
|
async fn gc(&self, _now_ms: u64) -> Result<ManagedTempGcReport, ManagedTempError> {
|
|
Ok(ManagedTempGcReport { deleted: 1 })
|
|
}
|
|
|
|
async fn purge_all(&self) -> Result<ManagedTempGcReport, ManagedTempError> {
|
|
let deleted = self.leases.lock().unwrap().len();
|
|
self.leases.lock().unwrap().clear();
|
|
Ok(ManagedTempGcReport { deleted })
|
|
}
|
|
}
|
|
|
|
fn project() -> Project {
|
|
Project::new(
|
|
ProjectId::from_uuid(Uuid::from_u128(7)),
|
|
"demo",
|
|
ProjectPath::new("/tmp/managed-temp-provider-project").unwrap(),
|
|
RemoteRef::local(),
|
|
1_000,
|
|
)
|
|
.unwrap()
|
|
}
|
|
|
|
fn provider(store: Arc<FakeStore>) -> AppManagedTempToolProvider {
|
|
let store_port = store as Arc<dyn ManagedTempStore>;
|
|
AppManagedTempToolProvider {
|
|
allocate: Arc::new(AllocateManagedTemp::new(
|
|
Arc::clone(&store_port),
|
|
Arc::new(FixedIds),
|
|
Arc::new(FixedClock),
|
|
)),
|
|
resolve: Arc::new(ResolveManagedTemp::new(
|
|
Arc::clone(&store_port),
|
|
Arc::new(FixedClock),
|
|
)),
|
|
release: Arc::new(ReleaseManagedTemp::new(
|
|
Arc::clone(&store_port),
|
|
Arc::new(FixedClock),
|
|
)),
|
|
gc: Arc::new(GcManagedTemp::new(store_port, Arc::new(FixedClock))),
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn managed_temp_allocate_tool_uses_current_project_and_generated_handle() {
|
|
let store = Arc::new(FakeStore::default());
|
|
let project = project();
|
|
let result = provider(Arc::clone(&store))
|
|
.handle_managed_temp_tool(
|
|
&project,
|
|
"agent",
|
|
"idea_managed_temp_allocate",
|
|
json!({ "kind": "directory", "purpose": "scratch" }),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(result["lease"]["ownerProjectId"], project.id.to_string());
|
|
assert_eq!(
|
|
result["lease"]["handle"],
|
|
"00000000-0000-0000-0000-000000000189"
|
|
);
|
|
assert_eq!(
|
|
result["lease"]["path"],
|
|
"/managed/00000000-0000-0000-0000-000000000189"
|
|
);
|
|
let allocations = store.allocations.lock().unwrap();
|
|
assert_eq!(allocations.len(), 1);
|
|
assert_eq!(allocations[0].owner_project_id, Some(project.id));
|
|
assert_eq!(
|
|
allocations[0].expires_at_ms,
|
|
Some(50_000 + MANAGED_TEMP_DEFAULT_TTL_MS)
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn managed_temp_resolve_tool_rejects_path_like_handle() {
|
|
let err = provider(Arc::new(FakeStore::default()))
|
|
.handle_managed_temp_tool(
|
|
&project(),
|
|
"agent",
|
|
"idea_managed_temp_resolve",
|
|
json!({ "handle": "../escape" }),
|
|
)
|
|
.await
|
|
.unwrap_err();
|
|
|
|
assert_eq!(err.code, "invalid");
|
|
}
|
|
}
|