Files
crank/crates/crank-registry/src/postgres/agent.rs
T
github-ops 1230d3a547
Deploy / deploy (push) Successful in 2m31s
CI / Rust Checks (push) Failing after 5m12s
CI / UI Checks (push) Has been skipped
CI / Frontend E2E (push) Has been skipped
CI / Deployment Manifests (push) Has been skipped
chore: publish clean community baseline
2026-06-19 13:48:47 +00:00

602 lines
19 KiB
Rust

use super::*;
impl PostgresRegistry {
pub async fn list_agents(
&self,
workspace_id: &WorkspaceId,
) -> Result<Vec<AgentSummary>, RegistryError> {
let rows = sqlx::query!(
"select
id,
workspace_id,
slug,
display_name,
description,
status,
current_draft_version,
latest_published_version,
created_at as \"created_at!: time::OffsetDateTime\",
updated_at as \"updated_at!: time::OffsetDateTime\",
published_at as \"published_at: time::OffsetDateTime\"
from agents
where workspace_id = $1
order by slug asc",
workspace_id.as_str(),
)
.fetch_all(&self.pool)
.await?;
rows.into_iter()
.map(|row| {
build_agent_summary(
row.id,
row.workspace_id,
row.slug,
row.display_name,
row.description,
row.status,
row.current_draft_version,
row.latest_published_version,
row.created_at,
row.updated_at,
row.published_at,
)
})
.collect()
}
pub async fn get_agent_summary(
&self,
workspace_id: &WorkspaceId,
agent_id: &AgentId,
) -> Result<Option<AgentSummary>, RegistryError> {
let row = sqlx::query!(
"select
id,
workspace_id,
slug,
display_name,
description,
status,
current_draft_version,
latest_published_version,
created_at as \"created_at!: time::OffsetDateTime\",
updated_at as \"updated_at!: time::OffsetDateTime\",
published_at as \"published_at: time::OffsetDateTime\"
from agents
where workspace_id = $1 and id = $2",
workspace_id.as_str(),
agent_id.as_str(),
)
.fetch_optional(&self.pool)
.await?;
row.map(|row| {
build_agent_summary(
row.id,
row.workspace_id,
row.slug,
row.display_name,
row.description,
row.status,
row.current_draft_version,
row.latest_published_version,
row.created_at,
row.updated_at,
row.published_at,
)
})
.transpose()
}
pub async fn create_agent(&self, request: CreateAgentRequest<'_>) -> Result<(), RegistryError> {
let mut tx = self.pool.begin().await?;
sqlx::query(
"insert into agents (
id,
workspace_id,
slug,
display_name,
description,
status,
current_draft_version,
latest_published_version,
created_at,
updated_at,
published_at
) values (
$1, $2, $3, $4, $5, $6, $7, $8, $9::timestamptz, $10::timestamptz, $11::timestamptz
)",
)
.bind(request.agent.id.as_str())
.bind(request.agent.workspace_id.as_str())
.bind(&request.agent.slug)
.bind(&request.agent.display_name)
.bind(&request.agent.description)
.bind(serialize_enum_text(&request.agent.status, "status")?)
.bind(to_db_version(request.agent.current_draft_version))
.bind(request.agent.latest_published_version.map(to_db_version))
.bind(request.agent.created_at)
.bind(request.agent.updated_at)
.bind(request.agent.published_at)
.execute(&mut *tx)
.await?;
insert_agent_version_row(&mut tx, request.version).await?;
replace_agent_bindings_rows(
&mut tx,
&request.agent.id,
request.version.version,
request.bindings,
)
.await?;
tx.commit().await?;
Ok(())
}
pub async fn create_agent_draft_version(
&self,
request: CreateAgentDraftVersionRequest<'_>,
) -> Result<(), RegistryError> {
let Some(summary) = self
.get_agent_summary(request.workspace_id, request.agent_id)
.await?
else {
return Err(RegistryError::AgentNotFound {
agent_id: request.agent_id.as_str().to_owned(),
});
};
let expected = summary.current_draft_version + 1;
if request.version.version != expected {
return Err(RegistryError::InvalidAgentVersionSequence {
agent_id: request.agent_id.as_str().to_owned(),
expected,
actual: request.version.version,
});
}
let mut tx = self.pool.begin().await?;
insert_agent_version_row(&mut tx, request.version).await?;
replace_agent_bindings_rows(
&mut tx,
request.agent_id,
request.version.version,
request.bindings,
)
.await?;
sqlx::query(
"update agents
set current_draft_version = $3,
updated_at = $4::timestamptz
where workspace_id = $1 and id = $2",
)
.bind(request.workspace_id.as_str())
.bind(request.agent_id.as_str())
.bind(to_db_version(request.version.version))
.bind(request.updated_at)
.execute(&mut *tx)
.await?;
tx.commit().await?;
Ok(())
}
pub async fn get_agent_version(
&self,
workspace_id: &WorkspaceId,
agent_id: &AgentId,
version: u32,
) -> Result<Option<AgentVersionRecord>, RegistryError> {
let summary = match self.get_agent_summary(workspace_id, agent_id).await? {
Some(value) => value,
None => return Ok(None),
};
let row = sqlx::query!(
"select
agent_id,
version,
status,
instructions_json,
tool_selection_policy_json,
created_at as \"created_at!: time::OffsetDateTime\"
from agent_versions
where agent_id = $1 and version = $2",
agent_id.as_str(),
to_db_version(version),
)
.fetch_optional(&self.pool)
.await?;
let Some(row) = row else {
return Ok(None);
};
let bindings = self.list_agent_bindings(agent_id, version).await?;
Ok(Some(build_agent_version_record(
&summary,
row.version,
row.status,
row.instructions_json,
row.tool_selection_policy_json,
row.created_at,
bindings,
)?))
}
pub async fn save_agent_bindings(
&self,
request: SaveAgentBindingsRequest<'_>,
) -> Result<(), RegistryError> {
if self
.get_agent_summary(request.workspace_id, request.agent_id)
.await?
.is_none()
{
return Err(RegistryError::AgentNotFound {
agent_id: request.agent_id.as_str().to_owned(),
});
}
let mut tx = self.pool.begin().await?;
replace_agent_bindings_rows(
&mut tx,
request.agent_id,
request.agent_version,
request.bindings,
)
.await?;
tx.commit().await?;
Ok(())
}
pub async fn update_agent_summary(
&self,
workspace_id: &WorkspaceId,
agent_id: &AgentId,
slug: &str,
display_name: &str,
description: &str,
updated_at: &time::OffsetDateTime,
) -> Result<(), RegistryError> {
let result = sqlx::query(
"update agents
set slug = $3,
display_name = $4,
description = $5,
updated_at = $6::timestamptz
where workspace_id = $1 and id = $2",
)
.bind(workspace_id.as_str())
.bind(agent_id.as_str())
.bind(slug)
.bind(display_name)
.bind(description)
.bind(updated_at)
.execute(&self.pool)
.await?;
if result.rows_affected() == 0 {
return Err(RegistryError::AgentNotFound {
agent_id: agent_id.as_str().to_owned(),
});
}
Ok(())
}
pub async fn delete_agent(
&self,
workspace_id: &WorkspaceId,
agent_id: &AgentId,
) -> Result<(), RegistryError> {
let result = sqlx::query("delete from agents where workspace_id = $1 and id = $2")
.bind(workspace_id.as_str())
.bind(agent_id.as_str())
.execute(&self.pool)
.await?;
if result.rows_affected() == 0 {
return Err(RegistryError::AgentNotFound {
agent_id: agent_id.as_str().to_owned(),
});
}
Ok(())
}
pub async fn publish_agent(
&self,
request: PublishAgentRequest<'_>,
) -> Result<(), RegistryError> {
if self
.get_agent_version(request.workspace_id, request.agent_id, request.version)
.await?
.is_none()
{
return Err(RegistryError::AgentNotFound {
agent_id: request.agent_id.as_str().to_owned(),
});
}
let mut tx = self.pool.begin().await?;
sqlx::query(
"insert into published_agents (
agent_id,
version,
published_at,
published_by
) values ($1, $2, $3::timestamptz, $4)
on conflict(agent_id) do update set
version = excluded.version,
published_at = excluded.published_at,
published_by = excluded.published_by",
)
.bind(request.agent_id.as_str())
.bind(to_db_version(request.version))
.bind(request.published_at)
.bind(request.published_by)
.execute(&mut *tx)
.await?;
sqlx::query(
"update agent_versions
set status = $1
where agent_id = $2 and version = $3",
)
.bind(serialize_enum_text(&AgentStatus::Published, "status")?)
.bind(request.agent_id.as_str())
.bind(to_db_version(request.version))
.execute(&mut *tx)
.await?;
sqlx::query(
"update agents
set status = $1,
latest_published_version = $2,
published_at = $3::timestamptz,
updated_at = $4::timestamptz
where id = $5 and workspace_id = $6",
)
.bind(serialize_enum_text(&AgentStatus::Published, "status")?)
.bind(to_db_version(request.version))
.bind(request.published_at)
.bind(request.published_at)
.bind(request.agent_id.as_str())
.bind(request.workspace_id.as_str())
.execute(&mut *tx)
.await?;
tx.commit().await?;
Ok(())
}
pub async fn unpublish_agent(
&self,
workspace_id: &WorkspaceId,
agent_id: &AgentId,
updated_at: &time::OffsetDateTime,
) -> Result<(), RegistryError> {
let mut tx = self.pool.begin().await?;
sqlx::query("delete from published_agents where agent_id = $1")
.bind(agent_id.as_str())
.execute(&mut *tx)
.await?;
sqlx::query(
"update agent_versions
set status = $1
where agent_id = $2
and version = (
select current_draft_version
from agents
where id = $2 and workspace_id = $3
)",
)
.bind(serialize_enum_text(&AgentStatus::Draft, "status")?)
.bind(agent_id.as_str())
.bind(workspace_id.as_str())
.execute(&mut *tx)
.await?;
let result = sqlx::query(
"update agents
set status = $1,
published_at = null,
updated_at = $2::timestamptz
where id = $3 and workspace_id = $4",
)
.bind(serialize_enum_text(&AgentStatus::Draft, "status")?)
.bind(updated_at)
.bind(agent_id.as_str())
.bind(workspace_id.as_str())
.execute(&mut *tx)
.await?;
if result.rows_affected() == 0 {
return Err(RegistryError::AgentNotFound {
agent_id: agent_id.as_str().to_owned(),
});
}
tx.commit().await?;
Ok(())
}
pub async fn archive_agent(
&self,
workspace_id: &WorkspaceId,
agent_id: &AgentId,
updated_at: &time::OffsetDateTime,
) -> Result<(), RegistryError> {
let mut tx = self.pool.begin().await?;
sqlx::query("delete from published_agents where agent_id = $1")
.bind(agent_id.as_str())
.execute(&mut *tx)
.await?;
sqlx::query(
"update agent_versions
set status = $1
where agent_id = $2
and version = (
select current_draft_version
from agents
where id = $2 and workspace_id = $3
)",
)
.bind(serialize_enum_text(&AgentStatus::Archived, "status")?)
.bind(agent_id.as_str())
.bind(workspace_id.as_str())
.execute(&mut *tx)
.await?;
let result = sqlx::query(
"update agents
set status = $1,
published_at = null,
updated_at = $2::timestamptz
where id = $3 and workspace_id = $4",
)
.bind(serialize_enum_text(&AgentStatus::Archived, "status")?)
.bind(updated_at)
.bind(agent_id.as_str())
.bind(workspace_id.as_str())
.execute(&mut *tx)
.await?;
if result.rows_affected() == 0 {
return Err(RegistryError::AgentNotFound {
agent_id: agent_id.as_str().to_owned(),
});
}
tx.commit().await?;
Ok(())
}
pub async fn get_published_agent_tools_by_slug(
&self,
workspace_slug: &str,
agent_slug: &str,
) -> Result<Vec<PublishedAgentTool>, RegistryError> {
let rows = sqlx::query!(
"select
w.id as workspace_id,
w.slug as workspace_slug,
a.id as agent_id,
a.slug as agent_slug,
b.tool_name,
b.tool_title,
coalesce(b.tool_description_override, ov.tool_description_json->>'description') as \"tool_description!\",
o.id,
o.name,
o.display_name,
o.category,
o.protocol,
o.security_level,
o.created_at as \"operation_created_at!: time::OffsetDateTime\",
o.updated_at as \"operation_updated_at!: time::OffsetDateTime\",
o.published_at as \"operation_published_at: time::OffsetDateTime\",
ov.version,
ov.status,
ov.target_json,
ov.input_schema_json,
ov.output_schema_json,
ov.input_mapping_json,
ov.output_mapping_json,
ov.execution_config_json,
ov.tool_description_json,
ov.samples_json,
ov.generated_draft_json,
ov.config_export_json,
ov.change_note,
ov.created_at as \"created_at!: time::OffsetDateTime\",
ov.created_by
from workspaces w
join agents a on a.workspace_id = w.id
join published_agents pa on pa.agent_id = a.id
join agent_operation_bindings b on b.agent_id = a.id and b.agent_version = pa.version
join operation_versions ov on ov.operation_id = b.operation_id and ov.version = b.operation_version
join operations o on o.id = ov.operation_id and o.workspace_id = w.id
where w.slug = $1 and a.slug = $2 and b.enabled = true
order by b.tool_name asc",
workspace_slug,
agent_slug,
)
.fetch_all(&self.pool)
.await?;
if rows.is_empty() {
let exists = sqlx::query!(
"select 1 as \"present!\"
from workspaces w
join agents a on a.workspace_id = w.id
join published_agents pa on pa.agent_id = a.id
where w.slug = $1 and a.slug = $2",
workspace_slug,
agent_slug,
)
.fetch_optional(&self.pool)
.await?;
if exists.is_none() {
return Err(RegistryError::PublishedAgentNotFound {
workspace_slug: workspace_slug.to_owned(),
agent_slug: agent_slug.to_owned(),
});
}
}
rows.into_iter()
.map(|row| {
let workspace_id = row.workspace_id.clone();
build_published_agent_tool(
row.workspace_id,
row.workspace_slug,
row.agent_id,
row.agent_slug,
row.tool_name,
row.tool_title,
row.tool_description,
build_operation_version_record(
row.id,
workspace_id,
row.name,
row.display_name,
row.category,
row.protocol,
row.security_level,
row.operation_created_at,
row.operation_updated_at,
row.operation_published_at,
row.version,
row.status,
row.target_json,
row.input_schema_json,
row.output_schema_json,
row.input_mapping_json,
row.output_mapping_json,
row.execution_config_json,
row.tool_description_json,
row.samples_json,
row.generated_draft_json,
row.config_export_json,
row.change_note,
row.created_at,
row.created_by,
)?
.snapshot,
)
})
.collect()
}
}