diff --git a/TASKS.md b/TASKS.md index 603ebb3..0d0620b 100644 --- a/TASKS.md +++ b/TASKS.md @@ -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 diff --git a/apps/admin-api/src/error.rs b/apps/admin-api/src/error.rs index 3ffcde9..f9b288e 100644 --- a/apps/admin-api/src/error.rs +++ b/apps/admin-api/src/error.rs @@ -142,6 +142,12 @@ impl From 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 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")) } diff --git a/crates/crank-registry/src/error.rs b/crates/crank-registry/src/error.rs index 2df780f..2db2998 100644 --- a/crates/crank-registry/src/error.rs +++ b/crates/crank-registry/src/error.rs @@ -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")] diff --git a/crates/crank-registry/src/lib.rs b/crates/crank-registry/src/lib.rs index 773ecc3..6e9d592 100644 --- a/crates/crank-registry/src/lib.rs +++ b/crates/crank-registry/src/lib.rs @@ -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; diff --git a/crates/crank-registry/src/migrations.rs b/crates/crank-registry/src/migrations.rs index d26b8b2..a590b93 100644 --- a/crates/crank-registry/src/migrations.rs +++ b/crates/crank-registry/src/migrations.rs @@ -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(()) } diff --git a/crates/crank-registry/src/model.rs b/crates/crank-registry/src/model.rs index d578871..a6146cf 100644 --- a/crates/crank-registry/src/model.rs +++ b/crates/crank-registry/src/model.rs @@ -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; +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct Page { + pub items: Vec, + 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, + pub mode: Option, + 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, + pub limit: u32, +} + #[derive(Clone, Debug, PartialEq, Eq)] pub struct SaveSampleMetadataRequest<'a> { pub sample: &'a OperationSampleMetadata, diff --git a/crates/crank-registry/src/postgres.rs b/crates/crank-registry/src/postgres.rs index 1344383..00c21e3 100644 --- a/crates/crank-registry/src/postgres.rs +++ b/crates/crank-registry/src/postgres.rs @@ -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 { + 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, 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 { + 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, 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::, _>>()?; + let total = items.len() as u64; + + Ok(Page { items, total }) + } + + pub async fn delete_expired_stream_sessions(&self, now: &str) -> Result { + 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 { + 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, 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 { + 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, 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::, _>>()?; + let total = items.len() as u64; + + Ok(Page { items, total }) + } + + pub async fn delete_expired_async_jobs(&self, now: &str) -> Result { + 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 { }) } +fn map_stream_session(row: &PgRow) -> Result { + Ok(StreamSession { + id: StreamSessionId::new(row.try_get::("id")?), + workspace_id: WorkspaceId::new(row.try_get::("workspace_id")?), + agent_id: row + .try_get::, _>("agent_id")? + .map(AgentId::new), + operation_id: OperationId::new(row.try_get::("operation_id")?), + protocol: deserialize_enum_text(&row.try_get::("protocol")?, "protocol")?, + mode: deserialize_enum_text(&row.try_get::("mode")?, "mode")?, + status: deserialize_enum_text(&row.try_get::("status")?, "status")?, + cursor: row + .try_get::>, _>("cursor_json")? + .map(|value| value.0), + state: row.try_get::, _>("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 { + Ok(AsyncJobHandle { + id: AsyncJobId::new(row.try_get::("id")?), + workspace_id: WorkspaceId::new(row.try_get::("workspace_id")?), + agent_id: row + .try_get::, _>("agent_id")? + .map(AgentId::new), + operation_id: OperationId::new(row.try_get::("operation_id")?), + status: deserialize_enum_text(&row.try_get::("status")?, "status")?, + progress: row.try_get::, _>("progress_json")?.0, + result: row + .try_get::>, _>("result_json")? + .map(|value| value.0), + error: row + .try_get::>, _>("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 { Ok(InvocationLogRecord { log: InvocationLog { @@ -3724,6 +4247,59 @@ fn from_db_version(value: i32, field: &'static str) -> Result 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::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, + } + } }