feat: add streaming state storage
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
@@ -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"))
|
||||
}
|
||||
|
||||
@@ -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")]
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user