use serde::{Deserialize, Serialize};
use std::collections::HashSet;
use tirea_state::State;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct CompactBoundary {
pub covers_through_message_id: String,
pub summary: String,
pub original_token_count: usize,
pub created_at_ms: u64,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ArtifactRef {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub message_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub tool_call_id: Option<String>,
pub label: String,
pub summary: String,
pub original_size: usize,
pub original_token_count: usize,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize, State)]
#[tirea(path = "__context", action = "ContextAction", scope = "thread")]
pub struct ContextState {
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub boundaries: Vec<CompactBoundary>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub artifact_refs: Vec<ArtifactRef>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ContextAction {
AddBoundary(CompactBoundary),
AddArtifact(ArtifactRef),
PruneArtifacts {
#[serde(default, skip_serializing_if = "Vec::is_empty")]
message_ids: Vec<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
tool_call_ids: Vec<String>,
},
}
impl ContextState {
pub(super) fn reduce(&mut self, action: ContextAction) {
match action {
ContextAction::AddBoundary(boundary) => {
self.boundaries.clear();
self.boundaries.push(boundary);
}
ContextAction::AddArtifact(artifact) => {
if let Some(existing) = self
.artifact_refs
.iter_mut()
.find(|existing| artifact_identity_matches(existing, &artifact))
{
*existing = artifact;
} else {
self.artifact_refs.push(artifact);
}
}
ContextAction::PruneArtifacts {
message_ids,
tool_call_ids,
} => {
if message_ids.is_empty() && tool_call_ids.is_empty() {
return;
}
let message_ids: HashSet<String> = message_ids.into_iter().collect();
let tool_call_ids: HashSet<String> = tool_call_ids.into_iter().collect();
self.artifact_refs.retain(|artifact| {
let message_match = artifact
.message_id
.as_ref()
.is_some_and(|id| message_ids.contains(id));
let tool_match = artifact
.tool_call_id
.as_ref()
.is_some_and(|id| tool_call_ids.contains(id));
!(message_match || tool_match)
});
}
}
}
pub(super) fn latest_boundary(&self) -> Option<&CompactBoundary> {
self.boundaries.last()
}
}
fn artifact_identity_matches(existing: &ArtifactRef, candidate: &ArtifactRef) -> bool {
existing.message_id.is_some() && existing.message_id == candidate.message_id
|| existing.tool_call_id.is_some() && existing.tool_call_id == candidate.tool_call_id
}