use super::Record;
use serde_json::{Value, json};
impl Record {
pub(super) fn next_sequence(&mut self, id: &str) -> anyhow::Result<u64> {
let conversation = self.conversation(id)?;
conversation.seq = conversation
.seq
.checked_add(1)
.ok_or_else(|| anyhow::anyhow!("conversation_sequence_exhausted"))?;
conversation.meta["last_seq"] = json!(conversation.seq);
Ok(conversation.seq)
}
pub(super) fn resequence_completed(&mut self) -> anyhow::Result<()> {
let items = std::mem::take(&mut self.completed);
self.completed_bytes = 0;
for mut item in items {
let id = item["conversation"]
.as_str()
.ok_or_else(|| anyhow::anyhow!("record_conversation_invalid"))?
.to_owned();
item["seq"] = json!(self.next_sequence(&id)?);
self.complete(item)?;
}
Ok(())
}
pub(super) fn measured(&self, value: &Value) -> anyhow::Result<usize> {
let bytes = serde_json::to_vec(value)?.len();
let limit = self
.seat
.harness
.content_bytes
.checked_mul(self.seat.harness.write_queue)
.ok_or_else(|| anyhow::anyhow!("record_buffer_bound_overflow"))?;
anyhow::ensure!(bytes <= limit, "record_buffer_bytes_exceeded");
#[cfg(test)]
self.measured_bytes
.set(self.measured_bytes.get().saturating_add(bytes));
Ok(bytes)
}
pub(super) fn complete(&mut self, item: Value) -> anyhow::Result<()> {
self.completed_bytes = self
.completed_bytes
.checked_add(self.measured(&item)?)
.ok_or_else(|| anyhow::anyhow!("record_buffer_bound_overflow"))?;
self.completed.push(item);
Ok(())
}
pub(super) fn message(&mut self, item: &Value) -> anyhow::Result<()> {
if item["complete"] != true {
self.mark_live(item)?;
self.send(json!({"item":item}))?;
}
let id = item["conversation"]
.as_str()
.ok_or_else(|| anyhow::anyhow!("record_conversation_invalid"))?
.to_owned();
let key = (
id,
format!("{}:{}:{}", item["kind"], item["role"], item["message_id"]),
);
let mut message = if let Some(message) = self.messages.remove(&key) {
self.messages_bytes = self
.messages_bytes
.checked_sub(self.measured(&message)?)
.ok_or_else(|| anyhow::anyhow!("record_buffer_accounting_invalid"))?;
message
} else {
json!({"text":""})
};
let previous = message["text"].as_str().unwrap_or_default().to_owned();
let truncated = message["truncated"] == true;
if let Some(fields) = item.as_object() {
for (key, value) in fields {
if key != "text" {
message[key] = value.clone();
}
}
}
if let Some(text) = item.get("text").and_then(Value::as_str) {
message["text"] = json!(if item["replace"] == true {
text.to_owned()
} else {
previous + text
});
}
if truncated {
message["truncated"] = json!(true);
}
super::super::stream::fit(&mut message, self.seat.harness.content_bytes)?;
if item["complete"] == true {
message["replace"] = json!(true);
self.complete(message)?;
} else {
self.messages_bytes = self
.messages_bytes
.checked_add(self.measured(&message)?)
.ok_or_else(|| anyhow::anyhow!("record_buffer_bound_overflow"))?;
self.messages.insert(key, message);
}
Ok(())
}
pub(super) fn tool(&mut self, item: &Value) -> anyhow::Result<()> {
let key = (
item["conversation"]
.as_str()
.ok_or_else(|| anyhow::anyhow!("record_conversation_invalid"))?
.to_owned(),
item["tool_call_id"]
.as_str()
.ok_or_else(|| anyhow::anyhow!("record_tool_id_invalid"))?
.to_owned(),
);
let mut tool = if let Some(tool) = self.tools.remove(&key) {
self.tools_bytes = self
.tools_bytes
.checked_sub(self.measured(&tool)?)
.ok_or_else(|| anyhow::anyhow!("record_buffer_accounting_invalid"))?;
tool
} else {
json!({})
};
if let Some(fields) = item.as_object() {
for (key, value) in fields {
tool[key] = value.clone();
}
}
super::super::stream::fit(&mut tool, self.seat.harness.content_bytes)?;
let status: Option<super::super::wire::ToolStatus> = tool
.get("status")
.filter(|status| !status.is_null())
.cloned()
.map(serde_json::from_value)
.transpose()
.map_err(|_| anyhow::anyhow!("record_tool_status_invalid"))?;
if status.is_some_and(super::super::wire::ToolStatus::terminal) {
self.complete(tool)?;
} else {
self.mark_live(item)?;
self.send(json!({"item":item}))?;
self.tools_bytes = self
.tools_bytes
.checked_add(self.measured(&tool)?)
.ok_or_else(|| anyhow::anyhow!("record_buffer_bound_overflow"))?;
self.tools.insert(key, tool);
}
Ok(())
}
pub(super) fn finish_messages(&mut self, id: &str) -> anyhow::Result<()> {
let keys: Vec<_> = self
.messages
.keys()
.filter(|(conversation, _)| conversation == id)
.cloned()
.collect();
for key in keys {
if let Some(mut item) = self.messages.remove(&key) {
self.messages_bytes = self
.messages_bytes
.checked_sub(self.measured(&item)?)
.ok_or_else(|| anyhow::anyhow!("record_buffer_accounting_invalid"))?;
item["seq"] = json!(self.next_sequence(id)?);
item["complete"] = json!(true);
item["replace"] = json!(true);
super::super::stream::fit(&mut item, self.seat.harness.content_bytes)?;
self.complete(item)?;
}
}
let keys: Vec<_> = self
.tools
.keys()
.filter(|(conversation, _)| conversation == id)
.cloned()
.collect();
for key in keys {
if let Some(mut item) = self.tools.remove(&key) {
self.tools_bytes = self
.tools_bytes
.checked_sub(self.measured(&item)?)
.ok_or_else(|| anyhow::anyhow!("record_buffer_accounting_invalid"))?;
item["seq"] = json!(self.next_sequence(id)?);
item["status"] = json!(super::super::wire::ToolStatus::Interrupted);
super::super::stream::fit(&mut item, self.seat.harness.content_bytes)?;
self.complete(item)?;
}
}
Ok(())
}
pub(super) fn check_bound(&self) -> anyhow::Result<()> {
let bytes = self
.completed_bytes
.checked_add(self.messages_bytes)
.and_then(|n| n.checked_add(self.tools_bytes))
.and_then(|n| n.checked_add(self.batch_bytes))
.ok_or_else(|| anyhow::anyhow!("record_buffer_bound_overflow"))?;
let limit = self
.seat
.harness
.content_bytes
.checked_mul(self.seat.harness.write_queue)
.ok_or_else(|| anyhow::anyhow!("record_buffer_bound_overflow"))?;
anyhow::ensure!(bytes <= limit, "record_buffer_bytes_exceeded");
Ok(())
}
pub(super) fn mark_live(&mut self, item: &Value) -> anyhow::Result<()> {
let id = item["conversation"]
.as_str()
.ok_or_else(|| anyhow::anyhow!("record_conversation_invalid"))?;
let seq = item["seq"]
.as_u64()
.ok_or_else(|| anyhow::anyhow!("record_sequence_invalid"))?;
let conversation = self
.conversations
.get_mut(id)
.ok_or_else(|| anyhow::anyhow!("record_conversation_invalid"))?;
conversation.live_seq = conversation.live_seq.max(seq);
Ok(())
}
}