478 lines
16 KiB
Rust
478 lines
16 KiB
Rust
use super::*;
|
|
use sha2::{Digest, Sha256};
|
|
|
|
impl PostgresRegistry {
|
|
pub async fn create_approval_request(
|
|
&self,
|
|
request: CreateApprovalRequest<'_>,
|
|
) -> Result<ApprovalRequestRecord, RegistryError> {
|
|
let fingerprint = approval_request_fingerprint(&request.approval.request_payload)?;
|
|
let mut transaction = self.pool.begin().await?;
|
|
sqlx::query(
|
|
"update approval_requests
|
|
set status = 'expired'
|
|
where agent_id = $1
|
|
and operation_id = $2
|
|
and operation_version = $3
|
|
and request_fingerprint = $4
|
|
and status = 'pending'
|
|
and expires_at <= $5",
|
|
)
|
|
.bind(request.approval.agent_id.as_str())
|
|
.bind(request.approval.operation_id.as_str())
|
|
.bind(to_db_version(request.approval.operation_version))
|
|
.bind(&fingerprint)
|
|
.bind(request.approval.created_at)
|
|
.execute(&mut *transaction)
|
|
.await?;
|
|
let row = sqlx::query(
|
|
"insert into approval_requests (
|
|
id,
|
|
workspace_id,
|
|
agent_id,
|
|
operation_id,
|
|
operation_version,
|
|
status,
|
|
risk_level,
|
|
request_payload_json,
|
|
response_payload_json,
|
|
created_at,
|
|
expires_at,
|
|
decided_at,
|
|
decided_by_key_id,
|
|
decision_note,
|
|
request_fingerprint
|
|
) values (
|
|
$1, $2, $3, $4, $5, $6, $7, $8,
|
|
$9, $10::timestamptz, $11::timestamptz, $12::timestamptz,
|
|
$13, $14, $15
|
|
)
|
|
on conflict (agent_id, operation_id, operation_version, request_fingerprint)
|
|
where status = 'pending' and request_fingerprint is not null
|
|
do update set request_fingerprint = excluded.request_fingerprint
|
|
returning
|
|
id, workspace_id, agent_id, operation_id, operation_version,
|
|
status, risk_level, request_payload_json, response_payload_json,
|
|
created_at, expires_at, decided_at, decided_by_key_id, decision_note",
|
|
)
|
|
.bind(request.approval.id.as_str())
|
|
.bind(request.approval.workspace_id.as_str())
|
|
.bind(request.approval.agent_id.as_str())
|
|
.bind(request.approval.operation_id.as_str())
|
|
.bind(to_db_version(request.approval.operation_version))
|
|
.bind(serialize_enum_text(
|
|
&request.approval.status,
|
|
"approval_status",
|
|
)?)
|
|
.bind(serialize_enum_text(
|
|
&request.approval.risk_level,
|
|
"approval_risk_level",
|
|
)?)
|
|
.bind(Json(&request.approval.request_payload))
|
|
.bind(request.approval.response_payload.as_ref().map(Json))
|
|
.bind(request.approval.created_at)
|
|
.bind(request.approval.expires_at)
|
|
.bind(request.approval.decided_at)
|
|
.bind(
|
|
request
|
|
.approval
|
|
.decided_by_key_id
|
|
.as_ref()
|
|
.map(PlatformApiKeyId::as_str),
|
|
)
|
|
.bind(request.approval.decision_note.as_deref())
|
|
.bind(fingerprint)
|
|
.fetch_one(&mut *transaction)
|
|
.await?;
|
|
transaction.commit().await?;
|
|
|
|
map_approval_request_row(row)
|
|
}
|
|
|
|
pub async fn list_pending_approval_requests_for_agent(
|
|
&self,
|
|
workspace_id: &WorkspaceId,
|
|
agent_id: &AgentId,
|
|
) -> Result<Vec<ApprovalRequestRecord>, RegistryError> {
|
|
let rows = sqlx::query(
|
|
"select
|
|
id,
|
|
workspace_id,
|
|
agent_id,
|
|
operation_id,
|
|
operation_version,
|
|
status,
|
|
risk_level,
|
|
request_payload_json,
|
|
response_payload_json,
|
|
created_at,
|
|
expires_at,
|
|
decided_at,
|
|
decided_by_key_id,
|
|
decision_note
|
|
from approval_requests
|
|
where workspace_id = $1
|
|
and agent_id = $2
|
|
and status = 'pending'
|
|
and expires_at > now()
|
|
order by created_at asc",
|
|
)
|
|
.bind(workspace_id.as_str())
|
|
.bind(agent_id.as_str())
|
|
.fetch_all(&self.pool)
|
|
.await?;
|
|
|
|
rows.into_iter().map(map_approval_request_row).collect()
|
|
}
|
|
|
|
pub async fn get_approval_request_for_agent(
|
|
&self,
|
|
workspace_id: &WorkspaceId,
|
|
agent_id: &AgentId,
|
|
approval_id: &ApprovalRequestId,
|
|
) -> Result<Option<ApprovalRequestRecord>, RegistryError> {
|
|
let row = sqlx::query(
|
|
"select
|
|
id,
|
|
workspace_id,
|
|
agent_id,
|
|
operation_id,
|
|
operation_version,
|
|
status,
|
|
risk_level,
|
|
request_payload_json,
|
|
response_payload_json,
|
|
created_at,
|
|
expires_at,
|
|
decided_at,
|
|
decided_by_key_id,
|
|
decision_note
|
|
from approval_requests
|
|
where workspace_id = $1
|
|
and agent_id = $2
|
|
and id = $3
|
|
limit 1",
|
|
)
|
|
.bind(workspace_id.as_str())
|
|
.bind(agent_id.as_str())
|
|
.bind(approval_id.as_str())
|
|
.fetch_optional(&self.pool)
|
|
.await?;
|
|
|
|
row.map(map_approval_request_row).transpose()
|
|
}
|
|
|
|
pub async fn list_approval_requests(
|
|
&self,
|
|
query: ListApprovalRequestsQuery<'_>,
|
|
) -> Result<Vec<ApprovalRequestRecord>, RegistryError> {
|
|
let status = query
|
|
.status
|
|
.map(|status| serialize_enum_text(&status, "approval_status"))
|
|
.transpose()?;
|
|
let rows = sqlx::query(
|
|
"select
|
|
id,
|
|
workspace_id,
|
|
agent_id,
|
|
operation_id,
|
|
operation_version,
|
|
status,
|
|
risk_level,
|
|
request_payload_json,
|
|
response_payload_json,
|
|
created_at,
|
|
expires_at,
|
|
decided_at,
|
|
decided_by_key_id,
|
|
decision_note
|
|
from approval_requests
|
|
where workspace_id = $1
|
|
and ($2::text is null or status = $2)
|
|
order by created_at desc
|
|
limit $3",
|
|
)
|
|
.bind(query.workspace_id.as_str())
|
|
.bind(status.as_deref())
|
|
.bind(i64::from(query.limit))
|
|
.fetch_all(&self.pool)
|
|
.await?;
|
|
|
|
rows.into_iter().map(map_approval_request_row).collect()
|
|
}
|
|
|
|
pub async fn get_approval_request(
|
|
&self,
|
|
workspace_id: &WorkspaceId,
|
|
approval_id: &ApprovalRequestId,
|
|
) -> Result<Option<ApprovalRequestRecord>, RegistryError> {
|
|
let row = sqlx::query(
|
|
"select
|
|
id,
|
|
workspace_id,
|
|
agent_id,
|
|
operation_id,
|
|
operation_version,
|
|
status,
|
|
risk_level,
|
|
request_payload_json,
|
|
response_payload_json,
|
|
created_at,
|
|
expires_at,
|
|
decided_at,
|
|
decided_by_key_id,
|
|
decision_note
|
|
from approval_requests
|
|
where workspace_id = $1
|
|
and id = $2
|
|
limit 1",
|
|
)
|
|
.bind(workspace_id.as_str())
|
|
.bind(approval_id.as_str())
|
|
.fetch_optional(&self.pool)
|
|
.await?;
|
|
|
|
row.map(map_approval_request_row).transpose()
|
|
}
|
|
|
|
pub async fn decide_approval_request(
|
|
&self,
|
|
request: DecideApprovalRequest<'_>,
|
|
) -> Result<Option<ApprovalRequestRecord>, RegistryError> {
|
|
let row = sqlx::query(
|
|
"update approval_requests
|
|
set status = $1,
|
|
response_payload_json = $2,
|
|
decided_at = $3::timestamptz,
|
|
decided_by_key_id = $4,
|
|
decision_note = $5
|
|
where workspace_id = $6
|
|
and agent_id = $7
|
|
and id = $8
|
|
and status = 'pending'
|
|
and expires_at > $3::timestamptz
|
|
returning
|
|
id,
|
|
workspace_id,
|
|
agent_id,
|
|
operation_id,
|
|
operation_version,
|
|
status,
|
|
risk_level,
|
|
request_payload_json,
|
|
response_payload_json,
|
|
created_at,
|
|
expires_at,
|
|
decided_at,
|
|
decided_by_key_id,
|
|
decision_note",
|
|
)
|
|
.bind(serialize_enum_text(&request.status, "approval_status")?)
|
|
.bind(request.response_payload.as_ref().map(Json))
|
|
.bind(request.decided_at)
|
|
.bind(request.decided_by_key_id.as_str())
|
|
.bind(request.decision_note)
|
|
.bind(request.workspace_id.as_str())
|
|
.bind(request.agent_id.as_str())
|
|
.bind(request.approval_id.as_str())
|
|
.fetch_optional(&self.pool)
|
|
.await?;
|
|
|
|
row.map(map_approval_request_row).transpose()
|
|
}
|
|
|
|
pub async fn finish_approval_request(
|
|
&self,
|
|
request: FinishApprovalRequest<'_>,
|
|
) -> Result<Option<ApprovalRequestRecord>, RegistryError> {
|
|
let row = sqlx::query(
|
|
"update approval_requests
|
|
set status = $1,
|
|
response_payload_json = $2,
|
|
decision_note = coalesce($3, decision_note)
|
|
where workspace_id = $4
|
|
and agent_id = $5
|
|
and id = $6
|
|
and status = 'executing'
|
|
returning
|
|
id,
|
|
workspace_id,
|
|
agent_id,
|
|
operation_id,
|
|
operation_version,
|
|
status,
|
|
risk_level,
|
|
request_payload_json,
|
|
response_payload_json,
|
|
created_at,
|
|
expires_at,
|
|
decided_at,
|
|
decided_by_key_id,
|
|
decision_note",
|
|
)
|
|
.bind(serialize_enum_text(&request.status, "approval_status")?)
|
|
.bind(request.response_payload.as_ref().map(Json))
|
|
.bind(request.decision_note)
|
|
.bind(request.workspace_id.as_str())
|
|
.bind(request.agent_id.as_str())
|
|
.bind(request.approval_id.as_str())
|
|
.fetch_optional(&self.pool)
|
|
.await?;
|
|
|
|
row.map(map_approval_request_row).transpose()
|
|
}
|
|
|
|
pub async fn claim_approval_request(
|
|
&self,
|
|
workspace_id: &WorkspaceId,
|
|
agent_id: &AgentId,
|
|
approval_id: &ApprovalRequestId,
|
|
started_at: OffsetDateTime,
|
|
) -> Result<Option<ApprovalRequestRecord>, RegistryError> {
|
|
let row = sqlx::query(
|
|
"update approval_requests
|
|
set status = 'executing',
|
|
execution_started_at = $1,
|
|
execution_attempts = execution_attempts + 1
|
|
where workspace_id = $2
|
|
and agent_id = $3
|
|
and id = $4
|
|
and status = 'approved'
|
|
returning
|
|
id, workspace_id, agent_id, operation_id, operation_version,
|
|
status, risk_level, request_payload_json, response_payload_json,
|
|
created_at, expires_at, decided_at, decided_by_key_id, decision_note",
|
|
)
|
|
.bind(started_at)
|
|
.bind(workspace_id.as_str())
|
|
.bind(agent_id.as_str())
|
|
.bind(approval_id.as_str())
|
|
.fetch_optional(&self.pool)
|
|
.await?;
|
|
|
|
row.map(map_approval_request_row).transpose()
|
|
}
|
|
|
|
pub async fn claim_next_recoverable_approval_request(
|
|
&self,
|
|
started_at: OffsetDateTime,
|
|
approved_before: OffsetDateTime,
|
|
stale_before: OffsetDateTime,
|
|
) -> Result<Option<ApprovalRequestRecord>, RegistryError> {
|
|
let row = sqlx::query(
|
|
"with candidate as (
|
|
select id
|
|
from approval_requests
|
|
where (status = 'approved' and decided_at <= $1)
|
|
or (status = 'executing' and execution_started_at < $2)
|
|
order by decided_at asc nulls last, created_at asc
|
|
for update skip locked
|
|
limit 1
|
|
)
|
|
update approval_requests as approval
|
|
set status = 'executing',
|
|
execution_started_at = $3,
|
|
execution_attempts = approval.execution_attempts + 1
|
|
from candidate
|
|
where approval.id = candidate.id
|
|
returning
|
|
approval.id, approval.workspace_id, approval.agent_id,
|
|
approval.operation_id, approval.operation_version, approval.status,
|
|
approval.risk_level, approval.request_payload_json,
|
|
approval.response_payload_json, approval.created_at, approval.expires_at,
|
|
approval.decided_at, approval.decided_by_key_id, approval.decision_note",
|
|
)
|
|
.bind(approved_before)
|
|
.bind(stale_before)
|
|
.bind(started_at)
|
|
.fetch_optional(&self.pool)
|
|
.await?;
|
|
|
|
row.map(map_approval_request_row).transpose()
|
|
}
|
|
|
|
pub async fn expire_approval_request(
|
|
&self,
|
|
request: ExpireApprovalRequest<'_>,
|
|
) -> Result<Option<ApprovalRequestRecord>, RegistryError> {
|
|
let row = sqlx::query(
|
|
"update approval_requests
|
|
set status = 'expired'
|
|
where workspace_id = $1
|
|
and agent_id = $2
|
|
and id = $3
|
|
and status = 'pending'
|
|
and expires_at <= $4::timestamptz
|
|
returning
|
|
id,
|
|
workspace_id,
|
|
agent_id,
|
|
operation_id,
|
|
operation_version,
|
|
status,
|
|
risk_level,
|
|
request_payload_json,
|
|
response_payload_json,
|
|
created_at,
|
|
expires_at,
|
|
decided_at,
|
|
decided_by_key_id,
|
|
decision_note",
|
|
)
|
|
.bind(request.workspace_id.as_str())
|
|
.bind(request.agent_id.as_str())
|
|
.bind(request.approval_id.as_str())
|
|
.bind(request.expired_at)
|
|
.fetch_optional(&self.pool)
|
|
.await?;
|
|
|
|
row.map(map_approval_request_row).transpose()
|
|
}
|
|
}
|
|
|
|
fn approval_request_fingerprint(payload: &Value) -> Result<String, RegistryError> {
|
|
let canonical = canonical_json(payload);
|
|
let encoded = serde_json::to_vec(&canonical)?;
|
|
Ok(format!("{:x}", Sha256::digest(encoded)))
|
|
}
|
|
|
|
fn canonical_json(value: &Value) -> Value {
|
|
match value {
|
|
Value::Object(object) => {
|
|
let sorted = object
|
|
.iter()
|
|
.map(|(key, value)| (key.clone(), canonical_json(value)))
|
|
.collect::<std::collections::BTreeMap<_, _>>();
|
|
Value::Object(sorted.into_iter().collect())
|
|
}
|
|
Value::Array(items) => Value::Array(items.iter().map(canonical_json).collect()),
|
|
_ => value.clone(),
|
|
}
|
|
}
|
|
|
|
fn map_approval_request_row(row: PgRow) -> Result<ApprovalRequestRecord, RegistryError> {
|
|
Ok(ApprovalRequestRecord {
|
|
approval: ApprovalRequest {
|
|
id: ApprovalRequestId::new(row.get::<String, _>("id")),
|
|
workspace_id: WorkspaceId::new(row.get::<String, _>("workspace_id")),
|
|
agent_id: AgentId::new(row.get::<String, _>("agent_id")),
|
|
operation_id: OperationId::new(row.get::<String, _>("operation_id")),
|
|
operation_version: from_db_version(row.get("operation_version"), "operation_version")?,
|
|
status: deserialize_enum_text(&row.get::<String, _>("status"), "approval_status")?,
|
|
risk_level: deserialize_enum_text(
|
|
&row.get::<String, _>("risk_level"),
|
|
"approval_risk_level",
|
|
)?,
|
|
request_payload: row.get::<Value, _>("request_payload_json"),
|
|
response_payload: row.get::<Option<Value>, _>("response_payload_json"),
|
|
created_at: row.get("created_at"),
|
|
expires_at: row.get("expires_at"),
|
|
decided_at: row.get("decided_at"),
|
|
decided_by_key_id: row
|
|
.get::<Option<String>, _>("decided_by_key_id")
|
|
.map(PlatformApiKeyId::new),
|
|
decision_note: row.get("decision_note"),
|
|
},
|
|
})
|
|
}
|