use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use crate::acp::protocol::ClientCapabilities;
use crate::tools::AcpClientFsState;
use tokio::sync::Mutex;
use std::io::Write;
use serde::{Deserialize, Serialize};
use crate::acp::bridge::sha256_first_12;
use crate::acp::permission::PendingPermissionStore;
use crate::message::Message;
use crate::runtime::AgentRuntime;
pub type SessionId = String;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionMetadata {
pub session_id: SessionId,
pub cwd: PathBuf,
pub turn: u64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub system_prompt: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub mcp_servers: Option<serde_json::Value>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SavedMessage {
pub id: String,
#[serde(flatten)]
pub message: Message,
#[serde(default)]
pub compressible: bool,
}
impl SavedMessage {
pub fn from_message(msg: Message) -> Self {
let id = sha256_first_12(&msg.content);
Self {
id,
message: msg,
compressible: false,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompactionHint {
pub turn_index: usize,
pub compressible: bool,
}
#[derive(Debug, Clone)]
pub struct SummarizedContext {
pub summary: String,
pub message_count: usize,
}
pub struct AcpSession {
pub runtime: AgentRuntime,
pub cwd: PathBuf,
pub turn: u64,
pub session_id: SessionId,
pub transcript: Vec<Message>,
pub cancel_token: Arc<tokio_util::sync::CancellationToken>,
pub system_prompt: Option<String>,
pub mcp_servers: Option<serde_json::Value>,
pub client_capabilities: Option<ClientCapabilities>,
pub acp_state: Arc<Mutex<AcpClientFsState>>,
pub mcp_child_pids: Vec<u32>,
}
impl AcpSession {
pub fn refresh_cancel_token(&mut self) -> Arc<tokio_util::sync::CancellationToken> {
let token = Arc::new(tokio_util::sync::CancellationToken::new());
self.runtime.set_interrupt_token((*token).clone());
self.cancel_token = token.clone();
token
}
pub async fn kill_mcp_children(&mut self) {
for pid in std::mem::take(&mut self.mcp_child_pids) {
#[cfg(unix)]
{
use tokio::time::timeout;
let grace = crate::mcp::DEFAULT_KILL_GRACE_PERIOD;
tracing::info!(server_pid = pid, "AcpSession: killing MCP child process");
tracing::info!(server_pid = pid, "kill-gracefully: sending SIGTERM");
let _ = std::process::Command::new("kill")
.args(["-TERM", &pid.to_string()])
.status();
let waited = timeout(grace, async {
loop {
let status = std::process::Command::new("kill")
.args(["-0", &pid.to_string()])
.status()
.ok();
match status {
Some(s) if s.success() => {
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
}
_ => {
break;
}
}
}
})
.await;
if waited.is_err() {
tracing::warn!(
server_pid = pid,
grace_ms = grace.as_millis(),
"kill-gracefully: SIGTERM timeout, sending SIGKILL"
);
let _ = std::process::Command::new("kill")
.args(["-9", &pid.to_string()])
.status();
} else {
tracing::info!(
server_pid = pid,
"kill-gracefully: process exited gracefully"
);
}
let _ = std::process::Command::new("wait")
.args([&pid.to_string()])
.status();
}
#[cfg(not(unix))]
{
let _ = pid;
tracing::info!("AcpSession: killing MCP child (non-Unix)");
let _ = std::process::Command::new("taskkill")
.args(["/F", "/PID", &pid.to_string()])
.status();
}
}
}
}
pub struct AcpSessionManager {
sessions: HashMap<SessionId, AcpSession>,
next_id: AtomicU64,
persistence_dir: Option<PathBuf>,
permission_store: PendingPermissionStore,
client_capabilities: Option<ClientCapabilities>,
}
impl Default for AcpSessionManager {
fn default() -> Self {
Self::new()
}
}
impl AcpSessionManager {
pub fn new() -> Self {
Self {
sessions: HashMap::new(),
next_id: AtomicU64::new(1),
persistence_dir: None,
permission_store: PendingPermissionStore::new(),
client_capabilities: None,
}
}
pub fn set_client_capabilities(&mut self, caps: Option<ClientCapabilities>) {
self.client_capabilities = caps;
}
pub fn get_client_capabilities(&self) -> Option<&ClientCapabilities> {
self.client_capabilities.as_ref()
}
pub fn with_persistence_dir(mut self, dir: PathBuf) -> Self {
self.persistence_dir = Some(dir);
self
}
pub fn persistence_dir(&self) -> Option<&Path> {
self.persistence_dir.as_deref()
}
pub fn next_session_id(&self) -> SessionId {
let n = self.next_id.fetch_add(1, Ordering::Relaxed);
format!("acp-sess-{n}")
}
pub fn insert_with_id(&mut self, session_id: SessionId, session: AcpSession) {
if let Some(pdir) = &self.persistence_dir {
let session_dir = pdir.join(&session_id);
let _ = std::fs::create_dir_all(&session_dir);
let _ = self.save_metadata_to_disk(&session_id, &session);
}
self.sessions.insert(session_id, session);
}
pub fn get_mut(&mut self, sid: &str) -> Option<&mut AcpSession> {
self.sessions.get_mut(sid)
}
pub fn remove(&mut self, sid: &str) -> Option<AcpSession> {
self.sessions.remove(sid)
}
pub fn all_turns(&self) -> Vec<(String, u64)> {
self.sessions
.iter()
.map(|(k, v)| (k.clone(), v.turn))
.collect()
}
pub fn get(&self, sid: &str) -> Option<&AcpSession> {
self.sessions.get(sid)
}
pub fn contains(&self, sid: &str) -> bool {
self.sessions.contains_key(sid)
}
pub fn permission_store(&mut self) -> &mut PendingPermissionStore {
&mut self.permission_store
}
pub fn session_dir(&self, sid: &str) -> Option<PathBuf> {
self.persistence_dir
.as_ref()
.map(|p| p.join(sanitize_session_id(sid)))
}
pub fn save_transcript(&self, sid: &str) -> std::io::Result<()> {
let session = match self.sessions.get(sid) {
Some(s) => s,
None => {
return Err(std::io::Error::new(
std::io::ErrorKind::NotFound,
"session not found",
))
}
};
let session_dir = match self.session_dir(sid) {
Some(d) => d,
None => return Ok(()), };
let _ = std::fs::create_dir_all(&session_dir);
let path = session_dir.join("transcript.jsonl");
let mut file = std::fs::File::create(&path)?;
for msg in &session.transcript {
let saved = SavedMessage::from_message(msg.clone());
let line = serde_json::to_string(&saved)?;
writeln!(&mut file, "{line}")?;
}
Ok(())
}
pub fn load_transcript_from_disk(&self, sid: &str) -> std::io::Result<Vec<SavedMessage>> {
let session_dir = self.session_dir(sid).ok_or_else(|| {
std::io::Error::new(std::io::ErrorKind::NotFound, "no persistence dir")
})?;
let path = session_dir.join("transcript.jsonl");
let content = std::fs::read_to_string(&path)?;
let mut messages = Vec::new();
for line in content.lines() {
let trimmed = line.trim();
if trimmed.is_empty() {
continue;
}
let saved: SavedMessage = serde_json::from_str(trimmed).map_err(|e| {
std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!("corrupt transcript: {e}"),
)
})?;
messages.push(saved);
}
Ok(messages)
}
fn save_metadata_to_disk(&self, sid: &str, session: &AcpSession) -> std::io::Result<()> {
let session_dir = match self.session_dir(sid) {
Some(d) => d,
None => return Ok(()),
};
let _ = std::fs::create_dir_all(&session_dir);
let metadata = SessionMetadata {
session_id: sid.to_string(),
cwd: session.cwd.clone(),
turn: session.turn,
system_prompt: session.system_prompt.clone(),
mcp_servers: session.mcp_servers.clone(),
};
let path = session_dir.join("metadata.json");
let json = serde_json::to_string_pretty(&metadata)?;
std::fs::write(&path, json)?;
Ok(())
}
pub fn save_metadata(&self, sid: &str) -> std::io::Result<()> {
let session = match self.sessions.get(sid) {
Some(s) => s,
None => {
return Err(std::io::Error::new(
std::io::ErrorKind::NotFound,
"session not found",
))
}
};
self.save_metadata_to_disk(sid, session)
}
pub fn load_metadata_from_disk(&self, sid: &str) -> std::io::Result<SessionMetadata> {
let session_dir = self.session_dir(sid).ok_or_else(|| {
std::io::Error::new(std::io::ErrorKind::NotFound, "no persistence dir")
})?;
let path = session_dir.join("metadata.json");
let content = std::fs::read_to_string(&path)?;
let metadata: SessionMetadata = serde_json::from_str(&content).map_err(|e| {
std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!("corrupt metadata: {e}"),
)
})?;
Ok(metadata)
}
pub fn session_exists_on_disk(&self, sid: &str) -> bool {
self.session_dir(sid)
.map(|d| d.join("metadata.json").exists())
.unwrap_or(false)
}
}
pub fn summarize_transcript(messages: &[Message]) -> SummarizedContext {
let mut accomplishments: Vec<String> = Vec::new();
let mut blockers: Vec<String> = Vec::new();
let mut in_flight: Option<String> = None;
for (i, msg) in messages.iter().enumerate() {
let content = msg.content.trim();
match msg.role {
crate::message::Role::User => {
if i == messages.len() - 1
|| messages[messages.len() - 1].role == crate::message::Role::User
{
in_flight = Some(content.to_string());
} else {
let mut goal = content.chars().take(80).collect::<String>();
if content.len() > 80 {
goal.push_str("...");
}
accomplishments.push(format!("user request: \"{goal}\""));
}
}
crate::message::Role::Assistant => {
if !content.is_empty() {
let mut snippet = content.chars().take(80).collect::<String>();
if content.len() > 80 {
snippet.push_str("...");
}
accomplishments.push(format!("assistant responded: \"{snippet}\""));
}
}
crate::message::Role::Tool => {
if content.contains("error")
|| content.contains("Error")
|| content.contains("failed")
{
let mut snippet = content.chars().take(60).collect::<String>();
if content.len() > 60 {
snippet.push_str("...");
}
blockers.push(format!("tool error: \"{snippet}\""));
} else if !content.is_empty() {
let mut snippet = content.chars().take(60).collect::<String>();
if content.len() > 60 {
snippet.push_str("...");
}
accomplishments.push(format!("tool result: \"{snippet}\""));
}
}
crate::message::Role::System => {
}
}
}
let mut summary_parts: Vec<String> = Vec::new();
if accomplishments.is_empty() {
summary_parts.push("No accomplishments recorded yet.".to_string());
} else {
let n = accomplishments.len();
let sample: Vec<&str> = accomplishments.iter().take(3).map(|s| s.as_str()).collect();
summary_parts.push(format!(
"The session has {n} recorded action(s): {}.",
sample.join("; ")
));
if accomplishments.len() > 3 {
summary_parts.push(format!("And {} more action(s).", accomplishments.len() - 3));
}
}
match in_flight {
Some(ref text) => summary_parts.push(format!("The user's current request is: \"{text}\".")),
None => summary_parts.push("No active user request.".to_string()),
}
if blockers.is_empty() {
summary_parts.push("No blockers detected.".to_string());
} else {
let sample: Vec<&str> = blockers.iter().take(2).map(|s| s.as_str()).collect();
summary_parts.push(format!(
"The session encountered {} blocker(s): {}.",
blockers.len(),
sample.join("; ")
));
}
SummarizedContext {
summary: summary_parts.join(" "),
message_count: messages.len(),
}
}
fn sanitize_session_id(sid: &str) -> String {
sid.replace(['/', '\\', ':', ' ', '\0'], "_")
.replace("..", "__")
}
pub fn compute_compaction_hints(
messages: &[SavedMessage],
recency_threshold: usize,
) -> Vec<CompactionHint> {
let total = messages.len();
messages
.iter()
.enumerate()
.map(|(i, _msg)| {
let compressible = i < total.saturating_sub(recency_threshold);
CompactionHint {
turn_index: i,
compressible,
}
})
.collect()
}
pub fn apply_compaction_hints(messages: &mut [SavedMessage], hints: &[CompactionHint]) {
for hint in hints {
if let Some(msg) = messages.get_mut(hint.turn_index) {
msg.compressible = hint.compressible;
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::message::Message;
use std::io::Write;
#[test]
fn new_manager_is_empty() {
let mgr = AcpSessionManager::new();
assert!(mgr.get("nonexistent").is_none());
}
#[test]
fn next_session_id_is_monotonic() {
let mgr = AcpSessionManager::new();
let a = mgr.next_session_id();
let b = mgr.next_session_id();
assert_ne!(a, b);
assert!(a.contains("acp-sess-"));
assert!(b.contains("acp-sess-"));
}
#[test]
fn insert_and_get() {
let mut mgr = AcpSessionManager::new();
let _sid = mgr.next_session_id();
assert!(mgr.remove("nonexistent").is_none());
}
#[test]
fn contains_works() {
let mgr = AcpSessionManager::new();
assert!(!mgr.contains("nonexistent"));
}
#[test]
fn save_and_load_transcript_roundtrip() {
let tmp = tempfile::TempDir::new().unwrap();
let pdir = tmp.path().to_path_buf();
let _mgr = AcpSessionManager::new();
let _ = pdir;
let session_dir = pdir.join("test-sess-1");
std::fs::create_dir_all(&session_dir).unwrap();
let messages = vec![
Message::user("Hello"),
Message::assistant("Hi there!"),
Message::user("What is the weather?"),
];
let path = session_dir.join("transcript.jsonl");
let mut file = std::fs::File::create(&path).unwrap();
for msg in &messages {
let saved = SavedMessage::from_message(msg.clone());
let line = serde_json::to_string(&saved).unwrap();
writeln!(file, "{line}").unwrap();
}
drop(file);
let content = std::fs::read_to_string(&path).unwrap();
let loaded: Vec<SavedMessage> = content
.lines()
.filter(|l| !l.trim().is_empty())
.map(|l| serde_json::from_str(l).unwrap())
.collect();
assert_eq!(loaded.len(), 3);
assert_eq!(loaded[0].message.content, "Hello");
assert_eq!(loaded[1].message.content, "Hi there!");
assert_eq!(loaded[2].message.content, "What is the weather?");
assert_eq!(loaded[0].id, sha256_first_12("Hello"));
assert_eq!(loaded[1].id, sha256_first_12("Hi there!"));
assert_eq!(loaded[2].id, sha256_first_12("What is the weather?"));
}
#[test]
fn saved_message_id_matches_content_hash() {
let msg = Message::user("test content");
let saved = SavedMessage::from_message(msg.clone());
assert_eq!(saved.id, sha256_first_12("test content"));
}
#[test]
fn same_content_same_id() {
let a = SavedMessage::from_message(Message::user("hello"));
let b = SavedMessage::from_message(Message::user("hello"));
assert_eq!(a.id, b.id);
}
#[test]
fn different_content_different_id() {
let a = SavedMessage::from_message(Message::user("hello"));
let b = SavedMessage::from_message(Message::user("world"));
assert_ne!(a.id, b.id);
}
#[test]
fn save_and_load_metadata_roundtrip() {
let tmp = tempfile::TempDir::new().unwrap();
let session_dir = tmp.path().join("test-sess-meta");
std::fs::create_dir_all(&session_dir).unwrap();
let metadata = SessionMetadata {
session_id: "test-sess-meta".to_string(),
cwd: PathBuf::from("/tmp"),
turn: 5,
system_prompt: Some("Be helpful.".to_string()),
mcp_servers: None,
};
let path = session_dir.join("metadata.json");
let json = serde_json::to_string_pretty(&metadata).unwrap();
std::fs::write(&path, json).unwrap();
let content = std::fs::read_to_string(&path).unwrap();
let loaded: SessionMetadata = serde_json::from_str(&content).unwrap();
assert_eq!(loaded.session_id, "test-sess-meta");
assert_eq!(loaded.turn, 5);
assert_eq!(loaded.system_prompt, Some("Be helpful.".to_string()));
}
#[test]
fn corrupt_transcript_returns_error() {
let tmp = tempfile::TempDir::new().unwrap();
let session_dir = tmp.path().join("bad-sess");
std::fs::create_dir_all(&session_dir).unwrap();
let path = session_dir.join("transcript.jsonl");
std::fs::write(&path, "not valid json\n").unwrap();
let _mgr = AcpSessionManager::new();
let content = std::fs::read_to_string(&path).unwrap();
let result: Result<SavedMessage, _> = serde_json::from_str(content.trim());
assert!(result.is_err());
}
#[test]
fn summarize_empty_transcript() {
let ctx = summarize_transcript(&[]);
assert!(ctx.summary.contains("No accomplishments"));
assert_eq!(ctx.message_count, 0);
}
#[test]
fn summarize_with_messages() {
let messages = vec![
Message::user("What is the capital of France?"),
Message::assistant("The capital of France is Paris."),
Message::user("Tell me more about Paris."),
];
let ctx = summarize_transcript(&messages);
assert!(ctx.summary.contains("capital of France"));
assert!(ctx.summary.contains("Tell me more"));
assert_eq!(ctx.message_count, 3);
}
#[test]
fn summarize_with_tool_errors_detects_blockers() {
let messages = vec![
Message::user("List the files"),
Message::assistant("Let me check the directory."),
Message::tool_result("tc-1", "error: permission denied"),
];
let ctx = summarize_transcript(&messages);
assert!(ctx.summary.contains("blocker"));
assert_eq!(ctx.message_count, 3);
}
#[test]
fn compaction_hints_all_compressible_when_threshold_zero() {
let mut messages = Vec::new();
for i in 0..5 {
let msg = SavedMessage::from_message(Message::user(format!("msg {i}")));
messages.push(msg);
}
let hints = compute_compaction_hints(&messages, 0);
assert_eq!(hints.len(), 5);
for hint in &hints {
assert!(
hint.compressible,
"all should be compressible with threshold 0"
);
}
}
#[test]
fn compaction_hints_recency_protected() {
let mut messages = Vec::new();
for i in 0..10 {
let msg = SavedMessage::from_message(Message::user(format!("msg {i}")));
messages.push(msg);
}
let hints = compute_compaction_hints(&messages, 2);
assert_eq!(hints.len(), 10);
for (i, hint) in hints.iter().enumerate().take(8) {
assert!(hint.compressible, "msg {i} should be compressible");
}
assert!(!hints[8].compressible, "msg 8 should NOT be compressible");
assert!(!hints[9].compressible, "msg 9 should NOT be compressible");
}
#[test]
fn compaction_hints_configurable_threshold() {
let mut messages = Vec::new();
for i in 0..10 {
let msg = SavedMessage::from_message(Message::user(format!("msg {i}")));
messages.push(msg);
}
let hints = compute_compaction_hints(&messages, 5);
for (i, hint) in hints.iter().enumerate().take(5) {
assert!(hint.compressible, "msg {i} should be compressible");
}
for (i, hint) in hints.iter().enumerate().skip(5) {
assert!(!hint.compressible, "msg {i} should NOT be compressible");
}
}
#[test]
fn apply_compaction_hints_sets_compressible() {
let mut messages = vec![
SavedMessage::from_message(Message::user("old")),
SavedMessage::from_message(Message::user("new")),
];
let hints = vec![
CompactionHint {
turn_index: 0,
compressible: true,
},
CompactionHint {
turn_index: 1,
compressible: false,
},
];
apply_compaction_hints(&mut messages, &hints);
assert!(messages[0].compressible);
assert!(!messages[1].compressible);
}
#[test]
fn cancel_token_is_arc_type() {
let _type_check: Arc<tokio_util::sync::CancellationToken> =
Arc::new(tokio_util::sync::CancellationToken::new());
}
#[test]
fn acp_session_field_cancel_token_is_arc() {
fn assert_arc_cancel_token(t: &Arc<tokio_util::sync::CancellationToken>) {
assert!(!t.is_cancelled());
}
let token = Arc::new(tokio_util::sync::CancellationToken::new());
assert_arc_cancel_token(&token);
}
#[tokio::test]
async fn drop_acp_session_does_not_cancel_token() {
use crate::llm::MockProvider;
let llm = Arc::new(MockProvider::new(vec![]));
let runtime = AgentRuntime::builder()
.llm(llm)
.build()
.expect("build runtime");
let shared_token = Arc::new(tokio_util::sync::CancellationToken::new());
let clone_for_drop = shared_token.clone();
{
let _session = AcpSession {
runtime,
cwd: PathBuf::from("/tmp"),
turn: 0,
session_id: "test".to_string(),
transcript: vec![],
cancel_token: clone_for_drop,
system_prompt: None,
mcp_servers: None,
client_capabilities: None,
acp_state: Arc::new(Mutex::new(AcpClientFsState::new())),
mcp_child_pids: Vec::new(),
};
}
assert!(
!shared_token.is_cancelled(),
"dropping AcpSession must not cancel token if other Arc clones exist"
);
}
}