feat: add streaming ui configuration

2026-04-06 12:04:49 +03:00
parent fdd0a45124
commit b40daf4f54
20 changed files with 2254 additions and 20 deletions
@@ -7,26 +7,27 @@ use base64::{
 };
 use crank_adapter_grpc::test_support as grpc_test_support;
 use crank_core::{
-    Agent, AgentId, AgentOperationBinding, AgentStatus, AgentVersion, AuthConfig, AuthKind,
-    AuthProfile, AuthProfileId, ConfigExport, ExportMode, GeneratedDraft, GeneratedDraftStatus,
-    InvitationId, InvitationStatus, InvitationToken, InvocationLevel, InvocationLog,
-    InvocationLogId, InvocationSource, InvocationStatus, MembershipRole, OperationId,
-    OperationStatus, PlatformApiKey, PlatformApiKeyId, PlatformApiKeyScope, PlatformApiKeyStatus,
-    Protocol, SampleId, Samples, Secret, SecretId, SecretKind, SecretStatus, Target, UsagePeriod,
-    UserId, UserSessionId, Workspace, WorkspaceId, WorkspaceStatus,
+    Agent, AgentId, AgentOperationBinding, AgentStatus, AgentVersion, AggregationMode,
+    AsyncJobHandle, AuthConfig, AuthKind, AuthProfile, AuthProfileId, ConfigExport, ExecutionMode,
+    ExportMode, GeneratedDraft, GeneratedDraftStatus, InvitationId, InvitationStatus,
+    InvitationToken, InvocationLevel, InvocationLog, InvocationLogId, InvocationSource,
+    InvocationStatus, JobStatus, MembershipRole, OperationId, OperationStatus, PlatformApiKey,
+    PlatformApiKeyId, PlatformApiKeyScope, PlatformApiKeyStatus, Protocol, SampleId, Samples,
+    Secret, SecretId, SecretKind, SecretStatus, StreamSession, StreamStatus, Target,
+    TransportBehavior, UsagePeriod, UserId, UserSessionId, Workspace, WorkspaceId, WorkspaceStatus,
 };
 use crank_mapping::{JsonPathRoot, MappingSet, infer_mapping_from_samples};
 use crank_proto::{ProtoService, services_from_descriptor_set_bytes};
 use crank_registry::{
-    AgentSummary, AgentVersionRecord, CreateAgentRequest, CreateInvitationRequest,
+    AgentSummary, AgentVersionRecord, AsyncJobFilter, CreateAgentRequest, CreateInvitationRequest,
    CreateInvocationLogRequest, CreatePlatformApiKeyRequest, CreateSecretRequest,
    CreateVersionRequest, CreateWorkspaceRequest, InvitationRecord, InvocationLogRecord,
    ListInvocationLogsQuery, MembershipRecord, OperationAgentRef, OperationSampleMetadata,
-    OperationSummary, OperationUsageSummary, OperationVersionRecord, PlatformApiKeyRecord,
+    OperationSummary, OperationUsageSummary, OperationVersionRecord, Page, PlatformApiKeyRecord,
    PostgresRegistry, PublishAgentRequest, PublishRequest, RegistryOperation, RotateSecretRequest,
    SampleKind, SaveAgentBindingsRequest, SaveAuthProfileRequest, SaveDescriptorMetadataRequest,
-    SaveSampleMetadataRequest, UpdateWorkspaceRequest, UsageAgentBreakdown, UsageBucket,
-    UsageOperationBreakdown, UsageQuery, UsageSummary, UsageTimelinePoint,
+    SaveSampleMetadataRequest, StreamSessionFilter, UpdateWorkspaceRequest, UsageAgentBreakdown,
+    UsageBucket, UsageOperationBreakdown, UsageQuery, UsageSummary, UsageTimelinePoint,
    WorkspaceMembershipRecord, WorkspaceRecord,
 };
 use crank_runtime::{PreparedRequest, RuntimeError, RuntimeExecutor, RuntimeOperation};
@@ -303,6 +304,78 @@ pub struct UsageOverviewResponse {
    pub agents: Vec<UsageAgentBreakdown>,
 }

