use super::startup::RuntimeContext;
use crate::{
config::EffectiveConfig,
sessions::{BoundedReadError, Session, SessionManager, SessionWriterBusy, chat},
};
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use std::{
ops::Range,
sync::{Arc, Mutex},
thread::JoinHandle,
};
const PAGE_BYTES: usize = 32_768;
const PAGE_ENTRIES: usize = 100;
const CHUNK_BYTES: usize = 16_384;
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct OpenParams {
session_id: String,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct PageParams {
session_id: String,
revision: String,
cursor: Option<String>,
}
pub(super) enum Action {
Create,
Open(String),
Page {
session_id: String,
revision: String,
cursor: Option<String>,
},
}
impl Action {
pub(super) fn parse(
method: &str,
params: &serde_json::value::RawValue,
) -> Result<Self, &'static str> {
if !params.get().starts_with('{') {
return Err("invalid_params");
}
match method {
"session.create" if super::empty_params(params) => Ok(Self::Create),
"session.open" => {
let params: OpenParams =
serde_json::from_str(params.get()).map_err(|_| "invalid_params")?;
validate_id(¶ms.session_id)?;
Ok(Self::Open(params.session_id))
}
"history.page" => {
let params: PageParams =
serde_json::from_str(params.get()).map_err(|_| "invalid_params")?;
validate_id(¶ms.session_id)?;
if !super::valid_id(¶ms.revision)
|| params
.cursor
.as_deref()
.is_some_and(|c| !super::valid_id(c))
{
return Err("invalid_params");
}
Ok(Self::Page {
session_id: params.session_id,
revision: params.revision,
cursor: params.cursor,
})
}
_ => Err("invalid_params"),
}
}
}
fn validate_id(id: &str) -> Result<(), &'static str> {
if id.len() > 64 || crate::sessions::validate_session_id(id.to_owned()).is_err() {
Err("invalid_params")
} else {
Ok(())
}
}
#[derive(Clone, Serialize)]
struct Entry {
message_id: String,
kind: &'static str,
text: String,
partial: bool,
part: usize,
last_part: bool,
}
pub(super) struct Snapshot {
revision: String,
entries: Vec<Entry>,
pages: Vec<Range<usize>>,
fingerprint: ([u8; 32], usize),
complete: bool,
invalidated: bool,
generation: u64,
file_present: bool,
}
impl Snapshot {
pub(super) fn read(session: &Session) -> Result<Self, &'static str> {
let (read, stats) = session
.read_events_tolerant_bounded_with_stats(
crate::context::REPLAY_JSONL_MAX_LINES,
crate::context::REPLAY_JSONL_MAX_BYTES,
)
.map_err(|_| "history_limit")?;
if !read.diagnostics.is_empty() || read.events.iter().any(|e| e.session_id != session.id())
{
return Err("session_invalid");
}
let mut snapshot = Self::from_events(
&read.events,
(stats.content_digest, stats.bytes_read),
session.replay_generation(),
)?;
snapshot.file_present = session.path().exists();
Ok(snapshot)
}
pub(super) fn input_identity(&self) -> Option<(&str, &str)> {
self.entries
.iter()
.rev()
.find(|e| e.kind == "user")
.map(|e| (e.message_id.as_str(), self.revision.as_str()))
}
fn from_events(
events: &[crate::sessions::SessionEvent],
fingerprint: ([u8; 32], usize),
generation: u64,
) -> Result<Self, &'static str> {
let (messages, complete) = chat::project(events)?;
let mut snapshot = Self {
revision: uuid::Uuid::new_v4().to_string(),
entries: Vec::new(),
pages: Vec::new(),
fingerprint,
complete,
invalidated: false,
generation,
file_present: fingerprint.1 > 0,
};
let mut total_bytes = 2;
let mut page_bytes = 512;
let mut page_start = 0;
for message in messages {
let text = crate::output::sanitize_display_text(&message.text);
let message = chat::ChatMessage { text, ..message };
let message_id = uuid::Uuid::new_v4().to_string();
let mut start = 0;
let mut encoded_bytes = 0;
let mut part = 0;
for (offset, character) in message.text.char_indices() {
let cost = serde_json::to_string(&character.to_string())
.map_err(|_| "content_limit")?
.len()
- 2;
if offset - start + character.len_utf8() > CHUNK_BYTES
|| encoded_bytes + cost > 12_000
{
snapshot.push(
Entry {
message_id: message_id.clone(),
kind: message.kind,
text: message.text[start..offset].to_owned(),
partial: message.partial,
part,
last_part: false,
},
&mut total_bytes,
&mut page_bytes,
&mut page_start,
)?;
start = offset;
encoded_bytes = 0;
part += 1;
}
encoded_bytes += cost;
}
snapshot.push(
Entry {
message_id,
kind: message.kind,
text: message.text[start..].to_owned(),
partial: message.partial,
part,
last_part: true,
},
&mut total_bytes,
&mut page_bytes,
&mut page_start,
)?;
}
snapshot.pages.push(page_start..snapshot.entries.len());
Ok(snapshot)
}
fn push(
&mut self,
entry: Entry,
total: &mut usize,
page_bytes: &mut usize,
page_start: &mut usize,
) -> Result<(), &'static str> {
let bytes = serde_json::to_vec(&entry)
.map_err(|_| "content_limit")?
.len()
+ 1;
*total += bytes;
if *total > chat::CHAT_MAX_BYTES || self.entries.len() == chat::CHAT_MAX_ENTRIES {
return Err("history_limit");
}
if bytes + 512 > PAGE_BYTES {
return Err("content_limit");
}
if *page_bytes + bytes > PAGE_BYTES || self.entries.len() - *page_start == PAGE_ENTRIES {
self.pages.push(*page_start..self.entries.len());
*page_start = self.entries.len();
*page_bytes = 512;
}
*page_bytes += bytes;
self.entries.push(entry);
Ok(())
}
fn page(
&mut self,
session: &Session,
revision: &str,
cursor: Option<&str>,
) -> Result<Value, &'static str> {
if self.invalidated || revision != self.revision {
return Err("stale_revision");
}
let page = match cursor {
None => 0,
Some(cursor) => cursor
.strip_prefix(&format!("{}.", self.revision))
.and_then(|index| index.parse::<usize>().ok())
.ok_or("stale_revision")?,
};
let range = self.pages.get(page).ok_or("stale_revision")?;
let fingerprint =
match session.content_fingerprint_bounded(crate::context::REPLAY_JSONL_MAX_BYTES) {
Ok(fingerprint) => fingerprint,
Err(_) => {
self.invalidated = true;
return Err("stale_revision");
}
};
if fingerprint != self.fingerprint
|| session.replay_generation() != self.generation
|| (self.file_present && std::fs::symlink_metadata(session.path()).is_err())
{
self.invalidated = true;
return Err("stale_revision");
}
let last = page + 1 == self.pages.len();
Ok(
json!({ "session_id": session.id(), "revision": self.revision,
"entries": self.entries[range.clone()],
"next_cursor": (!last).then(|| format!("{}.{next}", self.revision, next = page + 1)),
"complete": last && self.complete,
}),
)
}
}
pub(super) struct Attachment {
pub(super) session: Session,
pub(super) config: EffectiveConfig,
pub(super) task_scope: Option<crate::sessions::task_scope::TaskScope>,
snapshot: Snapshot,
notices: Vec<&'static str>,
durable: bool,
}
impl Attachment {
pub(super) fn validate_for_turn(&mut self) -> Result<(), &'static str> {
let revision = self.snapshot.revision.clone();
self.snapshot.page(&self.session, &revision, None)?;
let read = self
.session
.read_events_tolerant_bounded(
crate::context::REPLAY_JSONL_MAX_LINES,
crate::context::REPLAY_JSONL_MAX_BYTES,
)
.map_err(|_| "history_limit")?;
let scope = crate::sessions::task_scope::from_events(self.session.id(), &read.events)
.map_err(|_| "session_invalid")?;
if scope != self.task_scope {
return Err("session_invalid");
}
Ok(())
}
pub(super) fn refresh_history(&mut self) -> Result<(), &'static str> {
self.snapshot.invalidated = true;
self.durable = self
.session
.history_bytes()
.map_err(|_| "session_invalid")?
> 0;
self.snapshot = Snapshot::read(&self.session)?;
Ok(())
}
pub(super) fn revision(&self) -> &str {
&self.snapshot.revision
}
fn result(&self) -> Value {
json!({"session_id": self.session.id(), "durable": self.durable,
"revision": self.snapshot.revision, "notices": self.notices})
}
}
impl RuntimeContext {
pub(super) fn effective_config(&self) -> &EffectiveConfig {
self.attachment
.as_ref()
.map(|a| &a.config)
.unwrap_or(&self.launch.config)
}
fn session_action(
&mut self,
action: Action,
cancellation: &crate::cancellation::AgentCancellation,
) -> Result<Value, &'static str> {
cancellation.check().map_err(|_| "operation_cancelled")?;
match action {
Action::Page {
session_id,
revision,
cursor,
} => {
let attachment = self.attachment.as_mut().ok_or("session_missing")?;
if session_id != attachment.session.id() {
return Err("stale_session");
}
attachment
.snapshot
.page(&attachment.session, &revision, cursor.as_deref())
}
Action::Create => {
let session = SessionManager::new(self.launch.config.paths.sessions.clone())
.create()
.map_err(|_| "session_invalid")?;
self.attach(session, false, cancellation)
}
Action::Open(id) => {
if let Some(attachment) = &self.attachment
&& attachment.session.id() == id
{
return self.attach(
attachment.session.clone(),
attachment.durable,
cancellation,
);
}
let manager = SessionManager::new(self.launch.config.paths.sessions.clone());
let session = manager.open(&id).map_err(|_| "session_invalid")?;
match std::fs::symlink_metadata(session.path()) {
Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
return Err("session_missing");
}
Err(_) => return Err("session_invalid"),
Ok(_) => {}
}
let session = manager.open_existing(id).map_err(|_| "session_invalid")?;
self.attach(session, true, cancellation)
}
}
}
fn attach(
&mut self,
session: Session,
durable: bool,
cancellation: &crate::cancellation::AgentCancellation,
) -> Result<Value, &'static str> {
let session = session.admit_standalone_writer().map_err(|error| {
if error.downcast_ref::<SessionWriterBusy>().is_some() {
"session_busy"
} else {
"session_invalid"
}
})?;
let durable = match std::fs::symlink_metadata(session.path()) {
Ok(_) => true,
Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
if durable {
return Err("session_missing");
}
false
}
Err(_) => return Err("session_invalid"),
};
let (read, stats) = session
.read_events_tolerant_bounded_with_stats(
crate::context::REPLAY_JSONL_MAX_LINES,
crate::context::REPLAY_JSONL_MAX_BYTES,
)
.map_err(|error| {
if error.downcast_ref::<BoundedReadError>().is_some() {
"history_limit"
} else {
"session_invalid"
}
})?;
cancellation.check().map_err(|_| "operation_cancelled")?;
if !read.diagnostics.is_empty() || read.events.iter().any(|e| e.session_id != session.id())
{
return Err("session_invalid");
}
let mut notices = Vec::new();
let mut verified_cwd = false;
let mut recorded_cwds = std::collections::HashSet::new();
for event in &read.events {
cancellation.check().map_err(|_| "operation_cancelled")?;
if event.cwd.as_os_str().is_empty() || !recorded_cwds.insert(&event.cwd) {
continue;
}
if !event.cwd.is_absolute()
|| event.cwd.canonicalize().map_err(|_| "cwd_mismatch")? != self.launch.cwd
{
return Err("cwd_mismatch");
}
verified_cwd = true;
}
if durable && !verified_cwd {
notices.push("cwd_unverified");
}
let task_scope = crate::sessions::task_scope::from_events(session.id(), &read.events)
.map_err(|_| "session_invalid")?;
let mut saved = None;
for event in &read.events {
crate::sessions::preferences::observe(&mut saved, event);
}
let config = if let Some(saved) = saved {
let restored = crate::sessions::preferences::restore_with_settings(
&self.launch.config,
saved,
&self.launch.settings,
);
if !restored.selection_restored {
notices.push("saved_selection_unavailable");
} else if restored.notice.is_some() {
notices.push("saved_selection_notice");
}
restored.config
} else {
self.launch.config.clone()
};
let fingerprint = (stats.content_digest, stats.bytes_read);
let mut snapshot =
Snapshot::from_events(&read.events, fingerprint, session.replay_generation())?;
snapshot.file_present = durable;
if session
.content_fingerprint_bounded(crate::context::REPLAY_JSONL_MAX_BYTES)
.map_err(|_| "session_invalid")?
!= fingerprint
{
return Err("stale_revision");
}
if durable && std::fs::symlink_metadata(session.path()).is_err() {
return Err("session_missing");
}
cancellation.check().map_err(|_| "operation_cancelled")?;
let attachment = Attachment {
session,
config,
task_scope,
snapshot,
notices,
durable,
};
let result = attachment.result();
self.attachment = Some(attachment);
Ok(result)
}
}
pub(super) struct SessionWorker {
worker: Option<JoinHandle<(RuntimeContext, Result<Value, &'static str>)>>,
cancellation: crate::cancellation::AgentCancellationHandle,
}
impl SessionWorker {
pub(super) fn start(
runtime: RuntimeContext,
action: Action,
) -> Result<Self, (Box<RuntimeContext>, &'static str)> {
let slot = Arc::new(Mutex::new(Some(runtime)));
let owned = Arc::clone(&slot);
let (token, cancellation) = crate::cancellation::AgentCancellation::default().child_token();
match std::thread::Builder::new()
.name("local-agent-session".into())
.spawn(move || {
let mut runtime = owned
.lock()
.unwrap_or_else(|e| e.into_inner())
.take()
.expect("one session worker");
let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
runtime.session_action(action, &token)
}))
.unwrap_or(Err("session_invalid"));
(runtime, result)
}) {
Ok(worker) => Ok(Self {
worker: Some(worker),
cancellation,
}),
Err(_) => Err((
Box::new(
slot.lock()
.unwrap_or_else(|e| e.into_inner())
.take()
.expect("unstarted worker"),
),
"session_invalid",
)),
}
}
pub(super) fn poll(&mut self) -> Option<(RuntimeContext, Result<Value, &'static str>)> {
if !self.worker.as_ref()?.is_finished() {
return None;
}
self.worker.take()?.join().ok()
}
pub(super) fn cancel(&self) {
self.cancellation.cancel();
}
pub(super) fn finish(mut self) {
if let Some(worker) = self.worker.take() {
drop(worker.join());
}
}
}