crank/crates/crank-core/src/stream_session.rs

use serde::{Deserialize, Serialize};
use serde_json::Value;
use time::OffsetDateTime;

use crate::{
    ids::{AgentId, AsyncJobId, OperationId, StreamSessionId, WorkspaceId},
    protocol::Protocol,
    streaming::ExecutionMode,
};

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum StreamStatus {
    Created,
    Running,
    Stopped,
    Failed,
    Expired,
}

impl StreamStatus {
    pub fn is_terminal(self) -> bool {
        matches!(self, Self::Stopped | Self::Failed | Self::Expired)
    }
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct StreamSession {
    pub id: StreamSessionId,
    pub workspace_id: WorkspaceId,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub agent_id: Option<AgentId>,
    pub operation_id: OperationId,
    pub protocol: Protocol,
    pub mode: ExecutionMode,
    pub status: StreamStatus,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub cursor: Option<Value>,
    pub state: Value,
    #[serde(with = "time::serde::rfc3339")]
    pub expires_at: OffsetDateTime,
    #[serde(skip_serializing_if = "Option::is_none")]
    #[serde(with = "time::serde::rfc3339::option")]
    pub last_poll_at: Option<OffsetDateTime>,
    #[serde(with = "time::serde::rfc3339")]
    pub created_at: OffsetDateTime,
    #[serde(skip_serializing_if = "Option::is_none")]
    #[serde(with = "time::serde::rfc3339::option")]
    pub closed_at: Option<OffsetDateTime>,
}

impl StreamSession {
    pub fn is_expired(&self, now: OffsetDateTime) -> bool {
        self.expires_at <= now
    }

    pub fn can_poll(&self, now: OffsetDateTime) -> bool {
        !self.status.is_terminal() && !self.is_expired(now)
    }

    pub fn remaining_poll_delay_ms(&self, now: OffsetDateTime, poll_interval_ms: u64) -> u64 {
        let Some(last_poll_at) = self.last_poll_at else {
            return 0;
        };

        let elapsed_ms = (now - last_poll_at).whole_milliseconds().max(0) as u64;
        poll_interval_ms.saturating_sub(elapsed_ms)
    }

    pub fn mark_polled(&mut self, now: OffsetDateTime) {
        self.last_poll_at = Some(now);
    }

