Polish community UI copy and cleanup
This commit is contained in:
@@ -7,10 +7,9 @@ mod postgres;
|
||||
pub use error::RegistryError;
|
||||
pub use ext::{ExtensionMigration, RegistryExtension, apply_extension_migrations};
|
||||
pub use model::{
|
||||
AgentSummary, AgentVersionRecord, AsyncJobFilter, AsyncJobRecord, AuthUserRecord,
|
||||
CreateAgentDraftVersionRequest, CreateAgentRequest, CreateAsyncJobRequest,
|
||||
CreateInvitationRequest, CreateInvocationLogRequest, CreatePlatformApiKeyRequest,
|
||||
CreateSecretRequest, CreateStreamSessionRequest, CreateVersionRequest, CreateWorkspaceRequest,
|
||||
AgentSummary, AgentVersionRecord, AuthUserRecord, CreateAgentDraftVersionRequest,
|
||||
CreateAgentRequest, CreateInvitationRequest, CreateInvocationLogRequest,
|
||||
CreatePlatformApiKeyRequest, CreateSecretRequest, CreateVersionRequest, CreateWorkspaceRequest,
|
||||
CreateYamlImportJobRequest, DescriptorKind, DescriptorMetadata, InvitationRecord,
|
||||
InvocationLogRecord, ListInvocationLogsQuery, MembershipRecord, OperationAgentRef,
|
||||
OperationSampleMetadata, OperationSummary, OperationUsageSummary, OperationVersionRecord, Page,
|
||||
@@ -18,10 +17,9 @@ pub use model::{
|
||||
RegistryOperation, RotateSecretRequest, SampleKind, SaveAgentBindingsRequest,
|
||||
SaveAuthProfileRequest, SaveDescriptorMetadataRequest, SaveSampleMetadataRequest,
|
||||
SaveWorkspaceUpstreamRequest, SecretRecord, SecretVersionRecord, SessionRecord,
|
||||
StreamSessionFilter, StreamSessionRecord, UpdateAsyncJobStatusRequest,
|
||||
UpdateStreamSessionStateRequest, UpdateWorkspaceRequest, UsageAgentBreakdown, UsageBucket,
|
||||
UsageOperationBreakdown, UsageQuery, UsageRollupRecord, UsageSummary, UsageTimelinePoint,
|
||||
WorkspaceMembershipRecord, WorkspaceRecord, WorkspaceUpstream, WorkspaceUpstreamId,
|
||||
YamlImportJob, YamlImportJobCompletion, YamlImportJobId, YamlImportJobStatus,
|
||||
UpdateWorkspaceRequest, UsageAgentBreakdown, UsageBucket, UsageOperationBreakdown, UsageQuery,
|
||||
UsageRollupRecord, UsageSummary, UsageTimelinePoint, WorkspaceMembershipRecord,
|
||||
WorkspaceRecord, WorkspaceUpstream, WorkspaceUpstreamId, YamlImportJob,
|
||||
YamlImportJobCompletion, YamlImportJobId, YamlImportJobStatus,
|
||||
};
|
||||
pub use postgres::{PostgresPoolConfig, PostgresPoolConfigError, PostgresRegistry};
|
||||
|
||||
@@ -583,77 +583,5 @@ 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,
|
||||
last_poll_at timestamptz null,
|
||||
created_at timestamptz not null,
|
||||
updated_at timestamptz not null,
|
||||
finished_at timestamptz null
|
||||
)",
|
||||
)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
|
||||
query("alter table async_jobs add column if not exists last_poll_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,10 +1,9 @@
|
||||
use crank_core::{
|
||||
Agent, AgentId, AgentOperationBinding, AgentStatus, AgentVersion, AsyncJobHandle, AsyncJobId,
|
||||
AuthProfile, DescriptorId, ExecutionMode, ExportMode, InvitationToken, InvocationLevel,
|
||||
InvocationLog, InvocationSource, JobStatus, MembershipRole, Operation, OperationId,
|
||||
OperationSecurityLevel, OperationStatus, PlatformApiKey, Protocol, SampleId, Secret, SecretId,
|
||||
SecretVersion, StreamSession, StreamSessionId, StreamStatus, UsagePeriod, UsageRollup, User,
|
||||
UserSessionId, Workspace, WorkspaceId,
|
||||
Agent, AgentId, AgentOperationBinding, AgentStatus, AgentVersion, AuthProfile, DescriptorId,
|
||||
ExportMode, InvitationToken, InvocationLevel, InvocationLog, InvocationSource, MembershipRole,
|
||||
Operation, OperationId, OperationSecurityLevel, OperationStatus, PlatformApiKey, Protocol,
|
||||
SampleId, Secret, SecretId, SecretVersion, UsagePeriod, UsageRollup, User, UserSessionId,
|
||||
Workspace, WorkspaceId,
|
||||
};
|
||||
use crank_mapping::MappingSet;
|
||||
use crank_schema::Schema;
|
||||
@@ -117,16 +116,6 @@ pub struct WorkspaceUpstream {
|
||||
pub updated_at: OffsetDateTime,
|
||||
}
|
||||
|
||||
#[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,
|
||||
@@ -487,60 +476,6 @@ 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 OffsetDateTime>,
|
||||
pub last_poll_at: Option<&'a OffsetDateTime>,
|
||||
pub closed_at: Option<&'a OffsetDateTime>,
|
||||
}
|
||||
|
||||
#[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 OffsetDateTime>,
|
||||
pub updated_at: &'a OffsetDateTime,
|
||||
pub finished_at: Option<&'a OffsetDateTime>,
|
||||
}
|
||||
|
||||
#[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,
|
||||
|
||||
@@ -4,18 +4,15 @@ mod auth;
|
||||
mod observability;
|
||||
mod operation;
|
||||
mod secret;
|
||||
mod stream;
|
||||
mod upstream;
|
||||
mod workspace;
|
||||
mod yaml_import;
|
||||
|
||||
use crank_core::{
|
||||
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,
|
||||
AgentId, AgentOperationBinding, AgentStatus, AgentVersion, AuthProfile, HttpMethod,
|
||||
InvitationId, InvitationToken, InvocationLog, InvocationLogId, MembershipRole, OperationId,
|
||||
OperationStatus, PlatformApiKey, PlatformApiKeyId, PlatformApiKeyStatus, Secret, SecretId,
|
||||
SecretVersion, Target, UsageRollup, User, UserId, UserSessionId, Workspace, WorkspaceId,
|
||||
};
|
||||
use serde::{Serialize, de::DeserializeOwned};
|
||||
use serde_json::Value;
|
||||
@@ -32,18 +29,16 @@ use crate::{
|
||||
error::RegistryError,
|
||||
migrations,
|
||||
model::{
|
||||
AgentSummary, AgentVersionRecord, AsyncJobFilter, AuthUserRecord,
|
||||
CreateAgentDraftVersionRequest, CreateAgentRequest, CreateAsyncJobRequest,
|
||||
CreateInvitationRequest, CreateInvocationLogRequest, CreatePlatformApiKeyRequest,
|
||||
CreateSecretRequest, CreateStreamSessionRequest, CreateVersionRequest,
|
||||
AgentSummary, AgentVersionRecord, AuthUserRecord, CreateAgentDraftVersionRequest,
|
||||
CreateAgentRequest, CreateInvitationRequest, CreateInvocationLogRequest,
|
||||
CreatePlatformApiKeyRequest, CreateSecretRequest, CreateVersionRequest,
|
||||
CreateWorkspaceRequest, CreateYamlImportJobRequest, DescriptorMetadata, InvitationRecord,
|
||||
InvocationLogRecord, ListInvocationLogsQuery, MembershipRecord, OperationAgentRef,
|
||||
OperationSampleMetadata, OperationSummary, OperationUsageSummary, OperationVersionRecord,
|
||||
Page, PlatformApiKeyRecord, PublishAgentRequest, PublishRequest, PublishedAgentTool,
|
||||
PlatformApiKeyRecord, PublishAgentRequest, PublishRequest, PublishedAgentTool,
|
||||
RegistryOperation, RotateSecretRequest, SaveAgentBindingsRequest, SaveAuthProfileRequest,
|
||||
SaveDescriptorMetadataRequest, SaveSampleMetadataRequest, SaveWorkspaceUpstreamRequest,
|
||||
SecretRecord, SecretVersionRecord, SessionRecord, StreamSessionFilter,
|
||||
UpdateAsyncJobStatusRequest, UpdateStreamSessionStateRequest, UpdateWorkspaceRequest,
|
||||
SecretRecord, SecretVersionRecord, SessionRecord, UpdateWorkspaceRequest,
|
||||
UsageAgentBreakdown, UsageOperationBreakdown, UsageQuery, UsageRollupRecord, UsageSummary,
|
||||
UsageTimelinePoint, WorkspaceMembershipRecord, WorkspaceRecord, WorkspaceUpstream,
|
||||
YamlImportJob, YamlImportJobCompletion, YamlImportJobId, YamlImportJobStatus,
|
||||
@@ -520,52 +515,6 @@ fn map_user_update_error(error: sqlx::Error, user_id: &UserId, email: &str) -> R
|
||||
}
|
||||
}
|
||||
|
||||
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")?,
|
||||
last_poll_at: row.try_get("last_poll_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 {
|
||||
@@ -627,22 +576,6 @@ fn target_summary(target: &Target) -> (String, String) {
|
||||
}
|
||||
.to_owned(),
|
||||
),
|
||||
Target::Graphql(graphql) => (
|
||||
graphql.endpoint.clone(),
|
||||
match graphql.operation_type {
|
||||
GraphqlOperationType::Query => "QUERY",
|
||||
GraphqlOperationType::Mutation => "MUTATION",
|
||||
}
|
||||
.to_owned(),
|
||||
),
|
||||
Target::Grpc(grpc) => (grpc.server_addr.clone(), grpc.method.clone()),
|
||||
Target::Websocket(websocket) => (websocket.url.clone(), "SUBSCRIBE".to_owned()),
|
||||
Target::Soap(soap) => (
|
||||
soap.endpoint_override
|
||||
.clone()
|
||||
.unwrap_or_else(|| format!("wsdl://{}", soap.service_name)),
|
||||
soap.operation_name.clone(),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1061,59 +994,6 @@ 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 })
|
||||
}
|
||||
@@ -1137,13 +1017,12 @@ mod tests {
|
||||
|
||||
use crank_core::{
|
||||
AgentId, AgentOperationBinding, AgentStatus, AgentVersion, ApiKeyHeaderAuthConfig,
|
||||
AsyncJobHandle, AuthConfig, AuthKind, AuthProfile, ConfigExport, ExecutionConfig,
|
||||
ExportMode, GeneratedDraft, GeneratedDraftStatus, HttpMethod, InvocationLog, JobStatus,
|
||||
MembershipRole, OperationId, OperationSecurityLevel, OperationStatus, PlatformApiKey,
|
||||
PlatformApiKeyId, PlatformApiKeyScope, PlatformApiKeyStatus, Protocol, RestTarget,
|
||||
RetryPolicy, Samples, SecretId, StreamSession, StreamSessionId, StreamStatus, Target,
|
||||
ToolDescription, ToolExample, User, UserId, UserSessionId, WizardState, Workspace,
|
||||
WorkspaceId,
|
||||
AuthConfig, AuthKind, AuthProfile, ConfigExport, ExecutionConfig, ExportMode,
|
||||
GeneratedDraft, GeneratedDraftStatus, HttpMethod, InvocationLog, MembershipRole,
|
||||
OperationId, OperationSecurityLevel, OperationStatus, PlatformApiKey, PlatformApiKeyId,
|
||||
PlatformApiKeyScope, PlatformApiKeyStatus, Protocol, RestTarget, RetryPolicy, Samples,
|
||||
SecretId, Target, ToolDescription, ToolExample, User, UserId, UserSessionId, WizardState,
|
||||
Workspace, WorkspaceId,
|
||||
};
|
||||
use crank_mapping::{MappingRule, MappingSet};
|
||||
use crank_schema::{Schema, SchemaKind};
|
||||
@@ -1154,14 +1033,12 @@ mod tests {
|
||||
use crate::{
|
||||
PostgresRegistry, RegistryError,
|
||||
model::{
|
||||
AsyncJobFilter, CreateAgentRequest, CreateAsyncJobRequest, CreateInvocationLogRequest,
|
||||
CreatePlatformApiKeyRequest, CreateStreamSessionRequest, CreateVersionRequest,
|
||||
CreateWorkspaceRequest, CreateYamlImportJobRequest, DescriptorKind, DescriptorMetadata,
|
||||
OperationSampleMetadata, PlatformApiKeyRecord, PublishAgentRequest, PublishRequest,
|
||||
RegistryOperation, SampleKind, SaveAuthProfileRequest, SaveDescriptorMetadataRequest,
|
||||
SaveSampleMetadataRequest, StreamSessionFilter, UpdateAsyncJobStatusRequest,
|
||||
UpdateStreamSessionStateRequest, WorkspaceRecord, YamlImportJobCompletion,
|
||||
YamlImportJobId, YamlImportJobStatus,
|
||||
CreateAgentRequest, CreateInvocationLogRequest, CreatePlatformApiKeyRequest,
|
||||
CreateVersionRequest, CreateWorkspaceRequest, CreateYamlImportJobRequest,
|
||||
DescriptorKind, DescriptorMetadata, OperationSampleMetadata, PlatformApiKeyRecord,
|
||||
PublishAgentRequest, PublishRequest, RegistryOperation, SampleKind,
|
||||
SaveAuthProfileRequest, SaveDescriptorMetadataRequest, SaveSampleMetadataRequest,
|
||||
WorkspaceRecord, YamlImportJobCompletion, YamlImportJobId, YamlImportJobStatus,
|
||||
},
|
||||
};
|
||||
|
||||
@@ -2032,284 +1909,6 @@ 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, ×tamp("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 touched = registry
|
||||
.touch_stream_session_poll(&session.id, ×tamp("2026-04-06T12:00:20Z"))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
touched.last_poll_at,
|
||||
Some(timestamp("2026-04-06T12:00:20Z"))
|
||||
);
|
||||
|
||||
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(×tamp("2026-04-06T12:10:00Z")),
|
||||
last_poll_at: Some(×tamp("2026-04-06T12:01:00Z")),
|
||||
closed_at: Some(×tamp("2026-04-06T12:01:00Z")),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(updated.status, StreamStatus::Failed);
|
||||
assert_eq!(updated.closed_at, Some(timestamp("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: timestamp("2026-04-06T11:00:00Z"),
|
||||
..expired_session
|
||||
},
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let deleted = registry
|
||||
.delete_expired_stream_sessions(×tamp("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, ×tamp("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(×tamp("2026-04-06T12:15:00Z")),
|
||||
updated_at: ×tamp("2026-04-06T12:01:00Z"),
|
||||
finished_at: Some(×tamp("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 touched = registry
|
||||
.touch_async_job_poll(&job.id, ×tamp("2026-04-06T12:02:00Z"))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
touched.last_poll_at,
|
||||
Some(timestamp("2026-04-06T12:02:00Z"))
|
||||
);
|
||||
|
||||
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(timestamp("2026-04-06T11:00:00Z")),
|
||||
..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(×tamp("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: ×tamp("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,
|
||||
@@ -2454,9 +2053,7 @@ mod tests {
|
||||
retry_policy: Some(RetryPolicy { max_attempts: 3 }),
|
||||
auth_profile_ref: Some("auth_crank".into()),
|
||||
headers: BTreeMap::from([("X-Request-Id".to_owned(), "static".to_owned())]),
|
||||
protocol_options: None,
|
||||
response_cache: None,
|
||||
streaming: None,
|
||||
},
|
||||
tool_description: ToolDescription {
|
||||
title: "Create lead".to_owned(),
|
||||
@@ -2469,8 +2066,6 @@ mod tests {
|
||||
samples: Some(Samples {
|
||||
input_json_sample_ref: Some("sample_input".into()),
|
||||
output_json_sample_ref: Some("sample_output".into()),
|
||||
proto_file_ref: None,
|
||||
descriptor_ref: None,
|
||||
}),
|
||||
generated_draft: Some(GeneratedDraft {
|
||||
status: GeneratedDraftStatus::Available,
|
||||
@@ -2496,25 +2091,6 @@ 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: timestamp("2026-04-06T12:05:00Z"),
|
||||
last_poll_at: None,
|
||||
created_at: timestamp("2026-04-06T12:00:00Z"),
|
||||
closed_at: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn test_agent(id: &str, status: AgentStatus) -> crank_core::Agent {
|
||||
crank_core::Agent {
|
||||
id: AgentId::new(id),
|
||||
@@ -2575,22 +2151,4 @@ mod tests {
|
||||
created_at: timestamp(created_at),
|
||||
}
|
||||
}
|
||||
|
||||
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(timestamp("2026-04-06T12:05:00Z")),
|
||||
last_poll_at: None,
|
||||
created_at: timestamp("2026-04-06T12:00:00Z"),
|
||||
updated_at: timestamp("2026-04-06T12:00:00Z"),
|
||||
finished_at: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,634 +0,0 @@
|
||||
use super::*;
|
||||
use time::OffsetDateTime;
|
||||
|
||||
impl PostgresRegistry {
|
||||
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)
|
||||
.bind(request.session.created_at)
|
||||
.bind(request.session.closed_at)
|
||||
.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,
|
||||
expires_at as \"expires_at!: OffsetDateTime\",
|
||||
last_poll_at as \"last_poll_at: OffsetDateTime\",
|
||||
created_at as \"created_at!: OffsetDateTime\",
|
||||
closed_at as \"closed_at: OffsetDateTime\"
|
||||
from stream_sessions
|
||||
where id = $1",
|
||||
id.as_str(),
|
||||
)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
|
||||
row.map(|row| {
|
||||
Ok(StreamSession {
|
||||
id: StreamSessionId::new(row.id),
|
||||
workspace_id: WorkspaceId::new(row.workspace_id),
|
||||
agent_id: row.agent_id.map(AgentId::new),
|
||||
operation_id: OperationId::new(row.operation_id),
|
||||
protocol: deserialize_enum_text(&row.protocol, "protocol")?,
|
||||
mode: deserialize_enum_text(&row.mode, "mode")?,
|
||||
status: deserialize_enum_text(&row.status, "status")?,
|
||||
cursor: row.cursor_json,
|
||||
state: row.state_json,
|
||||
expires_at: row.expires_at,
|
||||
last_poll_at: row.last_poll_at,
|
||||
created_at: row.created_at,
|
||||
closed_at: row.closed_at,
|
||||
})
|
||||
})
|
||||
.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,
|
||||
expires_at,
|
||||
last_poll_at,
|
||||
created_at,
|
||||
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: &OffsetDateTime,
|
||||
) -> 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 touch_stream_session_poll(
|
||||
&self,
|
||||
id: &StreamSessionId,
|
||||
now: &OffsetDateTime,
|
||||
) -> Result<StreamSession, RegistryError> {
|
||||
let row = sqlx::query(
|
||||
"update stream_sessions
|
||||
set last_poll_at = $2::timestamptz
|
||||
where id = $1
|
||||
returning
|
||||
id,
|
||||
workspace_id,
|
||||
agent_id,
|
||||
operation_id,
|
||||
protocol,
|
||||
mode,
|
||||
status,
|
||||
cursor_json,
|
||||
state_json,
|
||||
expires_at,
|
||||
last_poll_at,
|
||||
created_at,
|
||||
closed_at",
|
||||
)
|
||||
.bind(id.as_str())
|
||||
.bind(now)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
|
||||
match row.as_ref() {
|
||||
Some(row) => map_stream_session(row),
|
||||
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,
|
||||
expires_at as \"expires_at!: OffsetDateTime\",
|
||||
last_poll_at as \"last_poll_at: OffsetDateTime\",
|
||||
created_at as \"created_at!: OffsetDateTime\",
|
||||
closed_at as \"closed_at: OffsetDateTime\"
|
||||
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",
|
||||
filter.workspace_id.as_str(),
|
||||
filter.agent_id.map(|value| value.as_str()),
|
||||
filter.operation_id.map(|value| value.as_str()),
|
||||
status,
|
||||
mode,
|
||||
limit,
|
||||
)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
let items = rows
|
||||
.into_iter()
|
||||
.map(|row| {
|
||||
Ok(StreamSession {
|
||||
id: StreamSessionId::new(row.id),
|
||||
workspace_id: WorkspaceId::new(row.workspace_id),
|
||||
agent_id: row.agent_id.map(AgentId::new),
|
||||
operation_id: OperationId::new(row.operation_id),
|
||||
protocol: deserialize_enum_text(&row.protocol, "protocol")?,
|
||||
mode: deserialize_enum_text(&row.mode, "mode")?,
|
||||
status: deserialize_enum_text(&row.status, "status")?,
|
||||
cursor: row.cursor_json,
|
||||
state: row.state_json,
|
||||
expires_at: row.expires_at,
|
||||
last_poll_at: row.last_poll_at,
|
||||
created_at: row.created_at,
|
||||
closed_at: row.closed_at,
|
||||
})
|
||||
})
|
||||
.collect::<Result<Vec<_>, RegistryError>>()?;
|
||||
let total = items.len() as u64;
|
||||
|
||||
Ok(Page { items, total })
|
||||
}
|
||||
|
||||
pub async fn delete_expired_stream_sessions(
|
||||
&self,
|
||||
now: &OffsetDateTime,
|
||||
) -> 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,
|
||||
last_poll_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, $13::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)
|
||||
.bind(request.job.last_poll_at)
|
||||
.bind(request.job.created_at)
|
||||
.bind(request.job.updated_at)
|
||||
.bind(request.job.finished_at)
|
||||
.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,
|
||||
expires_at as \"expires_at: OffsetDateTime\",
|
||||
last_poll_at as \"last_poll_at: OffsetDateTime\",
|
||||
created_at as \"created_at!: OffsetDateTime\",
|
||||
updated_at as \"updated_at!: OffsetDateTime\",
|
||||
finished_at as \"finished_at: OffsetDateTime\"
|
||||
from async_jobs
|
||||
where id = $1",
|
||||
id.as_str(),
|
||||
)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
|
||||
row.map(|row| {
|
||||
Ok(AsyncJobHandle {
|
||||
id: AsyncJobId::new(row.id),
|
||||
workspace_id: WorkspaceId::new(row.workspace_id),
|
||||
agent_id: row.agent_id.map(AgentId::new),
|
||||
operation_id: OperationId::new(row.operation_id),
|
||||
status: deserialize_enum_text(&row.status, "status")?,
|
||||
progress: row.progress_json,
|
||||
result: row.result_json,
|
||||
error: row.error_json,
|
||||
expires_at: row.expires_at,
|
||||
last_poll_at: row.last_poll_at,
|
||||
created_at: row.created_at,
|
||||
updated_at: row.updated_at,
|
||||
finished_at: row.finished_at,
|
||||
})
|
||||
})
|
||||
.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,
|
||||
expires_at,
|
||||
last_poll_at,
|
||||
created_at,
|
||||
updated_at,
|
||||
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 touch_async_job_poll(
|
||||
&self,
|
||||
id: &AsyncJobId,
|
||||
now: &OffsetDateTime,
|
||||
) -> Result<AsyncJobHandle, RegistryError> {
|
||||
let row = sqlx::query(
|
||||
"update async_jobs
|
||||
set last_poll_at = $2::timestamptz
|
||||
where id = $1
|
||||
returning
|
||||
id,
|
||||
workspace_id,
|
||||
agent_id,
|
||||
operation_id,
|
||||
status,
|
||||
progress_json,
|
||||
result_json,
|
||||
error_json,
|
||||
expires_at,
|
||||
last_poll_at,
|
||||
created_at,
|
||||
updated_at,
|
||||
finished_at",
|
||||
)
|
||||
.bind(id.as_str())
|
||||
.bind(now)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
|
||||
match row.as_ref() {
|
||||
Some(row) => map_async_job(row),
|
||||
None => Err(RegistryError::AsyncJobNotFound {
|
||||
job_id: id.as_str().to_owned(),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn cancel_async_job(
|
||||
&self,
|
||||
id: &AsyncJobId,
|
||||
now: &OffsetDateTime,
|
||||
) -> 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,
|
||||
expires_at as \"expires_at: OffsetDateTime\",
|
||||
last_poll_at as \"last_poll_at: OffsetDateTime\",
|
||||
created_at as \"created_at!: OffsetDateTime\",
|
||||
updated_at as \"updated_at!: OffsetDateTime\",
|
||||
finished_at as \"finished_at: OffsetDateTime\"
|
||||
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",
|
||||
filter.workspace_id.as_str(),
|
||||
filter.agent_id.map(|value| value.as_str()),
|
||||
filter.operation_id.map(|value| value.as_str()),
|
||||
status,
|
||||
limit,
|
||||
)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
let items = rows
|
||||
.into_iter()
|
||||
.map(|row| {
|
||||
Ok(AsyncJobHandle {
|
||||
id: AsyncJobId::new(row.id),
|
||||
workspace_id: WorkspaceId::new(row.workspace_id),
|
||||
agent_id: row.agent_id.map(AgentId::new),
|
||||
operation_id: OperationId::new(row.operation_id),
|
||||
status: deserialize_enum_text(&row.status, "status")?,
|
||||
progress: row.progress_json,
|
||||
result: row.result_json,
|
||||
error: row.error_json,
|
||||
expires_at: row.expires_at,
|
||||
last_poll_at: row.last_poll_at,
|
||||
created_at: row.created_at,
|
||||
updated_at: row.updated_at,
|
||||
finished_at: row.finished_at,
|
||||
})
|
||||
})
|
||||
.collect::<Result<Vec<_>, RegistryError>>()?;
|
||||
let total = items.len() as u64;
|
||||
|
||||
Ok(Page { items, total })
|
||||
}
|
||||
|
||||
pub async fn delete_expired_async_jobs(
|
||||
&self,
|
||||
now: &OffsetDateTime,
|
||||
) -> 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())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user