Files
crank/apps/admin-api/src/service/agents.rs
T

745 lines
27 KiB
Rust

use std::collections::BTreeMap;
use crank_core::{
Agent, AgentId, AgentOperationBinding, AgentStatus, AgentVersion, OperationId, SearchableTool,
ToolAccessMode, ToolSelectionPolicy, UsagePeriod, WorkspaceId, search_tool_catalog,
};
use crank_registry::{
AgentStateExpectation, AgentVersionRecord, CreateAgentDraftVersionRequest, CreateAgentRequest,
PublishAgentRequest, SaveAgentCatalogConfigRequest, UpdateAgentSummaryRequest, UsageBucket,
UsageQuery,
};
use serde_json::json;
use sha2::{Digest, Sha256};
use time::{OffsetDateTime, format_description::well_known::Rfc3339};
use tracing::{info, instrument};
use crate::{
error::ApiError,
service::{
AdminService, AgentCatalogPayload, AgentMutationResult, AgentPayload, AgentSummaryView,
CreatedAgentResponse, PublishAgentResponse, ToolSearchPreviewPayload, UpdateAgentPayload,
agent_mcp_endpoint, format_timestamp, map_agent_summary_view, new_prefixed_id, now_string,
today_start_utc,
},
};
impl AdminService {
pub fn agent_state_etag(agent: &AgentSummaryView) -> String {
let policy =
serde_json::to_string(&agent.tool_selection_policy).unwrap_or_else(|_| "{}".to_owned());
let operation_ids = agent.operation_ids.join(",");
let material = format!(
"agent-state-v2\0{}\0{}\0{}\0{}\0{}\0{}\0{}\0{:?}\0{:?}\0{}\0{}\0{}\0{}\0{}",
agent.workspace_id,
agent.id,
agent.slug,
agent.display_name,
agent.description,
agent.updated_at,
agent.current_draft_version,
agent.status,
agent.latest_published_version,
agent.catalog_revision,
agent.operation_count,
operation_ids,
policy,
agent.key_count
);
format!("\"{:x}\"", Sha256::digest(material.as_bytes()))
}
pub fn agent_state_expectation(
agent: &AgentSummaryView,
) -> Result<AgentStateExpectation, ApiError> {
let updated_at = OffsetDateTime::parse(&agent.updated_at, &Rfc3339)
.map_err(|_| ApiError::internal("invalid agent state timestamp"))?;
Ok(AgentStateExpectation {
status: agent.status,
current_draft_version: agent.current_draft_version,
latest_published_version: agent.latest_published_version,
catalog_revision: agent.catalog_revision,
updated_at,
})
}
#[instrument(skip(self, payload), fields(workspace_id = %workspace_id.as_str()))]
pub async fn preview_tool_search(
&self,
workspace_id: &WorkspaceId,
payload: ToolSearchPreviewPayload,
) -> Result<Vec<crank_core::ToolSearchMatch>, ApiError> {
self.ensure_workspace_exists(workspace_id).await?;
if payload.tool_selection_policy.mode != ToolAccessMode::Search {
return Err(ApiError::validation_with_context(
"tool search preview requires search mode",
json!({"field": "tool_selection_policy.mode"}),
));
}
let bindings = payload
.bindings
.iter()
.map(|binding| AgentOperationBinding {
operation_id: OperationId::new(binding.operation_id.clone()),
operation_version: binding.operation_version,
tool_name: binding.tool_name.clone(),
tool_title: binding.tool_title.clone(),
tool_description_override: binding.tool_description_override.clone(),
enabled: binding.enabled,
})
.collect::<Vec<_>>();
validate_tool_selection_policy(&payload.tool_selection_policy, &bindings)?;
let mut tools = Vec::new();
for binding in bindings.iter().filter(|binding| binding.enabled) {
let version = self
.registry
.get_operation_version(
workspace_id,
&binding.operation_id,
binding.operation_version,
)
.await?
.ok_or_else(|| {
ApiError::not_found_with_context(
format!("operation {} was not found", binding.operation_id.as_str()),
json!({"operation_id": binding.operation_id.as_str()}),
)
})?;
let groups = payload
.tool_selection_policy
.groups
.iter()
.filter(|group| {
group
.tool_names
.iter()
.any(|name| name == &binding.tool_name)
})
.collect::<Vec<_>>();
tools.push(SearchableTool {
name: binding.tool_name.clone(),
title: binding.tool_title.clone(),
description: binding
.tool_description_override
.clone()
.unwrap_or(version.snapshot.tool_description.description),
input_schema: serde_json::Value::Null,
group_ids: groups.iter().map(|group| group.id.clone()).collect(),
group_context: groups
.iter()
.map(|group| format!("{} {}", group.name, group.description))
.collect::<Vec<_>>()
.join(" "),
});
}
let max_results = payload.tool_selection_policy.search.max_results;
Ok(search_tool_catalog(
&tools,
&payload.query,
&payload.group_ids,
max_results,
))
}
#[instrument(skip(self))]
pub async fn list_agents(
&self,
workspace_id: &WorkspaceId,
) -> Result<Vec<AgentSummaryView>, ApiError> {
self.ensure_workspace_exists(workspace_id).await?;
let workspace = self.get_workspace(workspace_id).await?;
let summaries = self.registry.list_agents(workspace_id).await?;
let usage_start = today_start_utc()?;
let usage_end = now_string()?;
let usage = self
.registry
.list_usage_by_agent(UsageQuery {
workspace_id,
period: UsagePeriod::Last24Hours,
source: None,
created_after: &usage_start,
created_before: &usage_end,
bucket: UsageBucket::Hour,
})
.await?;
let calls_today = usage
.into_iter()
.map(|item| (item.agent_id.as_str().to_owned(), item.calls_total))
.collect::<BTreeMap<_, _>>();
let key_counts = self
.registry
.list_platform_api_keys(workspace_id)
.await?
.into_iter()
.fold(BTreeMap::new(), |mut counts, record| {
if let Some(agent_id) = record.api_key.agent_id {
*counts.entry(agent_id.as_str().to_owned()).or_insert(0usize) += 1;
}
counts
});
let mut items = Vec::with_capacity(summaries.len());
for summary in summaries {
let version = self
.get_agent_version(workspace_id, &summary.id, summary.current_draft_version)
.await?;
let operation_ids = version
.bindings
.iter()
.map(|binding| binding.operation_id.as_str().to_owned())
.collect::<Vec<_>>();
items.push(AgentSummaryView {
operation_count: operation_ids.len(),
operation_ids,
tool_selection_policy: version.snapshot.tool_selection_policy,
key_count: key_counts.get(summary.id.as_str()).copied().unwrap_or(0),
calls_today: calls_today.get(summary.id.as_str()).copied().unwrap_or(0),
mcp_endpoint: agent_mcp_endpoint(
workspace.workspace.slug.as_str(),
summary.slug.as_str(),
),
..map_agent_summary_view(summary)
});
}
Ok(items)
}
#[instrument(skip(self))]
pub async fn get_agent(
&self,
workspace_id: &WorkspaceId,
agent_id: &AgentId,
) -> Result<AgentSummaryView, ApiError> {
self.ensure_workspace_exists(workspace_id).await?;
let workspace = self.get_workspace(workspace_id).await?;
let summary = self
.registry
.get_agent_summary(workspace_id, agent_id)
.await?
.ok_or_else(|| {
ApiError::not_found_with_context(
format!("agent {} was not found", agent_id.as_str()),
json!({ "agent_id": agent_id.as_str() }),
)
})?;
let version = self
.get_agent_version(workspace_id, agent_id, summary.current_draft_version)
.await?;
let operation_ids = version
.bindings
.iter()
.map(|binding| binding.operation_id.as_str().to_owned())
.collect::<Vec<_>>();
let usage_start = today_start_utc()?;
let usage_end = now_string()?;
let usage = self
.registry
.get_usage_for_agent(
UsageQuery {
workspace_id,
period: UsagePeriod::Last24Hours,
source: None,
created_after: &usage_start,
created_before: &usage_end,
bucket: UsageBucket::Hour,
},
agent_id,
)
.await?;
let key_count = self
.registry
.list_platform_api_keys_for_agent(workspace_id, agent_id)
.await?
.len();
Ok(AgentSummaryView {
operation_count: operation_ids.len(),
operation_ids,
tool_selection_policy: version.snapshot.tool_selection_policy,
key_count,
calls_today: usage.map(|item| item.rollup.calls_total).unwrap_or(0),
mcp_endpoint: agent_mcp_endpoint(
workspace.workspace.slug.as_str(),
summary.slug.as_str(),
),
..map_agent_summary_view(summary)
})
}
#[instrument(skip(self))]
pub async fn get_agent_version(
&self,
workspace_id: &WorkspaceId,
agent_id: &AgentId,
version: u32,
) -> Result<AgentVersionRecord, ApiError> {
self.registry
.get_agent_version(workspace_id, agent_id, version)
.await?
.ok_or_else(|| {
ApiError::not_found_with_context(
format!(
"agent version {version} for {} was not found",
agent_id.as_str()
),
json!({
"agent_id": agent_id.as_str(),
"version": version,
}),
)
})
}
#[instrument(skip(self, payload), fields(workspace_id = %workspace_id.as_str(), agent_slug = %payload.slug))]
pub async fn create_agent(
&self,
workspace_id: &WorkspaceId,
payload: AgentPayload,
) -> Result<CreatedAgentResponse, ApiError> {
self.ensure_workspace_exists(workspace_id).await?;
if self
.find_agent_by_slug(workspace_id, &payload.slug)
.await?
.is_some()
{
return Err(ApiError::conflict_with_context(
format!("agent with slug {} already exists", payload.slug),
json!({ "slug": payload.slug }),
));
}
let now = OffsetDateTime::now_utc();
let agent_id = AgentId::new(new_prefixed_id("agent"));
let agent = Agent {
id: agent_id.clone(),
workspace_id: workspace_id.clone(),
slug: payload.slug,
display_name: payload.display_name,
description: payload.description,
status: AgentStatus::Draft,
current_draft_version: 1,
latest_published_version: None,
created_at: now,
updated_at: now,
published_at: None,
};
let version = AgentVersion {
agent_id: agent_id.clone(),
version: 1,
status: AgentStatus::Draft,
instructions: payload.instructions,
tool_selection_policy: payload.tool_selection_policy,
created_at: now,
};
self.registry
.create_agent(CreateAgentRequest {
agent: &agent,
version: &version,
bindings: &[],
})
.await?;
info!(
name: "admin.agent.created",
agent_id = %agent_id.as_str(),
version = 1,
"agent created"
);
Ok(CreatedAgentResponse {
agent_id: agent_id.as_str().to_owned(),
workspace_id: workspace_id.as_str().to_owned(),
version: 1,
status: AgentStatus::Draft,
})
}
#[instrument(skip(self, payload), fields(workspace_id = %workspace_id.as_str(), agent_id = %agent_id.as_str()))]
pub async fn update_agent(
&self,
workspace_id: &WorkspaceId,
agent_id: &AgentId,
payload: UpdateAgentPayload,
expected_state: Option<&AgentStateExpectation>,
) -> Result<AgentMutationResult, ApiError> {
let existing = self
.registry
.get_agent_summary(workspace_id, agent_id)
.await?
.ok_or_else(|| {
ApiError::not_found_with_context(
format!("agent {} was not found", agent_id.as_str()),
json!({ "agent_id": agent_id.as_str() }),
)
})?;
if existing.status == AgentStatus::Published
&& (payload.slug != existing.slug
|| payload.display_name != existing.display_name
|| payload.description != existing.description)
{
return Err(ApiError::conflict_with_context(
"published Agent summary is immutable; unpublish before editing",
json!({
"agent_id": agent_id.as_str(),
"error_code": "agent_published_summary_immutable",
"recovery": "unpublish_edit_publish"
}),
));
}
if payload.slug != existing.slug
&& self
.find_agent_by_slug(workspace_id, &payload.slug)
.await?
.is_some()
{
return Err(ApiError::conflict_with_context(
format!("agent with slug {} already exists", payload.slug),
json!({ "slug": payload.slug }),
));
}
let updated_at = OffsetDateTime::now_utc();
self.registry
.update_agent_summary(
workspace_id,
agent_id,
UpdateAgentSummaryRequest {
slug: &payload.slug,
display_name: &payload.display_name,
description: &payload.description,
updated_at: &updated_at,
expected_state,
},
)
.await?;
Ok(AgentMutationResult {
agent_id: agent_id.as_str().to_owned(),
workspace_id: workspace_id.as_str().to_owned(),
updated_at: format_timestamp(updated_at),
})
}
#[instrument(skip(self), fields(workspace_id = %workspace_id.as_str(), agent_id = %agent_id.as_str()))]
pub async fn delete_agent(
&self,
workspace_id: &WorkspaceId,
agent_id: &AgentId,
expected_state: Option<&AgentStateExpectation>,
) -> Result<AgentMutationResult, ApiError> {
let existing = self
.registry
.get_agent_summary(workspace_id, agent_id)
.await?
.ok_or_else(|| {
ApiError::not_found_with_context(
format!("agent {} was not found", agent_id.as_str()),
json!({ "agent_id": agent_id.as_str() }),
)
})?;
if existing.latest_published_version.is_some() {
return Err(ApiError::conflict_with_context(
"published Agent cannot be deleted; archive or unpublish it instead",
json!({
"agent_id": agent_id.as_str(),
"error_code": "agent_delete_forbidden",
"recovery": "archive_or_unpublish"
}),
));
}
self.registry
.delete_agent(workspace_id, agent_id, expected_state)
.await?;
Ok(AgentMutationResult {
agent_id: agent_id.as_str().to_owned(),
workspace_id: workspace_id.as_str().to_owned(),
updated_at: format_timestamp(existing.updated_at),
})
}
#[instrument(skip(self, payload), fields(workspace_id = %workspace_id.as_str(), agent_id = %agent_id.as_str()))]
pub async fn save_agent_bindings(
&self,
workspace_id: &WorkspaceId,
agent_id: &AgentId,
payload: AgentCatalogPayload,
expected_state: Option<&AgentStateExpectation>,
) -> Result<AgentVersionRecord, ApiError> {
let current_version = self
.ensure_editable_agent_version(workspace_id, agent_id)
.await?;
let (payload, requested_policy) = payload.into_parts();
let bindings = payload
.into_iter()
.map(|binding| AgentOperationBinding {
operation_id: OperationId::new(binding.operation_id),
operation_version: binding.operation_version,
tool_name: binding.tool_name,
tool_title: binding.tool_title,
tool_description_override: binding.tool_description_override,
enabled: binding.enabled,
})
.collect::<Vec<_>>();
self.validate_exact_published_agent_bindings(workspace_id, &bindings)
.await?;
let tool_selection_policy = requested_policy
.unwrap_or_else(|| current_version.snapshot.tool_selection_policy.clone());
validate_tool_selection_policy(&tool_selection_policy, &bindings)?;
self.registry
.save_agent_catalog_config(SaveAgentCatalogConfigRequest {
workspace_id,
agent_id,
agent_version: current_version.version,
bindings: &bindings,
tool_selection_policy: &tool_selection_policy,
expected_state,
})
.await?;
info!(
name: "admin.agent.bindings_saved",
agent_id = %agent_id.as_str(),
version = current_version.version,
binding_count = bindings.len(),
"agent bindings saved"
);
self.get_agent_version(workspace_id, agent_id, current_version.version)
.await
}
async fn ensure_editable_agent_version(
&self,
workspace_id: &WorkspaceId,
agent_id: &AgentId,
) -> Result<AgentVersionRecord, ApiError> {
let agent = self.get_agent(workspace_id, agent_id).await?;
let current = self
.get_agent_version(workspace_id, agent_id, agent.current_draft_version)
.await?;
if agent.latest_published_version != Some(agent.current_draft_version) {
return Ok(current);
}
let now = OffsetDateTime::now_utc();
let draft = AgentVersion {
agent_id: agent_id.clone(),
version: current.version + 1,
status: AgentStatus::Draft,
instructions: current.snapshot.instructions.clone(),
tool_selection_policy: current.snapshot.tool_selection_policy.clone(),
created_at: now,
};
self.registry
.create_agent_draft_version(CreateAgentDraftVersionRequest {
workspace_id,
agent_id,
version: &draft,
bindings: &current.bindings,
updated_at: &now,
})
.await?;
self.get_agent_version(workspace_id, agent_id, draft.version)
.await
}
#[instrument(skip(self), fields(workspace_id = %workspace_id.as_str(), agent_id = %agent_id.as_str(), version))]
pub async fn publish_agent(
&self,
workspace_id: &WorkspaceId,
agent_id: &AgentId,
version: u32,
expected_state: Option<&AgentStateExpectation>,
) -> Result<PublishAgentResponse, ApiError> {
let agent_version = self
.get_agent_version(workspace_id, agent_id, version)
.await?;
validate_tool_selection_policy(
&agent_version.snapshot.tool_selection_policy,
&agent_version.bindings,
)?;
if agent_version
.bindings
.iter()
.filter(|binding| binding.enabled)
.count()
== 0
{
return Err(ApiError::conflict_with_context(
"agent cannot be published without published enabled tools",
json!({
"agent_id": agent_id.as_str(),
"version": version,
"binding_count": agent_version.bindings.len()
}),
));
}
self.validate_exact_published_agent_bindings(workspace_id, &agent_version.bindings)
.await?;
let published_at = OffsetDateTime::now_utc();
self.registry
.publish_agent(PublishAgentRequest {
workspace_id,
agent_id,
version,
published_at: &published_at,
published_by: None,
expected_state,
})
.await?;
info!(
name: "admin.agent.published",
agent_id = %agent_id.as_str(),
version,
"agent published"
);
Ok(PublishAgentResponse {
agent_id: agent_id.as_str().to_owned(),
workspace_id: workspace_id.as_str().to_owned(),
published_version: version,
published_at: format_timestamp(published_at),
})
}
async fn validate_exact_published_agent_bindings(
&self,
workspace_id: &WorkspaceId,
bindings: &[AgentOperationBinding],
) -> Result<(), ApiError> {
for binding in bindings {
let Some(summary) = self
.registry
.get_operation_summary(workspace_id, &binding.operation_id)
.await?
else {
return Err(ApiError::not_found_with_context(
"operation was not found for Agent binding",
json!({
"error_code": "agent_binding_scope_denied",
"operation_id": binding.operation_id.as_str()
}),
));
};
if summary.status == crank_core::OperationStatus::Archived {
return Err(ApiError::unprocessable_with_context(
"archived operation cannot be added to an Agent binding",
json!({
"error_code": "agent_binding_archived_operation",
"operation_id": binding.operation_id.as_str()
}),
));
}
let Some(version) = self
.registry
.get_operation_version(
workspace_id,
&binding.operation_id,
binding.operation_version,
)
.await?
else {
return Err(ApiError::unprocessable_with_context(
"operation version is not published for Agent binding",
json!({
"error_code": "agent_binding_not_published",
"operation_id": binding.operation_id.as_str(),
"operation_version": binding.operation_version
}),
));
};
if !version.snapshot.is_published() {
return Err(ApiError::unprocessable_with_context(
"operation version is not published for Agent binding",
json!({
"error_code": "agent_binding_not_published",
"operation_id": binding.operation_id.as_str(),
"operation_version": binding.operation_version
}),
));
}
}
Ok(())
}
#[instrument(skip(self), fields(workspace_id = %workspace_id.as_str(), agent_id = %agent_id.as_str()))]
pub async fn unpublish_agent(
&self,
workspace_id: &WorkspaceId,
agent_id: &AgentId,
expected_state: Option<&AgentStateExpectation>,
) -> Result<AgentMutationResult, ApiError> {
self.ensure_workspace_exists(workspace_id).await?;
let updated_at = OffsetDateTime::now_utc();
self.registry
.unpublish_agent(workspace_id, agent_id, &updated_at, expected_state)
.await?;
info!(
name: "admin.agent.unpublished",
agent_id = %agent_id.as_str(),
"agent moved to draft"
);
Ok(AgentMutationResult {
agent_id: agent_id.as_str().to_owned(),
workspace_id: workspace_id.as_str().to_owned(),
updated_at: format_timestamp(updated_at),
})
}
#[instrument(skip(self), fields(workspace_id = %workspace_id.as_str(), agent_id = %agent_id.as_str()))]
pub async fn archive_agent(
&self,
workspace_id: &WorkspaceId,
agent_id: &AgentId,
expected_state: Option<&AgentStateExpectation>,
) -> Result<AgentMutationResult, ApiError> {
self.ensure_workspace_exists(workspace_id).await?;
let updated_at = OffsetDateTime::now_utc();
self.registry
.archive_agent(workspace_id, agent_id, &updated_at, expected_state)
.await?;
info!(
name: "admin.agent.archived",
agent_id = %agent_id.as_str(),
"agent archived"
);
Ok(AgentMutationResult {
agent_id: agent_id.as_str().to_owned(),
workspace_id: workspace_id.as_str().to_owned(),
updated_at: format_timestamp(updated_at),
})
}
}
fn validate_tool_selection_policy(
policy: &ToolSelectionPolicy,
bindings: &[AgentOperationBinding],
) -> Result<(), ApiError> {
policy
.validate_for_tools(
bindings
.iter()
.filter(|binding| binding.enabled)
.map(|binding| binding.tool_name.as_str()),
)
.map_err(|error| {
ApiError::validation_with_context(
error.to_string(),
json!({"field": "tool_selection_policy"}),
)
})
}