use async_trait::async_trait;
use starweaver_context::ResumableState;
use starweaver_core::{RunId, SessionId};
use starweaver_runtime::{AgentCheckpoint, AgentStreamRecord};
use crate::{
approval::{ApprovalRecord, DeferredToolRecord},
error::SessionStoreResult,
records::{
EnvironmentStateRef, RunRecord, RunStatus, SessionRecord, SessionStatus, StreamCursorRef,
},
resume::SessionResumeSnapshot,
trace::{CompactRunTrace, CompactSessionTrace},
};
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct SessionFilter {
pub status: Option<SessionStatus>,
pub profile: Option<String>,
pub workspace: Option<String>,
pub limit: Option<usize>,
}
#[async_trait]
pub trait SessionStore: Send + Sync {
async fn save_session(&self, session: SessionRecord) -> SessionStoreResult<()>;
async fn load_session(&self, session_id: &SessionId) -> SessionStoreResult<SessionRecord>;
async fn list_sessions(&self, filter: SessionFilter) -> SessionStoreResult<Vec<SessionRecord>>;
async fn update_session_status(
&self,
session_id: &SessionId,
status: SessionStatus,
) -> SessionStoreResult<()>;
async fn save_context_state(
&self,
session_id: &SessionId,
state: ResumableState,
) -> SessionStoreResult<()>;
async fn save_environment_state(
&self,
session_id: &SessionId,
environment_state: EnvironmentStateRef,
) -> SessionStoreResult<()>;
async fn append_run(&self, run: RunRecord) -> SessionStoreResult<()>;
async fn load_run(
&self,
session_id: &SessionId,
run_id: &RunId,
) -> SessionStoreResult<RunRecord>;
async fn list_runs(&self, session_id: &SessionId) -> SessionStoreResult<Vec<RunRecord>>;
async fn update_run_status(
&self,
session_id: &SessionId,
run_id: &RunId,
status: RunStatus,
output_preview: Option<String>,
) -> SessionStoreResult<()>;
async fn append_checkpoint(
&self,
session_id: &SessionId,
checkpoint: AgentCheckpoint,
) -> SessionStoreResult<()>;
async fn load_checkpoints(
&self,
session_id: &SessionId,
run_id: &RunId,
) -> SessionStoreResult<Vec<AgentCheckpoint>>;
async fn latest_checkpoint(
&self,
session_id: &SessionId,
run_id: &RunId,
) -> SessionStoreResult<Option<AgentCheckpoint>> {
let checkpoints = self.load_checkpoints(session_id, run_id).await?;
Ok(checkpoints.into_iter().last())
}
async fn append_stream_records(
&self,
session_id: &SessionId,
run_id: &RunId,
records: Vec<AgentStreamRecord>,
) -> SessionStoreResult<()>;
async fn replay_stream_records(
&self,
session_id: &SessionId,
run_id: &RunId,
) -> SessionStoreResult<Vec<AgentStreamRecord>>;
async fn replay_stream_records_after(
&self,
session_id: &SessionId,
run_id: &RunId,
after_sequence: Option<usize>,
) -> SessionStoreResult<Vec<AgentStreamRecord>> {
let records = self.replay_stream_records(session_id, run_id).await?;
Ok(records
.into_iter()
.filter(|record| after_sequence.is_none_or(|cursor| record.sequence > cursor))
.collect())
}
async fn save_stream_cursor(
&self,
session_id: &SessionId,
run_id: &RunId,
cursor: StreamCursorRef,
) -> SessionStoreResult<()>;
async fn append_approval(&self, approval: ApprovalRecord) -> SessionStoreResult<()>;
async fn load_approvals(
&self,
session_id: &SessionId,
run_id: &RunId,
) -> SessionStoreResult<Vec<ApprovalRecord>>;
async fn append_deferred_tool(&self, record: DeferredToolRecord) -> SessionStoreResult<()>;
async fn load_deferred_tools(
&self,
session_id: &SessionId,
run_id: &RunId,
) -> SessionStoreResult<Vec<DeferredToolRecord>>;
async fn resume_snapshot(
&self,
session_id: &SessionId,
run_id: &RunId,
) -> SessionStoreResult<SessionResumeSnapshot> {
let session = self.load_session(session_id).await?;
let run = self.load_run(session_id, run_id).await?;
let latest_checkpoint = self.latest_checkpoint(session_id, run_id).await?;
let after_sequence = latest_checkpoint
.as_ref()
.and_then(|checkpoint| checkpoint.resume.cursor.stream_cursor);
let stream_records = self
.replay_stream_records_after(session_id, run_id, after_sequence)
.await?;
let approvals = self.load_approvals(session_id, run_id).await?;
let deferred_tools = self.load_deferred_tools(session_id, run_id).await?;
let environment_state = run
.environment_state
.clone()
.or_else(|| session.environment_state.clone());
let mut stream_cursors = session.stream_cursors.clone();
stream_cursors.extend(run.stream_cursors.clone());
Ok(SessionResumeSnapshot {
state: session.state.clone(),
session,
run,
environment_state,
latest_checkpoint,
stream_records,
approvals,
deferred_tools,
stream_cursors,
})
}
async fn compact_run_trace(
&self,
session_id: &SessionId,
run_id: &RunId,
) -> SessionStoreResult<CompactRunTrace>;
async fn compact_session_trace(
&self,
session_id: &SessionId,
) -> SessionStoreResult<CompactSessionTrace>;
}