Усилить безопасность и надёжность выполнения операций
This commit is contained in:
@@ -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())
|
||||
}
|
||||
Reference in New Issue
Block a user