use std::{collections::BTreeSet, fs, path::Path, path::PathBuf, sync::Arc};
use chrono::Utc;
use serde::Serialize;
use serde_json::Value;
use starweaver_agent::ResumableState;
use starweaver_core::{ConversationId, RunId, SessionId};
use starweaver_environment::EnvironmentState;
use starweaver_model::{ModelMessage, ModelRequest, ModelRequestPart, ToolReturnPart};
use starweaver_runtime::{AgentStreamEvent, AgentStreamRecord};
use starweaver_session::{
ApprovalRecord, ApprovalStatus, DeferredToolRecord, EnvironmentStateRef, ExecutionStatus,
InputPart, RunRecord, RunStatus, SessionRecord, SessionSearchError, SessionSearchPage,
SessionSearchProvider, SessionSearchQuery, SessionSearchScope, SessionStatus,
SessionStoreError, StreamCursorRef,
};
use starweaver_storage::{LocalSessionSearchProvider, RunEvidenceCommit, SqliteStorage};
use starweaver_stream::{DisplayMessage, ReplayCursor, ReplayScope, ReplaySnapshot};
use uuid::Uuid;
use crate::{CliError, CliResult, config::CliConfig, error::io_error};
mod archive;
mod hitl;
mod replay;
mod session_store;
pub use archive::LocalStreamArchive;
use hitl::{
approval_tool_return, deferred_status_is_unresolved, deferred_tool_return,
existing_resume_tool_return_ids, latest_tool_call_order, pending_hitl_resume_error,
tool_return_control_flow,
};
pub use replay::DisplayReplayWindow;
pub use session_store::LocalSessionStore;
pub struct LocalStore {
storage: SqliteStorage,
file_store_path: PathBuf,
search_scope: SessionSearchScope,
}
pub struct RunArtifacts {
pub state: ResumableState,
pub environment_state: Option<EnvironmentState>,
pub raw_records: Vec<AgentStreamRecord>,
pub display_messages: Vec<DisplayMessage>,
pub display_snapshot: ReplaySnapshot,
pub approvals: Vec<ApprovalRecord>,
pub deferred_tools: Vec<DeferredToolRecord>,
pub status: RunStatus,
}
#[derive(Clone, Debug, Serialize)]
pub struct SessionSummary {
pub session_id: String,
pub title: Option<String>,
pub profile: Option<String>,
pub status: String,
pub head_run_id: Option<String>,
pub head_success_run_id: Option<String>,
pub active_run_id: Option<String>,
pub run_count: usize,
pub last_output_preview: Option<String>,
pub created_at: String,
pub updated_at: String,
}
#[derive(Clone, Debug, Serialize)]
pub struct RunSummary {
pub run_id: String,
pub sequence_no: usize,
pub status: String,
pub restore_from_run_id: Option<String>,
pub output_preview: Option<String>,
pub created_at: String,
pub updated_at: String,
}
#[derive(Clone, Debug, Default, Serialize)]
pub struct TrimReport {
pub sessions_scanned: usize,
pub runs_to_trim: usize,
pub runs_trimmed: usize,
pub bytes_reclaimed: u64,
pub dry_run: bool,
}
impl LocalStore {
pub fn open(config: &CliConfig) -> CliResult<Self> {
crate::config::ensure_config_dirs(config)?;
Ok(Self {
storage: SqliteStorage::open(&config.database_path).map_err(storage_error)?,
file_store_path: config.file_store_path.clone(),
search_scope: SessionSearchScope::local(
config.database_path.to_string_lossy().into_owned(),
),
})
}
pub fn create_session(
&mut self,
profile: &str,
title: Option<String>,
) -> CliResult<SessionRecord> {
self.storage
.create_session(Some(profile.to_string()), title)
.map_err(storage_error)
}
pub fn load_session(&self, session_id: &str) -> CliResult<SessionRecord> {
self.storage
.load_session(&SessionId::from_string(session_id))
.map_err(storage_error)
}
pub fn resolve_session_prefix(&self, value: &str) -> CliResult<String> {
self.storage
.resolve_session_prefix(value)
.map(|session_id| session_id.as_str().to_string())
.map_err(|error| match error {
SessionStoreError::NotFound(_) => CliError::NotFound(value.to_string()),
SessionStoreError::Failed(message) if message.contains("ambiguous") => {
CliError::Usage(message)
}
other => storage_error(other),
})
}
pub fn delete_session(&mut self, session_id: &str) -> CliResult<bool> {
let session_id = SessionId::from_string(session_id);
let deleted = self
.storage
.delete_session(&session_id)
.map_err(storage_error)?;
if deleted {
let path = self
.file_store_path
.join("sessions")
.join(session_id.as_str());
if path.exists() {
fs::remove_dir_all(&path).map_err(|error| io_error(&path, error))?;
}
}
Ok(deleted)
}
pub fn load_run(&self, session_id: &str, run_id: &str) -> CliResult<RunRecord> {
self.storage
.load_run(
&SessionId::from_string(session_id),
&RunId::from_string(run_id),
)
.map_err(storage_error)
}
pub fn latest_session(&self) -> CliResult<Option<SessionRecord>> {
Ok(self
.storage
.list_sessions()
.map_err(storage_error)?
.into_iter()
.find(|session| session.status == SessionStatus::Active))
}
pub fn append_run(
&mut self,
session_id: &str,
prompt: String,
restore_from_run_id: Option<String>,
profile: &str,
) -> CliResult<RunRecord> {
let session = self.load_session(session_id)?;
let mut run = RunRecord::new(session.session_id, RunId::new(), ConversationId::new());
run.restore_from_run_id = restore_from_run_id.map(RunId::from_string);
run.trigger_type = Some("cli".to_string());
run.profile = Some(profile.to_string());
run.input = vec![InputPart::text(prompt)];
self.storage.begin_run(run).map_err(storage_error)
}
pub fn complete_run(
&mut self,
run: &mut RunRecord,
output: String,
artifacts: RunArtifacts,
) -> CliResult<Vec<DisplayMessage>> {
let scope = ReplayScope::run(run.run_id.as_str());
let mut display_snapshot = artifacts.display_snapshot.clone();
if display_snapshot.scope.is_none() {
display_snapshot.scope = Some(scope.clone());
}
let raw_cursor = artifacts.raw_records.last().map(|record| {
StreamCursorRef::new(ReplayCursor::raw_runtime(scope.clone(), record.sequence))
});
let display_cursor = artifacts.display_messages.last().map(|message| {
StreamCursorRef::new(ReplayCursor::display(scope.clone(), message.sequence))
});
let environment_ref =
artifacts
.environment_state
.as_ref()
.map(|state| EnvironmentStateRef {
provider: state.provider_id.clone(),
reference: format!(
"sqlite:run_environment_records/{}/{}",
run.session_id.as_str(),
run.run_id.as_str()
),
revision: Some(format!("{}", state.files.len() + state.resources.len())),
metadata: state.metadata.clone(),
});
run.status = artifacts.status;
run.output_preview = Some(output);
run.updated_at = Utc::now();
run.environment_state.clone_from(&environment_ref);
run.stream_cursors = raw_cursor.into_iter().chain(display_cursor).collect();
let mut commit = RunEvidenceCommit::new(run.clone(), artifacts.state.clone());
commit.environment_state = artifacts
.environment_state
.as_ref()
.map(EnvironmentState::to_json);
commit.stream_records.clone_from(&artifacts.raw_records);
commit.approvals.clone_from(&artifacts.approvals);
commit.deferred_tools.clone_from(&artifacts.deferred_tools);
commit.stream_cursors.clone_from(&run.stream_cursors);
commit
.display_messages
.clone_from(&artifacts.display_messages);
commit.display_snapshot = Some(display_snapshot.clone());
*run = self
.storage
.commit_run_evidence(commit)
.map_err(storage_error)?;
self.write_run_blob(run, "raw.stream.json", &artifacts.raw_records)?;
self.write_run_blob(run, "display.compact.json", &display_snapshot)?;
self.write_run_blob(run, "context.state.json", &artifacts.state)?;
if let Some(environment_state) = artifacts.environment_state.as_ref() {
self.write_run_blob(run, "environment.state.json", &environment_state.to_json())?;
}
Ok(artifacts.display_messages)
}
pub fn fail_run(&mut self, run: &mut RunRecord, message: String) -> CliResult<()> {
self.fail_run_with_messages(run, message, &[])
}
pub fn fail_run_with_messages(
&mut self,
run: &mut RunRecord,
message: String,
messages: &[DisplayMessage],
) -> CliResult<()> {
let session = self.load_session(run.session_id.as_str())?;
run.status = RunStatus::Failed;
run.output_preview = Some(message);
run.updated_at = Utc::now();
let mut commit = RunEvidenceCommit::new(run.clone(), session.state);
commit.display_messages = messages.to_vec();
*run = self
.storage
.commit_run_evidence(commit)
.map_err(storage_error)?;
self.write_run_blob(run, "display.compact.json", &messages)?;
Ok(())
}
pub fn load_restore_state(
&self,
session_id: &str,
run_id: Option<&str>,
) -> CliResult<Option<ResumableState>> {
let Some(run_id) = run_id else {
return Ok(Some(self.load_session(session_id)?.state));
};
let mut state = self
.storage
.load_run_context(
&SessionId::from_string(session_id),
&RunId::from_string(run_id),
)
.map_err(storage_error)?;
if let Some(state) = state.as_mut() {
self.inject_resolved_hitl_tool_returns(session_id, run_id, state)?;
}
Ok(state)
}
fn inject_resolved_hitl_tool_returns(
&self,
session_id: &str,
run_id: &str,
state: &mut ResumableState,
) -> CliResult<()> {
let mut existing_returns = existing_resume_tool_return_ids(&state.message_history);
let tool_call_order = latest_tool_call_order(&state.message_history);
let latest_tool_call_ids = tool_call_order.iter().cloned().collect::<BTreeSet<_>>();
let approvals = self.list_approvals(Some(session_id), Some(run_id))?;
let deferred_tools = self.list_deferred_tools(Some(session_id), Some(run_id))?;
let pending_approvals = approvals
.iter()
.filter(|approval| {
approval.status == ApprovalStatus::Pending
&& !existing_returns.contains(&approval.action_id)
})
.map(|approval| approval.approval_id.clone())
.collect::<Vec<_>>();
let pending_deferred = deferred_tools
.iter()
.filter(|deferred| {
deferred_status_is_unresolved(deferred.status)
&& !existing_returns.contains(&deferred.tool_call_id)
})
.map(|deferred| deferred.deferred_id.clone())
.collect::<Vec<_>>();
if !pending_approvals.is_empty() || !pending_deferred.is_empty() {
return Err(pending_hitl_resume_error(
run_id,
&pending_approvals,
&pending_deferred,
));
}
let mut resolved = Vec::<(String, ModelRequestPart)>::new();
for tool_return in self.list_run_tool_returns(session_id, run_id)? {
if !latest_tool_call_ids.contains(&tool_return.tool_call_id)
|| tool_return_control_flow(&tool_return).is_some()
|| existing_returns.contains(&tool_return.tool_call_id)
{
continue;
}
existing_returns.insert(tool_return.tool_call_id.clone());
resolved.push((
tool_return.tool_call_id.clone(),
ModelRequestPart::ToolReturn(tool_return),
));
}
for approval in approvals {
if existing_returns.contains(&approval.action_id) {
continue;
}
if let Some(tool_return) = approval_tool_return(&approval) {
existing_returns.insert(approval.action_id.clone());
resolved.push((
approval.action_id.clone(),
ModelRequestPart::ToolReturn(tool_return),
));
}
}
for deferred in deferred_tools {
if existing_returns.contains(&deferred.tool_call_id) {
continue;
}
if let Some(tool_return) = deferred_tool_return(&deferred) {
existing_returns.insert(deferred.tool_call_id.clone());
resolved.push((
deferred.tool_call_id.clone(),
ModelRequestPart::ToolReturn(tool_return),
));
}
}
if resolved.is_empty() {
return Ok(());
}
resolved.sort_by_key(|(tool_call_id, _)| {
tool_call_order
.iter()
.position(|known| known == tool_call_id)
.unwrap_or(usize::MAX)
});
let mut metadata = serde_json::Map::new();
metadata.insert(
"starweaver.resume.hitl_results".to_string(),
serde_json::json!(true),
);
metadata.insert(
"starweaver.resume.source_run_id".to_string(),
serde_json::json!(run_id),
);
state
.message_history
.push(ModelMessage::Request(ModelRequest {
parts: resolved.into_iter().map(|(_, part)| part).collect(),
timestamp: Some(Utc::now()),
instructions: None,
run_id: Some(RunId::from_string(run_id)),
conversation_id: state.conversation_id.clone(),
metadata,
}));
Ok(())
}
fn list_run_tool_returns(
&self,
session_id: &str,
run_id: &str,
) -> CliResult<Vec<ToolReturnPart>> {
let records = self
.storage
.load_stream_records(
&SessionId::from_string(session_id),
&RunId::from_string(run_id),
)
.map_err(storage_error)?;
Ok(records
.into_iter()
.filter_map(|record| match record.event {
AgentStreamEvent::ToolReturn { tool_return, .. } => Some(tool_return),
_ => None,
})
.collect())
}
pub fn list_sessions(&self, limit: usize) -> CliResult<Vec<SessionSummary>> {
self.storage
.list_sessions()
.map_err(storage_error)?
.into_iter()
.take(limit)
.map(|session| {
let runs = self
.storage
.list_runs(&session.session_id)
.map_err(storage_error)?;
Ok(SessionSummary {
session_id: session.session_id.as_str().to_string(),
title: session.title,
profile: session.profile,
status: session_status_name(session.status).to_string(),
head_run_id: session.head_run_id.map(|id| id.as_str().to_string()),
head_success_run_id: session
.head_success_run_id
.map(|id| id.as_str().to_string()),
active_run_id: session.active_run_id.map(|id| id.as_str().to_string()),
run_count: runs.len(),
last_output_preview: runs.last().and_then(|run| run.output_preview.clone()),
created_at: session.created_at.to_rfc3339(),
updated_at: session.updated_at.to_rfc3339(),
})
})
.collect()
}
pub fn search_sessions(&self, query: SessionSearchQuery) -> CliResult<SessionSearchPage> {
let provider = LocalSessionSearchProvider::new(
Arc::new(self.storage.session_store()),
&self.search_scope,
)
.with_display_root(self.file_store_path.clone());
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.map_err(|error| CliError::Run(error.to_string()))?;
runtime
.block_on(provider.search(&self.search_scope, query))
.map_err(search_error)
}
pub fn list_runs(&self, session_id: &str, limit: usize) -> CliResult<Vec<RunSummary>> {
let mut runs = self
.storage
.list_runs(&SessionId::from_string(session_id))
.map_err(storage_error)?;
if runs.len() > limit {
runs.drain(..runs.len() - limit);
}
Ok(runs.into_iter().map(run_summary).collect())
}
pub fn replay_display(
&self,
session_id: &str,
run_id: Option<&str>,
after: Option<usize>,
) -> CliResult<Vec<DisplayMessage>> {
let session_id = SessionId::from_string(session_id);
let run_id = run_id.map(RunId::from_string);
self.storage
.load_display_messages(&session_id, run_id.as_ref(), after)
.map_err(storage_error)
}
pub fn trim(
&mut self,
sessions: Vec<String>,
keep_runs: usize,
dry_run: bool,
) -> CliResult<TrimReport> {
self.trim_with_age(sessions, keep_runs, None, dry_run)
}
pub fn trim_with_age(
&mut self,
sessions: Vec<String>,
keep_runs: usize,
older_than: Option<chrono::Duration>,
dry_run: bool,
) -> CliResult<TrimReport> {
let mut report = TrimReport {
sessions_scanned: sessions.len(),
dry_run,
..TrimReport::default()
};
let cutoff = older_than.map(|duration| Utc::now() - duration);
for session_id in sessions {
let session_id = SessionId::from_string(session_id);
let session = self
.storage
.load_session(&session_id)
.map_err(storage_error)?;
let runs = self.storage.list_runs(&session_id).map_err(storage_error)?;
let keep_from = runs.len().saturating_sub(keep_runs);
let candidates = runs
.into_iter()
.take(keep_from)
.filter(|run| session.active_run_id.as_ref() != Some(&run.run_id))
.filter(|run| cutoff.is_none_or(|cutoff| run.updated_at < cutoff))
.collect::<Vec<_>>();
report.runs_to_trim += candidates.len();
for run in &candidates {
report.bytes_reclaimed = report
.bytes_reclaimed
.saturating_add(self.run_file_bytes(session_id.as_str(), run.run_id.as_str())?);
}
if !dry_run && !candidates.is_empty() {
let run_ids = candidates
.iter()
.map(|run| run.run_id.clone())
.collect::<Vec<_>>();
report.runs_trimmed += self
.storage
.prune_runs(&session_id, &run_ids)
.map_err(storage_error)?;
for run_id in run_ids {
self.remove_run_files(session_id.as_str(), run_id.as_str())?;
}
}
}
Ok(report)
}
pub fn all_session_ids(&self) -> CliResult<Vec<String>> {
Ok(self
.storage
.list_sessions()
.map_err(storage_error)?
.into_iter()
.map(|session| session.session_id.as_str().to_string())
.collect())
}
fn write_run_blob<T: Serialize>(
&self,
run: &RunRecord,
name: &str,
value: &T,
) -> CliResult<()> {
let path = self
.file_store_path
.join("sessions")
.join(run.session_id.as_str())
.join("runs")
.join(run.run_id.as_str())
.join(name);
atomic_write_json(&path, value)
}
fn run_file_bytes(&self, session_id: &str, run_id: &str) -> CliResult<u64> {
directory_bytes(
&self
.file_store_path
.join("sessions")
.join(session_id)
.join("runs")
.join(run_id),
)
}
fn remove_run_files(&self, session_id: &str, run_id: &str) -> CliResult<()> {
let path = self
.file_store_path
.join("sessions")
.join(session_id)
.join("runs")
.join(run_id);
if path.exists() {
fs::remove_dir_all(&path).map_err(|error| io_error(&path, error))?;
}
Ok(())
}
pub fn list_approvals(
&self,
session_id: Option<&str>,
run_id: Option<&str>,
) -> CliResult<Vec<ApprovalRecord>> {
let session_id = session_id.map(SessionId::from_string);
let run_id = run_id.map(RunId::from_string);
self.storage
.list_approvals(session_id.as_ref(), run_id.as_ref())
.map_err(storage_error)
}
pub fn load_approval(&self, approval_id: &str) -> CliResult<ApprovalRecord> {
self.storage
.load_approval(approval_id)
.map_err(storage_error)
}
pub fn decide_approval(
&mut self,
approval_id: &str,
status: ApprovalStatus,
reason: Option<String>,
) -> CliResult<ApprovalRecord> {
self.storage
.decide_approval(
approval_id,
status,
Some("starweaver-cli".to_string()),
reason,
)
.map_err(storage_error)
}
pub fn list_deferred_tools(
&self,
session_id: Option<&str>,
run_id: Option<&str>,
) -> CliResult<Vec<DeferredToolRecord>> {
let session_id = session_id.map(SessionId::from_string);
let run_id = run_id.map(RunId::from_string);
self.storage
.list_deferred_tools(session_id.as_ref(), run_id.as_ref())
.map_err(storage_error)
}
pub fn load_deferred_tool(&self, deferred_id: &str) -> CliResult<DeferredToolRecord> {
self.storage
.load_deferred_tool(deferred_id)
.map_err(storage_error)
}
pub fn complete_deferred_tool(
&mut self,
deferred_id: &str,
response: Value,
) -> CliResult<DeferredToolRecord> {
self.storage
.resolve_deferred_tool(deferred_id, ExecutionStatus::Completed, response)
.map_err(storage_error)
}
pub fn fail_deferred_tool(
&mut self,
deferred_id: &str,
error: &str,
) -> CliResult<DeferredToolRecord> {
self.storage
.resolve_deferred_tool(
deferred_id,
ExecutionStatus::Failed,
serde_json::json!({"error": error}),
)
.map_err(storage_error)
}
}
fn run_summary(run: RunRecord) -> RunSummary {
RunSummary {
run_id: run.run_id.as_str().to_string(),
sequence_no: run.sequence_no,
status: run_status_name(run.status).to_string(),
restore_from_run_id: run
.restore_from_run_id
.map(|run_id| run_id.as_str().to_string()),
output_preview: run.output_preview,
created_at: run.created_at.to_rfc3339(),
updated_at: run.updated_at.to_rfc3339(),
}
}
const fn run_status_name(status: RunStatus) -> &'static str {
status.as_str()
}
const fn session_status_name(status: SessionStatus) -> &'static str {
match status {
SessionStatus::Active => "active",
SessionStatus::Archived => "archived",
SessionStatus::Failed => "failed",
SessionStatus::Deleted => "deleted",
}
}
fn storage_error(error: SessionStoreError) -> CliError {
match error {
SessionStoreError::NotFound(value) => CliError::NotFound(value),
other => CliError::Storage(other.to_string()),
}
}
fn search_error(error: SessionSearchError) -> CliError {
match error {
SessionSearchError::InvalidQuery(message) | SessionSearchError::InvalidCursor(message) => {
CliError::Usage(message)
}
SessionSearchError::Unsupported(message) => CliError::Unsupported(message),
SessionSearchError::Unavailable(message) | SessionSearchError::Failed(message) => {
CliError::Storage(message)
}
SessionSearchError::PermissionDenied => {
CliError::Storage("session search permission denied".to_string())
}
}
}
fn atomic_write_json<T: Serialize>(path: &Path, value: &T) -> CliResult<()> {
let parent = path
.parent()
.ok_or_else(|| CliError::Storage("missing parent path".to_string()))?;
fs::create_dir_all(parent).map_err(|error| io_error(parent, error))?;
let temp = path.with_extension(format!("{}.tmp", Uuid::new_v4()));
fs::write(&temp, serde_json::to_vec_pretty(value)?).map_err(|error| io_error(&temp, error))?;
fs::rename(&temp, path).map_err(|error| io_error(path, error))
}
fn directory_bytes(path: &Path) -> CliResult<u64> {
if !path.exists() {
return Ok(0);
}
let mut total = 0_u64;
for entry in fs::read_dir(path).map_err(|error| io_error(path, error))? {
let entry = entry.map_err(|error| io_error(path, error))?;
let entry_path = entry.path();
let metadata = entry
.metadata()
.map_err(|error| io_error(&entry_path, error))?;
if metadata.is_dir() {
total = total.saturating_add(directory_bytes(&entry_path)?);
} else {
total = total.saturating_add(metadata.len());
}
}
Ok(total)
}