clusterflux-public/crates/clusterflux-coordinator/src/service/protocol.rs

use std::collections::BTreeMap;

use clusterflux_core::{
    AgentId, AgentSignedRequest, ArtifactId, Authorization, Capability, CredentialKind,
    DataPlaneScope, Digest, DirectBulkTransferPlan, DownloadLink, EnvironmentRequirements,
    LimitKind, NodeCapabilities, NodeDescriptor, NodeEndpoint, NodeId, NodeSignedRequest,
    PanelEventKind, PanelState, Placement, ProcessId, ProjectId, ResourceLimits, SourcePreparation,
    SourceProviderKind, TaskBoundaryValue, TaskInstanceId, TaskJoinResult, TaskSpec, TenantId,
    UserId, VfsPath,
};
use serde::{Deserialize, Serialize};

mod responses;
pub use responses::*;

use crate::{AgentPublicKeyRecord, ProjectRecord};

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct TaskReplacementBundle {
    pub bundle_digest: Digest,
    pub wasm_module_base64: String,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum TaskFailureResolution {
    AcceptFailure,
    Cancel,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case", deny_unknown_fields)]
pub enum CoordinatorRequest {
    Ping,
    Authenticated {
        session_secret: String,
        request: AuthenticatedCoordinatorRequest,
    },
    AuthStatus {
        tenant: String,
        project: String,
        actor_user: String,
    },
    AdminStatus {
        tenant: String,
        actor_user: String,
        admin_proof: Digest,
        admin_nonce: String,
        issued_at_epoch_seconds: u64,
    },
    SuspendTenant {
        tenant: String,
        actor_user: String,
        target_tenant: String,
        admin_proof: Digest,
        admin_nonce: String,
        issued_at_epoch_seconds: u64,
    },
    CreateProject {
        tenant: String,
        actor_user: String,
        project: String,
        name: String,
    },
    SelectProject {
        tenant: String,
        actor_user: String,
        project: String,
    },
    ListProjects {
        tenant: String,
        actor_user: String,
    },
    RegisterAgentPublicKey {
        tenant: String,
        project: String,
        user: String,
        agent: String,
        public_key: String,
    },
    ListAgentPublicKeys {
        tenant: String,
        project: String,
        user: String,
    },
    RotateAgentPublicKey {
        tenant: String,
        project: String,
        user: String,
        agent: String,
        public_key: String,
    },
    RevokeAgentPublicKey {
        tenant: String,
        project: String,
        user: String,
        agent: String,
    },
    AttachNode {
        tenant: String,
        project: String,
        node: String,
        public_key: String,
    },
    CreateNodeEnrollmentGrant {
        tenant: String,
        project: String,
        actor_user: String,
        #[serde(default = "default_node_enrollment_ttl_seconds")]
        ttl_seconds: u64,
    },
    ExchangeNodeEnrollmentGrant {
        tenant: String,
        project: String,
        node: String,
        public_key: String,
        enrollment_grant: String,
    },
    NodeHeartbeat {
        node: String,
        #[serde(default)]
        node_signature: Option<NodeSignedRequest>,
    },
    SignedNode {
        node: String,
        node_signature: NodeSignedRequest,
        request: Box<CoordinatorRequest>,
    },
    ReportNodeCapabilities {
        tenant: String,
        project: String,
        node: String,
        capabilities: NodeCapabilities,
        cached_environment_digests: Vec<Digest>,
        #[serde(default)]
        dependency_cache_digests: Vec<Digest>,
        source_snapshots: Vec<Digest>,
        artifact_locations: Vec<String>,
        direct_connectivity: bool,
        online: bool,
    },
    ListNodeDescriptors {
        tenant: String,
        project: String,
        actor_user: String,
    },
    RevokeNodeCredential {
        tenant: String,
        project: String,
        actor_user: String,
        node: String,
    },
    ScheduleTask {
        tenant: String,
        project: String,
        environment: Option<EnvironmentRequirements>,
        environment_digest: Option<Digest>,
        required_capabilities: Vec<Capability>,
        #[serde(default)]
        dependency_cache: Option<Digest>,
        source_snapshot: Option<Digest>,
        required_artifacts: Vec<String>,
        prefer_node: Option<String>,
    },
    LaunchTask {
        tenant: String,
        project: String,
        #[serde(default)]
        actor_user: Option<String>,
        #[serde(default)]
        actor_agent: Option<String>,
        #[serde(default)]
        agent_public_key_fingerprint: Option<Digest>,
        #[serde(default)]
        agent_signature: Option<AgentSignedRequest>,
        task_spec: TaskSpec,
        #[serde(default)]
        wait_for_node: bool,
        artifact_path: String,
        wasm_module_base64: String,
    },
    LaunchChildTask {
        tenant: String,
        project: String,
        process: String,
        node: String,
        parent_task: String,
        task_spec: TaskSpec,
        #[serde(default)]
        wait_for_node: bool,
        artifact_path: String,
        wasm_module_base64: String,
    },
    JoinChildTask {
        tenant: String,
        project: String,
        process: String,
        node: String,
        parent_task: String,
        task: String,
    },
    PollTaskAssignment {
        tenant: String,
        project: String,
        node: String,
    },
    PollArtifactTransfer {
        tenant: String,
        project: String,
        node: String,
    },
    UploadArtifactTransferChunk {
        tenant: String,
        project: String,
        node: String,
        transfer_id: String,
        artifact: String,
        offset: u64,
        content_base64: String,
        chunk_digest: Digest,
        eof: bool,
    },
    FailArtifactTransfer {
        tenant: String,
        project: String,
        node: String,
        transfer_id: String,
        artifact: String,
        message: String,
    },
    RequestRendezvous {
        scope: DataPlaneScope,
        source: NodeEndpoint,
        destination: NodeEndpoint,
        direct_connectivity: bool,
        failure_reason: String,
    },
    RequestSourcePreparation {
        tenant: String,
        project: String,
        provider: SourceProviderKind,
    },
    CompleteSourcePreparation {
        tenant: String,
        project: String,
        node: String,
        provider: SourceProviderKind,
        source_snapshot: Digest,
    },
    StartProcess {
        tenant: String,
        project: String,
        #[serde(default)]
        actor_user: Option<String>,
        #[serde(default)]
        actor_agent: Option<String>,
        #[serde(default)]
        agent_public_key_fingerprint: Option<Digest>,
        #[serde(default)]
        agent_signature: Option<AgentSignedRequest>,
        process: String,
        #[serde(default)]
        restart: bool,
    },
    ReconnectNode {
        node: String,
        process: String,
        epoch: u64,
    },
    CancelTask {
        tenant: String,
        project: String,
        process: String,
        node: String,
        task: String,
    },
    CancelProcess {
        tenant: String,
        project: String,
        actor_user: String,
        process: String,
    },
    AbortProcess {
        tenant: String,
        project: String,
        actor_user: String,
        process: String,
    },
    ListProcesses {
        tenant: String,
        project: String,
        actor_user: String,
    },
    QuotaStatus {
        tenant: String,
        project: String,
        actor_user: String,
    },
    PollTaskControl {
        tenant: String,
        project: String,
        process: String,
        node: String,
        task: String,
    },
    RestartTask {
        tenant: String,
        project: String,
        actor_user: String,
        process: String,
        task: String,
        #[serde(default)]
        replacement_bundle: Option<TaskReplacementBundle>,
    },
    ResolveTaskFailure {
        tenant: String,
        project: String,
        actor_user: String,
        process: String,
        task: String,
        resolution: TaskFailureResolution,
    },
    DebugAttach {
        tenant: String,
        project: String,
        actor_user: String,
        process: String,
    },
    SetDebugBreakpoints {
        tenant: String,
        project: String,
        actor_user: String,
        process: String,
        probe_symbols: Vec<String>,
    },
    InspectDebugBreakpoints {
        tenant: String,
        project: String,
        actor_user: String,
        process: String,
    },
    CreateDebugEpoch {
        tenant: String,
        project: String,
        actor_user: String,
        process: String,
        stopped_task: String,
        reason: String,
    },
    ResumeDebugEpoch {
        tenant: String,
        project: String,
        actor_user: String,
        process: String,
        epoch: u64,
    },
    InspectDebugEpoch {
        tenant: String,
        project: String,
        actor_user: String,
        process: String,
        epoch: u64,
    },
    PollDebugCommand {
        tenant: String,
        project: String,
        process: String,
        node: String,
        task: String,
    },
    ReportDebugState {
        tenant: String,
        project: String,
        process: String,
        node: String,
        task: String,
        epoch: u64,
        state: DebugAcknowledgementState,
        #[serde(default)]
        stack_frames: Vec<String>,
        #[serde(default)]
        local_values: Vec<(String, String)>,
        #[serde(default)]
        task_args: Vec<(String, String)>,
        #[serde(default)]
        handles: Vec<(String, String)>,
        #[serde(default)]
        command_status: Option<String>,
        #[serde(default)]
        recent_output: Vec<String>,
        #[serde(default)]
        message: Option<String>,
    },
    ReportDebugProbeHit {
        tenant: String,
        project: String,
        process: String,
        node: String,
        task: String,
        probe_symbol: String,
    },
    ReportTaskLog {
        tenant: String,
        project: String,
        process: String,
        node: String,
        task: String,
        stdout_bytes: u64,
        stderr_bytes: u64,
        #[serde(default)]
        stdout_tail: String,
        #[serde(default)]
        stderr_tail: String,
        stdout_truncated: bool,
        stderr_truncated: bool,
        backpressured: bool,
    },
    ReportVfsMetadata {
        tenant: String,
        project: String,
        process: String,
        node: String,
        task: String,
        artifact_path: Option<String>,
        artifact_digest: Option<Digest>,
        artifact_size_bytes: Option<u64>,
        large_bytes_uploaded: bool,
    },
    TaskCompleted {
        tenant: String,
        project: String,
        process: String,
        node: String,
        task: String,
        #[serde(default)]
        terminal_state: Option<TaskTerminalState>,
        status_code: Option<i32>,
        stdout_bytes: u64,
        stderr_bytes: u64,
        #[serde(default)]
        stdout_tail: String,
        #[serde(default)]
        stderr_tail: String,
        #[serde(default)]
        stdout_truncated: bool,
        #[serde(default)]
        stderr_truncated: bool,
        artifact_path: Option<String>,
        artifact_digest: Option<Digest>,
        artifact_size_bytes: Option<u64>,
        #[serde(default)]
        result: Option<TaskBoundaryValue>,
    },
    ListTaskEvents {
        tenant: String,
        project: String,
        actor_user: String,
        #[serde(default)]
        process: Option<String>,
    },
    ListTaskSnapshots {
        tenant: String,
        project: String,
        actor_user: String,
        process: String,
    },
    JoinTask {
        tenant: String,
        project: String,
        actor_user: String,
        process: String,
        task: String,
    },
    RenderOperatorPanel {
        tenant: String,
        project: String,
        process: String,
        actor_user: String,
        max_download_bytes: u64,
        stopped: bool,
    },
    SubmitPanelEvent {
        tenant: String,
        project: String,
        process: String,
        widget_id: String,
        kind: PanelEventKind,
        max_events: u64,
    },
    CreateArtifactDownloadLink {
        tenant: String,
        project: String,
        actor_user: String,
        artifact: String,
        max_bytes: u64,
        #[serde(default = "default_download_ttl_seconds")]
        ttl_seconds: u64,
    },
    OpenArtifactDownloadStream {
        tenant: String,
        project: String,
        actor_user: String,
        artifact: String,
        max_bytes: u64,
        token_digest: Digest,
        chunk_bytes: u64,
    },
    RevokeArtifactDownloadLink {
        tenant: String,
        project: String,
        actor_user: String,
        artifact: String,
        token_digest: Digest,
    },
    ExportArtifactToNode {
        tenant: String,
        project: String,
        actor_user: String,
        artifact: String,
        receiver_node: String,
        direct_connectivity: bool,
        failure_reason: String,
    },
}

impl CoordinatorRequest {
    pub fn operation(&self) -> Result<String, String> {
        serde_json::to_value(self)
            .map_err(|err| format!("failed to encode coordinator request operation: {err}"))
            .map(|value| clusterflux_core::coordinator_payload_operation(&value))
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case", deny_unknown_fields)]
pub enum AuthenticatedCoordinatorRequest {
    AuthStatus,
    RevokeCliSession,
    CreateProject {
        project: String,
        name: String,
    },
    SelectProject {
        project: String,
    },
    ListProjects,
    RegisterAgentPublicKey {
        agent: String,
        public_key: String,
    },
    ListAgentPublicKeys,
    RotateAgentPublicKey {
        agent: String,
        public_key: String,
    },
    RevokeAgentPublicKey {
        agent: String,
    },
    CreateNodeEnrollmentGrant {
        #[serde(default = "default_node_enrollment_ttl_seconds")]
        ttl_seconds: u64,
    },
    ListNodeDescriptors,
    RevokeNodeCredential {
        node: String,
    },
    StartProcess {
        process: String,
        #[serde(default)]
        restart: bool,
    },
    ScheduleTask {
        environment: Option<EnvironmentRequirements>,
        environment_digest: Option<Digest>,
        required_capabilities: Vec<Capability>,
        #[serde(default)]
        dependency_cache: Option<Digest>,
        source_snapshot: Option<Digest>,
        required_artifacts: Vec<String>,
        prefer_node: Option<String>,
    },
    LaunchTask {
        task_spec: Box<TaskSpec>,
        #[serde(default)]
        wait_for_node: bool,
        artifact_path: String,
        wasm_module_base64: String,
    },
    CancelProcess {
        process: String,
    },
    AbortProcess {
        process: String,
    },
    ListProcesses,
    QuotaStatus,
    RestartTask {
        process: String,
        task: String,
        #[serde(default)]
        replacement_bundle: Option<TaskReplacementBundle>,
    },
    ResolveTaskFailure {
        process: String,
        task: String,
        resolution: TaskFailureResolution,
    },
    DebugAttach {
        process: String,
    },
    SetDebugBreakpoints {
        process: String,
        probe_symbols: Vec<String>,
    },
    InspectDebugBreakpoints {
        process: String,
    },
    CreateDebugEpoch {
        process: String,
        stopped_task: String,
        reason: String,
    },
    ResumeDebugEpoch {
        process: String,
        epoch: u64,
    },
    InspectDebugEpoch {
        process: String,
        epoch: u64,
    },
    ListTaskEvents {
        #[serde(default)]
        process: Option<String>,
    },
    ListTaskSnapshots {
        process: String,
    },
    JoinTask {
        process: String,
        task: String,
    },
    CreateArtifactDownloadLink {
        artifact: String,
        max_bytes: u64,
        #[serde(default = "default_download_ttl_seconds")]
        ttl_seconds: u64,
    },
    OpenArtifactDownloadStream {
        artifact: String,
        max_bytes: u64,
        token_digest: Digest,
        chunk_bytes: u64,
    },
    RevokeArtifactDownloadLink {
        artifact: String,
        token_digest: Digest,
    },
    ExportArtifactToNode {
        artifact: String,
        receiver_node: String,
        direct_connectivity: bool,
        failure_reason: String,
    },
}

fn default_download_ttl_seconds() -> u64 {
    900
}

fn default_node_enrollment_ttl_seconds() -> u64 {
    900
}