    pub fn mark_closed(&mut self, now: OffsetDateTime) {
        self.status = StreamStatus::Stopped;
        self.last_poll_at = Some(now);
        self.closed_at = Some(now);
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum JobStatus {
    Created,
    Running,
    Completed,
    Failed,
    Cancelled,
    Expired,
}

impl JobStatus {
    pub fn is_terminal(self) -> bool {
        matches!(
            self,
            Self::Completed | Self::Failed | Self::Cancelled | Self::Expired
        )
    }
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AsyncJobHandle {
    pub id: AsyncJobId,
    pub workspace_id: WorkspaceId,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub agent_id: Option<AgentId>,
    pub operation_id: OperationId,
    pub status: JobStatus,
    pub progress: Value,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub result: Option<Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    #[serde(with = "time::serde::rfc3339::option")]
    pub expires_at: Option<OffsetDateTime>,
    #[serde(skip_serializing_if = "Option::is_none")]
    #[serde(with = "time::serde::rfc3339::option")]
    pub last_poll_at: Option<OffsetDateTime>,
    #[serde(with = "time::serde::rfc3339")]
    pub created_at: OffsetDateTime,
    #[serde(with = "time::serde::rfc3339")]
    pub updated_at: OffsetDateTime,
    #[serde(skip_serializing_if = "Option::is_none")]
    #[serde(with = "time::serde::rfc3339::option")]
    pub finished_at: Option<OffsetDateTime>,
}

impl AsyncJobHandle {
    pub fn is_finished(&self) -> bool {
        self.status.is_terminal()
    }

    pub fn can_cancel(&self) -> bool {
        matches!(self.status, JobStatus::Created | JobStatus::Running)
    }

    pub fn remaining_poll_delay_ms(&self, now: OffsetDateTime, poll_interval_ms: u64) -> u64 {
        let Some(last_poll_at) = self.last_poll_at else {
            return 0;
        };

        let elapsed_ms = (now - last_poll_at).whole_milliseconds().max(0) as u64;
        poll_interval_ms.saturating_sub(elapsed_ms)
    }

    pub fn mark_finished(&mut self, now: OffsetDateTime, result: Value) {
        self.status = JobStatus::Completed;
        self.result = Some(result);
        self.error = None;
        self.updated_at = now;
        self.finished_at = Some(now);
    }

    pub fn mark_failed(&mut self, now: OffsetDateTime, error: Value) {
        self.status = JobStatus::Failed;
        self.error = Some(error);
        self.updated_at = now;
        self.finished_at = Some(now);
    }
}

#[cfg(test)]
mod tests {
    use serde_json::json;
    use time::{OffsetDateTime, format_description::well_known::Rfc3339};

    use super::{AsyncJobHandle, JobStatus, StreamSession, StreamStatus};
    use crate::{
        ids::{AsyncJobId, OperationId, StreamSessionId, WorkspaceId},
        protocol::Protocol,
        streaming::ExecutionMode,
    };

    fn timestamp(value: &str) -> OffsetDateTime {
        OffsetDateTime::parse(value, &Rfc3339).unwrap()
    }

    #[test]
    fn stream_session_tracks_poll_and_close_transitions() {
        let mut session = StreamSession {
            id: StreamSessionId::new("stream_01"),
            workspace_id: WorkspaceId::new("ws_01"),
            agent_id: None,
            operation_id: OperationId::new("op_01"),
            protocol: Protocol::Rest,
            mode: ExecutionMode::Session,
            status: StreamStatus::Running,
            cursor: None,
            state: json!({"cursor":"abc"}),
            expires_at: timestamp("2026-04-06T12:05:00Z"),
            last_poll_at: None,
            created_at: timestamp("2026-04-06T12:00:00Z"),
            closed_at: None,
        };

        assert!(session.can_poll(timestamp("2026-04-06T12:01:00Z")));

        session.mark_polled(timestamp("2026-04-06T12:01:00Z"));
        session.mark_closed(timestamp("2026-04-06T12:02:00Z"));

        assert_eq!(session.status, StreamStatus::Stopped);
        assert_eq!(session.closed_at, Some(timestamp("2026-04-06T12:02:00Z")));
        assert!(!session.can_poll(timestamp("2026-04-06T12:03:00Z")));
    }

    #[test]
    fn first_poll_has_no_required_delay() {
        let session = StreamSession {
            id: StreamSessionId::new("stream_01"),
            workspace_id: WorkspaceId::new("ws_01"),
            agent_id: None,
            operation_id: OperationId::new("op_01"),
            protocol: Protocol::Rest,
            mode: ExecutionMode::Session,
            status: StreamStatus::Running,
            cursor: None,
            state: json!({"cursor":"abc"}),
            expires_at: timestamp("2026-04-06T12:05:00Z"),
            last_poll_at: None,
            created_at: timestamp("2026-04-06T12:00:00Z"),
            closed_at: None,
        };

        assert_eq!(
            session.remaining_poll_delay_ms(timestamp("2026-04-06T12:00:00.100Z"), 250),
            0
        );
    }

    #[test]
    fn rapid_repeat_poll_reports_remaining_delay() {
        let session = StreamSession {
            id: StreamSessionId::new("stream_01"),
            workspace_id: WorkspaceId::new("ws_01"),
            agent_id: None,
            operation_id: OperationId::new("op_01"),
            protocol: Protocol::Rest,
            mode: ExecutionMode::Session,
            status: StreamStatus::Running,
            cursor: None,
            state: json!({"cursor":"abc"}),
            expires_at: timestamp("2026-04-06T12:05:00Z"),
            last_poll_at: Some(timestamp("2026-04-06T12:01:00Z")),
            created_at: timestamp("2026-04-06T12:00:00Z"),
            closed_at: None,
        };

        assert_eq!(
            session.remaining_poll_delay_ms(timestamp("2026-04-06T12:01:00.100Z"), 250),
            150
        );
    }

    #[test]
    fn async_job_tracks_finish_and_failure() {
        let mut job = AsyncJobHandle {
            id: AsyncJobId::new("job_01"),
            workspace_id: WorkspaceId::new("ws_01"),
            agent_id: None,
            operation_id: OperationId::new("op_01"),
            status: JobStatus::Running,
            progress: json!({"percent": 60}),
            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,
        };

        assert!(job.can_cancel());

        job.mark_finished(timestamp("2026-04-06T12:01:00Z"), json!({"ok": true}));
        assert!(job.is_finished());
        assert_eq!(job.status, JobStatus::Completed);

        let mut failed_job = job.clone();
        failed_job.status = JobStatus::Running;
        failed_job.result = None;
        failed_job.finished_at = None;

        failed_job.mark_failed(
            timestamp("2026-04-06T12:02:00Z"),
            json!({"message": "boom"}),
        );
        assert_eq!(failed_job.status, JobStatus::Failed);
        assert_eq!(
            failed_job.finished_at,
            Some(timestamp("2026-04-06T12:02:00Z"))
        );
    }

    #[test]
    fn async_job_reports_remaining_poll_delay() {
        let job = AsyncJobHandle {
            id: AsyncJobId::new("job_01"),
            workspace_id: WorkspaceId::new("ws_01"),
            agent_id: None,
            operation_id: OperationId::new("op_01"),
            status: JobStatus::Running,
            progress: json!({"percent": 60}),
            result: None,
            error: None,
            expires_at: Some(timestamp("2026-04-06T12:05:00Z")),
            last_poll_at: Some(timestamp("2026-04-06T12:01:00Z")),
            created_at: timestamp("2026-04-06T12:00:00Z"),
            updated_at: timestamp("2026-04-06T12:00:00Z"),
            finished_at: None,
        };

        assert_eq!(
            job.remaining_poll_delay_ms(timestamp("2026-04-06T12:01:00.100Z"), 250),
            150
        );
    }
}