Усилить безопасность и надёжность выполнения операций
This commit is contained in:
@@ -18,9 +18,8 @@ use crank_core::{
|
||||
OperationApprovalMode, PlatformApiKeyScope, SecretId,
|
||||
};
|
||||
use crank_registry::{
|
||||
ApprovalRequestRecord, CreateApprovalRequest, CreateInvocationLogRequest,
|
||||
DecideApprovalRequest, ExpireApprovalRequest, FinishApprovalRequest, PostgresRegistry,
|
||||
PublishedAgentTool,
|
||||
CreateApprovalRequest, CreateInvocationLogRequest, DecideApprovalRequest,
|
||||
ExpireApprovalRequest, PostgresRegistry, PublishedAgentTool,
|
||||
};
|
||||
use crank_runtime::{
|
||||
RequestRateLimiter, ResolvedAuth, RuntimeError, RuntimeExecutionRequest, RuntimeExecutor,
|
||||
@@ -30,13 +29,14 @@ use futures_util::stream;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Value, json};
|
||||
use time::OffsetDateTime;
|
||||
use tracing::info;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::{
|
||||
access::{
|
||||
credential_allows_security_level, require_approval_access, require_machine_access,
|
||||
serialize_machine_access_mode, serialize_security_level,
|
||||
},
|
||||
approval_execution::{execute_approved_request, spawn_approval_recovery},
|
||||
auth::{SharedMachineCredentialVerifier, VerifiedMachineCredential},
|
||||
catalog::PublishedToolCatalog,
|
||||
jsonrpc::{
|
||||
@@ -66,8 +66,8 @@ const TRANSPORT_SESSION_TTL_MS: u64 = 86_400_000;
|
||||
#[derive(Clone)]
|
||||
pub(super) struct AppState {
|
||||
pub(super) registry: PostgresRegistry,
|
||||
catalog: PublishedToolCatalog,
|
||||
runtime: RuntimeExecutor,
|
||||
pub(super) catalog: PublishedToolCatalog,
|
||||
pub(super) runtime: RuntimeExecutor,
|
||||
pub(super) api_rate_limiter: RequestRateLimiter,
|
||||
secret_crypto: SecretCrypto,
|
||||
sessions: SharedSessionStore,
|
||||
@@ -132,6 +132,59 @@ pub fn build_app(
|
||||
coordination_store: Arc<dyn CoordinationStateStore>,
|
||||
sessions: SharedSessionStore,
|
||||
credential_verifier: SharedMachineCredentialVerifier,
|
||||
) -> Router {
|
||||
build_app_inner(
|
||||
registry,
|
||||
refresh_interval,
|
||||
public_base_url,
|
||||
secret_crypto,
|
||||
runtime,
|
||||
api_rate_limiter,
|
||||
coordination_store,
|
||||
sessions,
|
||||
credential_verifier,
|
||||
false,
|
||||
)
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn build_app_with_background_workers(
|
||||
registry: PostgresRegistry,
|
||||
refresh_interval: Duration,
|
||||
public_base_url: Option<String>,
|
||||
secret_crypto: SecretCrypto,
|
||||
runtime: RuntimeExecutor,
|
||||
api_rate_limiter: RequestRateLimiter,
|
||||
coordination_store: Arc<dyn CoordinationStateStore>,
|
||||
sessions: SharedSessionStore,
|
||||
credential_verifier: SharedMachineCredentialVerifier,
|
||||
) -> Router {
|
||||
build_app_inner(
|
||||
registry,
|
||||
refresh_interval,
|
||||
public_base_url,
|
||||
secret_crypto,
|
||||
runtime,
|
||||
api_rate_limiter,
|
||||
coordination_store,
|
||||
sessions,
|
||||
credential_verifier,
|
||||
true,
|
||||
)
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn build_app_inner(
|
||||
registry: PostgresRegistry,
|
||||
refresh_interval: Duration,
|
||||
public_base_url: Option<String>,
|
||||
secret_crypto: SecretCrypto,
|
||||
runtime: RuntimeExecutor,
|
||||
api_rate_limiter: RequestRateLimiter,
|
||||
coordination_store: Arc<dyn CoordinationStateStore>,
|
||||
sessions: SharedSessionStore,
|
||||
credential_verifier: SharedMachineCredentialVerifier,
|
||||
start_background_workers: bool,
|
||||
) -> Router {
|
||||
let state = Arc::new(AppState {
|
||||
registry: registry.clone(),
|
||||
@@ -143,6 +196,9 @@ pub fn build_app(
|
||||
credential_verifier,
|
||||
allowed_origins: AllowedOrigins::new(public_base_url),
|
||||
});
|
||||
if start_background_workers {
|
||||
spawn_approval_recovery(Arc::clone(&state));
|
||||
}
|
||||
|
||||
Router::new()
|
||||
.route("/health", get(health))
|
||||
@@ -315,7 +371,21 @@ async fn decide_approval_request(
|
||||
.await
|
||||
{
|
||||
Ok(Some(record)) if status == ApprovalRequestStatus::Approved => {
|
||||
match execute_approved_request(&state, &agent_path, record).await {
|
||||
let claimed = match state
|
||||
.registry
|
||||
.claim_approval_request(
|
||||
&record.approval.workspace_id,
|
||||
&record.approval.agent_id,
|
||||
&record.approval.id,
|
||||
OffsetDateTime::now_utc(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(claimed)) => claimed,
|
||||
Ok(None) => return StatusCode::CONFLICT.into_response(),
|
||||
Err(_) => return StatusCode::INTERNAL_SERVER_ERROR.into_response(),
|
||||
};
|
||||
match execute_approved_request(&state, &agent_path, claimed).await {
|
||||
Ok(record) => Json(json!(record)).into_response(),
|
||||
Err(response) => response,
|
||||
}
|
||||
@@ -408,103 +478,6 @@ async fn expire_approval_response(
|
||||
}
|
||||
}
|
||||
|
||||
async fn execute_approved_request(
|
||||
state: &Arc<AppState>,
|
||||
path: &AgentRoutePath,
|
||||
approval: ApprovalRequestRecord,
|
||||
) -> Result<ApprovalRequestRecord, Response> {
|
||||
let tools = state
|
||||
.catalog
|
||||
.list_tools(&path.workspace_slug, &path.agent_slug)
|
||||
.await
|
||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response())?;
|
||||
let Some(tool) = tools.into_iter().find(|tool| {
|
||||
tool.operation.id == approval.approval.operation_id
|
||||
&& tool.operation.version == approval.approval.operation_version
|
||||
}) else {
|
||||
return Err(StatusCode::NOT_FOUND.into_response());
|
||||
};
|
||||
|
||||
let operation = runtime_operation(&tool);
|
||||
let request_preview = build_request_preview(
|
||||
&state.runtime,
|
||||
&operation,
|
||||
&approval.approval.request_payload,
|
||||
);
|
||||
let started_at = Instant::now();
|
||||
let resolved_auth =
|
||||
resolve_operation_auth(state, &tool.workspace_id, &operation.execution_config).await;
|
||||
let result = match resolved_auth {
|
||||
Ok(resolved_auth) => {
|
||||
state
|
||||
.runtime
|
||||
.execute_request(
|
||||
RuntimeExecutionRequest::new(&operation, &approval.approval.request_payload)
|
||||
.with_optional_auth(resolved_auth.as_ref()),
|
||||
)
|
||||
.await
|
||||
}
|
||||
Err(error) => Err(error),
|
||||
};
|
||||
|
||||
let (status, response_payload, invocation_status, invocation_level, message, error_kind) =
|
||||
match result {
|
||||
Ok(output) => (
|
||||
ApprovalRequestStatus::Completed,
|
||||
output,
|
||||
InvocationStatus::Ok,
|
||||
InvocationLevel::Info,
|
||||
"approved tool call completed",
|
||||
None,
|
||||
),
|
||||
Err(error) => (
|
||||
ApprovalRequestStatus::Failed,
|
||||
json!({
|
||||
"error": {
|
||||
"code": runtime_error_code(&error),
|
||||
"message": error.to_string(),
|
||||
}
|
||||
}),
|
||||
InvocationStatus::Error,
|
||||
InvocationLevel::Error,
|
||||
"approved tool call failed",
|
||||
Some(runtime_error_code(&error)),
|
||||
),
|
||||
};
|
||||
|
||||
let _ = persist_invocation(
|
||||
state,
|
||||
&tool,
|
||||
InvocationRecord {
|
||||
request_id: Some(approval.approval.id.as_str()),
|
||||
tool_name: &tool.tool_name,
|
||||
status: invocation_status,
|
||||
level: invocation_level,
|
||||
message,
|
||||
status_code: None,
|
||||
error_kind,
|
||||
duration: started_at.elapsed(),
|
||||
request_preview,
|
||||
response_preview: response_payload.clone(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
state
|
||||
.registry
|
||||
.finish_approval_request(FinishApprovalRequest {
|
||||
workspace_id: &approval.approval.workspace_id,
|
||||
agent_id: &approval.approval.agent_id,
|
||||
approval_id: &approval.approval.id,
|
||||
status,
|
||||
response_payload: Some(response_payload),
|
||||
decision_note: None,
|
||||
})
|
||||
.await
|
||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response())?
|
||||
.ok_or_else(|| StatusCode::CONFLICT.into_response())
|
||||
}
|
||||
|
||||
async fn mcp_get(
|
||||
Path(path): Path<AgentRoutePath>,
|
||||
State(state): State<Arc<AppState>>,
|
||||
@@ -857,7 +830,7 @@ async fn handle_tool_call(
|
||||
.await
|
||||
}
|
||||
|
||||
async fn resolve_operation_auth(
|
||||
pub(super) async fn resolve_operation_auth(
|
||||
state: &Arc<AppState>,
|
||||
workspace_id: &crank_core::WorkspaceId,
|
||||
execution_config: &crank_core::ExecutionConfig,
|
||||
@@ -1004,7 +977,7 @@ async fn handle_base_tool_call(
|
||||
|
||||
match result {
|
||||
Ok(output) => {
|
||||
let _ = persist_invocation(
|
||||
if let Err(error) = persist_invocation(
|
||||
&state,
|
||||
&tool,
|
||||
InvocationRecord {
|
||||
@@ -1020,12 +993,15 @@ async fn handle_base_tool_call(
|
||||
response_preview: output.clone(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
.await
|
||||
{
|
||||
warn!(error = %error, "successful invocation log write failed");
|
||||
}
|
||||
|
||||
success_tool_response(message, response_mode, &session.protocol_version, output)
|
||||
}
|
||||
Err(error) => {
|
||||
let _ = persist_invocation(
|
||||
if let Err(log_error) = persist_invocation(
|
||||
&state,
|
||||
&tool,
|
||||
InvocationRecord {
|
||||
@@ -1041,7 +1017,10 @@ async fn handle_base_tool_call(
|
||||
response_preview: Value::Null,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
.await
|
||||
{
|
||||
warn!(error = %log_error, "failed invocation log write failed");
|
||||
}
|
||||
|
||||
tool_error_response(
|
||||
message,
|
||||
@@ -1146,17 +1125,22 @@ async fn maybe_create_custom_pending_approval(
|
||||
decision_note: None,
|
||||
};
|
||||
|
||||
if let Err(error) = state
|
||||
let persisted_approval = match state
|
||||
.registry
|
||||
.create_approval_request(CreateApprovalRequest {
|
||||
approval: &approval,
|
||||
})
|
||||
.await
|
||||
{
|
||||
return Some(internal_jsonrpc_error(message, error));
|
||||
}
|
||||
Ok(approval) => approval,
|
||||
Err(error) => return Some(internal_jsonrpc_error(message, error)),
|
||||
};
|
||||
let response_payload = persisted_approval
|
||||
.approval
|
||||
.response_payload
|
||||
.unwrap_or(response_payload);
|
||||
|
||||
let _ = persist_invocation(
|
||||
if let Err(error) = persist_invocation(
|
||||
state,
|
||||
tool,
|
||||
InvocationRecord {
|
||||
@@ -1172,7 +1156,10 @@ async fn maybe_create_custom_pending_approval(
|
||||
response_preview: response_payload.clone(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
.await
|
||||
{
|
||||
warn!(error = %error, "pending approval invocation log write failed");
|
||||
}
|
||||
|
||||
Some(success_tool_response(
|
||||
message,
|
||||
@@ -1404,7 +1391,7 @@ fn take_confirmation_token(arguments: &mut Value) -> Option<String> {
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
}
|
||||
|
||||
fn build_request_preview(
|
||||
pub(super) fn build_request_preview(
|
||||
runtime: &RuntimeExecutor,
|
||||
operation: &RuntimeOperation,
|
||||
arguments: &Value,
|
||||
@@ -1420,20 +1407,20 @@ fn build_request_preview(
|
||||
}
|
||||
}
|
||||
|
||||
struct InvocationRecord<'a> {
|
||||
request_id: Option<&'a str>,
|
||||
tool_name: &'a str,
|
||||
status: InvocationStatus,
|
||||
level: InvocationLevel,
|
||||
message: &'a str,
|
||||
status_code: Option<u16>,
|
||||
error_kind: Option<&'a str>,
|
||||
duration: Duration,
|
||||
request_preview: Value,
|
||||
response_preview: Value,
|
||||
pub(super) struct InvocationRecord<'a> {
|
||||
pub(super) request_id: Option<&'a str>,
|
||||
pub(super) tool_name: &'a str,
|
||||
pub(super) status: InvocationStatus,
|
||||
pub(super) level: InvocationLevel,
|
||||
pub(super) message: &'a str,
|
||||
pub(super) status_code: Option<u16>,
|
||||
pub(super) error_kind: Option<&'a str>,
|
||||
pub(super) duration: Duration,
|
||||
pub(super) request_preview: Value,
|
||||
pub(super) response_preview: Value,
|
||||
}
|
||||
|
||||
async fn persist_invocation(
|
||||
pub(super) async fn persist_invocation(
|
||||
state: &Arc<AppState>,
|
||||
tool: &PublishedAgentTool,
|
||||
record: InvocationRecord<'_>,
|
||||
@@ -1543,7 +1530,7 @@ fn resolve_generated_tool(
|
||||
None
|
||||
}
|
||||
|
||||
fn runtime_operation(tool: &PublishedAgentTool) -> RuntimeOperation {
|
||||
pub(super) fn runtime_operation(tool: &PublishedAgentTool) -> RuntimeOperation {
|
||||
let mut operation = RuntimeOperation::from(tool.operation.clone());
|
||||
operation.tool_name = tool.tool_name.clone();
|
||||
operation.tool_description.title = tool.tool_title.clone();
|
||||
|
||||
@@ -0,0 +1,237 @@
|
||||
use std::{sync::Arc, time::Instant};
|
||||
|
||||
use axum::{
|
||||
http::StatusCode,
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use crank_core::{ApprovalRequestStatus, InvocationLevel, InvocationSource, InvocationStatus};
|
||||
use crank_registry::{ApprovalRequestRecord, FinishApprovalRequest};
|
||||
use crank_runtime::{RuntimeExecutionRequest, RuntimeRequestContext};
|
||||
use serde_json::json;
|
||||
use time::OffsetDateTime;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::{
|
||||
app::{
|
||||
AgentRoutePath, AppState, InvocationRecord, build_request_preview, persist_invocation,
|
||||
resolve_operation_auth, runtime_operation,
|
||||
},
|
||||
tool_error::runtime_error_code,
|
||||
};
|
||||
|
||||
const RECOVERY_INTERVAL: std::time::Duration = std::time::Duration::from_secs(5);
|
||||
const RECOVERY_GRACE: time::Duration = time::Duration::seconds(5);
|
||||
const EXECUTION_LEASE: time::Duration = time::Duration::minutes(6);
|
||||
|
||||
pub(super) fn spawn_approval_recovery(state: Arc<AppState>) {
|
||||
tokio::spawn(async move {
|
||||
let mut interval = tokio::time::interval(RECOVERY_INTERVAL);
|
||||
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
||||
loop {
|
||||
interval.tick().await;
|
||||
recover_approved_requests(&state).await;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
async fn recover_approved_requests(state: &Arc<AppState>) {
|
||||
for _ in 0..32 {
|
||||
let now = OffsetDateTime::now_utc();
|
||||
let approval = match state
|
||||
.registry
|
||||
.claim_next_recoverable_approval_request(
|
||||
now,
|
||||
now - RECOVERY_GRACE,
|
||||
now - EXECUTION_LEASE,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(approval)) => approval,
|
||||
Ok(None) => break,
|
||||
Err(error) => {
|
||||
warn!(error = %error, "approval recovery query failed");
|
||||
break;
|
||||
}
|
||||
};
|
||||
let Some(path) = approval_agent_path(state, &approval).await else {
|
||||
continue;
|
||||
};
|
||||
if execute_approved_request(state, &path, approval)
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
warn!("recovered approval execution did not finish");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn approval_agent_path(
|
||||
state: &Arc<AppState>,
|
||||
approval: &ApprovalRequestRecord,
|
||||
) -> Option<AgentRoutePath> {
|
||||
let workspace = match state
|
||||
.registry
|
||||
.get_workspace(&approval.approval.workspace_id)
|
||||
.await
|
||||
{
|
||||
Ok(Some(workspace)) => workspace,
|
||||
Ok(None) => return None,
|
||||
Err(error) => {
|
||||
warn!(error = %error, "approval workspace lookup failed");
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let agent = match state
|
||||
.registry
|
||||
.get_agent_summary(&approval.approval.workspace_id, &approval.approval.agent_id)
|
||||
.await
|
||||
{
|
||||
Ok(Some(agent)) => agent,
|
||||
Ok(None) => return None,
|
||||
Err(error) => {
|
||||
warn!(error = %error, "approval agent lookup failed");
|
||||
return None;
|
||||
}
|
||||
};
|
||||
Some(AgentRoutePath {
|
||||
workspace_slug: workspace.workspace.slug,
|
||||
agent_slug: agent.slug,
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) async fn execute_approved_request(
|
||||
state: &Arc<AppState>,
|
||||
path: &AgentRoutePath,
|
||||
approval: ApprovalRequestRecord,
|
||||
) -> Result<ApprovalRequestRecord, Response> {
|
||||
let tools = state
|
||||
.catalog
|
||||
.list_tools(&path.workspace_slug, &path.agent_slug)
|
||||
.await
|
||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response())?;
|
||||
let Some(tool) = tools.into_iter().find(|tool| {
|
||||
tool.operation.id == approval.approval.operation_id
|
||||
&& tool.operation.version == approval.approval.operation_version
|
||||
}) else {
|
||||
return finish_unavailable_approval(state, &approval).await;
|
||||
};
|
||||
|
||||
let operation = runtime_operation(&tool);
|
||||
let request_preview = build_request_preview(
|
||||
&state.runtime,
|
||||
&operation,
|
||||
&approval.approval.request_payload,
|
||||
);
|
||||
let started_at = Instant::now();
|
||||
let runtime_request_context =
|
||||
RuntimeRequestContext::from_request_id(approval.approval.id.as_str().to_owned())
|
||||
.with_response_cache_scope(
|
||||
tool.workspace_id.as_str().to_owned(),
|
||||
tool.agent_id.as_str().to_owned(),
|
||||
)
|
||||
.with_metering_context(
|
||||
tool.workspace_id.clone(),
|
||||
Some(tool.agent_id.clone()),
|
||||
InvocationSource::AgentToolCall,
|
||||
)
|
||||
.with_approval_granted();
|
||||
let resolved_auth =
|
||||
resolve_operation_auth(state, &tool.workspace_id, &operation.execution_config).await;
|
||||
let result = match resolved_auth {
|
||||
Ok(resolved_auth) => {
|
||||
state
|
||||
.runtime
|
||||
.execute_request(
|
||||
RuntimeExecutionRequest::new(&operation, &approval.approval.request_payload)
|
||||
.with_optional_auth(resolved_auth.as_ref())
|
||||
.with_context(&runtime_request_context),
|
||||
)
|
||||
.await
|
||||
}
|
||||
Err(error) => Err(error),
|
||||
};
|
||||
|
||||
let (status, response_payload, invocation_status, invocation_level, message, error_kind) =
|
||||
match result {
|
||||
Ok(output) => (
|
||||
ApprovalRequestStatus::Completed,
|
||||
output,
|
||||
InvocationStatus::Ok,
|
||||
InvocationLevel::Info,
|
||||
"approved tool call completed",
|
||||
None,
|
||||
),
|
||||
Err(error) => (
|
||||
ApprovalRequestStatus::Failed,
|
||||
json!({
|
||||
"error": {
|
||||
"code": runtime_error_code(&error),
|
||||
"message": error.to_string(),
|
||||
}
|
||||
}),
|
||||
InvocationStatus::Error,
|
||||
InvocationLevel::Error,
|
||||
"approved tool call failed",
|
||||
Some(runtime_error_code(&error)),
|
||||
),
|
||||
};
|
||||
|
||||
if let Err(error) = persist_invocation(
|
||||
state,
|
||||
&tool,
|
||||
InvocationRecord {
|
||||
request_id: Some(approval.approval.id.as_str()),
|
||||
tool_name: &tool.tool_name,
|
||||
status: invocation_status,
|
||||
level: invocation_level,
|
||||
message,
|
||||
status_code: None,
|
||||
error_kind,
|
||||
duration: started_at.elapsed(),
|
||||
request_preview,
|
||||
response_preview: response_payload.clone(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
{
|
||||
warn!(error = %error, "approved invocation log write failed");
|
||||
}
|
||||
|
||||
state
|
||||
.registry
|
||||
.finish_approval_request(FinishApprovalRequest {
|
||||
workspace_id: &approval.approval.workspace_id,
|
||||
agent_id: &approval.approval.agent_id,
|
||||
approval_id: &approval.approval.id,
|
||||
status,
|
||||
response_payload: Some(response_payload),
|
||||
decision_note: None,
|
||||
})
|
||||
.await
|
||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response())?
|
||||
.ok_or_else(|| StatusCode::CONFLICT.into_response())
|
||||
}
|
||||
|
||||
async fn finish_unavailable_approval(
|
||||
state: &Arc<AppState>,
|
||||
approval: &ApprovalRequestRecord,
|
||||
) -> Result<ApprovalRequestRecord, Response> {
|
||||
state
|
||||
.registry
|
||||
.finish_approval_request(FinishApprovalRequest {
|
||||
workspace_id: &approval.approval.workspace_id,
|
||||
agent_id: &approval.approval.agent_id,
|
||||
approval_id: &approval.approval.id,
|
||||
status: ApprovalRequestStatus::Failed,
|
||||
response_payload: Some(json!({
|
||||
"error": {
|
||||
"code": "approved_operation_unavailable",
|
||||
"message": "the approved operation version is no longer published"
|
||||
}
|
||||
})),
|
||||
decision_note: None,
|
||||
})
|
||||
.await
|
||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response())?
|
||||
.ok_or_else(|| StatusCode::CONFLICT.into_response())
|
||||
}
|
||||
@@ -1,13 +1,13 @@
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
sync::Arc,
|
||||
time::{Duration, Instant},
|
||||
time::{Duration, Instant, SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
|
||||
use crank_core::{CacheScope, CoordinationStateStore, CoordinationStateValue};
|
||||
use crank_registry::{PostgresRegistry, PublishedAgentTool, RegistryError};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio::sync::RwLock;
|
||||
use tokio::sync::{Mutex, RwLock};
|
||||
use tracing::info;
|
||||
|
||||
#[derive(Clone)]
|
||||
@@ -16,6 +16,7 @@ pub struct PublishedToolCatalog {
|
||||
refresh_interval: Duration,
|
||||
coordination_store: Arc<dyn CoordinationStateStore>,
|
||||
cached: Arc<RwLock<HashMap<CatalogKey, CachedCatalog>>>,
|
||||
refresh_locks: Arc<Mutex<HashMap<CatalogKey, Arc<Mutex<()>>>>>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
|
||||
@@ -33,6 +34,7 @@ struct CachedCatalog {
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
struct CatalogSnapshot {
|
||||
tools: Vec<PublishedAgentTool>,
|
||||
generated_at_ms: u64,
|
||||
}
|
||||
|
||||
impl PublishedToolCatalog {
|
||||
@@ -46,6 +48,7 @@ impl PublishedToolCatalog {
|
||||
refresh_interval,
|
||||
coordination_store,
|
||||
cached: Arc::new(RwLock::new(HashMap::new())),
|
||||
refresh_locks: Arc::new(Mutex::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -81,12 +84,32 @@ impl PublishedToolCatalog {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if let Some(tools) = self.load_shared_snapshot(workspace_slug, agent_slug).await {
|
||||
let refresh_lock = {
|
||||
let mut locks = self.refresh_locks.lock().await;
|
||||
Arc::clone(
|
||||
locks
|
||||
.entry(key.clone())
|
||||
.or_insert_with(|| Arc::new(Mutex::new(()))),
|
||||
)
|
||||
};
|
||||
let _refresh_guard = refresh_lock.lock().await;
|
||||
let still_stale = {
|
||||
let guard = self.cached.read().await;
|
||||
match guard.get(&key).and_then(|entry| entry.loaded_at) {
|
||||
Some(loaded_at) => loaded_at.elapsed() >= self.refresh_interval,
|
||||
None => true,
|
||||
}
|
||||
};
|
||||
if !still_stale {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if let Some((tools, age)) = self.load_shared_snapshot(workspace_slug, agent_slug).await {
|
||||
let mut guard = self.cached.write().await;
|
||||
guard.insert(
|
||||
key,
|
||||
CachedCatalog {
|
||||
loaded_at: Some(Instant::now()),
|
||||
loaded_at: Instant::now().checked_sub(age),
|
||||
tools,
|
||||
},
|
||||
);
|
||||
@@ -136,7 +159,7 @@ impl PublishedToolCatalog {
|
||||
&self,
|
||||
workspace_slug: &str,
|
||||
agent_slug: &str,
|
||||
) -> Option<Vec<PublishedAgentTool>> {
|
||||
) -> Option<(Vec<PublishedAgentTool>, Duration)> {
|
||||
if self.refresh_interval.is_zero() {
|
||||
return None;
|
||||
}
|
||||
@@ -150,9 +173,9 @@ impl PublishedToolCatalog {
|
||||
Ok(value) => value?,
|
||||
Err(_) => return None,
|
||||
};
|
||||
serde_json::from_value::<CatalogSnapshot>(value.payload)
|
||||
.ok()
|
||||
.map(|snapshot| snapshot.tools)
|
||||
let snapshot = serde_json::from_value::<CatalogSnapshot>(value.payload).ok()?;
|
||||
let age = Duration::from_millis(now_unix_ms().saturating_sub(snapshot.generated_at_ms));
|
||||
(age < self.refresh_interval).then_some((snapshot.tools, age))
|
||||
}
|
||||
|
||||
async fn store_shared_snapshot(
|
||||
@@ -167,6 +190,7 @@ impl PublishedToolCatalog {
|
||||
|
||||
let payload = match serde_json::to_value(CatalogSnapshot {
|
||||
tools: tools.to_vec(),
|
||||
generated_at_ms: now_unix_ms(),
|
||||
}) {
|
||||
Ok(payload) => payload,
|
||||
Err(_) => return,
|
||||
@@ -184,6 +208,13 @@ impl PublishedToolCatalog {
|
||||
}
|
||||
}
|
||||
|
||||
fn now_unix_ms() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|duration| u64::try_from(duration.as_millis()).unwrap_or(u64::MAX))
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
impl CatalogKey {
|
||||
fn new(workspace_slug: &str, agent_slug: &str) -> Self {
|
||||
Self {
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
mod access;
|
||||
mod app;
|
||||
mod approval_execution;
|
||||
pub mod auth;
|
||||
pub mod catalog;
|
||||
pub mod jsonrpc;
|
||||
@@ -9,4 +10,4 @@ pub mod session;
|
||||
pub mod tool_error;
|
||||
mod transport;
|
||||
|
||||
pub use app::build_app;
|
||||
pub use app::{build_app, build_app_with_background_workers};
|
||||
|
||||
Reference in New Issue
Block a user