feat: expose streaming test run results

2026-04-07 18:30:08 +03:00
parent 35e9c8c754
commit 1a4f0ea6f3
9 changed files with 865 additions and 91 deletions
@@ -20,14 +20,15 @@ use crank_core::{
 use crank_mapping::{JsonPathRoot, MappingSet, infer_mapping_from_samples};
 use crank_proto::{ProtoService, services_from_descriptor_set_bytes};
 use crank_registry::{
-    AgentSummary, AgentVersionRecord, AsyncJobFilter, CreateAgentRequest, CreateInvitationRequest,
-    CreateInvocationLogRequest, CreatePlatformApiKeyRequest, CreateSecretRequest,
-    CreateVersionRequest, CreateWorkspaceRequest, InvitationRecord, InvocationLogRecord,
-    ListInvocationLogsQuery, MembershipRecord, OperationAgentRef, OperationSampleMetadata,
-    OperationSummary, OperationUsageSummary, OperationVersionRecord, Page, PlatformApiKeyRecord,
-    PostgresRegistry, PublishAgentRequest, PublishRequest, RegistryError, RegistryOperation,
-    RotateSecretRequest, SampleKind, SaveAgentBindingsRequest, SaveAuthProfileRequest,
-    SaveDescriptorMetadataRequest, SaveSampleMetadataRequest, StreamSessionFilter,
+    AgentSummary, AgentVersionRecord, AsyncJobFilter, CreateAgentRequest, CreateAsyncJobRequest,
+    CreateInvitationRequest, CreateInvocationLogRequest, CreatePlatformApiKeyRequest,
+    CreateSecretRequest, CreateStreamSessionRequest, CreateVersionRequest, CreateWorkspaceRequest,
+    InvitationRecord, InvocationLogRecord, ListInvocationLogsQuery, MembershipRecord,
+    OperationAgentRef, OperationSampleMetadata, OperationSummary, OperationUsageSummary,
+    OperationVersionRecord, Page, PlatformApiKeyRecord, PostgresRegistry, PublishAgentRequest,
+    PublishRequest, RegistryError, RegistryOperation, RotateSecretRequest, SampleKind,
+    SaveAgentBindingsRequest, SaveAuthProfileRequest, SaveDescriptorMetadataRequest,
+    SaveSampleMetadataRequest, StreamSessionFilter, UpdateAsyncJobStatusRequest,
    UpdateWorkspaceRequest, UsageAgentBreakdown, UsageBucket, UsageOperationBreakdown, UsageQuery,
    UsageSummary, UsageTimelinePoint, WorkspaceMembershipRecord, WorkspaceRecord,
 };
@@ -128,9 +129,123 @@ pub struct TestRunPayload {
 #[derive(Clone, Debug, Serialize)]
 pub struct TestRunResult {
    pub ok: bool,
+    pub mode: ExecutionMode,
    pub request_preview: Value,
    pub response_preview: Value,
    pub errors: Vec<Value>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub window: Option<WindowTestRunView>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub stream_session: Option<StreamSessionStartView>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub async_job: Option<AsyncJobStartView>,
+}
+
+#[derive(Clone, Debug, Serialize)]
+pub struct WindowTestRunView {
+    pub window_complete: bool,
+    pub truncated: bool,
+    pub has_more: bool,
+    pub cursor: Option<Value>,
+}
+
+#[derive(Clone, Debug, Serialize)]
+pub struct StreamSessionStartView {
+    pub session_id: String,
+    pub status: StreamStatus,
+    pub expires_at: String,
+    pub poll_after_ms: u64,
+    pub preview: Value,
+}
+
+#[derive(Clone, Debug, Serialize)]
+pub struct AsyncJobStartView {
+    pub job_id: String,
+    pub status: JobStatus,
+    pub progress: Value,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+struct StoredSessionState {
+    input: Value,
+    summary: Value,
+    items: Vec<Value>,
+    next_index: usize,
+    batch_size: usize,
+}
+
+enum TestRunOutcome {
+    Unary {
+        output: Value,
+    },
+    Window {
+        output: crank_runtime::WindowExecutionResult,
+    },
+    Session {
+        output: StreamSessionStartView,
+    },
+    AsyncJob {
+        output: AsyncJobStartView,
+    },
+}
+
+impl TestRunOutcome {
+    fn into_result_views(
+        self,
+    ) -> (
+        &'static str,
+        Value,
+        Option<WindowTestRunView>,
+        Option<StreamSessionStartView>,
+        Option<AsyncJobStartView>,
+    ) {
+        match self {
+            Self::Unary { output } => ("admin test run completed", output, None, None, None),
+            Self::Window { output } => (
+                "admin window test run completed",
+                json!({
+                    "summary": output.summary,
+                    "items": output.items,
+                    "cursor": output.cursor,
+                    "window_complete": output.window_complete,
+                    "truncated": output.truncated,
+                    "has_more": output.has_more,
+                }),
+                Some(WindowTestRunView {
+                    window_complete: output.window_complete,
+                    truncated: output.truncated,
+                    has_more: output.has_more,
+                    cursor: output.cursor,
+                }),
+                None,
+                None,
+            ),
+            Self::Session { output } => (
+                "admin session test run started",
+                json!({
+                    "session_id": output.session_id,
+                    "status": output.status,
+                    "expires_at": output.expires_at,
+                    "poll_after_ms": output.poll_after_ms,
+                    "preview": output.preview,
+                }),
+                None,
+                Some(output),
+                None,
+            ),
+            Self::AsyncJob { output } => (
+                "admin async job test run started",
+                json!({
+                    "job_id": output.job_id,
+                    "status": output.status,
+                    "progress": output.progress,
+                }),
+                None,
+                None,
+                Some(output),
+            ),
+        }
+    }
 }

 #[derive(Clone, Debug, Deserialize)]
@@ -1889,6 +2004,12 @@ impl AdminService {
            .get_operation_version(workspace_id, operation_id, payload.version)
            .await?;
        let runtime = RuntimeOperation::from(record.snapshot.clone());
+        let mode = runtime
+            .execution_config
+            .streaming
+            .as_ref()
+            .map(|streaming| streaming.mode)
+            .unwrap_or(ExecutionMode::Unary);
        let request_preview =
            match build_request_preview(&record.snapshot.input_mapping, &payload.input) {
                Ok(preview) => preview,
@@ -1910,11 +2031,15 @@ impl AdminService {
                    .await?;
                    return Ok(TestRunResult {
                        ok: false,
+                        mode,
                        request_preview: Value::Null,
                        response_preview: Value::Null,
                        errors: vec![crate::error::runtime_test_failure(&RuntimeError::Mapping(
                            error,
                        ))],
+                        window: None,
+                        stream_session: None,
+                        async_job: None,
                    });
                }
            };
@@ -1924,16 +2049,39 @@ impl AdminService {
            .await;
        let started_at = std::time::Instant::now();
        match match resolved_auth {
-            Ok(resolved_auth) => {
-                self.runtime
+            Ok(resolved_auth) => match mode {
+                ExecutionMode::Unary => self
+                    .runtime
                    .execute_with_auth(&runtime, &payload.input, resolved_auth.as_ref())
                    .await
-            }
+                    .map(|output| TestRunOutcome::Unary { output }),
+                ExecutionMode::Window => self
+                    .runtime
+                    .execute_window_with_auth(&runtime, &payload.input, resolved_auth.as_ref())
+                    .await
+                    .map(|output| TestRunOutcome::Window { output }),
+                ExecutionMode::Session => self
+                    .start_stream_session_test(
+                        workspace_id,
+                        &record.snapshot,
+                        &runtime,
+                        &payload.input,
+                        resolved_auth.as_ref(),
+                    )
+                    .await
+                    .map(|output| TestRunOutcome::Session { output }),
+                ExecutionMode::AsyncJob => self
+                    .start_async_job_test(workspace_id, &record.snapshot, &runtime, &payload.input)
+                    .await
+                    .map(|output| TestRunOutcome::AsyncJob { output }),
+            },
            Err(error) => Err(error),
        } {
-            Ok(output) => {
+            Ok(outcome) => {
                let duration_ms =
                    u64::try_from(started_at.elapsed().as_millis()).unwrap_or(u64::MAX);
+                let (message, response_preview, window, stream_session, async_job) =
+                    outcome.into_result_views();
                self.record_invocation(InvocationRecordRequest {
                    workspace_id,
                    agent_id: None,
@@ -1941,19 +2089,23 @@ impl AdminService {
                    source: InvocationSource::AdminTestRun,
                    level: InvocationLevel::Info,
                    status: InvocationStatus::Ok,
-                    message: "admin test run completed".to_owned(),
+                    message: message.to_owned(),
                    status_code: None,
                    error_kind: None,
                    duration_ms,
                    request_preview: request_preview.clone(),
-                    response_preview: output.clone(),
+                    response_preview: response_preview.clone(),
                })
                .await?;
                Ok(TestRunResult {
                    ok: true,
+                    mode,
                    request_preview,
-                    response_preview: output,
+                    response_preview,
                    errors: Vec::new(),
+                    window,
+                    stream_session,
+                    async_job,
                })
            }
            Err(error) => {
@@ -1976,9 +2128,13 @@ impl AdminService {
                .await?;
                Ok(TestRunResult {
                    ok: false,
+                    mode,
                    request_preview,
                    response_preview: Value::Null,
                    errors: vec![crate::error::runtime_test_failure(&error)],
+                    window: None,
+                    stream_session: None,
+                    async_job: None,
                })
            }
        }
@@ -2009,6 +2165,187 @@ impl AdminService {
            .map(Some)
    }

+    async fn start_stream_session_test(
+        &self,
+        workspace_id: &WorkspaceId,
+        operation: &RegistryOperation,
+        runtime: &RuntimeOperation,
+        input: &Value,
+        resolved_auth: Option<&ResolvedAuth>,
+    ) -> Result<StreamSessionStartView, RuntimeError> {
+        let streaming = runtime.execution_config.streaming.as_ref().ok_or_else(|| {
+            RuntimeError::MissingStreamingConfig {
+                operation_id: runtime.operation_id.as_str().to_owned(),
+            }
+        })?;
+        let seed = self
+            .runtime
+            .execute_session_seed_with_auth(runtime, input, resolved_auth)
+            .await?;
+        let batch_size = streaming.max_items.unwrap_or(10).max(1) as usize;
+        let preview_count = seed.items.len().min(batch_size);
+        let preview_items = seed.items[..preview_count].to_vec();
+        let next_index = preview_count;
+        let created_at = now_string().map_err(|error| RuntimeError::SecretCrypto {
+            details: error.to_string(),
+        })?;
+        let expires_at = add_millis(
+            &created_at,
+            streaming.max_session_lifetime_ms.unwrap_or(60_000),
+        )
+        .map_err(|error| RuntimeError::SecretCrypto {
+            details: error.to_string(),
+        })?;
+        let session = StreamSession {
+            id: crank_core::StreamSessionId::new(new_prefixed_id("sess")),
+            workspace_id: workspace_id.clone(),
+            agent_id: None,
+            operation_id: operation.id.clone(),
+            protocol: operation.protocol,
+            mode: ExecutionMode::Session,
+            status: StreamStatus::Running,
+            cursor: (next_index < seed.items.len()).then(|| json!(next_index)),
+            state: json!(StoredSessionState {
+                input: input.clone(),
+                summary: seed.summary.clone(),
+                items: seed.items.clone(),
+                next_index,
+                batch_size,
+            }),
+            expires_at: expires_at.clone(),
+            last_poll_at: Some(created_at.clone()),
+            created_at,
+            closed_at: None,
+        };
+        self.registry
+            .create_stream_session(CreateStreamSessionRequest { session: &session })
+            .await
+            .map_err(|error| RuntimeError::SecretCrypto {
+                details: error.to_string(),
+            })?;
+
+        Ok(StreamSessionStartView {
+            session_id: session.id.as_str().to_owned(),
+            status: session.status,
+            expires_at,
+            poll_after_ms: streaming.poll_interval_ms.unwrap_or(1_000),
+            preview: json!({
+                "summary": seed.summary,
+                "items": preview_items,
+            }),
+        })
+    }
+
+    async fn start_async_job_test(
+        &self,
+        workspace_id: &WorkspaceId,
+        operation: &RegistryOperation,
+        runtime: &RuntimeOperation,
+        input: &Value,
+    ) -> Result<AsyncJobStartView, RuntimeError> {
+        let created_at = now_string().map_err(|error| RuntimeError::SecretCrypto {
+            details: error.to_string(),
+        })?;
+        let expires_at =
+            add_millis(&created_at, 300_000).map_err(|error| RuntimeError::SecretCrypto {
+                details: error.to_string(),
+            })?;
+        let job = AsyncJobHandle {
+            id: crank_core::AsyncJobId::new(new_prefixed_id("job")),
+            workspace_id: workspace_id.clone(),
+            agent_id: None,
+            operation_id: operation.id.clone(),
+            status: JobStatus::Running,
+            progress: json!({ "pct": 0 }),
+            result: None,
+            error: None,
+            expires_at: Some(expires_at),
+            created_at: created_at.clone(),
+            updated_at: created_at.clone(),
+            finished_at: None,
+        };
+        self.registry
+            .create_async_job(CreateAsyncJobRequest { job: &job })
+            .await
+            .map_err(|error| RuntimeError::SecretCrypto {
+                details: error.to_string(),
+            })?;
+
+        let registry = self.registry.clone();
+        let secret_crypto = self.secret_crypto.clone();
+        let runtime_for_task = self.runtime.clone();
+        let workspace_for_task = workspace_id.clone();
+        let operation_for_task = runtime.clone();
+        let input_for_task = input.clone();
+        let job_id = job.id.clone();
+        tokio::spawn(async move {
+            let resolved_auth = resolve_runtime_auth_for_task(
+                &registry,
+                &secret_crypto,
+                &workspace_for_task,
+                &operation_for_task.execution_config,
+            )
+            .await;
+            let result = match resolved_auth {
+                Ok(resolved_auth) => {
+                    runtime_for_task
+                        .execute_with_auth(
+                            &operation_for_task,
+                            &input_for_task,
+                            resolved_auth.as_ref(),
+                        )
+                        .await
+                }
+                Err(error) => Err(error),
+            };
+            let finished_at = match now_string() {
+                Ok(value) => value,
+                Err(_) => return,
+            };
+            let _ = match result {
+                Ok(output) => {
+                    registry
+                        .update_async_job_status(UpdateAsyncJobStatusRequest {
+                            job_id: &job_id,
+                            current_status: JobStatus::Running,
+                            next_status: JobStatus::Completed,
+                            progress: &json!({ "pct": 100 }),
+                            result: Some(&output),
+                            error: None,
+                            expires_at: None,
+                            updated_at: &finished_at,
+                            finished_at: Some(&finished_at),
+                        })
+                        .await
+                }
+                Err(error) => {
+                    registry
+                        .update_async_job_status(UpdateAsyncJobStatusRequest {
+                            job_id: &job_id,
+                            current_status: JobStatus::Running,
+                            next_status: JobStatus::Failed,
+                            progress: &json!({ "pct": 100 }),
+                            result: None,
+                            error: Some(&json!({
+                                "code": runtime_error_code(&error),
+                                "message": error.to_string(),
+                            })),
+                            expires_at: None,
+                            updated_at: &finished_at,
+                            finished_at: Some(&finished_at),
+                        })
+                        .await
+                }
+            };
+        });
+
+        Ok(AsyncJobStartView {
+            job_id: job.id.as_str().to_owned(),
+            status: job.status,
+            progress: job.progress,
+        })
+    }
+
    async fn resolve_auth_profile(
        &self,
        workspace_id: &WorkspaceId,
@@ -4267,6 +4604,74 @@ fn now_string() -> Result<String, ApiError> {
        .map_err(|error| ApiError::internal(error.to_string()))
 }

+fn add_millis(timestamp: &str, millis: u64) -> Result<String, ApiError> {
+    let parsed = OffsetDateTime::parse(timestamp, &Rfc3339)
+        .map_err(|error| ApiError::internal(error.to_string()))?;
+    parsed
+        .checked_add(time::Duration::milliseconds(millis as i64))
+        .ok_or_else(|| ApiError::internal("failed to add millis to timestamp"))?
+        .format(&Rfc3339)
+        .map_err(|error| ApiError::internal(error.to_string()))
+}
+
+async fn resolve_runtime_auth_for_task(
+    registry: &PostgresRegistry,
+    secret_crypto: &SecretCrypto,
+    workspace_id: &WorkspaceId,
+    execution_config: &crank_core::ExecutionConfig,
+) -> Result<Option<ResolvedAuth>, RuntimeError> {
+    let Some(auth_profile_id) = execution_config.auth_profile_ref.as_ref() else {
+        return Ok(None);
+    };
+
+    let auth_profile = registry
+        .get_auth_profile(workspace_id, auth_profile_id)
+        .await
+        .map_err(|error| RuntimeError::SecretCrypto {
+            details: error.to_string(),
+        })?
+        .ok_or_else(|| RuntimeError::MissingAuthProfile {
+            auth_profile_id: auth_profile_id.as_str().to_owned(),
+        })?;
+
+    let mut secrets = BTreeMap::new();
+    let used_at = now_string().map_err(|error| RuntimeError::SecretCrypto {
+        details: error.to_string(),
+    })?;
+
+    for secret_id in auth_profile.config.secret_ids() {
+        let secret = registry
+            .get_secret(workspace_id, secret_id)
+            .await
+            .map_err(|error| RuntimeError::SecretCrypto {
+                details: error.to_string(),
+            })?
+            .ok_or_else(|| RuntimeError::MissingSecret {
+                secret_id: secret_id.as_str().to_owned(),
+            })?;
+        let version = registry
+            .get_current_secret_version(workspace_id, secret_id)
+            .await
+            .map_err(|error| RuntimeError::SecretCrypto {
+                details: error.to_string(),
+            })?
+            .ok_or_else(|| RuntimeError::MissingSecretVersion {
+                secret_id: secret_id.as_str().to_owned(),
+                version: secret.secret.current_version,
+            })?;
+        let plaintext = secret_crypto.decrypt(&version.secret_version.ciphertext)?;
+        registry
+            .touch_secret(workspace_id, secret_id, &used_at)
+            .await
+            .map_err(|error| RuntimeError::SecretCrypto {
+                details: error.to_string(),
+            })?;
+        secrets.insert(secret_id.clone(), plaintext);
+    }
+
+    ResolvedAuth::from_profile(&auth_profile, &secrets).map(Some)
+}
+
 fn default_invitation_expiry() -> Result<String, ApiError> {
    OffsetDateTime::now_utc()
        .checked_add(time::Duration::days(7))