feat: add streaming state storage
This commit is contained in:
@@ -2,20 +2,19 @@
|
|||||||
|
|
||||||
## Current
|
## Current
|
||||||
|
|
||||||
### `feat/streaming-core-model`
|
### `feat/stream-session-store`
|
||||||
|
|
||||||
Status: completed
|
Status: completed
|
||||||
|
|
||||||
DoD:
|
DoD:
|
||||||
- core domain types exist for window/session/job execution
|
- persistent tables exist for `stream_sessions` and `async_jobs`
|
||||||
- `ExecutionConfig` can carry `streaming`
|
- registry supports create/get/update/close/list/cleanup for both stores
|
||||||
- protocol-aware validation follows documented capability matrix
|
- invalid state transitions are rejected explicitly
|
||||||
- streaming types are exported from `crank-core`
|
- migration and storage tests cover lifecycle and cleanup
|
||||||
- core tests cover serde and validation rules
|
|
||||||
|
|
||||||
## Next
|
## Next
|
||||||
|
|
||||||
- `feat/stream-session-store`
|
- `feat/runtime-window-mode`
|
||||||
|
|
||||||
## Backlog
|
## Backlog
|
||||||
|
|
||||||
|
|||||||
@@ -142,6 +142,12 @@ impl From<RegistryError> for ApiError {
|
|||||||
RegistryError::SecretNotFound { secret_id } => {
|
RegistryError::SecretNotFound { secret_id } => {
|
||||||
Self::not_found(format!("secret {secret_id} was not found"))
|
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 } => {
|
RegistryError::InvocationLogNotFound { log_id } => {
|
||||||
Self::not_found(format!("invocation log {log_id} was not found"))
|
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(
|
RegistryError::SecretNameAlreadyExists { workspace_id, name } => Self::conflict(
|
||||||
format!("secret with name {name} already exists in workspace {workspace_id}"),
|
format!("secret with name {name} already exists in workspace {workspace_id}"),
|
||||||
),
|
),
|
||||||
|
RegistryError::InvalidStreamSessionTransition { .. }
|
||||||
|
| RegistryError::InvalidAsyncJobTransition { .. } => Self::conflict(value.to_string()),
|
||||||
RegistryError::UserEmailAlreadyExists { email } => {
|
RegistryError::UserEmailAlreadyExists { email } => {
|
||||||
Self::conflict(format!("user with email {email} already exists"))
|
Self::conflict(format!("user with email {email} already exists"))
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -25,6 +25,22 @@ pub enum RegistryError {
|
|||||||
PlatformApiKeyNotFound { key_id: String },
|
PlatformApiKeyNotFound { key_id: String },
|
||||||
#[error("secret {secret_id} was not found")]
|
#[error("secret {secret_id} was not found")]
|
||||||
SecretNotFound { secret_id: String },
|
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}")]
|
#[error("secret with name {name} already exists in workspace {workspace_id}")]
|
||||||
SecretNameAlreadyExists { workspace_id: String, name: String },
|
SecretNameAlreadyExists { workspace_id: String, name: String },
|
||||||
#[error("invocation log {log_id} was not found")]
|
#[error("invocation log {log_id} was not found")]
|
||||||
|
|||||||
@@ -5,17 +5,20 @@ mod postgres;
|
|||||||
|
|
||||||
pub use error::RegistryError;
|
pub use error::RegistryError;
|
||||||
pub use model::{
|
pub use model::{
|
||||||
AgentSummary, AgentVersionRecord, AuthUserRecord, CreateAgentRequest, CreateInvitationRequest,
|
AgentSummary, AgentVersionRecord, AsyncJobFilter, AsyncJobRecord, AuthUserRecord,
|
||||||
CreateInvocationLogRequest, CreatePlatformApiKeyRequest, CreateSecretRequest,
|
CreateAgentRequest, CreateAsyncJobRequest, CreateInvitationRequest, CreateInvocationLogRequest,
|
||||||
|
CreatePlatformApiKeyRequest, CreateSecretRequest, CreateStreamSessionRequest,
|
||||||
CreateVersionRequest, CreateWorkspaceRequest, CreateYamlImportJobRequest, DescriptorKind,
|
CreateVersionRequest, CreateWorkspaceRequest, CreateYamlImportJobRequest, DescriptorKind,
|
||||||
DescriptorMetadata, InvitationRecord, InvocationLogRecord, ListInvocationLogsQuery,
|
DescriptorMetadata, InvitationRecord, InvocationLogRecord, ListInvocationLogsQuery,
|
||||||
MembershipRecord, OperationAgentRef, OperationSampleMetadata, OperationSummary,
|
MembershipRecord, OperationAgentRef, OperationSampleMetadata, OperationSummary,
|
||||||
OperationUsageSummary, OperationVersionRecord, PlatformApiKeyRecord, PublishAgentRequest,
|
OperationUsageSummary, OperationVersionRecord, Page, PlatformApiKeyRecord, PublishAgentRequest,
|
||||||
PublishRequest, PublishedAgentTool, RegistryOperation, RotateSecretRequest, SampleKind,
|
PublishRequest, PublishedAgentTool, RegistryOperation, RotateSecretRequest, SampleKind,
|
||||||
SaveAgentBindingsRequest, SaveAuthProfileRequest, SaveDescriptorMetadataRequest,
|
SaveAgentBindingsRequest, SaveAuthProfileRequest, SaveDescriptorMetadataRequest,
|
||||||
SaveSampleMetadataRequest, SecretRecord, SecretVersionRecord, SessionRecord,
|
SaveSampleMetadataRequest, SecretRecord, SecretVersionRecord, SessionRecord,
|
||||||
UpdateWorkspaceRequest, UsageAgentBreakdown, UsageBucket, UsageOperationBreakdown, UsageQuery,
|
StreamSessionFilter, StreamSessionRecord, UpdateAsyncJobStatusRequest,
|
||||||
UsageRollupRecord, UsageSummary, UsageTimelinePoint, WorkspaceMembershipRecord,
|
UpdateStreamSessionStateRequest, UpdateWorkspaceRequest, UsageAgentBreakdown, UsageBucket,
|
||||||
WorkspaceRecord, YamlImportJob, YamlImportJobCompletion, YamlImportJobId, YamlImportJobStatus,
|
UsageOperationBreakdown, UsageQuery, UsageRollupRecord, UsageSummary, UsageTimelinePoint,
|
||||||
|
WorkspaceMembershipRecord, WorkspaceRecord, YamlImportJob, YamlImportJobCompletion,
|
||||||
|
YamlImportJobId, YamlImportJobStatus,
|
||||||
};
|
};
|
||||||
pub use postgres::PostgresRegistry;
|
pub use postgres::PostgresRegistry;
|
||||||
|
|||||||
@@ -513,5 +513,72 @@ pub async fn apply_postgres(pool: &PgPool) -> Result<(), sqlx::Error> {
|
|||||||
.execute(pool)
|
.execute(pool)
|
||||||
.await?;
|
.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(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,8 +1,10 @@
|
|||||||
use crank_core::{
|
use crank_core::{
|
||||||
Agent, AgentId, AgentOperationBinding, AgentStatus, AgentVersion, AuthProfile, DescriptorId,
|
Agent, AgentId, AgentOperationBinding, AgentStatus, AgentVersion, AsyncJobHandle, AsyncJobId,
|
||||||
ExportMode, InvitationToken, InvocationLevel, InvocationLog, InvocationSource, MembershipRole,
|
AuthProfile, DescriptorId, ExecutionMode, ExportMode, InvitationToken, InvocationLevel,
|
||||||
Operation, OperationId, OperationStatus, PlatformApiKey, Protocol, SampleId, Secret, SecretId,
|
InvocationLog, InvocationSource, JobStatus, MembershipRole, Operation, OperationId,
|
||||||
SecretVersion, UsagePeriod, UsageRollup, User, UserSessionId, Workspace, WorkspaceId,
|
OperationStatus, PlatformApiKey, Protocol, SampleId, Secret, SecretId, SecretVersion,
|
||||||
|
StreamSession, StreamSessionId, StreamStatus, UsagePeriod, UsageRollup, User, UserSessionId,
|
||||||
|
Workspace, WorkspaceId,
|
||||||
};
|
};
|
||||||
use crank_mapping::MappingSet;
|
use crank_mapping::MappingSet;
|
||||||
use crank_schema::Schema;
|
use crank_schema::Schema;
|
||||||
@@ -42,6 +44,12 @@ define_registry_id!(YamlImportJobId);
|
|||||||
|
|
||||||
pub type RegistryOperation = Operation<Schema, MappingSet>;
|
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)]
|
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||||
pub struct WorkspaceRecord {
|
pub struct WorkspaceRecord {
|
||||||
pub workspace: Workspace,
|
pub workspace: Workspace,
|
||||||
@@ -95,6 +103,16 @@ pub struct SecretVersionRecord {
|
|||||||
pub secret_version: SecretVersion,
|
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)]
|
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||||
pub struct InvocationLogRecord {
|
pub struct InvocationLogRecord {
|
||||||
pub log: InvocationLog,
|
pub log: InvocationLog,
|
||||||
@@ -432,6 +450,60 @@ pub struct RotateSecretRequest<'a> {
|
|||||||
pub created_by: Option<&'a crank_core::UserId>,
|
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)]
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||||
pub struct SaveSampleMetadataRequest<'a> {
|
pub struct SaveSampleMetadataRequest<'a> {
|
||||||
pub sample: &'a OperationSampleMetadata,
|
pub sample: &'a OperationSampleMetadata,
|
||||||
|
|||||||
@@ -1,8 +1,9 @@
|
|||||||
use crank_core::{
|
use crank_core::{
|
||||||
AgentId, AgentOperationBinding, AgentStatus, AgentVersion, AuthProfile, GraphqlOperationType,
|
AgentId, AgentOperationBinding, AgentStatus, AgentVersion, AsyncJobHandle, AsyncJobId,
|
||||||
HttpMethod, InvitationId, InvitationToken, InvocationLog, InvocationLogId, MembershipRole,
|
AuthProfile, GraphqlOperationType, HttpMethod, InvitationId, InvitationToken, InvocationLog,
|
||||||
OperationId, OperationStatus, PlatformApiKey, PlatformApiKeyId, PlatformApiKeyStatus, Secret,
|
InvocationLogId, JobStatus, MembershipRole, OperationId, OperationStatus, PlatformApiKey,
|
||||||
SecretId, SecretVersion, Target, UsageRollup, User, UserId, UserSessionId, Workspace,
|
PlatformApiKeyId, PlatformApiKeyStatus, Secret, SecretId, SecretVersion, StreamSession,
|
||||||
|
StreamSessionId, StreamStatus, Target, UsageRollup, User, UserId, UserSessionId, Workspace,
|
||||||
WorkspaceId,
|
WorkspaceId,
|
||||||
};
|
};
|
||||||
use serde::{Serialize, de::DeserializeOwned};
|
use serde::{Serialize, de::DeserializeOwned};
|
||||||
@@ -17,16 +18,18 @@ use crate::{
|
|||||||
error::RegistryError,
|
error::RegistryError,
|
||||||
migrations,
|
migrations,
|
||||||
model::{
|
model::{
|
||||||
AgentSummary, AgentVersionRecord, AuthUserRecord, CreateAgentRequest,
|
AgentSummary, AgentVersionRecord, AsyncJobFilter, AuthUserRecord, CreateAgentRequest,
|
||||||
CreateInvitationRequest, CreateInvocationLogRequest, CreatePlatformApiKeyRequest,
|
CreateAsyncJobRequest, CreateInvitationRequest, CreateInvocationLogRequest,
|
||||||
CreateSecretRequest, CreateVersionRequest, CreateWorkspaceRequest,
|
CreatePlatformApiKeyRequest, CreateSecretRequest, CreateStreamSessionRequest,
|
||||||
CreateYamlImportJobRequest, DescriptorMetadata, InvitationRecord, InvocationLogRecord,
|
CreateVersionRequest, CreateWorkspaceRequest, CreateYamlImportJobRequest,
|
||||||
ListInvocationLogsQuery, MembershipRecord, OperationAgentRef, OperationSampleMetadata,
|
DescriptorMetadata, InvitationRecord, InvocationLogRecord, ListInvocationLogsQuery,
|
||||||
OperationSummary, OperationUsageSummary, OperationVersionRecord, PlatformApiKeyRecord,
|
MembershipRecord, OperationAgentRef, OperationSampleMetadata, OperationSummary,
|
||||||
|
OperationUsageSummary, OperationVersionRecord, Page, PlatformApiKeyRecord,
|
||||||
PublishAgentRequest, PublishRequest, PublishedAgentTool, RegistryOperation,
|
PublishAgentRequest, PublishRequest, PublishedAgentTool, RegistryOperation,
|
||||||
RotateSecretRequest, SaveAgentBindingsRequest, SaveAuthProfileRequest,
|
RotateSecretRequest, SaveAgentBindingsRequest, SaveAuthProfileRequest,
|
||||||
SaveDescriptorMetadataRequest, SaveSampleMetadataRequest, SecretRecord,
|
SaveDescriptorMetadataRequest, SaveSampleMetadataRequest, SecretRecord,
|
||||||
SecretVersionRecord, SessionRecord, UpdateWorkspaceRequest, UsageAgentBreakdown,
|
SecretVersionRecord, SessionRecord, StreamSessionFilter, UpdateAsyncJobStatusRequest,
|
||||||
|
UpdateStreamSessionStateRequest, UpdateWorkspaceRequest, UsageAgentBreakdown,
|
||||||
UsageOperationBreakdown, UsageQuery, UsageRollupRecord, UsageSummary, UsageTimelinePoint,
|
UsageOperationBreakdown, UsageQuery, UsageRollupRecord, UsageSummary, UsageTimelinePoint,
|
||||||
WorkspaceMembershipRecord, WorkspaceRecord, YamlImportJob, YamlImportJobCompletion,
|
WorkspaceMembershipRecord, WorkspaceRecord, YamlImportJob, YamlImportJobCompletion,
|
||||||
YamlImportJobId, YamlImportJobStatus,
|
YamlImportJobId, YamlImportJobStatus,
|
||||||
@@ -860,6 +863,481 @@ impl PostgresRegistry {
|
|||||||
Ok(())
|
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(
|
pub async fn create_platform_api_key(
|
||||||
&self,
|
&self,
|
||||||
request: CreatePlatformApiKeyRequest<'_>,
|
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> {
|
fn map_invocation_log_record(row: &PgRow) -> Result<InvocationLogRecord, RegistryError> {
|
||||||
Ok(InvocationLogRecord {
|
Ok(InvocationLogRecord {
|
||||||
log: InvocationLog {
|
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> {
|
fn to_u64(value: i64, field: &'static str) -> Result<u64, RegistryError> {
|
||||||
u64::try_from(value).map_err(|_| RegistryError::InvalidNumericValue { field, value })
|
u64::try_from(value).map_err(|_| RegistryError::InvalidNumericValue { field, value })
|
||||||
}
|
}
|
||||||
@@ -3746,10 +4322,11 @@ mod tests {
|
|||||||
};
|
};
|
||||||
|
|
||||||
use crank_core::{
|
use crank_core::{
|
||||||
ApiKeyHeaderAuthConfig, AuthConfig, AuthKind, AuthProfile, ConfigExport, ExecutionConfig,
|
ApiKeyHeaderAuthConfig, AsyncJobHandle, AuthConfig, AuthKind, AuthProfile, ConfigExport,
|
||||||
ExportMode, GeneratedDraft, GeneratedDraftStatus, HttpMethod, OperationId, OperationStatus,
|
ExecutionConfig, ExportMode, GeneratedDraft, GeneratedDraftStatus, HttpMethod, JobStatus,
|
||||||
Protocol, RestTarget, RetryPolicy, Samples, SecretRef, Target, ToolDescription,
|
OperationId, OperationStatus, Protocol, RestTarget, RetryPolicy, Samples, SecretRef,
|
||||||
ToolExample, WorkspaceId,
|
StreamSession, StreamSessionId, StreamStatus, Target, ToolDescription, ToolExample,
|
||||||
|
WorkspaceId,
|
||||||
};
|
};
|
||||||
use crank_mapping::{MappingRule, MappingSet};
|
use crank_mapping::{MappingRule, MappingSet};
|
||||||
use crank_schema::{Schema, SchemaKind};
|
use crank_schema::{Schema, SchemaKind};
|
||||||
@@ -3759,9 +4336,11 @@ mod tests {
|
|||||||
use crate::{
|
use crate::{
|
||||||
PostgresRegistry, RegistryError,
|
PostgresRegistry, RegistryError,
|
||||||
model::{
|
model::{
|
||||||
|
AsyncJobFilter, CreateAsyncJobRequest, CreateStreamSessionRequest,
|
||||||
CreateVersionRequest, CreateYamlImportJobRequest, DescriptorKind, DescriptorMetadata,
|
CreateVersionRequest, CreateYamlImportJobRequest, DescriptorKind, DescriptorMetadata,
|
||||||
OperationSampleMetadata, PublishRequest, RegistryOperation, SampleKind,
|
OperationSampleMetadata, PublishRequest, RegistryOperation, SampleKind,
|
||||||
SaveAuthProfileRequest, SaveDescriptorMetadataRequest, SaveSampleMetadataRequest,
|
SaveAuthProfileRequest, SaveDescriptorMetadataRequest, SaveSampleMetadataRequest,
|
||||||
|
StreamSessionFilter, UpdateAsyncJobStatusRequest, UpdateStreamSessionStateRequest,
|
||||||
YamlImportJobCompletion, YamlImportJobId, YamlImportJobStatus,
|
YamlImportJobCompletion, YamlImportJobId, YamlImportJobStatus,
|
||||||
},
|
},
|
||||||
};
|
};
|
||||||
@@ -4003,6 +4582,265 @@ mod tests {
|
|||||||
database.cleanup().await;
|
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 {
|
struct TestDatabase {
|
||||||
admin_pool: PgPool,
|
admin_pool: PgPool,
|
||||||
database_url: String,
|
database_url: String,
|
||||||
@@ -4182,4 +5020,39 @@ mod tests {
|
|||||||
published_at: None,
|
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