198 lines
7.6 KiB
Rust
198 lines
7.6 KiB
Rust
//! A durable, self-compacting ("infinite") claude session.
|
|
|
|
use std::sync::Mutex;
|
|
|
|
use serde_json::Value;
|
|
|
|
use crate::{
|
|
Attach, Claude, CompactionPolicy, Config, Error, Result, SessionStore, Sink, Telemetry,
|
|
};
|
|
|
|
/// A named claude session that outlives the model's context window by
|
|
/// compacting itself. Bundles the three things a durable session needs:
|
|
///
|
|
/// - a **name** (the constant session title it resumes / creates under),
|
|
/// - a **store** ([`SessionStore`], to find the backing file so a resume vs.
|
|
/// create is decided without a wasted spawn), and
|
|
/// - a **policy** ([`CompactionPolicy`], deciding *when* to compact).
|
|
///
|
|
/// [`InfiniteSession::run`] keeps the session alive across turns:
|
|
///
|
|
/// - **resume-or-create** — resumes the titled session, creating it on first
|
|
/// use (or after its file was archived away);
|
|
/// - **reactive** — if a turn overflows ([`Error::PromptTooLong`]), it compacts
|
|
/// and retries the same prompt once;
|
|
/// - **proactive** — after a clean turn it consults the policy and, if due,
|
|
/// runs an optional checkpoint turn then compacts.
|
|
///
|
|
/// Resetting/archiving the session is intentionally *not* part of this type —
|
|
/// that stays with the caller.
|
|
pub struct InfiniteSession<P: CompactionPolicy> {
|
|
name: String,
|
|
store: SessionStore,
|
|
policy: P,
|
|
}
|
|
|
|
/// What [`InfiniteSession::run`] did, beyond streaming the turn to the sink.
|
|
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
|
pub struct Progress {
|
|
/// The turn resumed no existing session — a fresh one was created.
|
|
pub created: bool,
|
|
/// A compaction ran (reactively on overflow, or proactively per policy).
|
|
pub compacted: bool,
|
|
/// Everything parsed from the answering turn's stream (usage, cost,
|
|
/// context window, resolved model). The authoritative copy — finalised at
|
|
/// the turn's `result` event.
|
|
pub telemetry: Telemetry,
|
|
}
|
|
|
|
impl<P: CompactionPolicy> InfiniteSession<P> {
|
|
/// Build a durable session for `name`, backed by `store`, governed by
|
|
/// `policy`.
|
|
pub fn new(name: impl Into<String>, store: SessionStore, policy: P) -> Self {
|
|
Self {
|
|
name: name.into(),
|
|
store,
|
|
policy,
|
|
}
|
|
}
|
|
|
|
/// Run one turn, keeping the session infinite (see the type docs).
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Propagates any non-`SessionNotFound` [`Error`] from the underlying
|
|
/// runs. A reactive retry that *still* overflows surfaces as
|
|
/// [`Error::PromptTooLong`]; rate-limit / auth / hard failures propagate
|
|
/// unchanged for the caller to handle.
|
|
pub async fn run(&self, config: &Config, prompt: &str, sink: &impl Sink) -> Result<Progress> {
|
|
// Resolve resume-vs-create once, up front, so `created` reflects the
|
|
// session state at the START of the run: the reactive-retry path can't
|
|
// read it from the retry (the session exists by then).
|
|
let existed = self.store.find_by_title(&self.name).is_some();
|
|
let meter = TelemetrySink::new(sink);
|
|
let created = match self.attempt(config, prompt, &meter, existed).await {
|
|
Ok(created) => created,
|
|
Err(Error::PromptTooLong) => {
|
|
// The session is already past the window — no turn can run on
|
|
// it and the detail is gone (no checkpoint possible). Compact,
|
|
// then retry the same prompt once; the retry is the answering
|
|
// turn, so its telemetry is what we report. The failed attempt
|
|
// created/resumed the session, so the retry resumes it, and
|
|
// `created` still reflects the pre-run state.
|
|
self.compact(config, sink).await?;
|
|
let retry = TelemetrySink::new(sink);
|
|
self.attempt(config, prompt, &retry, true).await?;
|
|
return Ok(Progress {
|
|
created: !existed,
|
|
compacted: true,
|
|
telemetry: retry.snapshot(),
|
|
});
|
|
}
|
|
Err(other) => return Err(other),
|
|
};
|
|
let telemetry = meter.snapshot();
|
|
|
|
// Proactive: the turn completed on a healthy session. If the policy
|
|
// says it's due, checkpoint (best-effort) then compact.
|
|
if self.policy.should_compact(telemetry.usage()) {
|
|
if let Some(checkpoint) = self.policy.checkpoint_prompt() {
|
|
let _ = self.attempt(config, checkpoint, sink, true).await;
|
|
}
|
|
// Only claim a compaction if it actually succeeded — the flag feeds
|
|
// stats + the auto-reset watermark, so a failed best-effort
|
|
// `/compact` must not report that the context shrank.
|
|
let compacted = self.compact(config, sink).await.is_ok();
|
|
return Ok(Progress {
|
|
created,
|
|
compacted,
|
|
telemetry,
|
|
});
|
|
}
|
|
Ok(Progress {
|
|
created,
|
|
compacted: false,
|
|
telemetry,
|
|
})
|
|
}
|
|
|
|
/// Force a `/compact` on the session (e.g. operator-driven). Resume-only:
|
|
/// if the session doesn't exist there is nothing to compact, so a
|
|
/// [`Error::SessionNotFound`] is swallowed as a no-op `Ok`.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Propagates any error other than `SessionNotFound` from the compact run.
|
|
pub async fn compact(&self, config: &Config, sink: &impl Sink) -> Result<()> {
|
|
match Claude::run(config, &Attach::Resume(self.name.clone()), "/compact", sink).await {
|
|
Err(Error::SessionNotFound) => Ok(()),
|
|
other => other,
|
|
}
|
|
}
|
|
|
|
/// One resume-or-create turn. `existed` is the caller's up-front
|
|
/// resume-vs-create decision (whether the titled session was on disk before
|
|
/// the run) — passing it in rather than re-checking keeps `created`
|
|
/// reporting consistent across the reactive-retry path. Still self-heals if
|
|
/// the backing file vanished between the check and the run. Returns whether
|
|
/// a fresh session was created.
|
|
async fn attempt(
|
|
&self,
|
|
config: &Config,
|
|
prompt: &str,
|
|
sink: &impl Sink,
|
|
existed: bool,
|
|
) -> Result<bool> {
|
|
let attach = if existed {
|
|
Attach::Resume(self.name.clone())
|
|
} else {
|
|
Attach::Create(self.name.clone())
|
|
};
|
|
match Claude::run(config, &attach, prompt, sink).await {
|
|
Ok(()) => Ok(!existed),
|
|
// We thought it existed but the resume missed (raced an archive) —
|
|
// self-heal by creating.
|
|
Err(Error::SessionNotFound) if existed => {
|
|
Claude::run(config, &Attach::Create(self.name.clone()), prompt, sink).await?;
|
|
Ok(true)
|
|
}
|
|
Err(other) => Err(other),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// A [`Sink`] that forwards to an inner sink while accumulating the turn's
|
|
/// [`Telemetry`] from the stream. Cheap; the driver calls it synchronously
|
|
/// from one reader task, so the `Mutex` only satisfies `&self`.
|
|
struct TelemetrySink<'a, S: Sink> {
|
|
inner: &'a S,
|
|
telemetry: Mutex<Telemetry>,
|
|
}
|
|
|
|
impl<'a, S: Sink> TelemetrySink<'a, S> {
|
|
fn new(inner: &'a S) -> Self {
|
|
Self {
|
|
inner,
|
|
telemetry: Mutex::new(Telemetry::default()),
|
|
}
|
|
}
|
|
|
|
fn snapshot(&self) -> Telemetry {
|
|
self.telemetry.lock().unwrap().clone()
|
|
}
|
|
}
|
|
|
|
impl<S: Sink> Sink for TelemetrySink<'_, S> {
|
|
fn on_event(&self, event: &Value) {
|
|
self.telemetry.lock().unwrap().observe(event);
|
|
self.inner.on_event(event);
|
|
}
|
|
|
|
fn on_stdout_line(&self, line: &str) {
|
|
self.inner.on_stdout_line(line);
|
|
}
|
|
|
|
fn on_stderr_line(&self, line: &str) {
|
|
self.inner.on_stderr_line(line);
|
|
}
|
|
}
|