use std::collections::{BTreeMap, BTreeSet};
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{mpsc, Arc};
use std::time::{Duration, Instant};
use notify::{Event, RecommendedWatcher, RecursiveMode, Watcher};
use serde::Serialize;
use crate::{
DiscoveryQuery, HarnessCatalog, HarnessId, SessionDescriptor, SessionLocator, StorageLocator,
};
const RECONCILE_INTERVAL: Duration = Duration::from_secs(60);
const MAX_SUBSCRIPTION_ROWS: usize = 2_048;
const INVALIDATION_QUEUE_CAPACITY: usize = 1_024;
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize)]
pub struct SessionIndexKey {
pub harness: String,
pub session_id: String,
}
impl SessionIndexKey {
fn from_locator(locator: &SessionLocator) -> Self {
Self {
harness: locator.harness.as_str().to_string(),
session_id: locator.session_id.clone(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum SessionIndexChange {
Added {
descriptor: SessionDescriptor,
},
Updated {
descriptor: SessionDescriptor,
},
Removed {
key: SessionIndexKey,
},
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct SessionIndexDelta {
pub revision: u64,
pub changes: Vec<SessionIndexChange>,
}
pub(crate) struct SessionIndexSubscription {
query: DiscoveryQuery,
current: BTreeMap<SessionIndexKey, SessionDescriptor>,
revision: u64,
receiver: mpsc::Receiver<notify::Result<Event>>,
overflowed: Arc<AtomicBool>,
_watcher: RecommendedWatcher,
last_reconcile: Instant,
}
impl SessionIndexSubscription {
pub(crate) fn homes(&self) -> &crate::HarnessHomes {
&self.query.homes
}
pub(crate) fn open(
mut query: DiscoveryQuery,
) -> Result<(Self, Vec<SessionDescriptor>), String> {
validate_query(&query)?;
query.cursor = None;
query.limit = Some(query.limit.unwrap_or(100));
let initial = HarnessCatalog::new()
.discover_page(&query)
.map_err(|error| error.to_string())?
.sessions;
let current = descriptor_map(initial.iter().cloned());
let (sender, receiver) = mpsc::sync_channel(INVALIDATION_QUEUE_CAPACITY);
let overflowed = Arc::new(AtomicBool::new(false));
let callback_overflowed = Arc::clone(&overflowed);
let mut watcher = notify::recommended_watcher(move |event| {
if sender.try_send(event).is_err() {
callback_overflowed.store(true, Ordering::Release);
}
})
.map_err(|error| error.to_string())?;
for root in watch_roots(&query) {
if let Some(watched) = existing_watch_root(&root) {
watcher
.watch(&watched, RecursiveMode::Recursive)
.map_err(|error| format!("cannot watch {}: {error}", watched.display()))?;
}
}
Ok((
Self {
query,
current,
revision: 1,
receiver,
overflowed,
_watcher: watcher,
last_reconcile: Instant::now(),
},
initial,
))
}
pub(crate) fn poll(&mut self) -> Result<Option<SessionIndexDelta>, String> {
let mut paths = BTreeSet::new();
let mut reconcile = self.overflowed.swap(false, Ordering::AcqRel);
while let Ok(event) = self.receiver.try_recv() {
match event {
Ok(event) => paths.extend(event.paths),
Err(_) => reconcile = true,
}
}
if self.last_reconcile.elapsed() >= RECONCILE_INTERVAL {
reconcile = true;
}
if paths.is_empty() && !reconcile {
return Ok(None);
}
let before = self.current.clone();
if reconcile {
self.reconcile()?;
} else {
let mut needs_fill = false;
for path in paths {
needs_fill |= self.refresh_path(&path)?;
}
if needs_fill {
self.reconcile()?;
} else {
self.retain_page_limit();
}
}
let changes = diff_descriptors(&before, &self.current);
if changes.is_empty() {
return Ok(None);
}
self.revision = self.revision.saturating_add(1);
Ok(Some(SessionIndexDelta {
revision: self.revision,
changes,
}))
}
fn reconcile(&mut self) -> Result<(), String> {
self.last_reconcile = Instant::now();
let sessions = HarnessCatalog::new()
.discover_page(&self.query)
.map_err(|error| error.to_string())?
.sessions;
self.current = descriptor_map(sessions);
Ok(())
}
fn refresh_path(&mut self, path: &Path) -> Result<bool, String> {
if path.extension().and_then(|value| value.to_str()) != Some("jsonl") {
return Ok(false);
}
let event_path = normalized_path(path);
let known = self.current.iter().find_map(|(key, descriptor)| {
(normalized_path(descriptor.locator.storage.path()) == event_path)
.then(|| (key.clone(), descriptor.clone()))
});
let locator = match &known {
Some((_, descriptor)) => descriptor.locator.clone(),
None => match locator_for_path(&self.query, &event_path) {
Some(locator) => locator,
None => return Ok(false),
},
};
let refreshed = HarnessCatalog::new()
.refresh_file_descriptor(
&locator,
self.query.workspace.as_deref(),
self.query.include_topic_candidates,
)
.map_err(|error| error.to_string())?;
match (known, refreshed) {
(Some((old_key, _)), None) => {
self.current.remove(&old_key);
Ok(true)
}
(Some((old_key, _)), Some(descriptor)) => {
self.current.remove(&old_key);
self.current.insert(
SessionIndexKey::from_locator(&descriptor.locator),
descriptor,
);
Ok(false)
}
(None, Some(descriptor)) => {
self.current.insert(
SessionIndexKey::from_locator(&descriptor.locator),
descriptor,
);
Ok(false)
}
(None, None) => Ok(false),
}
}
fn retain_page_limit(&mut self) {
let limit = self.query.limit.unwrap_or(100);
let mut sessions = self.current.values().cloned().collect::<Vec<_>>();
sort_descriptors(&mut sessions);
sessions.truncate(limit);
self.current = descriptor_map(sessions);
}
}
pub(crate) fn validate_query(query: &DiscoveryQuery) -> Result<(), String> {
if query.cursor.is_some() {
return Err("sessions.index.subscribe does not accept a cursor".into());
}
let limit = query.limit.unwrap_or(100);
if limit == 0 || limit > MAX_SUBSCRIPTION_ROWS {
return Err(format!(
"sessions.index.subscribe limit must be between 1 and {MAX_SUBSCRIPTION_ROWS}"
));
}
if query.harnesses.is_empty()
|| query
.harnesses
.iter()
.any(|harness| !matches!(harness.as_str(), HarnessId::CLAUDE_CODE | HarnessId::CODEX))
{
return Err(
"sessions.index.subscribe currently requires explicit claude-code and/or codex harnesses"
.into(),
);
}
Ok(())
}
fn watch_roots(query: &DiscoveryQuery) -> BTreeSet<PathBuf> {
query
.harnesses
.iter()
.filter_map(|harness| match harness.as_str() {
HarnessId::CLAUDE_CODE => Some(query.homes.claude_code.clone()),
HarnessId::CODEX => Some(query.homes.codex.clone()),
_ => None,
})
.collect()
}
fn existing_watch_root(root: &Path) -> Option<PathBuf> {
if root.is_dir() {
return Some(root.to_path_buf());
}
root.parent()
.filter(|parent| parent.is_dir())
.map(Path::to_path_buf)
}
fn locator_for_path(query: &DiscoveryQuery, path: &Path) -> Option<SessionLocator> {
let claude_root = normalized_path(&query.homes.claude_code);
let codex_root = normalized_path(&query.homes.codex);
let harness = if query
.harnesses
.iter()
.any(|harness| harness.as_str() == HarnessId::CLAUDE_CODE)
&& path.starts_with(&claude_root)
{
if path
.components()
.any(|component| component.as_os_str() == "subagents")
{
return None;
}
HarnessId::CLAUDE_CODE
} else if query
.harnesses
.iter()
.any(|harness| harness.as_str() == HarnessId::CODEX)
&& path.starts_with(&codex_root)
{
HarnessId::CODEX
} else {
return None;
};
Some(SessionLocator {
harness: HarnessId::new(harness),
session_id: path
.file_stem()
.and_then(|value| value.to_str())
.unwrap_or("unknown")
.to_string(),
storage: StorageLocator::File {
path: path.to_path_buf(),
},
})
}
fn normalized_path(path: &Path) -> PathBuf {
fs::canonicalize(path).unwrap_or_else(|_| path.to_path_buf())
}
fn descriptor_map(
descriptors: impl IntoIterator<Item = SessionDescriptor>,
) -> BTreeMap<SessionIndexKey, SessionDescriptor> {
descriptors
.into_iter()
.map(|descriptor| {
(
SessionIndexKey::from_locator(&descriptor.locator),
descriptor,
)
})
.collect()
}
fn sort_descriptors(descriptors: &mut [SessionDescriptor]) {
descriptors.sort_by(|left, right| {
right
.updated_at_ms
.cmp(&left.updated_at_ms)
.then_with(|| left.locator.harness.cmp(&right.locator.harness))
.then_with(|| left.locator.session_id.cmp(&right.locator.session_id))
});
}
fn diff_descriptors(
before: &BTreeMap<SessionIndexKey, SessionDescriptor>,
after: &BTreeMap<SessionIndexKey, SessionDescriptor>,
) -> Vec<SessionIndexChange> {
let mut changes = Vec::new();
for (key, descriptor) in after {
match before.get(key) {
None => changes.push(SessionIndexChange::Added {
descriptor: descriptor.clone(),
}),
Some(previous) if previous != descriptor => {
changes.push(SessionIndexChange::Updated {
descriptor: descriptor.clone(),
});
}
Some(_) => {}
}
}
for key in before.keys() {
if !after.contains_key(key) {
changes.push(SessionIndexChange::Removed { key: key.clone() });
}
}
changes
}
#[cfg(test)]
mod tests {
use super::*;
fn descriptor(id: &str, updated_at_ms: u64) -> SessionDescriptor {
SessionDescriptor {
locator: SessionLocator {
harness: HarnessId::new(HarnessId::CODEX),
session_id: id.into(),
storage: StorageLocator::File {
path: PathBuf::from(format!("/{id}.jsonl")),
},
},
cwd: None,
title: None,
preview_candidates: Vec::new(),
latest_message_candidates: Vec::new(),
updated_at_ms: Some(updated_at_ms),
message_count: None,
model: None,
}
}
#[test]
fn index_delta_is_a_complete_deterministic_replacement_set() {
let before = descriptor_map([descriptor("removed", 1), descriptor("updated", 2)]);
let after = descriptor_map([descriptor("updated", 3), descriptor("added", 4)]);
let changes = diff_descriptors(&before, &after);
assert!(matches!(
&changes[0],
SessionIndexChange::Added { descriptor } if descriptor.locator.session_id == "added"
));
assert!(matches!(
&changes[1],
SessionIndexChange::Updated { descriptor } if descriptor.locator.session_id == "updated"
));
assert!(matches!(
&changes[2],
SessionIndexChange::Removed { key } if key.session_id == "removed"
));
}
}