426 lines
12 KiB
Rust
426 lines
12 KiB
Rust
//! Free functions for session persistence operations.
|
||
//!
|
||
//! These functions record and restore session state without owning a Worker.
|
||
//! The caller (typically Pod) holds the Worker directly and calls these
|
||
//! functions after state-mutating operations.
|
||
|
||
use crate::SessionId;
|
||
use crate::session_log::{self, EntryHash, HashedEntry, LogEntry, Outcome, SessionOrigin};
|
||
use crate::store::{Store, StoreError};
|
||
use llm_worker::llm_client::RequestConfig;
|
||
use llm_worker::llm_client::types::Item;
|
||
|
||
/// State snapshot for creating a SessionStart entry.
|
||
pub struct SessionStartState<'a> {
|
||
pub system_prompt: Option<&'a str>,
|
||
pub config: &'a RequestConfig,
|
||
pub history: &'a [Item],
|
||
}
|
||
|
||
/// Create a new session, writing the initial `SessionStart` entry.
|
||
///
|
||
/// Returns the new session ID and head hash.
|
||
pub async fn create_session(
|
||
store: &impl Store,
|
||
state: SessionStartState<'_>,
|
||
) -> Result<(SessionId, EntryHash), StoreError> {
|
||
let session_id = crate::new_session_id();
|
||
let hash = create_session_with_id(store, session_id, state).await?;
|
||
Ok((session_id, hash))
|
||
}
|
||
|
||
/// Write a fresh `SessionStart` entry using a pre-generated session ID.
|
||
///
|
||
/// Used by callers that need to reserve a session ID synchronously but
|
||
/// defer the initial log append (e.g. Pod, which resolves a templated
|
||
/// system prompt only at first turn). Returns the resulting head hash.
|
||
pub async fn create_session_with_id(
|
||
store: &impl Store,
|
||
session_id: SessionId,
|
||
state: SessionStartState<'_>,
|
||
) -> Result<EntryHash, StoreError> {
|
||
let entry = LogEntry::SessionStart {
|
||
ts: session_log::now_millis(),
|
||
system_prompt: state.system_prompt.map(String::from),
|
||
config: state.config.clone(),
|
||
history: state.history.to_vec(),
|
||
forked_from: None,
|
||
compacted_from: None,
|
||
};
|
||
let hash = session_log::compute_hash(None, &entry);
|
||
let hashed_entry = HashedEntry {
|
||
hash: hash.clone(),
|
||
prev_hash: None,
|
||
entry,
|
||
};
|
||
store.append(session_id, &hashed_entry).await?;
|
||
Ok(hash)
|
||
}
|
||
|
||
/// Create a compacted session from an existing one.
|
||
///
|
||
/// Records `compacted_from` provenance linking back to the source session.
|
||
/// Returns the new session ID and head hash.
|
||
pub async fn create_compacted_session(
|
||
store: &impl Store,
|
||
state: SessionStartState<'_>,
|
||
source_session_id: SessionId,
|
||
source_head_hash: EntryHash,
|
||
) -> Result<(SessionId, EntryHash), StoreError> {
|
||
let session_id = crate::new_session_id();
|
||
let entry = LogEntry::SessionStart {
|
||
ts: session_log::now_millis(),
|
||
system_prompt: state.system_prompt.map(String::from),
|
||
config: state.config.clone(),
|
||
history: state.history.to_vec(),
|
||
forked_from: None,
|
||
compacted_from: Some(SessionOrigin {
|
||
session_id: source_session_id,
|
||
at_hash: source_head_hash,
|
||
}),
|
||
};
|
||
let hash = session_log::compute_hash(None, &entry);
|
||
let hashed_entry = HashedEntry {
|
||
hash: hash.clone(),
|
||
prev_hash: None,
|
||
entry,
|
||
};
|
||
store.append(session_id, &hashed_entry).await?;
|
||
Ok((session_id, hash))
|
||
}
|
||
|
||
/// Restore session state from a stored log.
|
||
///
|
||
/// Returns the reconstructed state. The caller is responsible for
|
||
/// applying it to a Worker.
|
||
pub async fn restore(
|
||
store: &impl Store,
|
||
session_id: SessionId,
|
||
) -> Result<crate::session_log::RestoredState, StoreError> {
|
||
let entries = store.read_all(session_id).await?;
|
||
Ok(session_log::collect_state(&entries))
|
||
}
|
||
|
||
/// Check if the store's head still matches the expected head hash.
|
||
/// If not, auto-fork into a new session.
|
||
///
|
||
/// Updates `session_id` and `head_hash` in place when a fork occurs.
|
||
pub async fn ensure_head_or_fork(
|
||
store: &impl Store,
|
||
session_id: &mut SessionId,
|
||
head_hash: &mut Option<EntryHash>,
|
||
state: SessionStartState<'_>,
|
||
) -> Result<(), StoreError> {
|
||
let store_head = store.read_head_hash(*session_id).await?;
|
||
if store_head == *head_hash {
|
||
return Ok(());
|
||
}
|
||
let fork_id = crate::new_session_id();
|
||
let entry = LogEntry::SessionStart {
|
||
ts: session_log::now_millis(),
|
||
system_prompt: state.system_prompt.map(String::from),
|
||
config: state.config.clone(),
|
||
history: state.history.to_vec(),
|
||
forked_from: None,
|
||
compacted_from: None,
|
||
};
|
||
let hash = session_log::compute_hash(None, &entry);
|
||
let hashed_entry = HashedEntry {
|
||
hash: hash.clone(),
|
||
prev_hash: None,
|
||
entry,
|
||
};
|
||
store.create_session(fork_id, &[hashed_entry]).await?;
|
||
*session_id = fork_id;
|
||
*head_hash = Some(hash);
|
||
Ok(())
|
||
}
|
||
|
||
/// Log the history delta — new items added since the previous snapshot.
|
||
///
|
||
/// Classifies items into UserInput, AssistantItems, ToolResults, and
|
||
/// HookInjectedItems entries automatically.
|
||
pub async fn save_delta(
|
||
store: &impl Store,
|
||
session_id: SessionId,
|
||
head_hash: &mut Option<EntryHash>,
|
||
new_items: &[Item],
|
||
) -> Result<(), StoreError> {
|
||
if new_items.is_empty() {
|
||
return Ok(());
|
||
}
|
||
|
||
let ts = session_log::now_millis();
|
||
let mut i = 0;
|
||
|
||
while i < new_items.len() {
|
||
let item = &new_items[i];
|
||
if item.is_user_message() {
|
||
append_entry(
|
||
store,
|
||
session_id,
|
||
head_hash,
|
||
LogEntry::UserInput {
|
||
ts,
|
||
item: new_items[i].clone(),
|
||
},
|
||
)
|
||
.await?;
|
||
i += 1;
|
||
} else if item.is_tool_result() {
|
||
let start = i;
|
||
while i < new_items.len() && new_items[i].is_tool_result() {
|
||
i += 1;
|
||
}
|
||
append_entry(
|
||
store,
|
||
session_id,
|
||
head_hash,
|
||
LogEntry::ToolResults {
|
||
ts,
|
||
items: new_items[start..i].to_vec(),
|
||
},
|
||
)
|
||
.await?;
|
||
} else if item.is_assistant_message() || item.is_tool_call() || item.is_reasoning() {
|
||
let start = i;
|
||
while i < new_items.len()
|
||
&& (new_items[i].is_assistant_message()
|
||
|| new_items[i].is_tool_call()
|
||
|| new_items[i].is_reasoning())
|
||
{
|
||
i += 1;
|
||
}
|
||
append_entry(
|
||
store,
|
||
session_id,
|
||
head_hash,
|
||
LogEntry::AssistantItems {
|
||
ts,
|
||
items: new_items[start..i].to_vec(),
|
||
},
|
||
)
|
||
.await?;
|
||
} else {
|
||
append_entry(
|
||
store,
|
||
session_id,
|
||
head_hash,
|
||
LogEntry::HookInjectedItems {
|
||
ts,
|
||
items: vec![new_items[i].clone()],
|
||
},
|
||
)
|
||
.await?;
|
||
i += 1;
|
||
}
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// Log a TurnEnd entry.
|
||
pub async fn save_turn_end(
|
||
store: &impl Store,
|
||
session_id: SessionId,
|
||
head_hash: &mut Option<EntryHash>,
|
||
turn_count: usize,
|
||
) -> Result<(), StoreError> {
|
||
append_entry(
|
||
store,
|
||
session_id,
|
||
head_hash,
|
||
LogEntry::TurnEnd {
|
||
ts: session_log::now_millis(),
|
||
turn_count,
|
||
},
|
||
)
|
||
.await
|
||
}
|
||
|
||
/// Log a RunOutcome entry.
|
||
pub async fn save_outcome(
|
||
store: &impl Store,
|
||
session_id: SessionId,
|
||
head_hash: &mut Option<EntryHash>,
|
||
outcome: Outcome,
|
||
interrupted: bool,
|
||
) -> Result<(), StoreError> {
|
||
append_entry(
|
||
store,
|
||
session_id,
|
||
head_hash,
|
||
LogEntry::RunOutcome {
|
||
ts: session_log::now_millis(),
|
||
outcome,
|
||
interrupted,
|
||
},
|
||
)
|
||
.await
|
||
}
|
||
|
||
/// Log an `LlmUsage` entry — 1 LLM リクエスト分の Usage スナップショット。
|
||
///
|
||
/// `history_len` は送信時の `history.len()`。`input_total_tokens` は
|
||
/// その prefix をプロバイダが実測した占有量（プロンプト全長）で、
|
||
/// プロバイダ別の正規化（Anthropic では `input + cache_read + cache_creation`）を
|
||
/// 済ませた値を渡す。
|
||
pub async fn save_usage(
|
||
store: &impl Store,
|
||
session_id: SessionId,
|
||
head_hash: &mut Option<EntryHash>,
|
||
history_len: usize,
|
||
input_total_tokens: u64,
|
||
cache_read_tokens: u64,
|
||
cache_write_tokens: u64,
|
||
output_tokens: u64,
|
||
) -> Result<(), StoreError> {
|
||
append_entry(
|
||
store,
|
||
session_id,
|
||
head_hash,
|
||
LogEntry::LlmUsage {
|
||
ts: session_log::now_millis(),
|
||
history_len,
|
||
input_total_tokens,
|
||
cache_read_tokens,
|
||
cache_write_tokens,
|
||
output_tokens,
|
||
},
|
||
)
|
||
.await
|
||
}
|
||
|
||
/// Log a `Locked` entry (KV cache locked).
|
||
pub async fn save_cache_locked(
|
||
store: &impl Store,
|
||
session_id: SessionId,
|
||
head_hash: &mut Option<EntryHash>,
|
||
locked_prefix_len: usize,
|
||
) -> Result<(), StoreError> {
|
||
append_entry(
|
||
store,
|
||
session_id,
|
||
head_hash,
|
||
LogEntry::Locked {
|
||
ts: session_log::now_millis(),
|
||
locked_prefix_len,
|
||
},
|
||
)
|
||
.await
|
||
}
|
||
|
||
/// Log a `CacheUnlocked` entry.
|
||
pub async fn save_cache_unlocked(
|
||
store: &impl Store,
|
||
session_id: SessionId,
|
||
head_hash: &mut Option<EntryHash>,
|
||
) -> Result<(), StoreError> {
|
||
append_entry(
|
||
store,
|
||
session_id,
|
||
head_hash,
|
||
LogEntry::CacheUnlocked {
|
||
ts: session_log::now_millis(),
|
||
},
|
||
)
|
||
.await
|
||
}
|
||
|
||
/// Log a `ConfigChanged` entry.
|
||
pub async fn save_config_changed(
|
||
store: &impl Store,
|
||
session_id: SessionId,
|
||
head_hash: &mut Option<EntryHash>,
|
||
config: &RequestConfig,
|
||
) -> Result<(), StoreError> {
|
||
append_entry(
|
||
store,
|
||
session_id,
|
||
head_hash,
|
||
LogEntry::ConfigChanged {
|
||
ts: session_log::now_millis(),
|
||
config: config.clone(),
|
||
},
|
||
)
|
||
.await
|
||
}
|
||
|
||
/// Fork the current state into a new session.
|
||
pub async fn fork(
|
||
store: &impl Store,
|
||
state: SessionStartState<'_>,
|
||
) -> Result<SessionId, StoreError> {
|
||
let fork_id = crate::new_session_id();
|
||
let entry = LogEntry::SessionStart {
|
||
ts: session_log::now_millis(),
|
||
system_prompt: state.system_prompt.map(String::from),
|
||
config: state.config.clone(),
|
||
history: state.history.to_vec(),
|
||
forked_from: None,
|
||
compacted_from: None,
|
||
};
|
||
let hash = session_log::compute_hash(None, &entry);
|
||
let hashed_entry = HashedEntry {
|
||
hash,
|
||
prev_hash: None,
|
||
entry,
|
||
};
|
||
store.create_session(fork_id, &[hashed_entry]).await?;
|
||
Ok(fork_id)
|
||
}
|
||
|
||
/// Fork from an arbitrary point in a stored session's log.
|
||
pub async fn fork_at(
|
||
store: &impl Store,
|
||
source_id: SessionId,
|
||
at_hash: &EntryHash,
|
||
) -> Result<SessionId, StoreError> {
|
||
let entries = store.read_all(source_id).await?;
|
||
let cut = entries
|
||
.iter()
|
||
.position(|e| &e.hash == at_hash)
|
||
.map(|i| i + 1)
|
||
.unwrap_or(entries.len());
|
||
let state = session_log::collect_state(&entries[..cut]);
|
||
|
||
let fork_id = crate::new_session_id();
|
||
let entry = LogEntry::SessionStart {
|
||
ts: session_log::now_millis(),
|
||
system_prompt: state.system_prompt,
|
||
config: state.config,
|
||
history: state.history,
|
||
forked_from: Some(session_log::SessionOrigin {
|
||
session_id: source_id,
|
||
at_hash: at_hash.clone(),
|
||
}),
|
||
compacted_from: None,
|
||
};
|
||
let hash = session_log::compute_hash(None, &entry);
|
||
let hashed_entry = HashedEntry {
|
||
hash,
|
||
prev_hash: None,
|
||
entry,
|
||
};
|
||
store.create_session(fork_id, &[hashed_entry]).await?;
|
||
Ok(fork_id)
|
||
}
|
||
|
||
// ── Private helper ──────────────────────────────────────────────────────
|
||
|
||
async fn append_entry(
|
||
store: &impl Store,
|
||
session_id: SessionId,
|
||
head_hash: &mut Option<EntryHash>,
|
||
entry: LogEntry,
|
||
) -> Result<(), StoreError> {
|
||
let hash = session_log::compute_hash(head_hash.as_ref(), &entry);
|
||
let hashed_entry = HashedEntry {
|
||
hash: hash.clone(),
|
||
prev_hash: head_hash.clone(),
|
||
entry,
|
||
};
|
||
store.append(session_id, &hashed_entry).await?;
|
||
*head_hash = Some(hash);
|
||
Ok(())
|
||
}
|