feat: add streaming state storage

This commit is contained in:
a.tolmachev
2026-04-06 10:23:01 +03:00
parent 6a0381b8e5
commit 7e4f3d142e
7 changed files with 1070 additions and 32 deletions
+6 -7
View File
@@ -2,20 +2,19 @@
## Current
### `feat/streaming-core-model`
### `feat/stream-session-store`
Status: completed
DoD:
- core domain types exist for window/session/job execution
- `ExecutionConfig` can carry `streaming`
- protocol-aware validation follows documented capability matrix
- streaming types are exported from `crank-core`
- core tests cover serde and validation rules
- persistent tables exist for `stream_sessions` and `async_jobs`
- registry supports create/get/update/close/list/cleanup for both stores
- invalid state transitions are rejected explicitly
- migration and storage tests cover lifecycle and cleanup
## Next
- `feat/stream-session-store`
- `feat/runtime-window-mode`
## Backlog
+8
View File
@@ -142,6 +142,12 @@ impl From<RegistryError> for ApiError {
RegistryError::SecretNotFound { secret_id } => {
Self::not_found(format!("secret {secret_id} was not found"))
}
RegistryError::StreamSessionNotFound { session_id } => {
Self::not_found(format!("stream session {session_id} was not found"))
}
RegistryError::AsyncJobNotFound { job_id } => {
Self::not_found(format!("async job {job_id} was not found"))
}
RegistryError::InvocationLogNotFound { log_id } => {
Self::not_found(format!("invocation log {log_id} was not found"))
}
@@ -177,6 +183,8 @@ impl From<RegistryError> for ApiError {
RegistryError::SecretNameAlreadyExists { workspace_id, name } => Self::conflict(
format!("secret with name {name} already exists in workspace {workspace_id}"),
),
RegistryError::InvalidStreamSessionTransition { .. }
| RegistryError::InvalidAsyncJobTransition { .. } => Self::conflict(value.to_string()),
RegistryError::UserEmailAlreadyExists { email } => {
Self::conflict(format!("user with email {email} already exists"))
}
+16
View File
@@ -25,6 +25,22 @@ pub enum RegistryError {
PlatformApiKeyNotFound { key_id: String },
#[error("secret {secret_id} was not found")]
SecretNotFound { secret_id: String },
#[error("stream session {session_id} was not found")]
StreamSessionNotFound { session_id: String },
#[error("async job {job_id} was not found")]
AsyncJobNotFound { job_id: String },
#[error("invalid stream session transition for {session_id}: {from} -> {to}")]
InvalidStreamSessionTransition {
session_id: String,
from: String,
to: String,
},
#[error("invalid async job transition for {job_id}: {from} -> {to}")]
InvalidAsyncJobTransition {
job_id: String,
from: String,
to: String,
},
#[error("secret with name {name} already exists in workspace {workspace_id}")]
SecretNameAlreadyExists { workspace_id: String, name: String },
#[error("invocation log {log_id} was not found")]
+9 -6
View File
@@ -5,17 +5,20 @@ mod postgres;
pub use error::RegistryError;
pub use model::{
AgentSummary, AgentVersionRecord, AuthUserRecord, CreateAgentRequest, CreateInvitationRequest,
CreateInvocationLogRequest, CreatePlatformApiKeyRequest, CreateSecretRequest,
AgentSummary, AgentVersionRecord, AsyncJobFilter, AsyncJobRecord, AuthUserRecord,
CreateAgentRequest, CreateAsyncJobRequest, CreateInvitationRequest, CreateInvocationLogRequest,
CreatePlatformApiKeyRequest, CreateSecretRequest, CreateStreamSessionRequest,
CreateVersionRequest, CreateWorkspaceRequest, CreateYamlImportJobRequest, DescriptorKind,
DescriptorMetadata, InvitationRecord, InvocationLogRecord, ListInvocationLogsQuery,
MembershipRecord, OperationAgentRef, OperationSampleMetadata, OperationSummary,
OperationUsageSummary, OperationVersionRecord, PlatformApiKeyRecord, PublishAgentRequest,
OperationUsageSummary, OperationVersionRecord, Page, PlatformApiKeyRecord, PublishAgentRequest,
PublishRequest, PublishedAgentTool, RegistryOperation, RotateSecretRequest, SampleKind,
SaveAgentBindingsRequest, SaveAuthProfileRequest, SaveDescriptorMetadataRequest,
SaveSampleMetadataRequest, SecretRecord, SecretVersionRecord, SessionRecord,
UpdateWorkspaceRequest, UsageAgentBreakdown, UsageBucket, UsageOperationBreakdown, UsageQuery,
UsageRollupRecord, UsageSummary, UsageTimelinePoint, WorkspaceMembershipRecord,
WorkspaceRecord, YamlImportJob, YamlImportJobCompletion, YamlImportJobId, YamlImportJobStatus,
StreamSessionFilter, StreamSessionRecord, UpdateAsyncJobStatusRequest,
UpdateStreamSessionStateRequest, UpdateWorkspaceRequest, UsageAgentBreakdown, UsageBucket,
UsageOperationBreakdown, UsageQuery, UsageRollupRecord, UsageSummary, UsageTimelinePoint,
WorkspaceMembershipRecord, WorkspaceRecord, YamlImportJob, YamlImportJobCompletion,
YamlImportJobId, YamlImportJobStatus,
};
pub use postgres::PostgresRegistry;
+67
View File
@@ -513,5 +513,72 @@ pub async fn apply_postgres(pool: &PgPool) -> Result<(), sqlx::Error> {
.execute(pool)
.await?;
query(
"create table if not exists stream_sessions (
id text primary key,
workspace_id text not null references workspaces(id) on delete cascade,
agent_id text null references agents(id) on delete set null,
operation_id text not null references operations(id) on delete cascade,
protocol text not null,
mode text not null,
status text not null,
cursor_json jsonb null,
state_json jsonb not null,
expires_at timestamptz not null,
last_poll_at timestamptz null,
created_at timestamptz not null,
closed_at timestamptz null
)",
)
.execute(pool)
.await?;
query(
"create index if not exists stream_sessions_workspace_status_idx
on stream_sessions(workspace_id, status, created_at desc)",
)
.execute(pool)
.await?;
query(
"create index if not exists stream_sessions_expires_at_idx
on stream_sessions(expires_at)",
)
.execute(pool)
.await?;
query(
"create table if not exists async_jobs (
id text primary key,
workspace_id text not null references workspaces(id) on delete cascade,
agent_id text null references agents(id) on delete set null,
operation_id text not null references operations(id) on delete cascade,
status text not null,
progress_json jsonb not null,
result_json jsonb null,
error_json jsonb null,
expires_at timestamptz null,
created_at timestamptz not null,
updated_at timestamptz not null,
finished_at timestamptz null
)",
)
.execute(pool)
.await?;
query(
"create index if not exists async_jobs_workspace_status_idx
on async_jobs(workspace_id, status, updated_at desc)",
)
.execute(pool)
.await?;
query(
"create index if not exists async_jobs_expires_at_idx
on async_jobs(expires_at)",
)
.execute(pool)
.await?;
Ok(())
}
+76 -4
View File
@@ -1,8 +1,10 @@
use crank_core::{
Agent, AgentId, AgentOperationBinding, AgentStatus, AgentVersion, AuthProfile, DescriptorId,
ExportMode, InvitationToken, InvocationLevel, InvocationLog, InvocationSource, MembershipRole,
Operation, OperationId, OperationStatus, PlatformApiKey, Protocol, SampleId, Secret, SecretId,
SecretVersion, UsagePeriod, UsageRollup, User, UserSessionId, Workspace, WorkspaceId,
Agent, AgentId, AgentOperationBinding, AgentStatus, AgentVersion, AsyncJobHandle, AsyncJobId,
AuthProfile, DescriptorId, ExecutionMode, ExportMode, InvitationToken, InvocationLevel,
InvocationLog, InvocationSource, JobStatus, MembershipRole, Operation, OperationId,
OperationStatus, PlatformApiKey, Protocol, SampleId, Secret, SecretId, SecretVersion,
StreamSession, StreamSessionId, StreamStatus, UsagePeriod, UsageRollup, User, UserSessionId,
Workspace, WorkspaceId,
};
use crank_mapping::MappingSet;
use crank_schema::Schema;
@@ -42,6 +44,12 @@ define_registry_id!(YamlImportJobId);
pub type RegistryOperation = Operation<Schema, MappingSet>;
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct Page<T> {
pub items: Vec<T>,
pub total: u64,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct WorkspaceRecord {
pub workspace: Workspace,
@@ -95,6 +103,16 @@ pub struct SecretVersionRecord {
pub secret_version: SecretVersion,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct StreamSessionRecord {
pub session: StreamSession,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AsyncJobRecord {
pub job: AsyncJobHandle,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct InvocationLogRecord {
pub log: InvocationLog,
@@ -432,6 +450,60 @@ pub struct RotateSecretRequest<'a> {
pub created_by: Option<&'a crank_core::UserId>,
}
#[derive(Clone, Debug, PartialEq)]
pub struct CreateStreamSessionRequest<'a> {
pub session: &'a StreamSession,
}
#[derive(Clone, Debug, PartialEq)]
pub struct UpdateStreamSessionStateRequest<'a> {
pub session_id: &'a StreamSessionId,
pub current_status: StreamStatus,
pub next_status: StreamStatus,
pub cursor: Option<&'a Value>,
pub state: &'a Value,
pub expires_at: Option<&'a str>,
pub last_poll_at: Option<&'a str>,
pub closed_at: Option<&'a str>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct StreamSessionFilter<'a> {
pub workspace_id: &'a WorkspaceId,
pub agent_id: Option<&'a AgentId>,
pub operation_id: Option<&'a OperationId>,
pub status: Option<StreamStatus>,
pub mode: Option<ExecutionMode>,
pub limit: u32,
}
#[derive(Clone, Debug, PartialEq)]
pub struct CreateAsyncJobRequest<'a> {
pub job: &'a AsyncJobHandle,
}
#[derive(Clone, Debug, PartialEq)]
pub struct UpdateAsyncJobStatusRequest<'a> {
pub job_id: &'a AsyncJobId,
pub current_status: JobStatus,
pub next_status: JobStatus,
pub progress: &'a Value,
pub result: Option<&'a Value>,
pub error: Option<&'a Value>,
pub expires_at: Option<&'a str>,
pub updated_at: &'a str,
pub finished_at: Option<&'a str>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AsyncJobFilter<'a> {
pub workspace_id: &'a WorkspaceId,
pub agent_id: Option<&'a AgentId>,
pub operation_id: Option<&'a OperationId>,
pub status: Option<JobStatus>,
pub limit: u32,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SaveSampleMetadataRequest<'a> {
pub sample: &'a OperationSampleMetadata,
+888 -15
View File
@@ -1,8 +1,9 @@
use crank_core::{
AgentId, AgentOperationBinding, AgentStatus, AgentVersion, AuthProfile, GraphqlOperationType,
HttpMethod, InvitationId, InvitationToken, InvocationLog, InvocationLogId, MembershipRole,
OperationId, OperationStatus, PlatformApiKey, PlatformApiKeyId, PlatformApiKeyStatus, Secret,
SecretId, SecretVersion, Target, UsageRollup, User, UserId, UserSessionId, Workspace,
AgentId, AgentOperationBinding, AgentStatus, AgentVersion, AsyncJobHandle, AsyncJobId,
AuthProfile, GraphqlOperationType, HttpMethod, InvitationId, InvitationToken, InvocationLog,
InvocationLogId, JobStatus, MembershipRole, OperationId, OperationStatus, PlatformApiKey,
PlatformApiKeyId, PlatformApiKeyStatus, Secret, SecretId, SecretVersion, StreamSession,
StreamSessionId, StreamStatus, Target, UsageRollup, User, UserId, UserSessionId, Workspace,
WorkspaceId,
};
use serde::{Serialize, de::DeserializeOwned};
@@ -17,16 +18,18 @@ use crate::{
error::RegistryError,
migrations,
model::{
AgentSummary, AgentVersionRecord, AuthUserRecord, CreateAgentRequest,
CreateInvitationRequest, CreateInvocationLogRequest, CreatePlatformApiKeyRequest,
CreateSecretRequest, CreateVersionRequest, CreateWorkspaceRequest,
CreateYamlImportJobRequest, DescriptorMetadata, InvitationRecord, InvocationLogRecord,
ListInvocationLogsQuery, MembershipRecord, OperationAgentRef, OperationSampleMetadata,
OperationSummary, OperationUsageSummary, OperationVersionRecord, PlatformApiKeyRecord,
AgentSummary, AgentVersionRecord, AsyncJobFilter, AuthUserRecord, CreateAgentRequest,
CreateAsyncJobRequest, CreateInvitationRequest, CreateInvocationLogRequest,
CreatePlatformApiKeyRequest, CreateSecretRequest, CreateStreamSessionRequest,
CreateVersionRequest, CreateWorkspaceRequest, CreateYamlImportJobRequest,
DescriptorMetadata, InvitationRecord, InvocationLogRecord, ListInvocationLogsQuery,
MembershipRecord, OperationAgentRef, OperationSampleMetadata, OperationSummary,
OperationUsageSummary, OperationVersionRecord, Page, PlatformApiKeyRecord,
PublishAgentRequest, PublishRequest, PublishedAgentTool, RegistryOperation,
RotateSecretRequest, SaveAgentBindingsRequest, SaveAuthProfileRequest,
SaveDescriptorMetadataRequest, SaveSampleMetadataRequest, SecretRecord,
SecretVersionRecord, SessionRecord, UpdateWorkspaceRequest, UsageAgentBreakdown,
SecretVersionRecord, SessionRecord, StreamSessionFilter, UpdateAsyncJobStatusRequest,
UpdateStreamSessionStateRequest, UpdateWorkspaceRequest, UsageAgentBreakdown,
UsageOperationBreakdown, UsageQuery, UsageRollupRecord, UsageSummary, UsageTimelinePoint,
WorkspaceMembershipRecord, WorkspaceRecord, YamlImportJob, YamlImportJobCompletion,
YamlImportJobId, YamlImportJobStatus,
@@ -860,6 +863,481 @@ impl PostgresRegistry {
Ok(())
}
pub async fn create_stream_session(
&self,
request: CreateStreamSessionRequest<'_>,
) -> Result<StreamSession, RegistryError> {
sqlx::query(
"insert into stream_sessions (
id,
workspace_id,
agent_id,
operation_id,
protocol,
mode,
status,
cursor_json,
state_json,
expires_at,
last_poll_at,
created_at,
closed_at
) values (
$1, $2, $3, $4, $5, $6, $7, $8::jsonb, $9::jsonb, $10::timestamptz,
$11::timestamptz, $12::timestamptz, $13::timestamptz
)",
)
.bind(request.session.id.as_str())
.bind(request.session.workspace_id.as_str())
.bind(
request
.session
.agent_id
.as_ref()
.map(|value| value.as_str()),
)
.bind(request.session.operation_id.as_str())
.bind(serialize_enum_text(&request.session.protocol, "protocol")?)
.bind(serialize_enum_text(&request.session.mode, "mode")?)
.bind(serialize_enum_text(&request.session.status, "status")?)
.bind(request.session.cursor.clone().map(Json))
.bind(Json(request.session.state.clone()))
.bind(&request.session.expires_at)
.bind(request.session.last_poll_at.as_deref())
.bind(&request.session.created_at)
.bind(request.session.closed_at.as_deref())
.execute(&self.pool)
.await?;
Ok(request.session.clone())
}
pub async fn get_stream_session(
&self,
id: &StreamSessionId,
) -> Result<Option<StreamSession>, RegistryError> {
let row = sqlx::query(
"select
id,
workspace_id,
agent_id,
operation_id,
protocol,
mode,
status,
cursor_json,
state_json,
to_char(expires_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as expires_at,
to_char(last_poll_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as last_poll_at,
to_char(created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as created_at,
to_char(closed_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as closed_at
from stream_sessions
where id = $1",
)
.bind(id.as_str())
.fetch_optional(&self.pool)
.await?;
row.as_ref().map(map_stream_session).transpose()
}
pub async fn update_stream_session_state(
&self,
request: UpdateStreamSessionStateRequest<'_>,
) -> Result<StreamSession, RegistryError> {
validate_stream_session_transition(
request.session_id,
request.current_status,
request.next_status,
)?;
let row = sqlx::query(
"update stream_sessions
set status = $3,
cursor_json = $4::jsonb,
state_json = $5::jsonb,
expires_at = coalesce($6::timestamptz, expires_at),
last_poll_at = coalesce($7::timestamptz, last_poll_at),
closed_at = case
when $8::timestamptz is null then closed_at
else $8::timestamptz
end
where id = $1
and status = $2
returning
id,
workspace_id,
agent_id,
operation_id,
protocol,
mode,
status,
cursor_json,
state_json,
to_char(expires_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as expires_at,
to_char(last_poll_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as last_poll_at,
to_char(created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as created_at,
to_char(closed_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as closed_at",
)
.bind(request.session_id.as_str())
.bind(serialize_enum_text(&request.current_status, "status")?)
.bind(serialize_enum_text(&request.next_status, "status")?)
.bind(request.cursor.cloned().map(Json))
.bind(Json(request.state.clone()))
.bind(request.expires_at)
.bind(request.last_poll_at)
.bind(request.closed_at)
.fetch_optional(&self.pool)
.await?;
match row.as_ref() {
Some(row) => map_stream_session(row),
None => {
if self.get_stream_session(request.session_id).await?.is_some() {
Err(RegistryError::InvalidStreamSessionTransition {
session_id: request.session_id.as_str().to_owned(),
from: serialize_enum_text(&request.current_status, "status")?,
to: serialize_enum_text(&request.next_status, "status")?,
})
} else {
Err(RegistryError::StreamSessionNotFound {
session_id: request.session_id.as_str().to_owned(),
})
}
}
}
}
pub async fn close_stream_session(
&self,
id: &StreamSessionId,
now: &str,
) -> Result<(), RegistryError> {
let result = sqlx::query(
"update stream_sessions
set status = 'stopped',
last_poll_at = $2::timestamptz,
closed_at = $2::timestamptz
where id = $1
and status in ('created', 'running')",
)
.bind(id.as_str())
.bind(now)
.execute(&self.pool)
.await?;
if result.rows_affected() > 0 {
return Ok(());
}
match self.get_stream_session(id).await? {
Some(existing) => Err(RegistryError::InvalidStreamSessionTransition {
session_id: id.as_str().to_owned(),
from: serialize_enum_text(&existing.status, "status")?,
to: "stopped".to_owned(),
}),
None => Err(RegistryError::StreamSessionNotFound {
session_id: id.as_str().to_owned(),
}),
}
}
pub async fn list_stream_sessions(
&self,
filter: StreamSessionFilter<'_>,
) -> Result<Page<StreamSession>, RegistryError> {
let status = filter
.status
.as_ref()
.map(|value| serialize_enum_text(value, "status"))
.transpose()?;
let mode = filter
.mode
.as_ref()
.map(|value| serialize_enum_text(value, "mode"))
.transpose()?;
let limit = i64::from(filter.limit);
let rows = sqlx::query(
"select
id,
workspace_id,
agent_id,
operation_id,
protocol,
mode,
status,
cursor_json,
state_json,
to_char(expires_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as expires_at,
to_char(last_poll_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as last_poll_at,
to_char(created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as created_at,
to_char(closed_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as closed_at
from stream_sessions
where workspace_id = $1
and ($2::text is null or agent_id = $2)
and ($3::text is null or operation_id = $3)
and ($4::text is null or status = $4)
and ($5::text is null or mode = $5)
order by created_at desc
limit $6",
)
.bind(filter.workspace_id.as_str())
.bind(filter.agent_id.map(|value| value.as_str()))
.bind(filter.operation_id.map(|value| value.as_str()))
.bind(status)
.bind(mode)
.bind(limit)
.fetch_all(&self.pool)
.await?;
let items = rows
.iter()
.map(map_stream_session)
.collect::<Result<Vec<_>, _>>()?;
let total = items.len() as u64;
Ok(Page { items, total })
}
pub async fn delete_expired_stream_sessions(&self, now: &str) -> Result<u64, RegistryError> {
let result = sqlx::query(
"delete from stream_sessions
where expires_at <= $1::timestamptz",
)
.bind(now)
.execute(&self.pool)
.await?;
Ok(result.rows_affected())
}
pub async fn create_async_job(
&self,
request: CreateAsyncJobRequest<'_>,
) -> Result<AsyncJobHandle, RegistryError> {
sqlx::query(
"insert into async_jobs (
id,
workspace_id,
agent_id,
operation_id,
status,
progress_json,
result_json,
error_json,
expires_at,
created_at,
updated_at,
finished_at
) values (
$1, $2, $3, $4, $5, $6::jsonb, $7::jsonb, $8::jsonb, $9::timestamptz,
$10::timestamptz, $11::timestamptz, $12::timestamptz
)",
)
.bind(request.job.id.as_str())
.bind(request.job.workspace_id.as_str())
.bind(request.job.agent_id.as_ref().map(|value| value.as_str()))
.bind(request.job.operation_id.as_str())
.bind(serialize_enum_text(&request.job.status, "status")?)
.bind(Json(request.job.progress.clone()))
.bind(request.job.result.clone().map(Json))
.bind(request.job.error.clone().map(Json))
.bind(request.job.expires_at.as_deref())
.bind(&request.job.created_at)
.bind(&request.job.updated_at)
.bind(request.job.finished_at.as_deref())
.execute(&self.pool)
.await?;
Ok(request.job.clone())
}
pub async fn get_async_job(
&self,
id: &AsyncJobId,
) -> Result<Option<AsyncJobHandle>, RegistryError> {
let row = sqlx::query(
"select
id,
workspace_id,
agent_id,
operation_id,
status,
progress_json,
result_json,
error_json,
to_char(expires_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as expires_at,
to_char(created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as created_at,
to_char(updated_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as updated_at,
to_char(finished_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as finished_at
from async_jobs
where id = $1",
)
.bind(id.as_str())
.fetch_optional(&self.pool)
.await?;
row.as_ref().map(map_async_job).transpose()
}
pub async fn update_async_job_status(
&self,
request: UpdateAsyncJobStatusRequest<'_>,
) -> Result<AsyncJobHandle, RegistryError> {
validate_async_job_transition(request.job_id, request.current_status, request.next_status)?;
let row = sqlx::query(
"update async_jobs
set status = $3,
progress_json = $4::jsonb,
result_json = $5::jsonb,
error_json = $6::jsonb,
expires_at = coalesce($7::timestamptz, expires_at),
updated_at = $8::timestamptz,
finished_at = case
when $9::timestamptz is null then finished_at
else $9::timestamptz
end
where id = $1
and status = $2
returning
id,
workspace_id,
agent_id,
operation_id,
status,
progress_json,
result_json,
error_json,
to_char(expires_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as expires_at,
to_char(created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as created_at,
to_char(updated_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as updated_at,
to_char(finished_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as finished_at",
)
.bind(request.job_id.as_str())
.bind(serialize_enum_text(&request.current_status, "status")?)
.bind(serialize_enum_text(&request.next_status, "status")?)
.bind(Json(request.progress.clone()))
.bind(request.result.cloned().map(Json))
.bind(request.error.cloned().map(Json))
.bind(request.expires_at)
.bind(request.updated_at)
.bind(request.finished_at)
.fetch_optional(&self.pool)
.await?;
match row.as_ref() {
Some(row) => map_async_job(row),
None => {
if self.get_async_job(request.job_id).await?.is_some() {
Err(RegistryError::InvalidAsyncJobTransition {
job_id: request.job_id.as_str().to_owned(),
from: serialize_enum_text(&request.current_status, "status")?,
to: serialize_enum_text(&request.next_status, "status")?,
})
} else {
Err(RegistryError::AsyncJobNotFound {
job_id: request.job_id.as_str().to_owned(),
})
}
}
}
}
pub async fn cancel_async_job(&self, id: &AsyncJobId, now: &str) -> Result<(), RegistryError> {
let result = sqlx::query(
"update async_jobs
set status = 'cancelled',
updated_at = $2::timestamptz,
finished_at = $2::timestamptz
where id = $1
and status in ('created', 'running')",
)
.bind(id.as_str())
.bind(now)
.execute(&self.pool)
.await?;
if result.rows_affected() > 0 {
return Ok(());
}
match self.get_async_job(id).await? {
Some(existing) => Err(RegistryError::InvalidAsyncJobTransition {
job_id: id.as_str().to_owned(),
from: serialize_enum_text(&existing.status, "status")?,
to: "cancelled".to_owned(),
}),
None => Err(RegistryError::AsyncJobNotFound {
job_id: id.as_str().to_owned(),
}),
}
}
pub async fn list_async_jobs(
&self,
filter: AsyncJobFilter<'_>,
) -> Result<Page<AsyncJobHandle>, RegistryError> {
let status = filter
.status
.as_ref()
.map(|value| serialize_enum_text(value, "status"))
.transpose()?;
let limit = i64::from(filter.limit);
let rows = sqlx::query(
"select
id,
workspace_id,
agent_id,
operation_id,
status,
progress_json,
result_json,
error_json,
to_char(expires_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as expires_at,
to_char(created_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as created_at,
to_char(updated_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as updated_at,
to_char(finished_at at time zone 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS\"Z\"') as finished_at
from async_jobs
where workspace_id = $1
and ($2::text is null or agent_id = $2)
and ($3::text is null or operation_id = $3)
and ($4::text is null or status = $4)
order by updated_at desc
limit $5",
)
.bind(filter.workspace_id.as_str())
.bind(filter.agent_id.map(|value| value.as_str()))
.bind(filter.operation_id.map(|value| value.as_str()))
.bind(status)
.bind(limit)
.fetch_all(&self.pool)
.await?;
let items = rows
.iter()
.map(map_async_job)
.collect::<Result<Vec<_>, _>>()?;
let total = items.len() as u64;
Ok(Page { items, total })
}
pub async fn delete_expired_async_jobs(&self, now: &str) -> Result<u64, RegistryError> {
let result = sqlx::query(
"delete from async_jobs
where expires_at is not null
and expires_at <= $1::timestamptz",
)
.bind(now)
.execute(&self.pool)
.await?;
Ok(result.rows_affected())
}
pub async fn create_platform_api_key(
&self,
request: CreatePlatformApiKeyRequest<'_>,
@@ -3315,6 +3793,51 @@ fn map_secret_record(row: &PgRow) -> Result<SecretRecord, RegistryError> {
})
}
fn map_stream_session(row: &PgRow) -> Result<StreamSession, RegistryError> {
Ok(StreamSession {
id: StreamSessionId::new(row.try_get::<String, _>("id")?),
workspace_id: WorkspaceId::new(row.try_get::<String, _>("workspace_id")?),
agent_id: row
.try_get::<Option<String>, _>("agent_id")?
.map(AgentId::new),
operation_id: OperationId::new(row.try_get::<String, _>("operation_id")?),
protocol: deserialize_enum_text(&row.try_get::<String, _>("protocol")?, "protocol")?,
mode: deserialize_enum_text(&row.try_get::<String, _>("mode")?, "mode")?,
status: deserialize_enum_text(&row.try_get::<String, _>("status")?, "status")?,
cursor: row
.try_get::<Option<Json<Value>>, _>("cursor_json")?
.map(|value| value.0),
state: row.try_get::<Json<Value>, _>("state_json")?.0,
expires_at: row.try_get("expires_at")?,
last_poll_at: row.try_get("last_poll_at")?,
created_at: row.try_get("created_at")?,
closed_at: row.try_get("closed_at")?,
})
}
fn map_async_job(row: &PgRow) -> Result<AsyncJobHandle, RegistryError> {
Ok(AsyncJobHandle {
id: AsyncJobId::new(row.try_get::<String, _>("id")?),
workspace_id: WorkspaceId::new(row.try_get::<String, _>("workspace_id")?),
agent_id: row
.try_get::<Option<String>, _>("agent_id")?
.map(AgentId::new),
operation_id: OperationId::new(row.try_get::<String, _>("operation_id")?),
status: deserialize_enum_text(&row.try_get::<String, _>("status")?, "status")?,
progress: row.try_get::<Json<Value>, _>("progress_json")?.0,
result: row
.try_get::<Option<Json<Value>>, _>("result_json")?
.map(|value| value.0),
error: row
.try_get::<Option<Json<Value>>, _>("error_json")?
.map(|value| value.0),
expires_at: row.try_get("expires_at")?,
created_at: row.try_get("created_at")?,
updated_at: row.try_get("updated_at")?,
finished_at: row.try_get("finished_at")?,
})
}
fn map_invocation_log_record(row: &PgRow) -> Result<InvocationLogRecord, RegistryError> {
Ok(InvocationLogRecord {
log: InvocationLog {
@@ -3724,6 +4247,59 @@ fn from_db_version(value: i32, field: &'static str) -> Result<u32, RegistryError
})
}
fn validate_stream_session_transition(
session_id: &StreamSessionId,
from: StreamStatus,
to: StreamStatus,
) -> Result<(), RegistryError> {
let allowed = matches!(
(from, to),
(StreamStatus::Created, StreamStatus::Running)
| (StreamStatus::Running, StreamStatus::Running)
| (StreamStatus::Running, StreamStatus::Stopped)
| (StreamStatus::Running, StreamStatus::Expired)
| (StreamStatus::Running, StreamStatus::Failed)
);
if allowed {
Ok(())
} else {
Err(RegistryError::InvalidStreamSessionTransition {
session_id: session_id.as_str().to_owned(),
from: serialize_enum_text(&from, "status").unwrap_or_else(|_| "unknown".to_owned()),
to: serialize_enum_text(&to, "status").unwrap_or_else(|_| "unknown".to_owned()),
})
}
}
fn validate_async_job_transition(
job_id: &AsyncJobId,
from: JobStatus,
to: JobStatus,
) -> Result<(), RegistryError> {
let allowed = matches!(
(from, to),
(JobStatus::Created, JobStatus::Running)
| (JobStatus::Running, JobStatus::Running)
| (JobStatus::Running, JobStatus::Completed)
| (JobStatus::Running, JobStatus::Failed)
| (JobStatus::Running, JobStatus::Cancelled)
| (JobStatus::Completed, JobStatus::Expired)
| (JobStatus::Failed, JobStatus::Expired)
| (JobStatus::Cancelled, JobStatus::Expired)
);
if allowed {
Ok(())
} else {
Err(RegistryError::InvalidAsyncJobTransition {
job_id: job_id.as_str().to_owned(),
from: serialize_enum_text(&from, "status").unwrap_or_else(|_| "unknown".to_owned()),
to: serialize_enum_text(&to, "status").unwrap_or_else(|_| "unknown".to_owned()),
})
}
}
fn to_u64(value: i64, field: &'static str) -> Result<u64, RegistryError> {
u64::try_from(value).map_err(|_| RegistryError::InvalidNumericValue { field, value })
}
@@ -3746,10 +4322,11 @@ mod tests {
};
use crank_core::{
ApiKeyHeaderAuthConfig, AuthConfig, AuthKind, AuthProfile, ConfigExport, ExecutionConfig,
ExportMode, GeneratedDraft, GeneratedDraftStatus, HttpMethod, OperationId, OperationStatus,
Protocol, RestTarget, RetryPolicy, Samples, SecretRef, Target, ToolDescription,
ToolExample, WorkspaceId,
ApiKeyHeaderAuthConfig, AsyncJobHandle, AuthConfig, AuthKind, AuthProfile, ConfigExport,
ExecutionConfig, ExportMode, GeneratedDraft, GeneratedDraftStatus, HttpMethod, JobStatus,
OperationId, OperationStatus, Protocol, RestTarget, RetryPolicy, Samples, SecretRef,
StreamSession, StreamSessionId, StreamStatus, Target, ToolDescription, ToolExample,
WorkspaceId,
};
use crank_mapping::{MappingRule, MappingSet};
use crank_schema::{Schema, SchemaKind};
@@ -3759,9 +4336,11 @@ mod tests {
use crate::{
PostgresRegistry, RegistryError,
model::{
AsyncJobFilter, CreateAsyncJobRequest, CreateStreamSessionRequest,
CreateVersionRequest, CreateYamlImportJobRequest, DescriptorKind, DescriptorMetadata,
OperationSampleMetadata, PublishRequest, RegistryOperation, SampleKind,
SaveAuthProfileRequest, SaveDescriptorMetadataRequest, SaveSampleMetadataRequest,
StreamSessionFilter, UpdateAsyncJobStatusRequest, UpdateStreamSessionStateRequest,
YamlImportJobCompletion, YamlImportJobId, YamlImportJobStatus,
},
};
@@ -4003,6 +4582,265 @@ mod tests {
database.cleanup().await;
}
#[tokio::test]
async fn manages_stream_sessions_with_transitions_and_cleanup() {
let database = TestDatabase::new().await;
let registry = database.registry().await;
let operation = test_operation("op_stream_01", 1, OperationStatus::Draft);
let session = test_stream_session("stream_01", "op_stream_01", StreamStatus::Running);
registry
.create_operation(&test_workspace_id(), &operation, None)
.await
.unwrap();
registry
.create_stream_session(CreateStreamSessionRequest { session: &session })
.await
.unwrap();
let loaded = registry
.get_stream_session(&session.id)
.await
.unwrap()
.unwrap();
assert_eq!(loaded.status, StreamStatus::Running);
let closable_session =
test_stream_session("stream_closable", "op_stream_01", StreamStatus::Running);
registry
.create_stream_session(CreateStreamSessionRequest {
session: &closable_session,
})
.await
.unwrap();
registry
.close_stream_session(&closable_session.id, "2026-04-06T12:00:30Z")
.await
.unwrap();
let closed = registry
.get_stream_session(&closable_session.id)
.await
.unwrap()
.unwrap();
assert_eq!(closed.status, StreamStatus::Stopped);
let updated = registry
.update_stream_session_state(UpdateStreamSessionStateRequest {
session_id: &session.id,
current_status: StreamStatus::Running,
next_status: StreamStatus::Failed,
cursor: Some(&json!({"cursor":"next"})),
state: &json!({"phase":"failed"}),
expires_at: Some("2026-04-06T12:10:00Z"),
last_poll_at: Some("2026-04-06T12:01:00Z"),
closed_at: Some("2026-04-06T12:01:00Z"),
})
.await
.unwrap();
assert_eq!(updated.status, StreamStatus::Failed);
assert_eq!(updated.closed_at.as_deref(), Some("2026-04-06T12:01:00Z"));
let page = registry
.list_stream_sessions(StreamSessionFilter {
workspace_id: &test_workspace_id(),
agent_id: None,
operation_id: None,
status: Some(StreamStatus::Failed),
mode: Some(crank_core::ExecutionMode::Session),
limit: 10,
})
.await
.unwrap();
assert_eq!(page.total, 1);
assert_eq!(page.items[0].id, session.id);
let expired_session =
test_stream_session("stream_expired", "op_stream_01", StreamStatus::Running);
registry
.create_stream_session(CreateStreamSessionRequest {
session: &StreamSession {
id: expired_session.id.clone(),
expires_at: "2026-04-06T11:00:00Z".to_owned(),
..expired_session
},
})
.await
.unwrap();
let deleted = registry
.delete_expired_stream_sessions("2026-04-06T12:00:00Z")
.await
.unwrap();
assert_eq!(deleted, 1);
database.cleanup().await;
}
#[tokio::test]
async fn rejects_invalid_stream_session_transition() {
let database = TestDatabase::new().await;
let registry = database.registry().await;
let operation = test_operation("op_stream_02", 1, OperationStatus::Draft);
let session = test_stream_session("stream_invalid", "op_stream_02", StreamStatus::Stopped);
registry
.create_operation(&test_workspace_id(), &operation, None)
.await
.unwrap();
registry
.create_stream_session(CreateStreamSessionRequest { session: &session })
.await
.unwrap();
let error = registry
.update_stream_session_state(UpdateStreamSessionStateRequest {
session_id: &session.id,
current_status: StreamStatus::Stopped,
next_status: StreamStatus::Running,
cursor: None,
state: &json!({"phase":"resume"}),
expires_at: None,
last_poll_at: None,
closed_at: None,
})
.await
.unwrap_err();
assert!(matches!(
error,
RegistryError::InvalidStreamSessionTransition { .. }
));
database.cleanup().await;
}
#[tokio::test]
async fn manages_async_jobs_with_transitions_and_cleanup() {
let database = TestDatabase::new().await;
let registry = database.registry().await;
let operation = test_operation("op_job_01", 1, OperationStatus::Draft);
let job = test_async_job("job_01", "op_job_01", JobStatus::Running);
registry
.create_operation(&test_workspace_id(), &operation, None)
.await
.unwrap();
registry
.create_async_job(CreateAsyncJobRequest { job: &job })
.await
.unwrap();
let cancellable_job = test_async_job("job_cancel", "op_job_01", JobStatus::Running);
registry
.create_async_job(CreateAsyncJobRequest {
job: &cancellable_job,
})
.await
.unwrap();
registry
.cancel_async_job(&cancellable_job.id, "2026-04-06T12:00:30Z")
.await
.unwrap();
let cancelled = registry
.get_async_job(&cancellable_job.id)
.await
.unwrap()
.unwrap();
assert_eq!(cancelled.status, JobStatus::Cancelled);
let updated = registry
.update_async_job_status(UpdateAsyncJobStatusRequest {
job_id: &job.id,
current_status: JobStatus::Running,
next_status: JobStatus::Completed,
progress: &json!({"percent":100}),
result: Some(&json!({"ok":true})),
error: None,
expires_at: Some("2026-04-06T12:15:00Z"),
updated_at: "2026-04-06T12:01:00Z",
finished_at: Some("2026-04-06T12:01:00Z"),
})
.await
.unwrap();
assert_eq!(updated.status, JobStatus::Completed);
assert_eq!(updated.result, Some(json!({"ok":true})));
let loaded = registry.get_async_job(&job.id).await.unwrap().unwrap();
assert_eq!(loaded.status, JobStatus::Completed);
let page = registry
.list_async_jobs(AsyncJobFilter {
workspace_id: &test_workspace_id(),
agent_id: None,
operation_id: None,
status: Some(JobStatus::Completed),
limit: 10,
})
.await
.unwrap();
assert_eq!(page.total, 1);
let expired_job = AsyncJobHandle {
id: crank_core::AsyncJobId::new("job_expired"),
expires_at: Some("2026-04-06T11:00:00Z".to_owned()),
..test_async_job("job_expired", "op_job_01", JobStatus::Cancelled)
};
registry
.create_async_job(CreateAsyncJobRequest { job: &expired_job })
.await
.unwrap();
let deleted = registry
.delete_expired_async_jobs("2026-04-06T12:00:00Z")
.await
.unwrap();
assert_eq!(deleted, 1);
database.cleanup().await;
}
#[tokio::test]
async fn rejects_invalid_async_job_transition() {
let database = TestDatabase::new().await;
let registry = database.registry().await;
let operation = test_operation("op_job_02", 1, OperationStatus::Draft);
let job = test_async_job("job_invalid", "op_job_02", JobStatus::Completed);
registry
.create_operation(&test_workspace_id(), &operation, None)
.await
.unwrap();
registry
.create_async_job(CreateAsyncJobRequest { job: &job })
.await
.unwrap();
let error = registry
.update_async_job_status(UpdateAsyncJobStatusRequest {
job_id: &job.id,
current_status: JobStatus::Completed,
next_status: JobStatus::Running,
progress: &json!({"percent":50}),
result: None,
error: None,
expires_at: None,
updated_at: "2026-04-06T12:01:00Z",
finished_at: None,
})
.await
.unwrap_err();
assert!(matches!(
error,
RegistryError::InvalidAsyncJobTransition { .. }
));
database.cleanup().await;
}
struct TestDatabase {
admin_pool: PgPool,
database_url: String,
@@ -4182,4 +5020,39 @@ mod tests {
published_at: None,
}
}
fn test_stream_session(id: &str, operation_id: &str, status: StreamStatus) -> StreamSession {
StreamSession {
id: StreamSessionId::new(id),
workspace_id: test_workspace_id(),
agent_id: None,
operation_id: OperationId::new(operation_id),
protocol: Protocol::Rest,
mode: crank_core::ExecutionMode::Session,
status,
cursor: Some(json!({"cursor":"initial"})),
state: json!({"phase":"running"}),
expires_at: "2026-04-06T12:05:00Z".to_owned(),
last_poll_at: None,
created_at: "2026-04-06T12:00:00Z".to_owned(),
closed_at: None,
}
}
fn test_async_job(id: &str, operation_id: &str, status: JobStatus) -> AsyncJobHandle {
AsyncJobHandle {
id: crank_core::AsyncJobId::new(id),
workspace_id: test_workspace_id(),
agent_id: None,
operation_id: OperationId::new(operation_id),
status,
progress: json!({"percent": 25}),
result: None,
error: None,
expires_at: Some("2026-04-06T12:05:00Z".to_owned()),
created_at: "2026-04-06T12:00:00Z".to_owned(),
updated_at: "2026-04-06T12:00:00Z".to_owned(),
finished_at: None,
}
}
}