+#[derive(Clone, Debug, Serialize)]
+pub struct ProtocolCapabilityView {
+    pub protocol: Protocol,
+    pub supports_execution_modes: Vec<ExecutionMode>,
+    pub supports_transport_behaviors: Vec<TransportBehavior>,
+    pub supports_auth_kinds: Vec<String>,
+    pub supports_upload_artifacts: Vec<String>,
+    pub supports_cursor_path: bool,
+    pub supports_done_path: bool,
+    pub supports_aggregation_mode: Vec<AggregationMode>,
+}
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct StreamSessionsQuery {
+    pub operation_id: Option<String>,
+    pub agent_id: Option<String>,
+    pub status: Option<StreamStatus>,
+    pub mode: Option<ExecutionMode>,
+    pub page: Option<u32>,
+    pub page_size: Option<u32>,
+}
+
+#[derive(Clone, Debug, Serialize)]
+pub struct StreamSessionSummaryView {
+    pub id: String,
+    pub operation_id: String,
+    pub agent_id: Option<String>,
+    pub protocol: Protocol,
+    pub mode: ExecutionMode,
+    pub status: StreamStatus,
+    pub expires_at: String,
+    pub last_poll_at: Option<String>,
+    pub created_at: String,
+    pub closed_at: Option<String>,
+}
+
+#[derive(Clone, Debug, Serialize)]
+pub struct StreamSessionDetailView {
+    #[serde(flatten)]
+    pub summary: StreamSessionSummaryView,
+    pub cursor: Option<Value>,
+}
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct AsyncJobsQuery {
+    pub operation_id: Option<String>,
+    pub agent_id: Option<String>,
+    pub status: Option<JobStatus>,
+    pub page: Option<u32>,
+    pub page_size: Option<u32>,
+}
+
+#[derive(Clone, Debug, Serialize)]
+pub struct AsyncJobSummaryView {
+    pub id: String,
+    pub operation_id: String,
+    pub agent_id: Option<String>,
+    pub status: JobStatus,
+    pub progress: Value,
+    pub created_at: String,
+    pub updated_at: String,
+    pub finished_at: Option<String>,
+    pub expires_at: Option<String>,
+}
+
+#[derive(Clone, Debug, Serialize)]
+pub struct AsyncJobDetailView {
+    #[serde(flatten)]
+    pub summary: AsyncJobSummaryView,
+    pub error: Option<Value>,
+}
+
 #[derive(Clone, Debug, Deserialize)]
 pub struct GenerateDraftPayload {
    #[serde(default)]
@@ -1273,6 +1346,185 @@ impl AdminService {
            })
    }

