use std::path::PathBuf;
use std::sync::{Arc, Mutex, PoisonError};
use wasmtime::component::Linker;
use crate::bindings::hook::murmur::conversation::read::MessagePage;
use crate::workdir_fs::WorkdirRoot;
pub(crate) const CONVERSATION_IFACE_VERSIONED: &str = "murmur:conversation/read@0.2.0";
pub(crate) struct ConversationState {
root: Option<PathBuf>,
workdir: WorkdirRoot,
context_id: Mutex<Option<String>>,
cache: Mutex<crate::conversation::RecordCache>,
}
impl ConversationState {
pub(crate) fn new(root: Option<PathBuf>, workdir: &WorkdirRoot) -> Self {
Self {
root,
workdir: workdir.clone(),
context_id: Mutex::new(None),
cache: Mutex::new(crate::conversation::RecordCache::default()),
}
}
pub(crate) fn set_context(&self, context_id: Option<String>) {
*self.lock() = context_id;
}
pub(crate) fn read_messages(
&self,
cursor: Option<String>,
limit: u32,
) -> Result<MessagePage, String> {
let context_id = self.lock().clone();
let path = match (self.root.as_deref(), context_id.as_deref()) {
(Some(root), Some(context_id))
if crate::state_store::validate_store_name(context_id).is_ok() =>
{
Some(root.join(context_id).join("conversation.jsonl"))
}
_ => None,
};
crate::conversation::page(
&mut self.cache(),
path.as_deref(),
&self.workdir,
cursor.as_deref(),
limit,
)
}
#[cfg(test)]
pub(crate) fn record_reads(&self) -> u32 {
self.cache().reads()
}
fn lock(&self) -> std::sync::MutexGuard<'_, Option<String>> {
self.context_id
.lock()
.unwrap_or_else(PoisonError::into_inner)
}
fn cache(&self) -> std::sync::MutexGuard<'_, crate::conversation::RecordCache> {
self.cache.lock().unwrap_or_else(PoisonError::into_inner)
}
}
pub(crate) fn add_conversation_to_linker<T: 'static>(
linker: &mut Linker<T>,
state: Option<Arc<ConversationState>>,
) -> Result<(), String> {
let mut instance = linker
.instance(CONVERSATION_IFACE_VERSIONED)
.map_err(|e| format!("failed to define {CONVERSATION_IFACE_VERSIONED}: {e}"))?;
instance
.func_wrap(
"read-messages",
move |_store: wasmtime::StoreContextMut<'_, T>,
(cursor, limit): (Option<String>, u32)| {
Ok((state.as_ref().map_or_else(
|| Err("not-granted".to_string()),
|state| state.read_messages(cursor, limit),
),))
},
)
.map_err(|e| {
format!("failed to register {CONVERSATION_IFACE_VERSIONED}#read-messages: {e}")
})?;
Ok(())
}
pub(crate) fn add_remote_conversation_to_linker<T: 'static>(
linker: &mut Linker<T>,
link: Arc<crate::guest_worker::RemoteLink>,
) -> Result<(), String> {
let mut instance = linker
.instance(CONVERSATION_IFACE_VERSIONED)
.map_err(|e| format!("failed to define {CONVERSATION_IFACE_VERSIONED}: {e}"))?;
instance
.func_wrap_async(
"read-messages",
move |_store: wasmtime::StoreContextMut<'_, T>,
(cursor, limit): (Option<String>, u32)| {
let link = Arc::clone(&link);
let fut: Box<
dyn std::future::Future<
Output = wasmtime::Result<(Result<MessagePage, String>,)>,
> + Send,
> = Box::new(async move { Ok((link.conversation(cursor, limit).await,)) });
fut
},
)
.map_err(|e| {
format!("failed to register {CONVERSATION_IFACE_VERSIONED}#read-messages: {e}")
})?;
Ok(())
}
#[cfg(test)]
pub(crate) mod test_support {
use super::CONVERSATION_IFACE_VERSIONED;
use wasmtime::component::Component;
fn preamble() -> String {
format!(
r#" (import "murmur:hook/lifecycle@0.9.0" (instance $lct
(type $o (option string))
(type $ci (enum "replace-context" "seed-context"))
(export "context-insertion" (type (eq $ci)))
(type $oci (option 2))
(type $m (record
(field "role" string)
(field "content" string)
(field "id" $o)
(field "source-id" $o)
(field "inserted-by" $oci)))
(export "message" (type (eq $m)))
))
(import "{CONVERSATION_IFACE_VERSIONED}" (instance $conv
(type $optstr (option string))
(type $ci (enum "replace-context" "seed-context"))
(export "context-insertion" (type (eq $ci)))
(type $optci (option 2))
(type $message (record
(field "role" string)
(field "content" string)
(field "id" $optstr)
(field "source-id" $optstr)
(field "inserted-by" $optci)))
(export "message" (type (eq $message)))
(type $msglist (list 5))
(type $page (record
(field "messages" $msglist)
(field "next-cursor" $optstr)
(field "total" u32)))
(export "message-page" (type (eq $page)))
(type $ret (result 8 (error string)))
(export "read-messages"
(func (param "cursor" $optstr) (param "limit" u32) (result $ret)))
))
(alias export $conv "read-messages" (func $read))
(core module $libc
(memory (export "memory") 4)
(global $bump (mut i32) (i32.const 32768))
(func (export "realloc") (param i32 i32 i32 i32) (result i32)
(local $p i32)
(local.set $p (i32.and (i32.add (global.get $bump) (i32.const 7)) (i32.const -8)))
(global.set $bump (i32.add (local.get $p) (i32.add (local.get 3) (i32.const 8))))
(local.get $p))
)
(core instance $li (instantiate $libc))
(alias core export $li "memory" (core memory $mem))
(alias core export $li "realloc" (core func $realloc))
(core func $read_l
(canon lower (func $read) (memory $mem) (realloc $realloc) string-encoding=utf8))"#
)
}
pub(crate) const REPORT_SEP: char = '\u{1f}';
const CORE_HELPERS: &str = r#"
(import "libc" "memory" (memory 4))
(import "conv" "read" (func $read (param i32 i32 i32 i32 i32)))
(global $cur (mut i32) (i32.const 0))
(func $put (param $b i32)
(i32.store8 (global.get $cur) (local.get $b))
(global.set $cur (i32.add (global.get $cur) (i32.const 1))))
(func $append (param $src i32) (param $n i32)
(memory.copy (global.get $cur) (local.get $src) (local.get $n))
(global.set $cur (i32.add (global.get $cur) (local.get $n))))
(func $dec (param $v i32)
(if (i32.ge_u (local.get $v) (i32.const 10))
(then (call $dec (i32.div_u (local.get $v) (i32.const 10)))))
(call $put (i32.add (i32.const 48) (i32.rem_u (local.get $v) (i32.const 10)))))
(func $emit_page (param $rp i32) (local $i i32) (local $rec i32) (local $n i32)
(if (i32.eqz (i32.load8_u (local.get $rp)))
(then
(call $put (i32.const 84)) (call $put (i32.const 61))
(call $dec (i32.load (i32.add (local.get $rp) (i32.const 24))))
(local.set $n (i32.load (i32.add (local.get $rp) (i32.const 8))))
(local.set $i (i32.const 0))
(block $done
(loop $next
(br_if $done (i32.ge_u (local.get $i) (local.get $n)))
(local.set $rec (i32.add (i32.load (i32.add (local.get $rp) (i32.const 4)))
(i32.mul (local.get $i) (i32.const 44))))
(call $put (i32.const 31))
;; id, when the option carries one
(if (i32.load8_u (i32.add (local.get $rec) (i32.const 16)))
(then (call $append (i32.load (i32.add (local.get $rec) (i32.const 20)))
(i32.load (i32.add (local.get $rec) (i32.const 24))))))
(call $put (i32.const 61))
(call $append (i32.load (local.get $rec)) (i32.load (i32.add (local.get $rec) (i32.const 4))))
(call $put (i32.const 61))
;; inserted-by: `r` for replace-context, `s` for seed-context, `-` for none
(if (i32.load8_u (i32.add (local.get $rec) (i32.const 40)))
(then (if (i32.load8_u (i32.add (local.get $rec) (i32.const 41)))
(then (call $put (i32.const 115)))
(else (call $put (i32.const 114)))))
(else (call $put (i32.const 45))))
(local.set $i (i32.add (local.get $i) (i32.const 1)))
(br $next)))
(call $put (i32.const 31))
(call $put (i32.const 78)) (call $put (i32.const 61))
(if (i32.load8_u (i32.add (local.get $rp) (i32.const 12)))
(then (call $append (i32.load (i32.add (local.get $rp) (i32.const 16)))
(i32.load (i32.add (local.get $rp) (i32.const 20)))))
(else (call $put (i32.const 45)))))
(else
(call $put (i32.const 33))
(call $append (i32.load (i32.add (local.get $rp) (i32.const 4)))
(i32.load (i32.add (local.get $rp) (i32.const 8)))))))
"#;
const LIFECYCLE_TYPES: &str = r#"
(type $context-insertion (enum "replace-context" "seed-context"))
(type $message (record
(field "role" string)
(field "content" string)
(field "id" (option string))
(field "source-id" (option string))
(field "inserted-by" (option $context-insertion))))
(type $tool-manifest (record (field "binary-name" string) (field "content" string)))
(type $hook-output (variant
(case "none")
(case "replace-context" (list $message))
(case "write-manifests" (list $tool-manifest))
(case "artifact" string)
(case "reopen-task" string)
(case "seed-context" (list $message))
(case "deny" string)))
(type $task-end-event (record
(field "task-id" string)
(field "exit-status" string)))
"#;
pub(crate) fn reader_double(engine: &wasmtime::Engine, limit: u32) -> Component {
let core_body = format!(
r#"
(func (export "ontaskend") (param i32 i32 i32 i32) (result i32)
(global.set $cur (i32.const 4096))
(call $read (i32.const 0) (i32.const 0) (i32.const 0) (i32.const {limit}) (i32.const 1024))
(call $emit_page (i32.const 1024))
(call $put (i32.const 124))
;; The second call reuses the first page's `next-cursor`, when it had one.
(if (i32.and (i32.eqz (i32.load8_u (i32.const 1024)))
(i32.load8_u (i32.const 1036)))
(then (call $read (i32.const 1) (i32.load (i32.const 1040)) (i32.load (i32.const 1044))
(i32.const {limit}) (i32.const 2048)))
(else (call $read (i32.const 0) (i32.const 0) (i32.const 0)
(i32.const {limit}) (i32.const 2048))))
(call $emit_page (i32.const 2048))
(i32.store (i32.const 128) (i32.const 0))
(i32.store (i32.const 132) (i32.const 4))
(i32.store (i32.const 136) (i32.const 4096))
(i32.store (i32.const 140) (i32.sub (global.get $cur) (i32.const 4096)))
(i32.const 128))
(func (export "noop"))
"#
);
let wat = format!(
"(component\n{preamble}\n\n (core module $m{CORE_HELPERS}{core_body} )\n \
(core instance $i (instantiate $m\n (with \"libc\" (instance $li))\n \
(with \"conv\" (instance (export \"read\" (func $read_l))))))\n\
{LIFECYCLE_TYPES}\n \
(type $ontaskend-ft (func (param \"event\" $task-end-event) \
(result (result $hook-output (error string)))))\n \
(func $ontaskend (type $ontaskend-ft)\n (canon lift (core func $i \"ontaskend\") \
(memory $mem) (realloc $realloc) string-encoding=utf8))\n \
(func $noop (canon lift (core func $i \"noop\")))\n \
(instance $lc\n (export \"context-insertion\" (type $context-insertion))\n (export \"message\" (type $message))\n \
(export \"tool-manifest\" (type $tool-manifest))\n \
(export \"hook-output\" (type $hook-output))\n \
(export \"task-end-event\" (type $task-end-event))\n \
(export \"on-task-end\" (func $ontaskend))\n{stubs} )\n \
(export \"murmur:hook/lifecycle@0.9.0\" (instance $lc))\n)",
preamble = preamble(),
stubs = [
"on-session-start",
"on-inference",
"on-tool-call",
"on-shell",
"on-compaction",
"on-session-end",
]
.iter()
.map(|name| format!(" (export \"{name}\" (func $noop))\n"))
.collect::<String>(),
);
let bytes = wat::parse_str(&wat).expect("conversation double WAT parses");
Component::new(engine, &bytes).expect("conversation double compiles")
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::path::Path;
#[test]
fn the_guest_double_compiles() {
let mut config = wasmtime::Config::new();
config.wasm_component_model(true);
let engine = wasmtime::Engine::new(&config).expect("engine builds");
test_support::reader_double(&engine, 2);
}
#[test]
fn an_importing_component_links_with_and_without_the_grant() {
let mut config = wasmtime::Config::new();
config.wasm_component_model(true);
let engine = wasmtime::Engine::new(&config).expect("engine builds");
let component = test_support::reader_double(&engine, 2);
let workdir = tempfile::tempdir().unwrap();
for state in [
None,
Some(Arc::new(ConversationState::new(
None,
&WorkdirRoot::open(workdir.path()).unwrap(),
))),
] {
let mut linker: Linker<()> = Linker::new(&engine);
add_conversation_to_linker(&mut linker, state).unwrap();
let mut store = wasmtime::Store::new(&engine, ());
linker
.instantiate(&mut store, &component)
.expect("the import is defined on every linker");
}
}
#[test]
fn no_record_and_no_context_both_read_as_an_empty_page() {
let workdir = tempfile::tempdir().unwrap();
let home = tempfile::tempdir().unwrap();
let no_record = ConversationState::new(None, &WorkdirRoot::open(workdir.path()).unwrap());
no_record.set_context(Some("ctx_1".to_string()));
let no_context = ConversationState::new(
Some(home.path().to_path_buf()),
&WorkdirRoot::open(workdir.path()).unwrap(),
);
for state in [&no_record, &no_context] {
let page = state.read_messages(None, 10).expect("an Ok empty page");
assert!(page.messages.is_empty());
assert_eq!(page.next_cursor, None);
assert_eq!(page.total, 0);
}
}
#[test]
fn a_read_follows_the_context_the_state_was_moved_to() {
use serde_json::json;
let workdir = tempfile::tempdir().unwrap();
let home = tempfile::tempdir().unwrap();
let root = home.path().join("conversations/capsule");
for (context_id, count) in [("ctx_a", 2), ("ctx_b", 3)] {
let dir = root.join(context_id);
std::fs::create_dir_all(&dir).unwrap();
let raw: String = (0..count)
.map(|n| {
format!(
"{}\n",
json!({"role": "user", "content": "m", "id": format!("{context_id}_{n}")})
)
})
.collect();
std::fs::write(dir.join("conversation.jsonl"), raw).unwrap();
}
let state = ConversationState::new(Some(root), &WorkdirRoot::open(workdir.path()).unwrap());
state.set_context(Some("ctx_a".to_string()));
let first = state.read_messages(None, 10).unwrap();
assert_eq!(first.total, 2);
state.set_context(Some("ctx_b".to_string()));
let second = state.read_messages(None, 10).unwrap();
assert_eq!(second.total, 3);
assert_eq!(
second
.messages
.iter()
.map(|m| m.id.clone().unwrap())
.collect::<Vec<_>>(),
vec!["ctx_b_2", "ctx_b_1", "ctx_b_0"],
);
assert_eq!(state.record_reads(), 2, "one read per record in scope");
}
#[test]
fn a_context_id_that_escapes_reads_as_no_record() {
let workdir = tempfile::tempdir().unwrap();
let home = tempfile::tempdir().unwrap();
let state = ConversationState::new(
Some(home.path().join("conversations/c")),
&WorkdirRoot::open(workdir.path()).unwrap(),
);
state.set_context(Some("../../escape".to_string()));
assert_eq!(state.read_messages(None, 10).unwrap().total, 0);
}
fn write_record_lines(root: &Path, context_id: &str, lines: &[serde_json::Value]) {
let dir = root.join(context_id);
std::fs::create_dir_all(&dir).unwrap();
let raw: String = lines.iter().map(|line| format!("{line}\n")).collect();
std::fs::write(dir.join("conversation.jsonl"), raw).unwrap();
}
#[test]
fn a_reader_sees_which_messages_a_hook_inserted() {
use serde_json::json;
use wasmtime::component::Val;
let workdir = tempfile::tempdir().unwrap();
let home = tempfile::tempdir().unwrap();
let root = home.path().join("conversations/capsule");
write_record_lines(
&root,
"ctx_marks",
&[
json!({"role": "user", "content": "a person's turn", "id": "msg_person"}),
json!({
"role": "user",
"content": "a summary",
"id": "msg_summary",
"inserted_by": "replace-context",
}),
json!({
"role": "user",
"content": "a seed",
"id": "msg_seed",
"inserted_by": "seed-context",
}),
],
);
let state = Arc::new(ConversationState::new(
Some(root),
&WorkdirRoot::open(workdir.path()).unwrap(),
));
state.set_context(Some("ctx_marks".to_string()));
let mut config = wasmtime::Config::new();
config.wasm_component_model(true);
let engine = wasmtime::Engine::new(&config).expect("engine builds");
let component = test_support::reader_double(&engine, 3);
let mut linker: Linker<()> = Linker::new(&engine);
add_conversation_to_linker(&mut linker, Some(state)).unwrap();
let mut store = wasmtime::Store::new(&engine, ());
let instance = linker.instantiate(&mut store, &component).unwrap();
let lifecycle = instance
.get_export_index(&mut store, None, "murmur:hook/lifecycle@0.9.0")
.expect("the double exports the lifecycle");
let on_task_end = instance
.get_export_index(&mut store, Some(&lifecycle), "on-task-end")
.and_then(|index| instance.get_func(&mut store, index))
.expect("the double exports on-task-end");
let mut results = [Val::Bool(false)];
on_task_end
.call(
&mut store,
&[Val::Record(vec![
("task-id".to_string(), Val::String("tsk_1".to_string())),
("exit-status".to_string(), Val::String("ok".to_string())),
])],
&mut results,
)
.expect("the read succeeds");
let Val::Result(Ok(Some(output))) = &results[0] else {
panic!("the double returns ok(hook-output): {:?}", results[0]);
};
let Val::Variant(case, Some(reason)) = output.as_ref() else {
panic!("the double returns reopen-task: {output:?}");
};
assert_eq!(case, "reopen-task");
let Val::String(report) = reason.as_ref() else {
panic!("reopen-task carries a string: {reason:?}");
};
let first_page = report.split('|').next().unwrap();
assert_eq!(
first_page,
[
"T=3",
"msg_seed=user=s",
"msg_summary=user=r",
"msg_person=user=-",
"N=-",
]
.join(&test_support::REPORT_SEP.to_string()),
);
}
#[test]
fn an_unrecognised_inserted_by_reads_as_none() {
use serde_json::json;
let workdir = tempfile::tempdir().unwrap();
let home = tempfile::tempdir().unwrap();
let root = home.path().join("conversations/capsule");
write_record_lines(
&root,
"ctx_odd",
&[
json!({"role": "user", "content": "m", "id": "msg_absent"}),
json!({"role": "user", "content": "m", "id": "msg_bogus", "inserted_by": "bogus"}),
json!({"role": "user", "content": "m", "id": "msg_number", "inserted_by": 42}),
],
);
let state = ConversationState::new(Some(root), &WorkdirRoot::open(workdir.path()).unwrap());
state.set_context(Some("ctx_odd".to_string()));
let page = state.read_messages(None, 10).expect("the read succeeds");
assert_eq!(page.total, 3);
assert_eq!(
page.messages
.iter()
.map(|m| (m.id.clone().unwrap(), m.inserted_by))
.collect::<Vec<_>>(),
vec![
("msg_number".to_string(), None),
("msg_bogus".to_string(), None),
("msg_absent".to_string(), None),
],
);
assert!(
!workdir.path().join("logs/bootstrap.log").exists(),
"an odd mark is not a defect worth reporting"
);
}
}