use crate::error::MemoryError;
use crate::ids::{FactId, RunId, ThreadId};
use crate::llm::Message;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "outcome", rename_all = "snake_case")]
pub enum ToolResult {
Ok {
value: serde_json::Value,
},
Err {
message: String,
},
}
impl ToolResult {
pub fn ok(value: serde_json::Value) -> Self {
Self::Ok { value }
}
pub fn err(message: impl Into<String>) -> Self {
Self::Err {
message: message.into(),
}
}
}
#[cfg(test)]
mod tool_result_factories_tests {
use super::*;
#[test]
fn ok_factory_builds_ok_variant() {
let v = serde_json::json!({"hit": true});
let r = ToolResult::ok(v.clone());
match r {
ToolResult::Ok { value } => assert_eq!(value, v),
_ => panic!("expected ToolResult::Ok variant"),
}
}
#[test]
fn err_factory_builds_err_variant() {
let r = ToolResult::err("boom");
match r {
ToolResult::Err { message } => assert_eq!(message, "boom"),
_ => panic!("expected ToolResult::Err variant"),
}
}
}
#[async_trait]
pub trait ShortTermMemory: Send + Sync {
async fn append(&self, thread: ThreadId, msg: Message) -> Result<(), MemoryError>;
async fn append_batch(
&self,
thread: ThreadId,
messages: Vec<Message>,
) -> Result<(), MemoryError> {
for msg in messages {
self.append(thread.clone(), msg).await?;
}
Ok(())
}
async fn load(&self, thread: ThreadId, max_tokens: usize) -> Result<Vec<Message>, MemoryError>;
async fn clear(&self, thread: ThreadId) -> Result<(), MemoryError>;
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum Scope {
Workspace(String),
Agent(String),
Global,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[non_exhaustive]
pub struct Fact {
pub text: String,
#[serde(default)]
pub metadata: serde_json::Value,
}
impl Fact {
pub fn new(text: impl Into<String>) -> Self {
Self {
text: text.into(),
metadata: serde_json::Value::Null,
}
}
pub fn with_metadata(mut self, metadata: serde_json::Value) -> Self {
self.metadata = metadata;
self
}
}
#[async_trait]
pub trait LongTermMemory: Send + Sync {
async fn remember(&self, scope: Scope, fact: Fact) -> Result<FactId, MemoryError>;
async fn recall(&self, scope: Scope, query: &str, k: usize) -> Result<Vec<Fact>, MemoryError>;
async fn forget(&self, id: FactId) -> Result<(), MemoryError>;
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[non_exhaustive]
pub enum Episode {
Started {
agent: String,
},
LlmCall {
tokens: u32,
latency_ms: u32,
#[serde(default)]
provider: Option<String>,
#[serde(default)]
model: Option<String>,
#[serde(default)]
prompt_tokens: Option<u32>,
#[serde(default)]
completion_tokens: Option<u32>,
},
ToolCall {
name: String,
args: serde_json::Value,
result: ToolResult,
},
BusPublish {
subject: String,
},
BusReceive {
subject: String,
},
BusCausalLink {
subject: String,
caused_by_run: String,
},
Completed,
Failed {
error: String,
},
SummaryCheckpoint {
input_message_count: u32,
summary_chars: u32,
latency_ms: u32,
tokens: u32,
},
Ops(serde_json::Value),
RunAttributed {
tenant_label: String,
},
RunOrigin {
parent_anchor: String,
},
MemoryRecall {
query: String,
#[serde(default)]
k: u32,
#[serde(default)]
returned_fact_ids: Vec<FactId>,
},
}
impl Episode {
pub fn llm_call(tokens: u32, latency_ms: u32) -> Self {
Episode::LlmCall {
tokens,
latency_ms,
provider: None,
model: None,
prompt_tokens: None,
completion_tokens: None,
}
}
}
#[derive(Debug, Clone, Default)]
pub struct RunFilter {
pub agent: Option<String>,
pub since: Option<DateTime<Utc>>,
pub until: Option<DateTime<Utc>>,
pub limit: Option<usize>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RunSummary {
pub run_id: RunId,
pub agent: String,
pub started_at: DateTime<Utc>,
pub finished_at: Option<DateTime<Utc>>,
pub episode_count: u32,
}
#[async_trait]
pub trait EpisodicMemory: Send + Sync {
async fn record(&self, run: RunId, event: Episode) -> Result<(), MemoryError>;
async fn replay(&self, run: RunId) -> Result<Vec<Episode>, MemoryError>;
async fn replay_many(&self, runs: &[RunId]) -> Result<Vec<(RunId, Vec<Episode>)>, MemoryError> {
let mut out = Vec::with_capacity(runs.len());
for &run in runs {
out.push((run, self.replay(run).await?));
}
Ok(out)
}
async fn list_runs(&self, filter: RunFilter) -> Result<Vec<RunSummary>, MemoryError>;
}
#[derive(Clone)]
pub struct MemoryHandles {
pub short_term: std::sync::Arc<dyn ShortTermMemory>,
pub long_term: std::sync::Arc<dyn LongTermMemory>,
pub episodic: std::sync::Arc<dyn EpisodicMemory>,
}
impl MemoryHandles {
pub fn new(
short_term: std::sync::Arc<dyn ShortTermMemory>,
long_term: std::sync::Arc<dyn LongTermMemory>,
episodic: std::sync::Arc<dyn EpisodicMemory>,
) -> Self {
Self {
short_term,
long_term,
episodic,
}
}
}
#[cfg(test)]
mod fact_ctor_tests {
use super::*;
#[test]
fn fact_new_defaults_metadata_null() {
let f = Fact::new("alice likes tea");
assert_eq!(f.text, "alice likes tea");
assert_eq!(f.metadata, serde_json::Value::Null);
let f2 = Fact::new("x").with_metadata(serde_json::json!({"k":"v"}));
assert_eq!(f2.metadata, serde_json::json!({"k":"v"}));
}
}
#[cfg(test)]
mod tests {
use super::*;
#[allow(dead_code)]
fn _assert_dyn_short(_: &dyn ShortTermMemory) {}
#[allow(dead_code)]
fn _assert_dyn_long(_: &dyn LongTermMemory) {}
#[allow(dead_code)]
fn _assert_dyn_episodic(_: &dyn EpisodicMemory) {}
fn kind_discriminant(episode: &Episode) -> &'static str {
match episode {
Episode::Started { .. } => "started",
Episode::LlmCall { .. } => "llm_call",
Episode::ToolCall { .. } => "tool_call",
Episode::BusPublish { .. } => "bus_publish",
Episode::BusReceive { .. } => "bus_receive",
Episode::BusCausalLink { .. } => "bus_causal_link",
Episode::Completed => "completed",
Episode::Failed { .. } => "failed",
Episode::SummaryCheckpoint { .. } => "summary_checkpoint",
Episode::Ops(_) => "ops",
Episode::RunAttributed { .. } => "run_attributed",
Episode::RunOrigin { .. } => "run_origin",
Episode::MemoryRecall { .. } => "memory_recall",
}
}
fn one_sample_per_variant() -> Vec<Episode> {
vec![
Episode::Started {
agent: String::new(),
},
Episode::LlmCall {
tokens: 0,
latency_ms: 0,
provider: None,
model: None,
prompt_tokens: None,
completion_tokens: None,
},
Episode::ToolCall {
name: String::new(),
args: serde_json::Value::Null,
result: ToolResult::Ok {
value: serde_json::Value::Null,
},
},
Episode::BusPublish {
subject: String::new(),
},
Episode::BusReceive {
subject: String::new(),
},
Episode::BusCausalLink {
subject: String::new(),
caused_by_run: String::new(),
},
Episode::Completed,
Episode::Failed {
error: String::new(),
},
Episode::SummaryCheckpoint {
input_message_count: 0,
summary_chars: 0,
latency_ms: 0,
tokens: 0,
},
Episode::Ops(serde_json::Value::Null),
Episode::RunAttributed {
tenant_label: String::new(),
},
Episode::RunOrigin {
parent_anchor: String::new(),
},
Episode::MemoryRecall {
query: String::new(),
k: 0,
returned_fact_ids: Vec::new(),
},
]
}
#[test]
fn episode_discriminants_match_published_envelope_schema() {
let mut discriminants: Vec<&str> = one_sample_per_variant()
.iter()
.map(kind_discriminant)
.collect();
discriminants.sort_unstable();
let schema_path = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.join("../../docs/schemas/runlog/episode.schema.json");
let text = std::fs::read_to_string(&schema_path)
.unwrap_or_else(|err| panic!("read {}: {err}", schema_path.display()));
let schema: serde_json::Value =
serde_json::from_str(&text).expect("envelope schema parses");
let mut published: Vec<&str> = schema["properties"]["kind"]["enum"]
.as_array()
.expect("envelope schema declares a kind enum")
.iter()
.map(|value| value.as_str().expect("kind enum is strings"))
.collect();
published.sort_unstable();
assert_eq!(
discriminants, published,
"Episode variants have drifted from the published envelope kind enum",
);
}
#[test]
fn legacy_llm_call_json_deserialises_with_none_for_new_fields() {
let legacy = serde_json::json!({
"LlmCall": {
"tokens": 42,
"latency_ms": 17
}
});
let ep: Episode = serde_json::from_value(legacy).expect("legacy LlmCall decodes");
match ep {
Episode::LlmCall {
tokens,
latency_ms,
provider,
model,
prompt_tokens,
completion_tokens,
} => {
assert_eq!(tokens, 42);
assert_eq!(latency_ms, 17);
assert!(provider.is_none());
assert!(model.is_none());
assert!(prompt_tokens.is_none());
assert!(completion_tokens.is_none());
}
other => panic!("expected LlmCall, got {other:?}"),
}
}
#[test]
fn llm_call_ctor_leaves_enrichment_fields_none() {
match Episode::llm_call(42, 17) {
Episode::LlmCall {
tokens,
latency_ms,
provider,
model,
prompt_tokens,
completion_tokens,
} => {
assert_eq!(tokens, 42);
assert_eq!(latency_ms, 17);
assert!(provider.is_none());
assert!(model.is_none());
assert!(prompt_tokens.is_none());
assert!(completion_tokens.is_none());
}
other => panic!("expected LlmCall, got {other:?}"),
}
}
#[test]
fn enriched_llm_call_round_trips() {
let original = Episode::LlmCall {
tokens: 60,
latency_ms: 17,
provider: Some("ollama".into()),
model: Some("qwen2.5:14b".into()),
prompt_tokens: Some(40),
completion_tokens: Some(20),
};
let json = serde_json::to_value(&original).expect("serialises");
let back: Episode = serde_json::from_value(json).expect("deserialises");
match back {
Episode::LlmCall {
provider, model, ..
} => {
assert_eq!(provider.as_deref(), Some("ollama"));
assert_eq!(model.as_deref(), Some("qwen2.5:14b"));
}
other => panic!("expected LlmCall, got {other:?}"),
}
}
#[test]
fn memory_recall_round_trips_and_defaults_legacy() {
let ep = Episode::MemoryRecall {
query: "q".into(),
k: 5,
returned_fact_ids: vec![FactId::new("fact_1")],
};
let json = serde_json::to_value(&ep).expect("serialises");
let back: Episode = serde_json::from_value(json).expect("deserialises");
match back {
Episode::MemoryRecall {
query,
k,
returned_fact_ids,
} => {
assert_eq!(query, "q");
assert_eq!(k, 5);
assert_eq!(returned_fact_ids, vec![FactId::new("fact_1")]);
}
other => panic!("expected MemoryRecall, got {other:?}"),
}
let legacy = serde_json::json!({ "MemoryRecall": { "query": "q" } });
let ep: Episode = serde_json::from_value(legacy).expect("legacy MemoryRecall decodes");
match ep {
Episode::MemoryRecall {
k,
returned_fact_ids,
..
} => {
assert_eq!(k, 0);
assert!(returned_fact_ids.is_empty());
}
other => panic!("expected MemoryRecall, got {other:?}"),
}
}
}