+    pub async fn list_protocol_capabilities(&self) -> Vec<ProtocolCapabilityView> {
+        [Protocol::Rest, Protocol::Graphql, Protocol::Grpc]
+            .into_iter()
+            .map(protocol_capability_view)
+            .collect()
+    }
+
+    #[instrument(skip(self))]
+    pub async fn list_stream_sessions(
+        &self,
+        workspace_id: &WorkspaceId,
+        query: StreamSessionsQuery,
+    ) -> Result<Page<StreamSessionSummaryView>, ApiError> {
+        self.ensure_workspace_exists(workspace_id).await?;
+        let operation_id = query.operation_id.as_deref().map(OperationId::new);
+        let agent_id = query.agent_id.as_deref().map(AgentId::new);
+        let limit = query.page_size.unwrap_or(20).max(1);
+        let page = self
+            .registry
+            .list_stream_sessions(StreamSessionFilter {
+                workspace_id,
+                agent_id: agent_id.as_ref(),
+                operation_id: operation_id.as_ref(),
+                status: query.status,
+                mode: query.mode,
+                limit,
+            })
+            .await?;
+
+        Ok(Page {
+            items: page
+                .items
+                .into_iter()
+                .map(stream_session_summary_view)
+                .collect(),
+            total: page.total,
+        })
+    }
+
+    #[instrument(skip(self))]
+    pub async fn get_stream_session(
+        &self,
+        workspace_id: &WorkspaceId,
+        session_id: &crank_core::StreamSessionId,
+    ) -> Result<StreamSessionDetailView, ApiError> {
+        self.ensure_workspace_exists(workspace_id).await?;
+        let session = self
+            .registry
+            .get_stream_session(session_id)
+            .await?
+            .ok_or_else(|| {
+                ApiError::not_found(format!(
+                    "stream session {} was not found",
+                    session_id.as_str()
+                ))
+            })?;
+        ensure_stream_session_workspace(&session, workspace_id)?;
+
+        Ok(stream_session_detail_view(session))
+    }
+
+    #[instrument(skip(self))]
+    pub async fn stop_stream_session(
+        &self,
+        workspace_id: &WorkspaceId,
+        session_id: &crank_core::StreamSessionId,
+    ) -> Result<StreamSessionDetailView, ApiError> {
+        self.ensure_workspace_exists(workspace_id).await?;
+        let session = self
+            .registry
+            .get_stream_session(session_id)
+            .await?
+            .ok_or_else(|| {
+                ApiError::not_found(format!(
+                    "stream session {} was not found",
+                    session_id.as_str()
+                ))
+            })?;
+        ensure_stream_session_workspace(&session, workspace_id)?;
+        self.registry
+            .close_stream_session(session_id, &now_string()?)
+            .await?;
+        let updated = self
+            .registry
+            .get_stream_session(session_id)
+            .await?
+            .ok_or_else(|| {
+                ApiError::not_found(format!(
+                    "stream session {} was not found",
+                    session_id.as_str()
+                ))
+            })?;
+
+        Ok(stream_session_detail_view(updated))
+    }
+
+    #[instrument(skip(self))]
+    pub async fn list_async_jobs(
+        &self,
+        workspace_id: &WorkspaceId,
+        query: AsyncJobsQuery,
+    ) -> Result<Page<AsyncJobSummaryView>, ApiError> {
+        self.ensure_workspace_exists(workspace_id).await?;
+        let operation_id = query.operation_id.as_deref().map(OperationId::new);
+        let agent_id = query.agent_id.as_deref().map(AgentId::new);
+        let limit = query.page_size.unwrap_or(20).max(1);
+        let page = self
+            .registry
+            .list_async_jobs(AsyncJobFilter {
+                workspace_id,
+                agent_id: agent_id.as_ref(),
+                operation_id: operation_id.as_ref(),
+                status: query.status,
+                limit,
+            })
+            .await?;
+
+        Ok(Page {
+            items: page.items.into_iter().map(async_job_summary_view).collect(),
+            total: page.total,
+        })
+    }
+
+    #[instrument(skip(self))]
+    pub async fn get_async_job(
+        &self,
+        workspace_id: &WorkspaceId,
+        job_id: &crank_core::AsyncJobId,
+    ) -> Result<AsyncJobDetailView, ApiError> {
+        self.ensure_workspace_exists(workspace_id).await?;
+        let job = self.registry.get_async_job(job_id).await?.ok_or_else(|| {
+            ApiError::not_found(format!("async job {} was not found", job_id.as_str()))
+        })?;
+        ensure_async_job_workspace(&job, workspace_id)?;
+
+        Ok(async_job_detail_view(job))
+    }
+
+    #[instrument(skip(self))]
+    pub async fn cancel_async_job(
+        &self,
+        workspace_id: &WorkspaceId,
+        job_id: &crank_core::AsyncJobId,
+    ) -> Result<AsyncJobDetailView, ApiError> {
+        self.ensure_workspace_exists(workspace_id).await?;
+        let job = self.registry.get_async_job(job_id).await?.ok_or_else(|| {
+            ApiError::not_found(format!("async job {} was not found", job_id.as_str()))
+        })?;
+        ensure_async_job_workspace(&job, workspace_id)?;
+        self.registry
+            .cancel_async_job(job_id, &now_string()?)
+            .await?;
+        let updated = self.registry.get_async_job(job_id).await?.ok_or_else(|| {
+            ApiError::not_found(format!("async job {} was not found", job_id.as_str()))
+        })?;
+
+        Ok(async_job_detail_view(updated))
+    }
+
+    #[instrument(skip(self))]
+    pub async fn get_async_job_result(
+        &self,
+        workspace_id: &WorkspaceId,
+        job_id: &crank_core::AsyncJobId,
+    ) -> Result<Value, ApiError> {
+        self.ensure_workspace_exists(workspace_id).await?;
+        let job = self.registry.get_async_job(job_id).await?.ok_or_else(|| {
+            ApiError::not_found(format!("async job {} was not found", job_id.as_str()))
+        })?;
+        ensure_async_job_workspace(&job, workspace_id)?;
+
+        match job.status {
+            JobStatus::Completed => Ok(job.result.unwrap_or(Value::Null)),
+            JobStatus::Failed => Err(ApiError::conflict("async job failed")),
+            JobStatus::Cancelled => Err(ApiError::conflict("async job was cancelled")),
+            _ => Err(ApiError::conflict("async job result is not ready")),
+        }
+    }
+
    pub async fn list_operations(
        &self,
        workspace_id: &WorkspaceId,
@@ -3762,6 +4014,127 @@ fn runtime_error_code(error: &RuntimeError) -> &'static str {
    }
 }

