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
@@ -28,6 +28,10 @@ use crate::{
            upload_descriptor_set, upload_input_json, upload_output_json, upload_proto_descriptor,
        },
        secrets::{create_secret, delete_secret, get_secret, list_secrets, rotate_secret},
+        streaming::{
+            cancel_async_job, get_async_job, get_async_job_result, get_stream_session,
+            list_async_jobs, list_protocol_capabilities, list_stream_sessions, stop_stream_session,
+        },
        workspaces::{create_workspace, get_workspace, list_workspaces, update_workspace},
    },
    state::AppState,
@@ -133,7 +137,18 @@ pub fn build_app(state: AppState) -> Router {
        .route("/logs/{log_id}", get(get_log))
        .route("/usage", get(get_usage))
        .route("/usage/operations/{operation_id}", get(get_operation_usage))
-        .route("/usage/agents/{agent_id}", get(get_agent_usage));
+        .route("/usage/agents/{agent_id}", get(get_agent_usage))
+        .route("/protocol-capabilities", get(list_protocol_capabilities))
+        .route("/stream-sessions", get(list_stream_sessions))
+        .route("/stream-sessions/{session_id}", get(get_stream_session))
+        .route(
+            "/stream-sessions/{session_id}/stop",
+            post(stop_stream_session),
+        )
+        .route("/async-jobs", get(list_async_jobs))
+        .route("/async-jobs/{job_id}", get(get_async_job))
+        .route("/async-jobs/{job_id}/cancel", post(cancel_async_job))
+        .route("/async-jobs/{job_id}/result", get(get_async_job_result));

    let workspace_root_router = Router::new()
        .route("/workspaces", get(list_workspaces).post(create_workspace))
@@ -5,6 +5,7 @@ pub mod auth_profiles;
 pub mod observability;
 pub mod operations;
 pub mod secrets;
+pub mod streaming;
 pub mod workspaces;

 use axum::Json;
@@ -0,0 +1,137 @@
+use axum::{
+    Json,
+    extract::{Path, Query, State},
+};
+use serde_json::{Value, json};
+
+use crate::{
+    error::ApiError,
+    routes::access::WorkspacePath,
+    service::{AsyncJobsQuery, StreamSessionsQuery},
+    state::AppState,
+};
+
+#[derive(serde::Deserialize)]
+pub struct WorkspaceStreamSessionPath {
+    pub workspace_id: String,
+    pub session_id: String,
+}
+
+#[derive(serde::Deserialize)]
+pub struct WorkspaceAsyncJobPath {
+    pub workspace_id: String,
+    pub job_id: String,
+}
+
+pub async fn list_protocol_capabilities(
+    Path(_path): Path<WorkspacePath>,
+    State(state): State<AppState>,
+) -> Result<Json<Value>, ApiError> {
+    Ok(Json(json!({
+        "items": state.service.list_protocol_capabilities().await
+    })))
+}
+
+pub async fn list_stream_sessions(
+    Path(path): Path<WorkspacePath>,
+    Query(query): Query<StreamSessionsQuery>,
+    State(state): State<AppState>,
+) -> Result<Json<Value>, ApiError> {
+    let page = state
+        .service
+        .list_stream_sessions(&path.workspace_id.as_str().into(), query.clone())
+        .await?;
+    Ok(Json(json!({
+        "items": page.items,
+        "page": query.page.unwrap_or(1),
+        "page_size": query.page_size.unwrap_or(20),
+        "total": page.total,
+    })))
+}
+
+pub async fn get_stream_session(
+    Path(path): Path<WorkspaceStreamSessionPath>,
+    State(state): State<AppState>,
+) -> Result<Json<Value>, ApiError> {
+    let session = state
+        .service
+        .get_stream_session(
+            &path.workspace_id.as_str().into(),
+            &path.session_id.as_str().into(),
+        )
+        .await?;
+    Ok(Json(json!(session)))
+}
+
+pub async fn stop_stream_session(
+    Path(path): Path<WorkspaceStreamSessionPath>,
+    State(state): State<AppState>,
+) -> Result<Json<Value>, ApiError> {
+    let session = state
+        .service
+        .stop_stream_session(
+            &path.workspace_id.as_str().into(),
+            &path.session_id.as_str().into(),
+        )
+        .await?;
+    Ok(Json(json!(session)))
+}
+
+pub async fn list_async_jobs(
+    Path(path): Path<WorkspacePath>,
+    Query(query): Query<AsyncJobsQuery>,
+    State(state): State<AppState>,
+) -> Result<Json<Value>, ApiError> {
+    let page = state
+        .service
+        .list_async_jobs(&path.workspace_id.as_str().into(), query.clone())
+        .await?;
+    Ok(Json(json!({
+        "items": page.items,
+        "page": query.page.unwrap_or(1),
+        "page_size": query.page_size.unwrap_or(20),
+        "total": page.total,
+    })))
+}
+
+pub async fn get_async_job(
+    Path(path): Path<WorkspaceAsyncJobPath>,
+    State(state): State<AppState>,
+) -> Result<Json<Value>, ApiError> {
+    let job = state
+        .service
+        .get_async_job(
+            &path.workspace_id.as_str().into(),
+            &path.job_id.as_str().into(),
+        )
+        .await?;
+    Ok(Json(json!(job)))
+}
+
+pub async fn cancel_async_job(
+    Path(path): Path<WorkspaceAsyncJobPath>,
+    State(state): State<AppState>,
+) -> Result<Json<Value>, ApiError> {
+    let job = state
+        .service
+        .cancel_async_job(
+            &path.workspace_id.as_str().into(),
+            &path.job_id.as_str().into(),
+        )
+        .await?;
+    Ok(Json(json!(job)))
+}
+
+pub async fn get_async_job_result(
+    Path(path): Path<WorkspaceAsyncJobPath>,
+    State(state): State<AppState>,
+) -> Result<Json<Value>, ApiError> {
+    let result = state
+        .service
+        .get_async_job_result(
+            &path.workspace_id.as_str().into(),
+            &path.job_id.as_str().into(),
+        )
+        .await?;
+    Ok(Json(json!(result)))
+}
@@ -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 {