602 lines
19 KiB
Rust
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()
|
|
}
|
|
}
|