+fn protocol_capability_view(protocol: Protocol) -> ProtocolCapabilityView {
+    let supports_execution_modes = [
+        ExecutionMode::Unary,
+        ExecutionMode::Window,
+        ExecutionMode::Session,
+        ExecutionMode::AsyncJob,
+    ]
+    .into_iter()
+    .filter(|mode| protocol.supports_execution_mode(*mode))
+    .collect();
+    let supports_transport_behaviors = [
+        TransportBehavior::RequestResponse,
+        TransportBehavior::ServerStream,
+        TransportBehavior::StatefulSession,
+        TransportBehavior::DeferredResult,
+    ]
+    .into_iter()
+    .filter(|behavior| protocol.supports_transport_behavior(*behavior))
+    .collect();
+    let supports_upload_artifacts = match protocol {
+        Protocol::Rest | Protocol::Graphql => Vec::new(),
+        Protocol::Grpc => vec!["proto".to_owned(), "descriptor_set".to_owned()],
+    };
+
+    ProtocolCapabilityView {
+        protocol,
+        supports_execution_modes,
+        supports_transport_behaviors,
+        supports_auth_kinds: vec![
+            "none".to_owned(),
+            "bearer".to_owned(),
+            "basic".to_owned(),
+            "api_key_header".to_owned(),
+            "api_key_query".to_owned(),
+        ],
+        supports_upload_artifacts,
+        supports_cursor_path: !matches!(protocol, Protocol::Graphql),
+        supports_done_path: !matches!(protocol, Protocol::Graphql),
+        supports_aggregation_mode: vec![
+            AggregationMode::RawItems,
+            AggregationMode::SummaryOnly,
+            AggregationMode::SummaryPlusSamples,
+            AggregationMode::Stats,
+            AggregationMode::LatestState,
+        ],
+    }
+}
+
+fn stream_session_summary_view(session: StreamSession) -> StreamSessionSummaryView {
+    StreamSessionSummaryView {
+        id: session.id.as_str().to_owned(),
+        operation_id: session.operation_id.as_str().to_owned(),
+        agent_id: session.agent_id.map(|value| value.as_str().to_owned()),
+        protocol: session.protocol,
+        mode: session.mode,
+        status: session.status,
+        expires_at: session.expires_at,
+        last_poll_at: session.last_poll_at,
+        created_at: session.created_at,
+        closed_at: session.closed_at,
+    }
+}
+
+fn stream_session_detail_view(session: StreamSession) -> StreamSessionDetailView {
+    StreamSessionDetailView {
+        cursor: session.cursor.clone(),
+        summary: stream_session_summary_view(session),
+    }
+}
+
+fn async_job_summary_view(job: AsyncJobHandle) -> AsyncJobSummaryView {
+    AsyncJobSummaryView {
+        id: job.id.as_str().to_owned(),
+        operation_id: job.operation_id.as_str().to_owned(),
+        agent_id: job.agent_id.map(|value| value.as_str().to_owned()),
+        status: job.status,
+        progress: job.progress,
+        created_at: job.created_at,
+        updated_at: job.updated_at,
+        finished_at: job.finished_at,
+        expires_at: job.expires_at,
+    }
+}
+
+fn async_job_detail_view(job: AsyncJobHandle) -> AsyncJobDetailView {
+    AsyncJobDetailView {
+        error: job.error.clone(),
+        summary: async_job_summary_view(job),
+    }
+}
+
+fn ensure_stream_session_workspace(
+    session: &StreamSession,
+    workspace_id: &WorkspaceId,
+) -> Result<(), ApiError> {
+    if &session.workspace_id == workspace_id {
+        return Ok(());
+    }
+
+    Err(ApiError::not_found(format!(
+        "stream session {} was not found in workspace {}",
+        session.id.as_str(),
+        workspace_id.as_str()
+    )))
+}
+
+fn ensure_async_job_workspace(
+    job: &AsyncJobHandle,
+    workspace_id: &WorkspaceId,
+) -> Result<(), ApiError> {
+    if &job.workspace_id == workspace_id {
+        return Ok(());
+    }
+
+    Err(ApiError::not_found(format!(
+        "async job {} was not found in workspace {}",
+        job.id.as_str(),
+        workspace_id.as_str()
+    )))
+}
+
 fn usage_window(period: UsagePeriod) -> Result<(UsagePeriod, String, UsageBucket), ApiError> {
    let now = OffsetDateTime::now_utc();
    let (start, bucket) = match period {