use serde::{Deserialize, Serialize};
use crate::catalog::LlmModel;
use crate::model_settings::ModelSettings;
use crate::reasoning::ReasoningEffort;
use crate::types::IsoString;
use super::{AssistantReply, ChatMessage, MessageId, ToolCallError, ToolCallResult, ToolDefinition};
#[doc = include_str!("docs/context.md")]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Context {
messages: Vec<ChatMessage>,
tools: Vec<ToolDefinition>,
#[serde(skip)]
reasoning_effort: ReasoningEffort,
#[serde(skip)]
model_settings: ModelSettings,
#[serde(skip)]
prompt_cache_key: Option<String>,
#[serde(skip)]
session_affinity_key: Option<String>,
}
impl Context {
pub fn new(messages: Vec<ChatMessage>, tools: Vec<ToolDefinition>) -> Self {
Self {
messages,
tools,
reasoning_effort: ReasoningEffort::Default,
model_settings: ModelSettings::default(),
prompt_cache_key: None,
session_affinity_key: None,
}
}
pub fn prompt_cache_key(&self) -> Option<&str> {
self.prompt_cache_key.as_deref()
}
pub fn set_prompt_cache_key(&mut self, key: Option<String>) {
self.prompt_cache_key = key;
}
pub fn session_affinity_key(&self) -> Option<&str> {
self.session_affinity_key.as_deref()
}
pub fn set_session_affinity_key(&mut self, key: Option<String>) {
self.session_affinity_key = key;
}
pub fn reasoning_effort(&self) -> ReasoningEffort {
self.reasoning_effort
}
pub fn set_reasoning_effort(&mut self, effort: ReasoningEffort) {
self.reasoning_effort = effort;
}
pub fn model_settings(&self) -> &ModelSettings {
&self.model_settings
}
pub fn set_model_settings(&mut self, settings: ModelSettings) {
self.model_settings = settings;
}
pub fn add_message(&mut self, message: ChatMessage) {
self.messages.push(message);
}
pub fn set_tools(&mut self, tools: Vec<ToolDefinition>) {
self.tools = tools;
}
pub fn set_system_content(&mut self, content: String) {
if let Some(ChatMessage::System { content: existing, .. }) = self.messages.first_mut() {
*existing = content;
} else {
self.messages.insert(0, ChatMessage::system(content));
}
}
pub fn system_content(&self) -> Option<&str> {
match self.messages.first() {
Some(ChatMessage::System { content, .. }) if !content.is_empty() => Some(content),
_ => None,
}
}
pub fn messages(&self) -> &Vec<ChatMessage> {
&self.messages
}
pub fn tools(&self) -> &Vec<ToolDefinition> {
&self.tools
}
pub fn message_count(&self) -> usize {
self.messages.len()
}
pub fn estimated_token_count(&self) -> u32 {
let message_bytes: usize = self.messages.iter().map(ChatMessage::estimated_bytes).sum();
let tool_bytes: usize = self
.tools
.iter()
.map(|tool| tool.name.len() + tool.description.len() + tool.parameters.to_string().len())
.sum();
let total_bytes = message_bytes + tool_bytes;
u32::try_from(total_bytes / 4).unwrap_or(u32::MAX)
}
pub fn push_assistant_turn(
&mut self,
message_id: MessageId,
model: Option<LlmModel>,
mut content: AssistantReply,
results: Vec<Result<ToolCallResult, ToolCallError>>,
) -> &ChatMessage {
content.pair_tool_calls(&results);
let reply_index = self.messages.len();
self.messages.push(ChatMessage::Assistant { message_id, model, content, timestamp: IsoString::now() });
self.messages.extend(results.into_iter().map(ChatMessage::ToolCallResult));
&self.messages[reply_index]
}
pub fn clear_conversation(&mut self) {
self.messages.retain(super::chat_message::ChatMessage::is_system);
}
pub fn replace_conversation(&mut self, messages: Vec<ChatMessage>) {
self.messages = self
.messages
.drain(..)
.filter(ChatMessage::is_system)
.chain(messages.into_iter().filter(|m| !m.is_system()))
.collect();
}
pub fn messages_for_summary(&self) -> Vec<&ChatMessage> {
self.messages.iter().filter(|msg| !msg.is_system()).collect()
}
pub fn with_compacted_summary(&self, message_id: MessageId, summary: &str) -> Context {
let system_messages: Vec<_> = self.messages.iter().filter(|msg| msg.is_system()).cloned().collect();
let non_system_count = self.messages.len() - system_messages.len();
let mut messages = system_messages;
if non_system_count > 0 {
messages.push(ChatMessage::Summary {
message_id,
content: summary.to_string(),
timestamp: IsoString::now(),
messages_compacted: non_system_count,
});
}
Context { messages, ..self.clone() }
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::AssistantContent;
use crate::testing::assistant_message;
fn create_test_context() -> Context {
let messages = vec![
ChatMessage::system("You are a helpful assistant."),
ChatMessage::user("Hello"),
assistant_message().text("Hi there!").build(),
ChatMessage::ToolCallResult(Ok(ToolCallResult {
id: "1".to_string(),
name: "tool1".to_string(),
arguments: "{}".to_string(),
result: "Result 1".to_string(),
})),
ChatMessage::ToolCallResult(Ok(ToolCallResult {
id: "2".to_string(),
name: "tool2".to_string(),
arguments: "{}".to_string(),
result: "Result 2".to_string(),
})),
ChatMessage::ToolCallResult(Ok(ToolCallResult {
id: "3".to_string(),
name: "tool3".to_string(),
arguments: "{}".to_string(),
result: "Result 3".to_string(),
})),
];
Context::new(messages, vec![])
}
#[test]
fn replace_conversation_preserves_system_message() {
let mut ctx = create_test_context();
ctx.replace_conversation(vec![ChatMessage::user("new")]);
assert_eq!(ctx.message_count(), 2);
assert!(ctx.messages()[0].is_system());
assert!(matches!(ctx.messages()[1], ChatMessage::User { .. }));
}
#[test]
fn replace_conversation_replaces_old_non_system_messages() {
let mut ctx = create_test_context();
ctx.replace_conversation(vec![assistant_message().text("replacement").build()]);
assert_eq!(ctx.message_count(), 2);
assert!(
ctx.messages()
.iter()
.all(|message| { !matches!(message, ChatMessage::User { .. } | ChatMessage::ToolCallResult(_)) })
);
assert!(
matches!(&ctx.messages()[1], ChatMessage::Assistant { content, .. } if content.text() == "replacement")
);
}
#[test]
fn replace_conversation_filters_incoming_system_messages() {
let mut ctx = create_test_context();
ctx.replace_conversation(vec![ChatMessage::system("wrong system"), ChatMessage::user("kept")]);
assert_eq!(ctx.message_count(), 2);
assert!(
matches!(ctx.messages()[0], ChatMessage::System { ref content, .. } if content == "You are a helpful assistant.")
);
assert!(matches!(ctx.messages()[1], ChatMessage::User { .. }));
}
#[test]
fn replace_conversation_does_not_change_tools() {
let tool = ToolDefinition::new("read_file", "Reads a file", serde_json::json!({}));
let mut ctx = Context::new(vec![ChatMessage::system("system")], vec![tool.clone()]);
ctx.replace_conversation(vec![ChatMessage::user("new")]);
assert_eq!(ctx.tools(), &vec![tool]);
}
#[test]
fn test_message_count() {
let ctx = create_test_context();
assert_eq!(ctx.message_count(), 6);
}
#[test]
fn test_with_compacted_summary_preserves_system_prompt() {
let ctx = create_test_context();
let compacted = ctx.with_compacted_summary(MessageId::new(), "This is a summary of previous conversation.");
assert_eq!(compacted.message_count(), 2);
assert!(compacted.messages()[0].is_system());
assert!(compacted.messages()[1].is_summary());
}
#[test]
fn test_with_compacted_summary_empty_context() {
let ctx = Context::new(vec![ChatMessage::system("System")], vec![]);
let compacted = ctx.with_compacted_summary(MessageId::new(), "Summary");
assert_eq!(compacted.message_count(), 1);
}
#[test]
fn test_messages_for_summary() {
let ctx = create_test_context();
let msgs = ctx.messages_for_summary();
assert_eq!(msgs.len(), 5);
assert!(msgs.iter().all(|m| !m.is_system()));
}
#[test]
fn test_prompt_cache_key_default_is_none() {
let ctx = create_test_context();
assert_eq!(ctx.prompt_cache_key(), None);
}
#[test]
fn test_prompt_cache_key_set_and_get() {
let mut ctx = create_test_context();
ctx.set_prompt_cache_key(Some("session-123".to_string()));
assert_eq!(ctx.prompt_cache_key(), Some("session-123"));
ctx.set_prompt_cache_key(None);
assert_eq!(ctx.prompt_cache_key(), None);
}
#[test]
fn test_prompt_cache_key_preserved_through_compaction() {
let mut ctx = create_test_context();
ctx.set_prompt_cache_key(Some("session-abc".to_string()));
let compacted = ctx.with_compacted_summary(MessageId::new(), "Summary");
assert_eq!(compacted.prompt_cache_key(), Some("session-abc"));
}
#[test]
fn session_affinity_key_is_runtime_metadata_preserved_by_compaction() {
let mut context = create_test_context();
assert_eq!(context.session_affinity_key(), None);
context.set_session_affinity_key(Some("conversation-123".to_string()));
assert_eq!(context.session_affinity_key(), Some("conversation-123"));
assert_eq!(
context.with_compacted_summary(MessageId::new(), "Summary").session_affinity_key(),
Some("conversation-123")
);
}
#[test]
fn test_reasoning_effort_initializes_to_default() {
let ctx = create_test_context();
assert_eq!(ctx.reasoning_effort(), ReasoningEffort::Default);
}
#[test]
fn test_reasoning_effort_set_and_get() {
let mut ctx = create_test_context();
ctx.set_reasoning_effort(crate::ReasoningEffort::High);
assert_eq!(ctx.reasoning_effort(), crate::ReasoningEffort::High);
ctx.set_reasoning_effort(ReasoningEffort::Default);
assert_eq!(ctx.reasoning_effort(), ReasoningEffort::Default);
}
#[test]
fn test_reasoning_effort_preserved_through_compaction() {
let mut ctx = create_test_context();
ctx.set_reasoning_effort(crate::ReasoningEffort::Disabled);
let compacted = ctx.with_compacted_summary(MessageId::new(), "Summary");
assert_eq!(compacted.reasoning_effort(), crate::ReasoningEffort::Disabled);
assert_eq!(ctx.clone().reasoning_effort(), ReasoningEffort::Disabled);
ctx.replace_conversation(vec![]);
assert_eq!(ctx.reasoning_effort(), ReasoningEffort::Disabled);
}
#[test]
fn test_estimated_token_count() {
use crate::ToolDefinition;
let ctx = create_test_context();
let base_estimate = ctx.estimated_token_count();
assert_eq!(base_estimate, 87 / 4);
let tool = ToolDefinition::new("read_file", "Reads a file", serde_json::json!({}));
let ctx_with_tools = Context::new(ctx.messages().clone(), vec![tool]);
let with_tools_estimate = ctx_with_tools.estimated_token_count();
assert_eq!(with_tools_estimate, (87 + 9 + 12 + 2) / 4);
assert!(with_tools_estimate > base_estimate);
}
#[test]
fn compaction_drops_encrypted_reasoning() {
let ctx = Context::new(
vec![
ChatMessage::user("Hello"),
assistant_message()
.model("anthropic:claude-opus-4-6")
.reasoning("thinking")
.encrypted_reasoning("r_test", "blob")
.text("I see.")
.build(),
],
vec![],
);
let compacted = ctx.with_compacted_summary(MessageId::new(), "Summary of conversation");
assert!(compacted.messages().iter().all(|message| encrypted_reasoning(message).is_empty()));
}
#[test]
fn assistant_turn_pairs_every_tool_call_with_a_result() {
let mut ctx = Context::new(vec![], vec![]);
let model: LlmModel = "anthropic:claude-opus-4-6".parse().unwrap();
let content = assistant_message()
.tool_call("call_1", "read", "{}")
.tool_call("call_cancelled", "read", "{}")
.build_reply();
let result = ToolCallResult {
id: "call_1".to_string(),
name: "read".to_string(),
arguments: "{}".to_string(),
result: "ok".to_string(),
};
let error = ToolCallError {
id: "call_retried".to_string(),
name: "bash".to_string(),
arguments: None,
error: "failed".to_string(),
};
let recorded = ctx
.push_assistant_turn(MessageId::new(), Some(model.clone()), content, vec![Ok(result), Err(error)])
.clone();
assert_eq!(ctx.messages()[0], recorded);
assert_eq!(ctx.message_count(), 3);
let ChatMessage::Assistant { model: author, content, .. } = recorded else {
panic!("expected assistant message")
};
let ids: Vec<_> = content.tool_calls().map(|call| call.id.as_str()).collect();
assert_eq!(ids, ["call_1", "call_retried"]);
assert_eq!(author, Some(model));
}
fn encrypted_reasoning(message: &ChatMessage) -> Vec<&str> {
let ChatMessage::Assistant { content, .. } = message else { return vec![] };
content
.blocks()
.iter()
.filter_map(|block| match block {
AssistantContent::Reasoning { encrypted: Some(encrypted), .. } => Some(encrypted.content.as_str()),
_ => None,
})
.collect()
}
}