#[cfg(test)]
use super::AgentOutputSink;
#[cfg(test)]
use super::AssistantTextSink;
use crate::sessions::{Session, SessionEventKind, try_record_session_event_batch};
use serde_json::json;
use std::path::Path;
#[cfg(test)]
pub(super) struct AssistantOnlySink<'a> {
pub(super) on_text_delta: &'a mut AssistantTextSink<'a>,
}
#[cfg(test)]
impl AgentOutputSink for AssistantOnlySink<'_> {
fn assistant_delta(&mut self, text: &str) -> anyhow::Result<()> {
(self.on_text_delta)(text)
}
fn tool_block(&mut self, _block: &str) -> anyhow::Result<()> {
Ok(())
}
}
pub(super) fn normalize_agent_identifier(value: String, fallback: &str) -> String {
let trimmed = value.trim();
if trimmed.is_empty() {
fallback.to_string()
} else {
trimmed.to_string()
}
}
pub(super) fn safe_error_message(error: &anyhow::Error) -> String {
error
.chain()
.next()
.map(ToString::to_string)
.unwrap_or_else(|| "unknown error".to_string())
}
const ASSISTANT_CHUNK_BATCH_BYTES: usize = 4 * 1024;
pub(super) struct AssistantChunkBatch<'a> {
session: Option<&'a Session>,
cwd: &'a Path,
chunks: Vec<String>,
bytes: usize,
}
impl<'a> AssistantChunkBatch<'a> {
pub(super) fn new(session: Option<&'a Session>, cwd: &'a Path) -> Self {
Self {
session,
cwd,
chunks: Vec::new(),
bytes: 0,
}
}
pub(super) fn push(&mut self, delta: &str) {
self.bytes = self.bytes.saturating_add(delta.len());
self.chunks.push(delta.to_string());
}
pub(super) fn flush_if_full(&mut self) -> anyhow::Result<()> {
if self.bytes >= ASSISTANT_CHUNK_BATCH_BYTES {
self.flush()?;
}
Ok(())
}
pub(super) fn flush(&mut self) -> anyhow::Result<()> {
if self.chunks.iter().all(|chunk| chunk.trim().is_empty()) {
self.chunks.clear();
self.bytes = 0;
return Ok(());
}
let chunks = std::mem::take(&mut self.chunks);
self.bytes = 0;
let text = chunks.concat();
try_record_session_event_batch(
self.session,
self.cwd,
SessionEventKind::AssistantChunk,
std::iter::once(json!({"text": text})),
)
.map(|_| ())
}
}
#[cfg(test)]
mod tests {
use super::*;
fn fail_flush_with_unavailable_session(batch: &mut AssistantChunkBatch<'_>, session: &Session) {
let backup = session.path().with_extension("backup");
std::fs::rename(session.path(), &backup).unwrap();
std::fs::create_dir(session.path()).unwrap();
assert!(batch.flush().is_err());
std::fs::remove_dir(session.path()).unwrap();
std::fs::rename(backup, session.path()).unwrap();
}
#[test]
fn failed_chunk_flush_stays_before_tools_and_terminal_text_repairs_missing_tail() {
use crate::{
agent::session_persistence::SessionPersistence, context::build_conversation_replay,
providers::ProviderConversationItem, sessions::TurnStatus,
};
for status in [
TurnStatus::Complete,
TurnStatus::Cancelled,
TurnStatus::Failed,
] {
let temp = tempfile::TempDir::new().unwrap();
let manager = crate::sessions::SessionManager::new(temp.path().join("sessions"));
let session = manager.create().unwrap();
let mut persistence = SessionPersistence::new(Some(&session), temp.path());
let mut sink = None;
persistence
.record_required(SessionEventKind::UserInput, json!({"text": "question"}))
.unwrap();
let mut batch = AssistantChunkBatch::new(Some(&session), temp.path());
batch.push("Before tools λ.");
fail_flush_with_unavailable_session(&mut batch, &session);
for (kind, payload) in [
(
SessionEventKind::AssistantOutput,
json!({"text": "Before tools λ."}),
),
(
SessionEventKind::ToolCall,
json!({"id": "call_1", "name": "read", "arguments": {}}),
),
(
SessionEventKind::ToolResult,
json!({"call_id": "call_1", "result": {"tool_name": "read", "success": true, "content": "ok"}}),
),
] {
persistence.try_record(kind, payload, &mut sink).unwrap();
}
batch.push("After tools");
batch.flush().unwrap();
batch.push(" accepted tail λ.");
fail_flush_with_unavailable_session(&mut batch, &session);
batch.flush().unwrap();
persistence
.record_terminal_status(
status,
"Before tools λ.\n\nAfter tools accepted tail λ.",
&mut sink,
)
.unwrap();
let events = session.read_events().unwrap();
let chunks: Vec<_> = events
.iter()
.filter(|event| event.kind() == Some(SessionEventKind::AssistantChunk))
.map(|event| event.payload["text"].as_str().unwrap())
.collect();
assert_eq!(chunks, ["After tools"]);
let replay = build_conversation_replay(Some(&session)).unwrap();
let ordered: Vec<_> = replay
.items
.iter()
.map(|item| match item {
ProviderConversationItem::Message(message) => message.content.clone(),
ProviderConversationItem::ResponseItem(item) => {
format!("call:{}", item["call_id"].as_str().unwrap())
}
ProviderConversationItem::ToolResult(result) => {
format!("result:{}", result.call_id)
}
other => panic!("unexpected replay item: {other:?}"),
})
.collect();
assert_eq!(
ordered,
[
"question",
"Before tools λ.",
"call:call_1",
"result:call_1",
"After tools accepted tail λ."
]
);
}
}
}