From 4f4d5d99479cd26d065a222cb28989d0c50e08dd Mon Sep 17 00:00:00 2001 From: Josh Field <10372036+HexaField@users.noreply.github.com> Date: Tue, 28 Jul 2026 19:15:38 +1000 Subject: [PATCH 1/6] feat(api/openai): tool/function calling on /v1 via kalosm constrained decoding Makes /v1/chat/completions a real OpenAI tool-calling endpoint for local models: request tools/tool_choice/parallel_tool_calls, response tool_calls + finish_reason "tool_calls" (oneshot + streaming). Local output is held to a schema-valid by compiling each tool's JSON-Schema into a kalosm ArcParser (openai_compat/tool_grammar.rs). Multi-turn tool results fold into prompt text; the no-tools path is unchanged; Cargo.toml unchanged (kalosm-sample via re-export). Co-Authored-By: Claude Opus 4.8 --- rust-executor/src/ai_service/mod.rs | 92 ++- rust-executor/src/api/openai_compat/chat.rs | 347 ++++++++-- rust-executor/src/api/openai_compat/mod.rs | 1 + .../src/api/openai_compat/tool_grammar.rs | 594 ++++++++++++++++++ rust-executor/src/api/openai_compat/types.rs | 104 ++- 5 files changed, 1061 insertions(+), 77 deletions(-) create mode 100644 rust-executor/src/api/openai_compat/tool_grammar.rs diff --git a/rust-executor/src/ai_service/mod.rs b/rust-executor/src/ai_service/mod.rs index 9e12c1a76..18554ff8e 100644 --- a/rust-executor/src/ai_service/mod.rs +++ b/rust-executor/src/ai_service/mod.rs @@ -118,10 +118,13 @@ struct LLMTaskSpawnRequest { } #[allow(dead_code)] -#[derive(Debug)] struct LLMTaskPromptRequest { pub task_id: String, pub prompt: String, + /// Optional decoding constraint (a tool-call grammar). `None` ⇒ normal + /// unconstrained generation, byte-for-byte the pre-tools behaviour. + /// Ignored on the remote path (upstream tool forwarding is a follow-up). + pub constraint: Option>, pub result_sender: oneshot::Sender>, } @@ -152,14 +155,38 @@ struct LLMTaskShutdownRequest { /// `done_sender` fires once the model has emitted its final token (or /// errored) and carries `PromptResult` for the closing chunk's `usage`. #[allow(dead_code)] -#[derive(Debug)] struct LLMTaskPromptStreamRequest { pub task_id: String, pub prompt: String, + /// See [`LLMTaskPromptRequest::constraint`]. + pub constraint: Option>, pub token_sender: mpsc::UnboundedSender, pub done_sender: oneshot::Sender>, } +// Manual `Debug` — these structs hold a non-`Debug` `ArcParser` constraint. +// `LLMTaskRequest` must stay `Debug` so that `SendError` +// converts into `anyhow::Error` via `?` at the channel send sites. +impl std::fmt::Debug for LLMTaskPromptRequest { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LLMTaskPromptRequest") + .field("task_id", &self.task_id) + .field("prompt", &self.prompt) + .field("constrained", &self.constraint.is_some()) + .finish() + } +} + +impl std::fmt::Debug for LLMTaskPromptStreamRequest { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LLMTaskPromptStreamRequest") + .field("task_id", &self.task_id) + .field("prompt", &self.prompt) + .field("constrained", &self.constraint.is_some()) + .finish() + } +} + #[allow(dead_code)] #[derive(Debug)] enum LLMTaskRequest { @@ -847,7 +874,28 @@ impl AIService { )); let result = rt.block_on(async { - task.run(prompt_request.prompt.clone()).all_text().await + match prompt_request.constraint.clone() { + // Tool-call grammar: constrain decoding and + // accumulate the (guaranteed on-grammar) text. + Some(parser) => { + use futures::StreamExt; + let mut stream = Box::pin( + task.run(prompt_request.prompt.clone()) + .with_constraints(parser), + ); + let mut acc = String::new(); + while let Some(token) = stream.next().await { + acc.push_str(&token); + } + acc + } + // No tools: unchanged unconstrained path. + None => { + task.run(prompt_request.prompt.clone()) + .all_text() + .await + } + } }); rt.block_on(publish_model_status( @@ -969,14 +1017,33 @@ impl AIService { // implements `Stream`; // polling it yields one token // chunk at a time. - let mut stream = - Box::pin(task.run(prompt_clone.clone())); let mut accumulated = String::new(); - while let Some(token) = stream.next().await { - accumulated.push_str(&token); - if token_sender.send(token).is_err() { - // consumer dropped — stop generating - break; + match stream_request.constraint.clone() { + // Tool-call grammar: constrained streaming. + Some(parser) => { + let mut stream = Box::pin( + task.run(prompt_clone.clone()) + .with_constraints(parser), + ); + while let Some(token) = stream.next().await { + accumulated.push_str(&token); + if token_sender.send(token).is_err() { + // consumer dropped — stop generating + break; + } + } + } + // No tools: unchanged unconstrained streaming. + None => { + let mut stream = + Box::pin(task.run(prompt_clone.clone())); + while let Some(token) = stream.next().await { + accumulated.push_str(&token); + if token_sender.send(token).is_err() { + // consumer dropped — stop generating + break; + } + } } } accumulated @@ -1205,6 +1272,7 @@ impl AIService { &self, model_id: String, messages: Vec<(String, String)>, + constraint: Option>, ) -> Result { let resolved = Self::replace_model_variables(&model_id)?; let (task, final_prompt) = Self::build_ephemeral_task(&resolved, messages); @@ -1235,6 +1303,7 @@ impl AIService { sender.send(LLMTaskRequest::Prompt(LLMTaskPromptRequest { task_id: task_id.clone(), prompt: final_prompt.clone(), + constraint, result_sender: prompt_tx, }))?; } @@ -1269,6 +1338,7 @@ impl AIService { &self, model_id: String, messages: Vec<(String, String)>, + constraint: Option>, ) -> Result<( mpsc::UnboundedReceiver, oneshot::Receiver>, @@ -1300,6 +1370,7 @@ impl AIService { sender.send(LLMTaskRequest::PromptStream(LLMTaskPromptStreamRequest { task_id: task_id.clone(), prompt: final_prompt, + constraint, token_sender: token_tx, done_sender: done_tx, }))?; @@ -1350,6 +1421,7 @@ impl AIService { task_id, prompt, result_sender, + constraint: None, }))?; } else { return Err(anyhow::anyhow!( diff --git a/rust-executor/src/api/openai_compat/chat.rs b/rust-executor/src/api/openai_compat/chat.rs index ad9027b01..0fafc6147 100644 --- a/rust-executor/src/api/openai_compat/chat.rs +++ b/rust-executor/src/api/openai_compat/chat.rs @@ -4,6 +4,17 @@ //! `AIService::prompt_messages{,_stream}` call. No DB-backed task is //! created; the model thread spawns the task in-memory for the duration //! of the call. +//! +//! ## Tool / function calling +//! +//! When the request carries `tools`, we inject the Hermes/Qwen +//! `…` block as a system message and, for +//! `tool_choice: "required"` or a named function, hand the local model a +//! grammar constraint (see [`super::tool_grammar`]) so decoding is forced +//! onto a well-formed `` block. `auto`/`none` generate freely; +//! any tool calls are recovered from the text afterwards. Assistant +//! `tool_calls` and `role:"tool"` results are folded back into prompt text +//! because the local chat template has no tool role. use std::convert::Infallible; use std::time::SystemTime; @@ -13,14 +24,17 @@ use axum::{ Json, }; use futures::Stream; +use kalosm::language::ArcParser; use uuid::Uuid; use super::errors::{OpenAIError, OpenAIResult}; use super::model_selector::resolve_model; +use super::tool_grammar::{self, ExtractedToolCall, ToolChoice}; use super::types::{ ChatChoice, ChatChunkChoice, ChatChunkDelta, ChatCompletionChunk, ChatCompletionRequest, - ChatCompletionResponse, ChatMessageContent, ChatResponseMessage, CompletionChoice, - CompletionRequest, CompletionResponse, Role, Usage, + ChatCompletionResponse, ChatMessage, ChatMessageContent, ChatResponseMessage, CompletionChoice, + CompletionRequest, CompletionResponse, FunctionCall, FunctionCallDelta, Role, ToolCall, + ToolCallDelta, ToolDef, Usage, }; use crate::agent::capabilities::{check_capability, AI_PROMPT_CAPABILITY}; use crate::ai_service::AIService; @@ -39,29 +53,61 @@ pub async fn chat_completions( // Resolve the OpenAI `model` string to an AD4M model_id. let model_id = resolve_model(&req.model, ModelType::Llm).await?; - let messages: Vec<(String, String)> = req - .messages - .iter() - .map(|m| { - ( - role_to_str(&m.role).to_string(), - m.content - .as_ref() - .map(ChatMessageContent::flatten_to_text) - .unwrap_or_default(), - ) - }) - .collect(); + let tools: Vec = req.tools.clone().unwrap_or_default(); + let has_tools = !tools.is_empty(); + let choice = tool_grammar::parse_tool_choice(&req.tool_choice, has_tools); + let parallel = req.parallel_tool_calls.unwrap_or(true); + // Tools are "active" (rendered + potentially constrained + parsed out) + // unless the caller explicitly disabled them with tool_choice: "none". + let tools_active = has_tools && choice != ToolChoice::None; + + // Assemble (role, content) pairs. Fold assistant tool calls and + // `role:"tool"` results into text, and prepend the tools system prompt. + let mut messages: Vec<(String, String)> = Vec::with_capacity(req.messages.len() + 1); + if tools_active { + messages.push(( + "system".to_string(), + tool_grammar::render_tools_system_prompt(&tools), + )); + } + for m in &req.messages { + messages.push(flatten_message(m)); + } + + // Constrain decoding for required / named choices (guarantees a + // well-formed call). Auto/None return `None` and generate freely. + let constraint = if tools_active { + tool_grammar::build_tool_call_parser(&tools, &choice, parallel) + } else { + None + }; if req.stream { - chat_stream(auth, req.model.clone(), model_id, messages).await + chat_stream( + auth, + req.model.clone(), + model_id, + messages, + constraint, + tools_active, + ) + .await } else { - chat_oneshot(auth, req.model.clone(), model_id, messages).await + chat_oneshot( + auth, + req.model.clone(), + model_id, + messages, + constraint, + tools_active, + ) + .await } } /// `POST /v1/completions` (legacy text-completion). Treats `prompt` as a -/// single user message with no system prompt. +/// single user message with no system prompt. Tools are not supported on +/// the legacy endpoint. pub async fn completions( auth: AuthContext, Json(req): Json, @@ -84,7 +130,7 @@ pub async fn completions( .await .map_err(|e| OpenAIError::internal(e.to_string()))?; let result = service - .prompt_messages(model_id, messages) + .prompt_messages(model_id, messages, None) .await .map_err(|e| OpenAIError::internal(e.to_string()))?; @@ -115,6 +161,8 @@ async fn chat_oneshot( requested_model: String, model_id: String, messages: Vec<(String, String)>, + constraint: Option>, + tools_active: bool, ) -> Result { if let Some(email) = user_email(&auth) { check_compute_credits(&email) @@ -125,7 +173,7 @@ async fn chat_oneshot( .await .map_err(|e| OpenAIError::internal(e.to_string()))?; let result = service - .prompt_messages(model_id, messages) + .prompt_messages(model_id, messages, constraint) .await .map_err(|e| OpenAIError::internal(e.to_string()))?; @@ -133,6 +181,32 @@ async fn chat_oneshot( let _ = bill_compute(&email, 1.0, "ai_prompt", Some("v1/chat/completions")); } + let tool_calls = if tools_active { + to_openai_tool_calls(tool_grammar::extract_tool_calls(&result.text)) + } else { + Vec::new() + }; + + let (message, finish_reason) = if tool_calls.is_empty() { + ( + ChatResponseMessage { + role: "assistant", + content: Some(result.text), + tool_calls: None, + }, + "stop", + ) + } else { + ( + ChatResponseMessage { + role: "assistant", + content: None, + tool_calls: Some(tool_calls), + }, + "tool_calls", + ) + }; + let body = ChatCompletionResponse { id: format!("chatcmpl-{}", Uuid::new_v4()), object: "chat.completion", @@ -140,11 +214,8 @@ async fn chat_oneshot( model: requested_model, choices: vec![ChatChoice { index: 0, - message: ChatResponseMessage { - role: "assistant", - content: result.text, - }, - finish_reason: "stop", + message, + finish_reason, }], usage: Usage { prompt_tokens: result.prompt_tokens as u64, @@ -160,6 +231,8 @@ async fn chat_stream( requested_model: String, model_id: String, messages: Vec<(String, String)>, + constraint: Option>, + tools_active: bool, ) -> Result { if let Some(email) = user_email(&auth) { check_compute_credits(&email) @@ -170,7 +243,7 @@ async fn chat_stream( .await .map_err(|e| OpenAIError::internal(e.to_string()))?; let (token_rx, done_rx) = service - .prompt_messages_stream(model_id, messages) + .prompt_messages_stream(model_id, messages, constraint) .await .map_err(|e| OpenAIError::internal(e.to_string()))?; @@ -197,6 +270,7 @@ async fn chat_stream( delta: ChatChunkDelta { role: Some("assistant"), content: None, + tool_calls: None, }, finish_reason: None, }], @@ -212,46 +286,111 @@ async fn chat_stream( async move { let mut token_rx = token_rx; - while let Some(token) = token_rx.recv().await { - let chunk = ChatCompletionChunk { - id: id.clone(), - object: "chat.completion.chunk", - created, - model: stream_model.clone(), - choices: vec![ChatChunkChoice { - index: 0, - delta: ChatChunkDelta { - role: None, - content: Some(token), - }, - finish_reason: None, - }], - }; - if event_tx - .send(Ok( - Event::default().data(serde_json::to_string(&chunk).unwrap()) - )) - .is_err() - { - return; + if tools_active { + // Tool-enabled requests buffer the full generation, then emit + // tool-call framing at the end. Constrained decoding streams + // the on-grammar `` text token-by-token, but we + // can only translate it into OpenAI `tool_calls[]` deltas once + // whole; likewise auto-mode output must be inspected as a + // whole to avoid leaking raw `` tags as content. + // This trades token streaming for correct tool structuring. + let mut accumulated = String::new(); + while let Some(token) = token_rx.recv().await { + accumulated.push_str(&token); } - } - // Final event with finish_reason. - let final_chunk = ChatCompletionChunk { - id: id.clone(), - object: "chat.completion.chunk", - created, - model: stream_model.clone(), - choices: vec![ChatChunkChoice { - index: 0, - delta: ChatChunkDelta::default(), - finish_reason: Some("stop"), - }], - }; - let _ = event_tx.send(Ok( - Event::default().data(serde_json::to_string(&final_chunk).unwrap()) - )); + let tool_calls = + to_openai_tool_calls(tool_grammar::extract_tool_calls(&accumulated)); + + if !tool_calls.is_empty() { + for (i, call) in tool_calls.into_iter().enumerate() { + let chunk = ChatCompletionChunk { + id: id.clone(), + object: "chat.completion.chunk", + created, + model: stream_model.clone(), + choices: vec![ChatChunkChoice { + index: 0, + delta: ChatChunkDelta { + role: None, + content: None, + tool_calls: Some(vec![ToolCallDelta { + index: i as u32, + id: Some(call.id), + kind: Some("function"), + function: Some(FunctionCallDelta { + name: Some(call.function.name), + arguments: Some(call.function.arguments), + }), + }]), + }, + finish_reason: None, + }], + }; + if event_tx + .send(Ok( + Event::default().data(serde_json::to_string(&chunk).unwrap()) + )) + .is_err() + { + return; + } + } + emit_final(&event_tx, &id, &stream_model, created, "tool_calls"); + } else { + // Model answered normally — emit the buffered text as one + // content delta, then a normal stop. + if !accumulated.is_empty() { + let chunk = ChatCompletionChunk { + id: id.clone(), + object: "chat.completion.chunk", + created, + model: stream_model.clone(), + choices: vec![ChatChunkChoice { + index: 0, + delta: ChatChunkDelta { + role: None, + content: Some(accumulated), + tool_calls: None, + }, + finish_reason: None, + }], + }; + let _ = event_tx.send(Ok( + Event::default().data(serde_json::to_string(&chunk).unwrap()) + )); + } + emit_final(&event_tx, &id, &stream_model, created, "stop"); + } + } else { + // Unchanged per-token streaming. + while let Some(token) = token_rx.recv().await { + let chunk = ChatCompletionChunk { + id: id.clone(), + object: "chat.completion.chunk", + created, + model: stream_model.clone(), + choices: vec![ChatChunkChoice { + index: 0, + delta: ChatChunkDelta { + role: None, + content: Some(token), + tool_calls: None, + }, + finish_reason: None, + }], + }; + if event_tx + .send(Ok( + Event::default().data(serde_json::to_string(&chunk).unwrap()) + )) + .is_err() + { + return; + } + } + emit_final(&event_tx, &id, &stream_model, created, "stop"); + } // Billing — best-effort, charged once per completed stream. // Per-token billing requires tokenizer-exact counts which the @@ -276,6 +415,86 @@ async fn chat_stream( Ok(Sse::new(stream).into_response()) } +/// Emit the final SSE chunk carrying an empty delta and `finish_reason`. +fn emit_final( + event_tx: &tokio::sync::mpsc::UnboundedSender>, + id: &str, + model: &str, + created: i64, + finish_reason: &'static str, +) { + let final_chunk = ChatCompletionChunk { + id: id.to_string(), + object: "chat.completion.chunk", + created, + model: model.to_string(), + choices: vec![ChatChunkChoice { + index: 0, + delta: ChatChunkDelta::default(), + finish_reason: Some(finish_reason), + }], + }; + let _ = event_tx.send(Ok( + Event::default().data(serde_json::to_string(&final_chunk).unwrap()) + )); +} + +/// Convert recovered tool calls into OpenAI response `tool_calls[]`, +/// minting a stable `call_…` id for each. +fn to_openai_tool_calls(extracted: Vec) -> Vec { + extracted + .into_iter() + .map(|c| ToolCall { + id: format!("call_{}", Uuid::new_v4()), + kind: "function".to_string(), + function: FunctionCall { + name: c.name, + arguments: c.arguments, + }, + }) + .collect() +} + +/// Flatten one inbound message to a `(role, text)` pair. Assistant tool +/// calls and `role:"tool"` results are rendered into text using the Qwen +/// `` / `` convention, since the local chat +/// template has no tool role. +fn flatten_message(m: &ChatMessage) -> (String, String) { + let base_text = m + .content + .as_ref() + .map(ChatMessageContent::flatten_to_text) + .unwrap_or_default(); + + match m.role { + // Tool result → a `` block in a user turn. + Role::Tool => ( + "user".to_string(), + format!("\n{}\n", base_text), + ), + // Assistant turn that called tools → re-render the calls so the model + // sees its own prior invocations. + Role::Assistant if m.tool_calls.as_ref().map_or(false, |c| !c.is_empty()) => { + let mut text = base_text; + if let Some(calls) = &m.tool_calls { + for call in calls { + if !text.is_empty() { + text.push('\n'); + } + let args = call.function.arguments.trim(); + let args = if args.is_empty() { "{}" } else { args }; + text.push_str(&format!( + "\n{{\"name\": \"{}\", \"arguments\": {}}}\n", + call.function.name, args + )); + } + } + ("assistant".to_string(), text) + } + _ => (role_to_str(&m.role).to_string(), base_text), + } +} + fn role_to_str(role: &Role) -> &'static str { match role { Role::System => "system", diff --git a/rust-executor/src/api/openai_compat/mod.rs b/rust-executor/src/api/openai_compat/mod.rs index 64b3c1973..754fad1c0 100644 --- a/rust-executor/src/api/openai_compat/mod.rs +++ b/rust-executor/src/api/openai_compat/mod.rs @@ -30,6 +30,7 @@ pub mod model_selector; pub mod models; pub mod realtime; pub mod router; +pub mod tool_grammar; pub mod tts_passthrough; pub mod types; diff --git a/rust-executor/src/api/openai_compat/tool_grammar.rs b/rust-executor/src/api/openai_compat/tool_grammar.rs new file mode 100644 index 000000000..6d4eafff1 --- /dev/null +++ b/rust-executor/src/api/openai_compat/tool_grammar.rs @@ -0,0 +1,594 @@ +//! Runtime compiler from OpenAI tool definitions (JSON-Schema `parameters`) +//! to a kalosm [`ArcParser`] that *constrains local-model decoding* to a +//! well-formed tool call. +//! +//! The target convention is Hermes / Qwen2.5-Instruct style: +//! +//! ```text +//! +//! {"name": "", "arguments": } +//! +//! ``` +//! +//! Parallel calls repeat the block separated by newlines. +//! +//! ## What is and isn't constrained +//! +//! * `tool_choice: "required"` / a named function → we return a parser that +//! forces the model to emit one (or, for `required` + `parallel`, several) +//! syntactically valid tool call(s). This is the guarantee: the emitted +//! `arguments` are always JSON that matches the declared schema shape. +//! * `tool_choice: "auto"` / `"none"` → we return `None`. The model +//! generates freely and any `` blocks it chooses to emit are +//! recovered afterwards with [`extract_tool_calls`]. Constraining +//! arbitrary prose is not expressible with the available parser +//! primitives (there is no "any free text" parser — `StringParser` only +//! matches a quoted JSON string), and forcing prose through one would +//! corrupt normal answers. +//! +//! ## Grammar construction +//! +//! Every schema node compiles to a uniform `ArcParser<()>` (output type +//! erased to `()` via `map_output`, then boxed). Uniformity is what lets +//! the compiler recurse over arbitrary nested schemas and fold N +//! alternatives with `.or(..)` without the combinator tuple types +//! exploding. We only care that the *text* is on-grammar; the structured +//! value is recovered separately by `serde_json` in [`extract_tool_calls`]. +//! +//! Only the JSON-Schema keywords that shape the value are read +//! (`type`, `properties`, `items`, `enum`); everything else (`minLength`, +//! `description`, `additionalProperties`, …) is ignored, which is the +//! "strip unsupported keywords" behaviour by construction. + +use std::borrow::Cow; + +use kalosm::language::{ + ArcParser, FloatParser, IntegerParser, LiteralParser, ParserExt, SeparatedParser, StringParser, +}; +use serde_json::Value; + +use super::types::ToolDef; + +/// Upper bound on a constrained JSON string value (characters). +const MAX_STRING_LEN: usize = 8192; +/// Upper bound on items in a constrained JSON array. +const MAX_ARRAY_ITEMS: usize = 64; +/// Upper bound on parallel tool calls in one turn. +const MAX_PARALLEL_CALLS: usize = 8; + +// --------------------------------------------------------------------------- +// tool_choice +// --------------------------------------------------------------------------- + +/// Resolved form of the request's `tool_choice` field. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ToolChoice { + /// Model decides whether to call a tool (OpenAI default when tools are + /// present). Not hard-constrained. + Auto, + /// Tools must not be called. + None, + /// The model must call at least one of the supplied tools. + Required, + /// The model must call exactly this function. + Named(String), +} + +/// Resolve the raw `tool_choice` JSON into a [`ToolChoice`]. Absent choice +/// defaults to `Auto` when tools are present, else `None`. Unknown values +/// degrade to `Auto` rather than rejecting the request. +pub fn parse_tool_choice(value: &Option, has_tools: bool) -> ToolChoice { + match value { + None => { + if has_tools { + ToolChoice::Auto + } else { + ToolChoice::None + } + } + Some(Value::String(s)) => match s.as_str() { + "none" => ToolChoice::None, + "required" => ToolChoice::Required, + // "auto" and anything unrecognised + _ => ToolChoice::Auto, + }, + Some(Value::Object(obj)) => obj + .get("function") + .and_then(|f| f.get("name")) + .and_then(Value::as_str) + .map(|name| ToolChoice::Named(name.to_string())) + .unwrap_or(ToolChoice::Auto), + Some(_) => ToolChoice::Auto, + } +} + +// --------------------------------------------------------------------------- +// System-prompt rendering (Hermes / Qwen convention) +// --------------------------------------------------------------------------- + +/// Render the `…` system-prompt block that tells the model +/// which functions it may call and in what format. Injected as a system +/// message; the local chat template renders it verbatim. +pub fn render_tools_system_prompt(tools: &[ToolDef]) -> String { + let mut out = String::from( + "# Tools\n\nYou may call one or more functions to assist with the user query.\n\n\ + You are provided with function signatures within XML tags:\n\n", + ); + for tool in tools { + if let Ok(json) = serde_json::to_string(tool) { + out.push_str(&json); + out.push('\n'); + } + } + out.push_str( + "\n\nFor each function call, return a json object with function name and \ + arguments within XML tags:\n\n\ + {\"name\": , \"arguments\": }\n", + ); + out +} + +// --------------------------------------------------------------------------- +// Grammar compiler +// --------------------------------------------------------------------------- + +/// Build a decoding constraint for the given tools + choice, or `None` when +/// the mode should generate freely (`auto` / `none`). +pub fn build_tool_call_parser( + tools: &[ToolDef], + choice: &ToolChoice, + parallel: bool, +) -> Option> { + let selected: Vec<&ToolDef> = match choice { + ToolChoice::Required => tools.iter().collect(), + ToolChoice::Named(name) => tools.iter().filter(|t| &t.function.name == name).collect(), + ToolChoice::Auto | ToolChoice::None => return None, + }; + if selected.is_empty() { + return None; + } + + let one = or_all( + selected + .iter() + .map(|tool| single_tool_call_parser(tool)) + .collect(), + ); + + let parser = if parallel && matches!(choice, ToolChoice::Required) { + // one ("\n" one){0,} — 1..=MAX blocks separated by newlines. + SeparatedParser::new(one, LiteralParser::new("\n"), 1..=MAX_PARALLEL_CALLS) + .map_output(|_| ()) + .boxed() + } else { + one + }; + + Some(parser) +} + +/// `\n{"name": "", "arguments": }\n` +fn single_tool_call_parser(tool: &ToolDef) -> ArcParser<()> { + let empty = Value::Object(serde_json::Map::new()); + let params = tool.function.parameters.as_ref().unwrap_or(&empty); + let args = object_parser(params); + + let prefix = lit(format!( + "\n{{\"name\": {}, \"arguments\": ", + json_string_literal(&tool.function.name) + )); + let suffix = lit("}\n"); + + seq2(seq2(prefix, args), suffix) +} + +/// Compile a JSON-Schema object node into an object parser. Emits every +/// declared property, in the schema's property order, joined by `", "`. +/// Objects with no `properties` compile to the empty object `{}`. +fn object_parser(schema: &Value) -> ArcParser<()> { + match schema.get("properties").and_then(Value::as_object) { + Some(props) if !props.is_empty() => { + let mut parts = props.iter().map(|(key, prop)| { + seq2( + lit(format!("{}: ", json_string_literal(key))), + value_parser(prop), + ) + }); + let mut body = parts.next().expect("properties non-empty"); + for part in parts { + body = seq2(seq2(body, lit(", ")), part); + } + seq2(seq2(lit("{"), body), lit("}")) + } + _ => lit("{}"), + } +} + +/// Compile a JSON-Schema value node into a value parser. +fn value_parser(schema: &Value) -> ArcParser<()> { + // Enumerations may omit `type`; a value must be one of the literals. + if let Some(values) = schema.get("enum").and_then(Value::as_array) { + let literals: Vec> = values + .iter() + .map(|v| lit(serde_json::to_string(v).unwrap_or_else(|_| "null".to_string()))) + .collect(); + if !literals.is_empty() { + return or_all(literals); + } + } + + // `type` may be a string or an array of strings (union); take the first. + let ty = match schema.get("type") { + Some(Value::String(s)) => Some(s.as_str()), + Some(Value::Array(a)) => a.iter().find_map(Value::as_str), + _ => None, + }; + + match ty { + Some("string") => string_value_parser(), + Some("integer") => IntegerParser::new(i128::MIN..=i128::MAX) + .map_output(|_| ()) + .boxed(), + Some("number") => FloatParser::new(f64::MIN..=f64::MAX).map_output(|_| ()).boxed(), + Some("boolean") => bool_parser(), + Some("null") => lit("null"), + Some("array") => array_parser(schema), + Some("object") => object_parser(schema), + // Unknown / missing type → permissive scalar. + _ => any_value_parser(), + } +} + +/// `[]` OR `[` item (", " item)* `]` +fn array_parser(schema: &Value) -> ArcParser<()> { + let item = match schema.get("items") { + Some(items) => value_parser(items), + None => any_value_parser(), + }; + let non_empty = seq2( + seq2( + lit("["), + SeparatedParser::new(item, LiteralParser::new(", "), 1..=MAX_ARRAY_ITEMS) + .map_output(|_| ()) + .boxed(), + ), + lit("]"), + ); + lit("[]").or(non_empty).boxed() +} + +/// A permissive scalar (string | number | boolean | null) for schema nodes +/// with no usable `type`. +fn any_value_parser() -> ArcParser<()> { + or_all(vec![ + string_value_parser(), + FloatParser::new(f64::MIN..=f64::MAX).map_output(|_| ()).boxed(), + bool_parser(), + lit("null"), + ]) +} + +fn string_value_parser() -> ArcParser<()> { + StringParser::new(0..=MAX_STRING_LEN) + .map_output(|_| ()) + .boxed() +} + +fn bool_parser() -> ArcParser<()> { + lit("true").or(lit("false")).boxed() +} + +// -- uniform ArcParser<()> combinator helpers ------------------------------ + +/// A literal, boxed to the uniform `ArcParser<()>` type. +fn lit(text: impl Into>) -> ArcParser<()> { + LiteralParser::new(text).boxed() +} + +/// `a` then `b`, output erased to `()`. +fn seq2(a: ArcParser<()>, b: ArcParser<()>) -> ArcParser<()> { + a.then(b).map_output(|_| ()).boxed() +} + +/// Fold a non-empty list of alternatives into a single `either` parser. +/// (An empty list degrades to a parser that matches the empty string, but +/// callers never pass one.) +fn or_all(parsers: Vec>) -> ArcParser<()> { + let mut iter = parsers.into_iter(); + let mut acc = match iter.next() { + Some(first) => first, + None => return lit(""), + }; + for parser in iter { + acc = acc.or(parser).boxed(); + } + acc +} + +/// JSON-encode `s` as a quoted string literal (including the surrounding +/// quotes and any necessary escaping). +fn json_string_literal(s: &str) -> String { + serde_json::to_string(s).unwrap_or_else(|_| format!("\"{}\"", s.replace('"', "\\\""))) +} + +// --------------------------------------------------------------------------- +// Extraction (model text → tool calls) +// --------------------------------------------------------------------------- + +/// A tool call recovered from model output. `arguments` is a JSON string +/// (object re-serialised), matching the OpenAI `function.arguments` shape. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ExtractedToolCall { + pub name: String, + pub arguments: String, +} + +/// Pull every `…` block out of `text` and parse the +/// inner JSON. Works for both constrained output (the text *is* the +/// block(s)) and auto-mode output (blocks embedded in prose). Falls back +/// to treating the whole trimmed text as a single bare JSON call. +pub fn extract_tool_calls(text: &str) -> Vec { + const OPEN: &str = ""; + const CLOSE: &str = ""; + + let mut calls = Vec::new(); + let mut rest = text; + while let Some(start) = rest.find(OPEN) { + let after = &rest[start + OPEN.len()..]; + let (block, next) = match after.find(CLOSE) { + Some(end) => (&after[..end], &after[end + CLOSE.len()..]), + None => (after, ""), + }; + if let Some(call) = parse_tool_call_json(block.trim()) { + calls.push(call); + } + rest = next; + } + + if calls.is_empty() { + if let Some(call) = parse_tool_call_json(text.trim()) { + calls.push(call); + } + } + calls +} + +fn parse_tool_call_json(candidate: &str) -> Option { + let value: Value = serde_json::from_str(candidate).ok()?; + let name = value.get("name")?.as_str()?.to_string(); + let arguments = match value.get("arguments") { + Some(Value::String(s)) => s.clone(), + Some(other) => other.to_string(), + None => "{}".to_string(), + }; + Some(ExtractedToolCall { name, arguments }) +} + +// --------------------------------------------------------------------------- +// Tests (no model required — pure grammar / parsing) +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + use crate::api::openai_compat::types::FunctionDef; + use kalosm::language::{CreateParserState, ParseStatus, Parser}; + use serde_json::json; + + /// Feed the whole input to the parser and report whether it reaches a + /// `Finished` state (i.e. the string is accepted by the grammar). + fn accepts(parser: &ArcParser<()>, input: &str) -> bool { + let state = parser.create_parser_state(); + matches!( + parser.parse(&state, input.as_bytes()), + Ok(ParseStatus::Finished { .. }) + ) + } + + fn tool(name: &str, parameters: Value) -> ToolDef { + ToolDef { + kind: "function".to_string(), + function: FunctionDef { + name: name.to_string(), + description: None, + parameters: Some(parameters), + }, + } + } + + fn obj(schema: Value) -> ArcParser<()> { + object_parser(&schema) + } + + #[test] + fn string_property_accepts_valid_rejects_wrong_type() { + let p = obj(json!({ + "type": "object", + "properties": { "location": { "type": "string" } } + })); + assert!(accepts(&p, r#"{"location": "NYC"}"#)); + // number where a string is required + assert!(!accepts(&p, r#"{"location": 5}"#)); + // undeclared extra property + assert!(!accepts(&p, r#"{"location": "NYC", "x": "y"}"#)); + } + + #[test] + fn integer_property() { + let p = obj(json!({ + "type": "object", + "properties": { "n": { "type": "integer" } } + })); + assert!(accepts(&p, r#"{"n": 42}"#)); + assert!(accepts(&p, r#"{"n": -7}"#)); + assert!(!accepts(&p, r#"{"n": "x"}"#)); + } + + #[test] + fn enum_property() { + let p = obj(json!({ + "type": "object", + "properties": { "unit": { "type": "string", "enum": ["celsius", "fahrenheit"] } } + })); + assert!(accepts(&p, r#"{"unit": "celsius"}"#)); + assert!(accepts(&p, r#"{"unit": "fahrenheit"}"#)); + assert!(!accepts(&p, r#"{"unit": "kelvin"}"#)); + } + + #[test] + fn boolean_property() { + let p = obj(json!({ + "type": "object", + "properties": { "b": { "type": "boolean" } } + })); + assert!(accepts(&p, r#"{"b": true}"#)); + assert!(accepts(&p, r#"{"b": false}"#)); + assert!(!accepts(&p, r#"{"b": 1}"#)); + } + + #[test] + fn array_property_including_empty() { + let p = obj(json!({ + "type": "object", + "properties": { "xs": { "type": "array", "items": { "type": "integer" } } } + })); + assert!(accepts(&p, r#"{"xs": [1, 2, 3]}"#)); + assert!(accepts(&p, r#"{"xs": []}"#)); + assert!(accepts(&p, r#"{"xs": [7]}"#)); + assert!(!accepts(&p, r#"{"xs": [1, "a"]}"#)); + } + + #[test] + fn nested_object_property() { + let p = obj(json!({ + "type": "object", + "properties": { + "loc": { + "type": "object", + "properties": { "city": { "type": "string" } } + } + } + })); + assert!(accepts(&p, r#"{"loc": {"city": "NYC"}}"#)); + assert!(!accepts(&p, r#"{"loc": {"city": 5}}"#)); + } + + #[test] + fn empty_parameters_compiles_to_empty_object() { + let p = obj(json!({ "type": "object" })); + assert!(accepts(&p, "{}")); + assert!(!accepts(&p, r#"{"x": 1}"#)); + } + + #[test] + fn required_tool_call_parser_matches_hermes_block() { + let weather = tool( + "get_weather", + json!({ + "type": "object", + "properties": { "location": { "type": "string" } } + }), + ); + let parser = build_tool_call_parser(&[weather], &ToolChoice::Required, false) + .expect("required ⇒ Some(parser)"); + + assert!(accepts( + &parser, + "\n{\"name\": \"get_weather\", \"arguments\": {\"location\": \"NYC\"}}\n" + )); + // wrong function name + assert!(!accepts( + &parser, + "\n{\"name\": \"other\", \"arguments\": {\"location\": \"NYC\"}}\n" + )); + // malformed arguments (number for a string field) + assert!(!accepts( + &parser, + "\n{\"name\": \"get_weather\", \"arguments\": {\"location\": 5}}\n" + )); + } + + #[test] + fn named_choice_restricts_to_one_tool() { + let a = tool("alpha", json!({ "type": "object", "properties": {} })); + let b = tool("beta", json!({ "type": "object", "properties": {} })); + let parser = build_tool_call_parser(&[a, b], &ToolChoice::Named("alpha".to_string()), true) + .expect("named ⇒ Some(parser)"); + assert!(accepts( + &parser, + "\n{\"name\": \"alpha\", \"arguments\": {}}\n" + )); + assert!(!accepts( + &parser, + "\n{\"name\": \"beta\", \"arguments\": {}}\n" + )); + } + + #[test] + fn auto_and_none_are_unconstrained() { + let weather = tool("get_weather", json!({ "type": "object", "properties": {} })); + assert!(build_tool_call_parser(&[weather.clone()], &ToolChoice::Auto, true).is_none()); + assert!(build_tool_call_parser(&[weather], &ToolChoice::None, true).is_none()); + } + + #[test] + fn parse_tool_choice_variants() { + assert_eq!(parse_tool_choice(&None, true), ToolChoice::Auto); + assert_eq!(parse_tool_choice(&None, false), ToolChoice::None); + assert_eq!(parse_tool_choice(&Some(json!("none")), true), ToolChoice::None); + assert_eq!( + parse_tool_choice(&Some(json!("required")), true), + ToolChoice::Required + ); + assert_eq!(parse_tool_choice(&Some(json!("auto")), true), ToolChoice::Auto); + assert_eq!( + parse_tool_choice( + &Some(json!({ "type": "function", "function": { "name": "foo" } })), + true + ), + ToolChoice::Named("foo".to_string()) + ); + } + + #[test] + fn extract_tool_calls_tagged_and_bare() { + // single tagged block + let one = extract_tool_calls( + "\n{\"name\": \"f\", \"arguments\": {\"a\": 1}}\n", + ); + assert_eq!(one.len(), 1); + assert_eq!(one[0].name, "f"); + assert_eq!(one[0].arguments, r#"{"a":1}"#); + + // two blocks (parallel) + let two = extract_tool_calls( + "\n{\"name\": \"a\", \"arguments\": {}}\n\n\ + \n{\"name\": \"b\", \"arguments\": {}}\n", + ); + assert_eq!(two.len(), 2); + assert_eq!(two[0].name, "a"); + assert_eq!(two[1].name, "b"); + + // bare JSON (no tags) + let bare = extract_tool_calls(r#"{"name": "g", "arguments": {"x": true}}"#); + assert_eq!(bare.len(), 1); + assert_eq!(bare[0].name, "g"); + assert_eq!(bare[0].arguments, r#"{"x":true}"#); + + // plain prose ⇒ nothing + assert!(extract_tool_calls("I cannot help with that.").is_empty()); + } + + #[test] + fn render_tools_prompt_contains_signatures() { + let weather = tool( + "get_weather", + json!({ "type": "object", "properties": { "location": { "type": "string" } } }), + ); + let prompt = render_tools_system_prompt(&[weather]); + assert!(prompt.contains("")); + assert!(prompt.contains("")); + assert!(prompt.contains("get_weather")); + assert!(prompt.contains("")); + } +} diff --git a/rust-executor/src/api/openai_compat/types.rs b/rust-executor/src/api/openai_compat/types.rs index 4878ac3fb..182cf3f5e 100644 --- a/rust-executor/src/api/openai_compat/types.rs +++ b/rust-executor/src/api/openai_compat/types.rs @@ -61,8 +61,8 @@ pub enum Role { System, User, Assistant, - /// We don't process tool/function messages today — accepted but - /// ignored in the prompt assembly so the request doesn't reject. + /// Tool-result / legacy-function messages. Folded into prompt text by + /// the chat handler (the local chat template has no tool role). Tool, Function, Developer, @@ -75,6 +75,16 @@ pub struct ChatMessage { pub content: Option, #[serde(default)] pub name: Option, + /// Present on `role:"assistant"` messages that called tools in a prior + /// turn. We render these back into the prompt text (the local chat + /// template has no tool role) so multi-turn tool conversations carry + /// the assistant's own calls. + #[serde(default)] + pub tool_calls: Option>, + /// Present on `role:"tool"` messages, linking a tool result to the + /// `id` of the assistant tool call it answers. + #[serde(default)] + pub tool_call_id: Option, } /// OpenAI accepts either a string or an array of content parts. For now we @@ -120,6 +130,55 @@ impl ChatMessageContent { } } +// --------------------------------------------------------------------------- +// Tools / function calling +// --------------------------------------------------------------------------- + +/// A tool definition supplied in the request `tools[]` array. Only +/// `type: "function"` is defined by the OpenAI spec today; we keep `kind` +/// permissive so unknown tool types don't reject the request. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ToolDef { + #[serde(rename = "type")] + pub kind: String, + pub function: FunctionDef, +} + +/// The function schema inside a [`ToolDef`]. `parameters` is a JSON-Schema +/// object describing the arguments. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct FunctionDef { + pub name: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub description: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub parameters: Option, +} + +/// A tool call — emitted by the assistant (response side) and echoed back +/// by the caller in a follow-up assistant message (request side). The +/// same struct serves both directions, so `kind` is an owned `String` +/// (defaulting to `"function"`) to stay `Deserialize`-able; serialized it +/// still reads `"type":"function"`. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ToolCall { + pub id: String, + #[serde(rename = "type", default = "default_tool_type")] + pub kind: String, + pub function: FunctionCall, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct FunctionCall { + pub name: String, + /// Arguments as a JSON *string* (per the OpenAI spec), not an object. + pub arguments: String, +} + +fn default_tool_type() -> String { + "function".to_string() +} + #[derive(Debug, Deserialize)] pub struct ChatCompletionRequest { pub model: String, @@ -138,6 +197,16 @@ pub struct ChatCompletionRequest { pub seed: Option, #[serde(default)] pub response_format: Option, + /// Tool definitions the model may call. Absent/empty ⇒ no tool calling. + #[serde(default)] + pub tools: Option>, + /// `"auto"` | `"none"` | `"required"` | `{"type":"function","function":{"name":…}}`. + #[serde(default)] + pub tool_choice: Option, + /// Whether the model may emit more than one tool call in a turn. + /// Defaults to `true` (OpenAI's default) when omitted. + #[serde(default)] + pub parallel_tool_calls: Option, #[serde(default)] pub user: Option, } @@ -162,7 +231,11 @@ pub struct ChatChoice { #[derive(Debug, Serialize)] pub struct ChatResponseMessage { pub role: &'static str, // "assistant" - pub content: String, + /// `None` (and omitted) when the turn is a tool call. + #[serde(skip_serializing_if = "Option::is_none")] + pub content: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub tool_calls: Option>, } // --------------------------------------------------------------------------- @@ -192,6 +265,31 @@ pub struct ChatChunkDelta { pub role: Option<&'static str>, #[serde(skip_serializing_if = "Option::is_none")] pub content: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub tool_calls: Option>, +} + +/// One entry in a streaming `delta.tool_calls[]`. `index` is the stable +/// key clients accumulate fragments by; `id`/`kind`/`function.name` arrive +/// on the first fragment for a call, `function.arguments` may stream in +/// pieces across chunks. +#[derive(Debug, Serialize)] +pub struct ToolCallDelta { + pub index: u32, + #[serde(skip_serializing_if = "Option::is_none")] + pub id: Option, + #[serde(rename = "type", skip_serializing_if = "Option::is_none")] + pub kind: Option<&'static str>, + #[serde(skip_serializing_if = "Option::is_none")] + pub function: Option, +} + +#[derive(Debug, Serialize)] +pub struct FunctionCallDelta { + #[serde(skip_serializing_if = "Option::is_none")] + pub name: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub arguments: Option, } // --------------------------------------------------------------------------- From d2a8179a4d441f98212f33880b64cc2ce049e6b8 Mon Sep 17 00:00:00 2001 From: Josh Field <10372036+HexaField@users.noreply.github.com> Date: Tue, 28 Jul 2026 19:46:36 +1000 Subject: [PATCH 2/6] feat(assistant): server-side AI-assistant run subsystem MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A Rust subsystem in rust-executor that runs AI-assistant turns server-side — persistent, concurrent, surviving client disconnect. Watches perspectives for new WE Message turns, runs the context->model->tools->persist loop reusing the /v1 tool-calling machinery in-process, streams by rewriting Message.content (auto-publishes to WE), and resumes interrupted runs from a durable RunState. - assistant_runtime/{entities,store,sdna,context,tools,run,registry,mod} - Subject classes match the WE contract (coasys/we#95): Assistant/Personality/Skill/McpServer in we-root, Thread/Message in the neighbourhood perspective, we:// predicates. - Built-in perspective/neighbourhood tools in-process; MCP client behind a trait seam (explicit error until rmcp client is wired — no silent stub). - Boots beside the MCP/REST subsystems, gated on config.enable_assistants (default on). Co-Authored-By: Claude Opus 4.8 --- .../src/assistant_runtime/context.rs | 203 ++++++++ .../src/assistant_runtime/entities.rs | 471 ++++++++++++++++++ rust-executor/src/assistant_runtime/mod.rs | 190 +++++++ .../src/assistant_runtime/registry.rs | 74 +++ rust-executor/src/assistant_runtime/run.rs | 417 ++++++++++++++++ rust-executor/src/assistant_runtime/sdna.rs | 227 +++++++++ rust-executor/src/assistant_runtime/store.rs | 241 +++++++++ rust-executor/src/assistant_runtime/tools.rs | 363 ++++++++++++++ rust-executor/src/config.rs | 4 + rust-executor/src/lib.rs | 18 + 10 files changed, 2208 insertions(+) create mode 100644 rust-executor/src/assistant_runtime/context.rs create mode 100644 rust-executor/src/assistant_runtime/entities.rs create mode 100644 rust-executor/src/assistant_runtime/mod.rs create mode 100644 rust-executor/src/assistant_runtime/registry.rs create mode 100644 rust-executor/src/assistant_runtime/run.rs create mode 100644 rust-executor/src/assistant_runtime/sdna.rs create mode 100644 rust-executor/src/assistant_runtime/store.rs create mode 100644 rust-executor/src/assistant_runtime/tools.rs diff --git a/rust-executor/src/assistant_runtime/context.rs b/rust-executor/src/assistant_runtime/context.rs new file mode 100644 index 000000000..2d375bb42 --- /dev/null +++ b/rust-executor/src/assistant_runtime/context.rs @@ -0,0 +1,203 @@ +//! Context assembly — turn an assistant's config + a thread's history into the +//! `(role, content)` message list the executor's `prompt_messages` consumes. +//! +//! The executor's `build_ephemeral_task` folds all `system` messages into the +//! system prompt, turns prior `user`/`assistant` pairs into few-shot examples, +//! and uses the **last** `user` message as the live prompt — it silently drops +//! any `tool`/`function` role. So this assembler: +//! +//! * emits the persona (assistant system prompt + personality bodies + skills) +//! and the Hermes/Qwen `` block as `system` messages, +//! * replays thread history as `user`/`assistant` turns, folding any `tool` +//! history into the preceding assistant text (nothing dropped), +//! * places the live user turn (which the run loop extends with a tool +//! transcript between iterations) as the final `user` message. +//! +//! Pure and unit-testable — no perspective or model required. + +use crate::api::openai_compat::tool_grammar::render_tools_system_prompt; +use crate::api::openai_compat::types::ToolDef; + +use super::entities::{Assistant, Message, Personality, Skill}; + +/// Assemble the message list for one model call. +/// +/// `history` is the thread's prior messages (ordered by `ts`), NOT including +/// the current user turn. `user_turn` is the live prompt text. +pub fn assemble_messages( + assistant: &Assistant, + personalities: &[Personality], + skills: &[Skill], + tools: &[ToolDef], + history: &[Message], + user_turn: &str, +) -> Vec<(String, String)> { + let mut out: Vec<(String, String)> = Vec::new(); + + if !assistant.system_prompt.trim().is_empty() { + out.push(("system".to_string(), assistant.system_prompt.clone())); + } + for p in personalities { + if !p.body.trim().is_empty() { + out.push(("system".to_string(), p.body.clone())); + } + } + if !skills.is_empty() { + out.push(("system".to_string(), render_skills(skills))); + } + if !tools.is_empty() { + out.push(("system".to_string(), render_tools_system_prompt(tools))); + } + + for m in history { + match m.role.as_str() { + "user" => out.push(("user".to_string(), m.content.clone())), + "assistant" => out.push(("assistant".to_string(), m.content.clone())), + "tool" => { + // Fold tool output into the preceding assistant turn so the + // executor doesn't drop it; otherwise surface it as user + // context. + let folded = format!("{}", m.content); + match out.last_mut() { + Some(last) if last.0 == "assistant" => { + last.1.push('\n'); + last.1.push_str(&folded); + } + _ => out.push(("user".to_string(), folded)), + } + } + _ => {} + } + } + + out.push(("user".to_string(), user_turn.to_string())); + out +} + +/// Render selected skills as a single `system` block. +fn render_skills(skills: &[Skill]) -> String { + let mut s = String::from("# Skills\n\nYou have the following skills available:\n"); + for sk in skills { + s.push_str(&format!("\n## {}\n", sk.name)); + if !sk.description.trim().is_empty() { + s.push_str(sk.description.trim()); + s.push('\n'); + } + if !sk.body.trim().is_empty() { + s.push_str(sk.body.trim()); + s.push('\n'); + } + } + s +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::api::openai_compat::types::{FunctionDef, ToolDef}; + use serde_json::json; + + fn assistant() -> Assistant { + Assistant { + id: "we://assistant/1".into(), + name: "Ada".into(), + model_id: "llama".into(), + system_prompt: "You are Ada.".into(), + ..Default::default() + } + } + + fn msg(role: &str, content: &str, ts: &str) -> Message { + Message { + id: format!("m-{ts}"), + thread_id: "t".into(), + role: role.into(), + content: content.into(), + tool_calls: String::new(), + ts: ts.into(), + status: "complete".into(), + } + } + + fn tool(name: &str) -> ToolDef { + ToolDef { + kind: "function".into(), + function: FunctionDef { + name: name.into(), + description: Some("desc".into()), + parameters: Some(json!({"type":"object","properties":{}})), + }, + } + } + + #[test] + fn persona_and_tools_lead_as_system_messages() { + let personalities = vec![Personality { + id: "p".into(), + name: "Witty".into(), + body: "Be witty.".into(), + }]; + let skills = vec![Skill { + id: "s".into(), + name: "Search".into(), + description: "find things".into(), + body: "use the index".into(), + }]; + let tools = vec![tool("get_weather")]; + + let out = assemble_messages(&assistant(), &personalities, &skills, &tools, &[], "hello"); + + // system: assistant prompt, personality body, skills block, tools block + assert_eq!(out[0], ("system".to_string(), "You are Ada.".to_string())); + assert_eq!(out[1], ("system".to_string(), "Be witty.".to_string())); + assert_eq!(out[2].0, "system"); + assert!(out[2].1.contains("# Skills")); + assert!(out[2].1.contains("Search")); + assert_eq!(out[3].0, "system"); + assert!(out[3].1.contains("")); + assert!(out[3].1.contains("get_weather")); + + // final message is the live user turn + assert_eq!(out.last().unwrap(), &("user".to_string(), "hello".to_string())); + } + + #[test] + fn history_replayed_as_user_assistant_turns() { + let history = vec![ + msg("user", "hi", "1"), + msg("assistant", "hello!", "2"), + ]; + let out = assemble_messages(&assistant(), &[], &[], &[], &history, "next question"); + // system(1) + user + assistant + final user + assert_eq!(out[0].0, "system"); + assert_eq!(out[1], ("user".to_string(), "hi".to_string())); + assert_eq!(out[2], ("assistant".to_string(), "hello!".to_string())); + assert_eq!(out[3], ("user".to_string(), "next question".to_string())); + } + + #[test] + fn tool_history_folds_into_assistant_turn() { + let history = vec![ + msg("user", "weather?", "1"), + msg("assistant", "let me check", "2"), + msg("tool", "sunny, 24C", "3"), + ]; + let out = assemble_messages(&assistant(), &[], &[], &[], &history, "thanks"); + // The tool message is folded into the assistant turn, not dropped. + let assistant_turn = out.iter().find(|(r, _)| r == "assistant").unwrap(); + assert!(assistant_turn.1.contains("let me check")); + assert!(assistant_turn.1.contains("sunny, 24C")); + // No standalone tool role leaks through. + assert!(!out.iter().any(|(r, _)| r == "tool")); + } + + #[test] + fn empty_system_prompt_is_omitted() { + let mut a = assistant(); + a.system_prompt = " ".into(); + let out = assemble_messages(&a, &[], &[], &[], &[], "hi"); + // Only the final user turn — no empty system message. + assert_eq!(out.len(), 1); + assert_eq!(out[0].0, "user"); + } +} diff --git a/rust-executor/src/assistant_runtime/entities.rs b/rust-executor/src/assistant_runtime/entities.rs new file mode 100644 index 000000000..eb0fb2440 --- /dev/null +++ b/rust-executor/src/assistant_runtime/entities.rs @@ -0,0 +1,471 @@ +//! WE entity model — subject classes materialised as `we://` links. +//! +//! Matches the WE contract (coasys/we#95): a subject instance is a base +//! expression URI carrying a flag link (`source --we://flag--> we://`) +//! plus one property link per field (`source --we://--> `). +//! HasMany relations are collection links (`source --we://--> target`). +//! +//! Personal config (`Assistant`/`Personality`/`Skill`/`McpServer`) lives in +//! the agent's we-root perspective; conversations (`Thread`/`Message`) live in +//! the neighbourhood perspective (fallback we-root). All predicates are the +//! literal `we://…` strings below so the links this subsystem reads and writes +//! are byte-for-byte the ones WE's `Ad4mModel` layer produces and consumes. +//! +//! Everything in this module is pure (no executor runtime) so it is unit +//! testable without a model or a live perspective — the perspective I/O that +//! turns these structs into links lives in [`super::store`]. + +use std::collections::HashMap; + +use ad4m_client::literal::{Literal, LiteralValue}; + +use crate::types::Link; + +// --------------------------------------------------------------------------- +// Predicates +// --------------------------------------------------------------------------- + +/// Instance flag predicate. `source --we://flag--> we://` marks a base +/// as an instance of a class. +pub const FLAG: &str = "we://flag"; + +// Class flag targets. +pub const CLASS_ASSISTANT: &str = "we://assistant"; +pub const CLASS_PERSONALITY: &str = "we://personality"; +pub const CLASS_SKILL: &str = "we://skill"; +pub const CLASS_MCP_SERVER: &str = "we://mcp_server"; +pub const CLASS_THREAD: &str = "we://thread"; +pub const CLASS_MESSAGE: &str = "we://message"; +pub const CLASS_RUN_STATE: &str = "we://run_state"; + +// Property predicates (snake_case of the field, prefixed `we://`). +pub const P_NAME: &str = "we://name"; +pub const P_MODEL_ID: &str = "we://model_id"; +pub const P_SYSTEM_PROMPT: &str = "we://system_prompt"; +pub const P_PERSONALITY_IDS: &str = "we://personality_ids"; +pub const P_SKILL_IDS: &str = "we://skill_ids"; +pub const P_MCP_SERVER_IDS: &str = "we://mcp_server_ids"; +pub const P_BODY: &str = "we://body"; +pub const P_DESCRIPTION: &str = "we://description"; +pub const P_TRANSPORT: &str = "we://transport"; +pub const P_URL: &str = "we://url"; +pub const P_COMMAND: &str = "we://command"; +pub const P_AUTH: &str = "we://auth"; +pub const P_TITLE: &str = "we://title"; +pub const P_ASSISTANT_ID: &str = "we://assistant_id"; +pub const P_CREATED_AT: &str = "we://created_at"; +pub const P_UPDATED_AT: &str = "we://updated_at"; +pub const P_THREAD_ID: &str = "we://thread_id"; +pub const P_ROLE: &str = "we://role"; +pub const P_CONTENT: &str = "we://content"; +pub const P_TOOL_CALLS: &str = "we://tool_calls"; +pub const P_TS: &str = "we://ts"; +pub const P_STATUS: &str = "we://status"; +pub const P_CURSOR: &str = "we://cursor"; +pub const P_PENDING_TOOL_CALL: &str = "we://pending_tool_call"; + +/// Thread → Message HasMany collection predicate. +pub const REL_MESSAGE: &str = "we://message"; + +// --------------------------------------------------------------------------- +// Literal encode / decode +// --------------------------------------------------------------------------- + +/// Encode a scalar string field as a `literal:string:…` URL, exactly as the +/// `Ad4mModel` property writer does. +pub fn encode_literal(value: &str) -> String { + Literal::from_string(value.to_string()) + .to_url() + .unwrap_or_else(|_| format!("literal:string:{}", value)) +} + +/// Decode a link target back to a plain string. `literal:…` URLs are decoded +/// to their inner value; any other target (a raw URI reference) is returned +/// verbatim. +pub fn decode_literal(target: &str) -> String { + match Literal::from_url(target.to_string()) { + Ok(lit) => match lit.get() { + Ok(LiteralValue::String(s)) => s, + Ok(other) => other.to_string(), + Err(_) => target.to_string(), + }, + Err(_) => target.to_string(), + } +} + +/// Parse a JSON-encoded `string[]` id array property (used for +/// `personalityIds`/`skillIds`/`mcpServerIds`). Tolerant: a missing or +/// malformed value yields an empty vector. +pub fn parse_id_array(value: Option<&String>) -> Vec { + value + .and_then(|s| serde_json::from_str::>(s).ok()) + .unwrap_or_default() +} + +// --------------------------------------------------------------------------- +// Link builders +// --------------------------------------------------------------------------- + +/// `source --we://flag--> we://` +pub fn flag_link(base: &str, class: &str) -> Link { + Link { + source: base.to_string(), + predicate: Some(FLAG.to_string()), + target: class.to_string(), + } +} + +/// `source --predicate--> literal:string:` +pub fn property_link(base: &str, predicate: &str, value: &str) -> Link { + Link { + source: base.to_string(), + predicate: Some(predicate.to_string()), + target: encode_literal(value), + } +} + +// --------------------------------------------------------------------------- +// Entities +// --------------------------------------------------------------------------- + +#[derive(Debug, Clone, Default, PartialEq)] +pub struct Assistant { + pub id: String, + pub name: String, + pub model_id: String, + pub system_prompt: String, + pub personality_ids: Vec, + pub skill_ids: Vec, + pub mcp_server_ids: Vec, +} + +impl Assistant { + pub fn from_props(id: String, p: &HashMap) -> Self { + Assistant { + id, + name: p.get(P_NAME).cloned().unwrap_or_default(), + model_id: p.get(P_MODEL_ID).cloned().unwrap_or_default(), + system_prompt: p.get(P_SYSTEM_PROMPT).cloned().unwrap_or_default(), + personality_ids: parse_id_array(p.get(P_PERSONALITY_IDS)), + skill_ids: parse_id_array(p.get(P_SKILL_IDS)), + mcp_server_ids: parse_id_array(p.get(P_MCP_SERVER_IDS)), + } + } +} + +#[derive(Debug, Clone, Default, PartialEq)] +pub struct Personality { + pub id: String, + pub name: String, + pub body: String, +} + +impl Personality { + pub fn from_props(id: String, p: &HashMap) -> Self { + Personality { + id, + name: p.get(P_NAME).cloned().unwrap_or_default(), + body: p.get(P_BODY).cloned().unwrap_or_default(), + } + } +} + +#[derive(Debug, Clone, Default, PartialEq)] +pub struct Skill { + pub id: String, + pub name: String, + pub description: String, + pub body: String, +} + +impl Skill { + pub fn from_props(id: String, p: &HashMap) -> Self { + Skill { + id, + name: p.get(P_NAME).cloned().unwrap_or_default(), + description: p.get(P_DESCRIPTION).cloned().unwrap_or_default(), + body: p.get(P_BODY).cloned().unwrap_or_default(), + } + } +} + +#[derive(Debug, Clone, Default, PartialEq)] +pub struct McpServer { + pub id: String, + pub name: String, + pub transport: String, + pub url: String, + pub command: String, + pub auth: String, +} + +impl McpServer { + pub fn from_props(id: String, p: &HashMap) -> Self { + McpServer { + id, + name: p.get(P_NAME).cloned().unwrap_or_default(), + transport: p.get(P_TRANSPORT).cloned().unwrap_or_default(), + url: p.get(P_URL).cloned().unwrap_or_default(), + command: p.get(P_COMMAND).cloned().unwrap_or_default(), + auth: p.get(P_AUTH).cloned().unwrap_or_default(), + } + } +} + +#[derive(Debug, Clone, Default, PartialEq)] +pub struct Thread { + pub id: String, + pub title: String, + pub assistant_id: String, + pub model_id: String, + pub created_at: String, + pub updated_at: String, +} + +impl Thread { + pub fn from_props(id: String, p: &HashMap) -> Self { + Thread { + id, + title: p.get(P_TITLE).cloned().unwrap_or_default(), + assistant_id: p.get(P_ASSISTANT_ID).cloned().unwrap_or_default(), + model_id: p.get(P_MODEL_ID).cloned().unwrap_or_default(), + created_at: p.get(P_CREATED_AT).cloned().unwrap_or_default(), + updated_at: p.get(P_UPDATED_AT).cloned().unwrap_or_default(), + } + } +} + +#[derive(Debug, Clone, Default, PartialEq)] +pub struct Message { + pub id: String, + pub thread_id: String, + pub role: String, + pub content: String, + /// JSON-encoded array of tool calls + results (may be empty). + pub tool_calls: String, + /// ISO-8601 timestamp. + pub ts: String, + /// `""` | `streaming` | `complete` | `error`. + pub status: String, +} + +impl Message { + pub fn from_props(id: String, p: &HashMap) -> Self { + Message { + id, + thread_id: p.get(P_THREAD_ID).cloned().unwrap_or_default(), + role: p.get(P_ROLE).cloned().unwrap_or_default(), + content: p.get(P_CONTENT).cloned().unwrap_or_default(), + tool_calls: p.get(P_TOOL_CALLS).cloned().unwrap_or_default(), + ts: p.get(P_TS).cloned().unwrap_or_default(), + status: p.get(P_STATUS).cloned().unwrap_or_default(), + } + } + + /// The links that materialise this message as a `Message` subject + /// instance (flag + one property link per set field). + pub fn to_links(&self) -> Vec { + let mut links = vec![ + flag_link(&self.id, CLASS_MESSAGE), + property_link(&self.id, P_THREAD_ID, &self.thread_id), + property_link(&self.id, P_ROLE, &self.role), + property_link(&self.id, P_CONTENT, &self.content), + property_link(&self.id, P_TS, &self.ts), + property_link(&self.id, P_STATUS, &self.status), + ]; + if !self.tool_calls.is_empty() { + links.push(property_link(&self.id, P_TOOL_CALLS, &self.tool_calls)); + } + links + } +} + +/// Per-run durability record. One instance per active run, keyed by thread. +#[derive(Debug, Clone, Default, PartialEq)] +pub struct RunState { + pub id: String, + pub thread_id: String, + /// `running` | `awaiting_tool` | `awaiting_model` | `done` | `error`. + pub status: String, + /// Progress marker (iteration index rendered as a string). + pub cursor: String, + /// JSON of an in-flight tool call awaiting execution (may be empty). + pub pending_tool_call: String, +} + +impl RunState { + pub fn from_props(id: String, p: &HashMap) -> Self { + RunState { + id, + thread_id: p.get(P_THREAD_ID).cloned().unwrap_or_default(), + status: p.get(P_STATUS).cloned().unwrap_or_default(), + cursor: p.get(P_CURSOR).cloned().unwrap_or_default(), + pending_tool_call: p.get(P_PENDING_TOOL_CALL).cloned().unwrap_or_default(), + } + } + + pub fn to_links(&self) -> Vec { + let mut links = vec![ + flag_link(&self.id, CLASS_RUN_STATE), + property_link(&self.id, P_THREAD_ID, &self.thread_id), + property_link(&self.id, P_STATUS, &self.status), + property_link(&self.id, P_CURSOR, &self.cursor), + ]; + if !self.pending_tool_call.is_empty() { + links.push(property_link(&self.id, P_PENDING_TOOL_CALL, &self.pending_tool_call)); + } + links + } + + /// A run that was interrupted mid-flight and should be re-enqueued on boot. + pub fn is_active(&self) -> bool { + matches!( + self.status.as_str(), + "running" | "awaiting_tool" | "awaiting_model" + ) + } +} + +// --------------------------------------------------------------------------- +// Trigger logic (pure — the core dedupe/resume decisions) +// --------------------------------------------------------------------------- + +/// Given a thread's messages ordered by `ts`, return the message that should +/// trigger a run: the trailing message iff it is a completed user turn. This +/// self-dedupes — once an assistant reply is appended, the trailing message is +/// no longer a completed user turn, so re-scanning is a no-op. +pub fn trigger_message(sorted: &[Message]) -> Option<&Message> { + match sorted.last() { + Some(m) if m.role == "user" && m.status == "complete" => Some(m), + _ => None, + } +} + +/// True when the trailing message is a half-written assistant reply — i.e. a +/// run died mid-stream and must be resumed. +pub fn needs_resume(sorted: &[Message]) -> bool { + matches!(sorted.last(), Some(m) if m.role == "assistant" && m.status == "streaming") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn literal_round_trip() { + for v in ["hello", "with spaces & symbols/:", "", "unicode ☃ 漢字"] { + let enc = encode_literal(v); + assert!(enc.starts_with("literal:"), "encoded = {enc}"); + assert_eq!(decode_literal(&enc), v); + } + } + + #[test] + fn decode_passes_through_raw_uris() { + assert_eq!(decode_literal("we://message/abc"), "we://message/abc"); + assert_eq!(decode_literal("did:key:xyz"), "did:key:xyz"); + } + + #[test] + fn id_array_parsing_is_tolerant() { + assert_eq!(parse_id_array(Some(&r#"["a","b"]"#.to_string())), vec!["a", "b"]); + assert!(parse_id_array(Some(&"not json".to_string())).is_empty()); + assert!(parse_id_array(None).is_empty()); + } + + #[test] + fn message_links_round_trip_through_props() { + let msg = Message { + id: "we://message/1".into(), + thread_id: "we://thread/9".into(), + role: "assistant".into(), + content: "hi there".into(), + tool_calls: r#"[{"name":"f"}]"#.into(), + ts: "2026-07-28T00:00:00Z".into(), + status: "complete".into(), + }; + // Rebuild the property map the way store::load_props would, decoding + // each link target. + let mut props = HashMap::new(); + for l in msg.to_links() { + if let Some(pred) = l.predicate { + if pred == FLAG { + continue; + } + props.insert(pred, decode_literal(&l.target)); + } + } + assert_eq!(Message::from_props(msg.id.clone(), &props), msg); + } + + #[test] + fn assistant_parses_id_arrays_and_scalars() { + let mut p = HashMap::new(); + p.insert(P_NAME.to_string(), "Ada".to_string()); + p.insert(P_MODEL_ID.to_string(), "llama".to_string()); + p.insert(P_SYSTEM_PROMPT.to_string(), "be helpful".to_string()); + p.insert(P_PERSONALITY_IDS.to_string(), r#"["p1"]"#.to_string()); + p.insert(P_SKILL_IDS.to_string(), r#"["s1","s2"]"#.to_string()); + let a = Assistant::from_props("we://assistant/1".into(), &p); + assert_eq!(a.name, "Ada"); + assert_eq!(a.model_id, "llama"); + assert_eq!(a.personality_ids, vec!["p1"]); + assert_eq!(a.skill_ids, vec!["s1", "s2"]); + assert!(a.mcp_server_ids.is_empty()); + } + + fn msg(role: &str, status: &str, ts: &str) -> Message { + Message { + id: format!("m-{ts}"), + thread_id: "t".into(), + role: role.into(), + content: String::new(), + tool_calls: String::new(), + ts: ts.into(), + status: status.into(), + } + } + + #[test] + fn trigger_fires_only_on_trailing_complete_user() { + let convo = vec![ + msg("user", "complete", "1"), + msg("assistant", "complete", "2"), + msg("user", "complete", "3"), + ]; + assert_eq!(trigger_message(&convo).map(|m| m.ts.as_str()), Some("3")); + + // Trailing assistant reply → no trigger (self-dedupe). + let answered = vec![msg("user", "complete", "1"), msg("assistant", "complete", "2")]; + assert!(trigger_message(&answered).is_none()); + + // Incomplete user turn → no trigger. + let partial = vec![msg("user", "streaming", "1")]; + assert!(trigger_message(&partial).is_none()); + + assert!(trigger_message(&[]).is_none()); + } + + #[test] + fn resume_detects_trailing_streaming_assistant() { + let interrupted = vec![msg("user", "complete", "1"), msg("assistant", "streaming", "2")]; + assert!(needs_resume(&interrupted)); + let done = vec![msg("user", "complete", "1"), msg("assistant", "complete", "2")]; + assert!(!needs_resume(&done)); + } + + #[test] + fn run_state_activity() { + let mut rs = RunState { + status: "running".into(), + ..Default::default() + }; + assert!(rs.is_active()); + rs.status = "awaiting_tool".into(); + assert!(rs.is_active()); + rs.status = "done".into(); + assert!(!rs.is_active()); + rs.status = "error".into(); + assert!(!rs.is_active()); + } +} diff --git a/rust-executor/src/assistant_runtime/mod.rs b/rust-executor/src/assistant_runtime/mod.rs new file mode 100644 index 000000000..c55542fe7 --- /dev/null +++ b/rust-executor/src/assistant_runtime/mod.rs @@ -0,0 +1,190 @@ +//! Server-side AI-assistant run subsystem. +//! +//! Runs assistant turns inside the executor — persistent, concurrent, and +//! surviving client disconnect. It watches perspectives for new completed user +//! messages (the WE `Message` subject class), runs the context → model → tools +//! → persist loop, and writes replies back into the perspective as `we://` +//! links so WE re-renders live. +//! +//! Boot: [`start`] is spawned as a subsystem thread beside the MCP + REST +//! servers in `lib.rs`, gated on `config.enable_assistants` (default on). It +//! resumes interrupted runs from durable `RunState`/message links, then watches +//! `PERSPECTIVE_LINK_ADDED_TOPIC` for new turns. +//! +//! Layout: +//! * [`entities`] — the WE subject-class model (pure link (de)serialisation). +//! * [`store`] — perspective read/write helpers. +//! * [`sdna`] — SDNA subject-class registration (bootstrap). +//! * [`context`] — context assembly for `prompt_messages` (pure). +//! * [`tools`] — built-in graph tools + the MCP-client follow-up seam. +//! * [`run`] — the per-thread loop. +//! * [`registry`] — one live run per thread. + +pub mod context; +pub mod entities; +pub mod registry; +pub mod run; +pub mod sdna; +pub mod store; +pub mod tools; + +use std::collections::{HashMap, HashSet}; +use std::sync::{Mutex, OnceLock}; + +use tokio::sync::broadcast::error::RecvError; + +use crate::agent::AgentContext; +use crate::perspectives::{all_perspectives, get_perspective}; +use crate::pubsub::{get_global_pubsub, PERSPECTIVE_LINK_ADDED_TOPIC}; +use crate::types::PerspectiveLinkWithOwner; + +use entities::{Message, CLASS_ASSISTANT, CLASS_MESSAGE, CLASS_THREAD}; +use registry::RunRegistry; +use run::RunInput; + +/// Perspectives for which SDNA subject classes have already been ensured this +/// process (avoids repeated SDNA writes). +static ENSURED: OnceLock>> = OnceLock::new(); + +fn ensured_set() -> &'static Mutex> { + ENSURED.get_or_init(|| Mutex::new(HashSet::new())) +} + +/// Entry point — never returns. Bootstraps, then watches for new turns. +pub async fn start() { + log::info!("assistant_runtime: starting"); + let registry = RunRegistry::new(); + + // Give the executor a moment to finish loading perspectives before the + // boot-time resume scan; the live watcher below catches everything after. + tokio::time::sleep(std::time::Duration::from_secs(5)).await; + + bootstrap(®istry).await; + + let pubsub = get_global_pubsub().await; + let mut rx = pubsub.subscribe(&PERSPECTIVE_LINK_ADDED_TOPIC).await; + log::info!("assistant_runtime: watching for new assistant turns"); + + loop { + match rx.recv().await { + Ok(payload) => { + if let Some(uuid) = we_perspective_of(&payload) { + let reg = registry.clone(); + tokio::spawn(async move { + scan_perspective(&uuid, ®).await; + }); + } + } + Err(RecvError::Lagged(n)) => { + log::warn!("assistant_runtime: lagged {n} link events"); + } + Err(RecvError::Closed) => { + log::warn!("assistant_runtime: link event channel closed; stopping watcher"); + break; + } + } + } +} + +/// Parse a `PERSPECTIVE_LINK_ADDED_TOPIC` payload and return the perspective +/// uuid iff the added link is a `we://` link (i.e. WE activity we should scan). +fn we_perspective_of(payload: &str) -> Option { + let evt: PerspectiveLinkWithOwner = serde_json::from_str(payload).ok()?; + let predicate = evt.link.data.predicate.as_deref().unwrap_or(""); + if predicate.starts_with("we://") { + Some(evt.perspective_uuid) + } else { + None + } +} + +/// On boot: register SDNA classes where WE is active, pick up unanswered user +/// turns, and resume interrupted runs from durable `RunState`. +async fn bootstrap(registry: &RunRegistry) { + for p in all_perspectives() { + let uuid = p.uuid.clone(); + maybe_ensure_classes(&uuid).await; + scan_perspective(&uuid, registry).await; + } + + // Explicitly resume runs whose durable RunState is still active but whose + // thread was left mid-stream (a trailing streaming assistant message, which + // the user-turn trigger alone would not re-fire). + for (perspective_uuid, rs) in store::active_run_states().await { + log::info!( + "assistant_runtime: resuming run for thread {} (status {})", + rs.thread_id, + rs.status + ); + registry + .start(RunInput { + perspective_uuid, + thread_id: rs.thread_id, + }) + .await; + } +} + +/// Scan a perspective's threads and start a run for any that has a trailing +/// completed user turn or a half-written (interrupted) assistant reply. +async fn scan_perspective(perspective_uuid: &str, registry: &RunRegistry) { + let Some(conv) = get_perspective(perspective_uuid) else { + return; + }; + + // Group every message by thread. + let mut by_thread: HashMap> = HashMap::new(); + for base in store::message_bases(&conv).await { + let msg = store::load_message(&conv, &base).await; + if msg.thread_id.is_empty() { + continue; + } + by_thread.entry(msg.thread_id.clone()).or_default().push(msg); + } + + for (thread_id, mut msgs) in by_thread { + msgs.sort_by(|a, b| a.ts.cmp(&b.ts)); + let should_run = + entities::trigger_message(&msgs).is_some() || entities::needs_resume(&msgs); + if should_run { + registry + .start(RunInput { + perspective_uuid: perspective_uuid.to_string(), + thread_id, + }) + .await; + } + } +} + +/// Register the WE SDNA subject classes in a perspective the first time we see +/// WE activity there. No-op if already ensured this process or if the classes +/// are already present (e.g. published by WE with SHACL). +async fn maybe_ensure_classes(perspective_uuid: &str) { + { + let guard = ensured_set().lock().unwrap(); + if guard.contains(perspective_uuid) { + return; + } + } + + let Some(mut p) = get_perspective(perspective_uuid) else { + return; + }; + + // Only touch perspectives that actually carry WE instances. + let has_we = !store::find_instances(&p, CLASS_ASSISTANT).await.is_empty() + || !store::find_instances(&p, CLASS_THREAD).await.is_empty() + || !store::find_instances(&p, CLASS_MESSAGE).await.is_empty(); + if !has_we { + return; + } + + let ctx = AgentContext::main_agent(); + sdna::ensure_subject_classes(&mut p, &ctx).await; + + ensured_set() + .lock() + .unwrap() + .insert(perspective_uuid.to_string()); +} diff --git a/rust-executor/src/assistant_runtime/registry.rs b/rust-executor/src/assistant_runtime/registry.rs new file mode 100644 index 000000000..ce4d5ee9a --- /dev/null +++ b/rust-executor/src/assistant_runtime/registry.rs @@ -0,0 +1,74 @@ +//! Run registry — one live run per thread, each on its own `tokio` task. +//! +//! Shaped like `ai_service`'s `llm_channel`: an `Arc>>` +//! guarding the live runs. [`RunRegistry::start`] dedupes under the lock so a +//! burst of link events (a user message is several links) can only ever spawn +//! one run per thread; a run removes itself from the map when it finishes. + +use std::collections::HashMap; +use std::sync::Arc; + +use tokio::sync::Mutex; +use tokio::task::JoinHandle; + +use super::run::{self, RunInput}; + +struct RunHandle { + handle: JoinHandle<()>, +} + +/// Cheaply-cloneable handle to the set of live runs, keyed by thread id. +#[derive(Clone)] +pub struct RunRegistry { + inner: Arc>>, +} + +impl RunRegistry { + pub fn new() -> Self { + Self { + inner: Arc::new(Mutex::new(HashMap::new())), + } + } + + /// Start a run for `input.thread_id` unless one is already live. Dedupe and + /// insertion happen under a single lock acquisition so concurrent callers + /// cannot both spawn. + pub async fn start(&self, input: RunInput) { + let thread_id = input.thread_id.clone(); + let mut guard = self.inner.lock().await; + + if let Some(existing) = guard.get(&thread_id) { + if !existing.handle.is_finished() { + log::debug!("assistant_runtime: run already active for thread {thread_id}"); + return; + } + } + + let registry = self.clone(); + let tid = thread_id.clone(); + let handle = tokio::spawn(async move { + if let Err(e) = run::run_thread(input).await { + log::error!("assistant_runtime: run_thread error for {tid}: {e}"); + } + registry.remove(&tid).await; + }); + + guard.insert(thread_id, RunHandle { handle }); + } + + async fn remove(&self, thread_id: &str) { + self.inner.lock().await.remove(thread_id); + } + + /// Number of currently-tracked runs (test/introspection helper). + #[cfg(test)] + pub async fn active_count(&self) -> usize { + self.inner.lock().await.len() + } +} + +impl Default for RunRegistry { + fn default() -> Self { + Self::new() + } +} diff --git a/rust-executor/src/assistant_runtime/run.rs b/rust-executor/src/assistant_runtime/run.rs new file mode 100644 index 000000000..06630fe20 --- /dev/null +++ b/rust-executor/src/assistant_runtime/run.rs @@ -0,0 +1,417 @@ +//! The per-thread assistant run loop. +//! +//! One `run_thread` call handles one user turn end-to-end: +//! +//! 1. resolve the thread → assistant → model → personalities / skills / MCP +//! servers (parsing the JSON id arrays and loading each record), +//! 2. create (or, on resume, reuse) the assistant reply `Message` in +//! `status: streaming`, +//! 3. loop: assemble context → `AIService::prompt_messages_stream` (streaming +//! tokens into `Message.content`) → `extract_tool_calls` → execute each tool +//! → fold the tool transcript back into the prompt → repeat until the model +//! answers without calling a tool, +//! 4. finalise `Message.status` (`complete`, or `error` on failure) and the +//! durable `RunState`. +//! +//! Streaming persists by rewriting the assistant `Message.content` link, which +//! auto-publishes on `PERSPECTIVE_LINK_ADDED_TOPIC` so WE's live query +//! re-renders — no separate token channel, matching WE. + +use anyhow::{anyhow, Result}; + +use kalosm::language::ArcParser; + +use crate::agent::AgentContext; +use crate::ai_service::AIService; +use crate::api::openai_compat::tool_grammar::{build_tool_call_parser, extract_tool_calls, ToolChoice}; + +use super::context::assemble_messages; +use super::entities::{ + Assistant, McpServer, Message, Personality, Skill, Thread, CLASS_ASSISTANT, CLASS_MCP_SERVER, + CLASS_THREAD, P_CONTENT, P_STATUS, P_TOOL_CALLS, +}; +use super::store; +use super::tools::{BuiltinTools, McpToolProvider, ToolProvider, ToolSet}; + +/// Upper bound on tool-call iterations per turn (defence against loops). +const MAX_ITERATIONS: usize = 8; +/// Rewrite the streaming `Message.content` link every N tokens. +const CONTENT_UPDATE_EVERY_TOKENS: usize = 24; + +/// What a run needs to start: the conversation perspective + the thread. +#[derive(Debug, Clone)] +pub struct RunInput { + pub perspective_uuid: String, + pub thread_id: String, +} + +/// Run one user turn for `input.thread_id`. Errors are handled internally +/// (assistant message + RunState marked `error`); the returned `Result` is for +/// the registry's logging only. +pub async fn run_thread(input: RunInput) -> Result<()> { + let RunInput { + perspective_uuid, + thread_id, + } = input; + + let mut conv = store::writable(&perspective_uuid) + .ok_or_else(|| anyhow!("Conversation perspective not found: {perspective_uuid}"))?; + + // Snapshot the thread's messages and split off the live user turn. + let msgs = store::thread_messages(&conv, &thread_id).await; + let Some(last_user_idx) = msgs.iter().rposition(|m| m.role == "user") else { + log::debug!("assistant_runtime: thread {thread_id} has no user message; skipping"); + return Ok(()); + }; + let user_turn = msgs[last_user_idx].content.clone(); + let history: Vec = msgs[..last_user_idx].to_vec(); + + // Reuse a trailing half-written assistant message (resume) or create one. + let reply_id = match msgs.get(last_user_idx + 1) { + Some(m) if m.role == "assistant" && m.status == "streaming" => m.id.clone(), + _ => { + let reply = Message { + id: format!("we://message/{}", uuid::Uuid::new_v4()), + thread_id: thread_id.clone(), + role: "assistant".to_string(), + content: String::new(), + tool_calls: String::new(), + ts: chrono::Utc::now().to_rfc3339(), + status: "streaming".to_string(), + }; + let ctx = AgentContext::main_agent(); + store::write_message(&mut conv, &reply, &ctx).await?; + reply.id + } + }; + + // Resolve the assistant configuration. On failure, mark the reply errored. + let config = match resolve_config(&conv, &thread_id).await { + Ok(c) => c, + Err(e) => { + log::error!("assistant_runtime: config resolution failed for thread {thread_id}: {e}"); + let _ = fail(&mut conv, &reply_id, &thread_id, &e.to_string()).await; + return Err(e); + } + }; + + let _ = store::upsert_run_state(&mut conv, &thread_id, "running", "0", "").await; + + match drive_loop(&mut conv, &reply_id, &thread_id, &config, &user_turn, &history).await { + Ok(()) => { + let _ = store::upsert_run_state(&mut conv, &thread_id, "done", "final", "").await; + Ok(()) + } + Err(e) => { + log::error!("assistant_runtime: run failed for thread {thread_id}: {e}"); + let _ = fail(&mut conv, &reply_id, &thread_id, &e.to_string()).await; + Err(e) + } + } +} + +/// Resolved config for a run. +struct RunConfig { + assistant: Assistant, + personalities: Vec, + skills: Vec, + mcp_servers: Vec, + model_id: String, +} + +async fn resolve_config( + conv: &crate::perspectives::perspective_instance::PerspectiveInstance, + thread_id: &str, +) -> Result { + // Thread instance: prefer the conversation perspective, else scan. + let thread_props = { + let local = store::load_props(conv, thread_id).await; + if local.contains_key(super::entities::P_ASSISTANT_ID) { + local + } else if let Some(p) = + store::find_perspective_with_instance(thread_id, CLASS_THREAD).await + { + store::load_props(&p, thread_id).await + } else { + local + } + }; + let thread = Thread::from_props(thread_id.to_string(), &thread_props); + if thread.assistant_id.is_empty() { + return Err(anyhow!("Thread {thread_id} has no assistantId")); + } + + // Assistant lives in the we-root perspective — locate it by instance. + let we_root = store::find_perspective_with_instance(&thread.assistant_id, CLASS_ASSISTANT) + .await + .ok_or_else(|| anyhow!("Assistant {} not found in any perspective", thread.assistant_id))?; + let assistant = Assistant::from_props( + thread.assistant_id.clone(), + &store::load_props(&we_root, &thread.assistant_id).await, + ); + + let mut personalities = Vec::new(); + for id in &assistant.personality_ids { + let props = store::load_props(&we_root, id).await; + if !props.is_empty() { + personalities.push(Personality::from_props(id.clone(), &props)); + } + } + let mut skills = Vec::new(); + for id in &assistant.skill_ids { + let props = store::load_props(&we_root, id).await; + if !props.is_empty() { + skills.push(Skill::from_props(id.clone(), &props)); + } + } + let mut mcp_servers = Vec::new(); + for id in &assistant.mcp_server_ids { + // MCP servers may be defined alongside the assistant or elsewhere. + let props = { + let local = store::load_props(&we_root, id).await; + if local.is_empty() { + match store::find_perspective_with_instance(id, CLASS_MCP_SERVER).await { + Some(p) => store::load_props(&p, id).await, + None => local, + } + } else { + local + } + }; + if !props.is_empty() { + mcp_servers.push(McpServer::from_props(id.clone(), &props)); + } + } + + let model_id = if !thread.model_id.is_empty() { + thread.model_id.clone() + } else if !assistant.model_id.is_empty() { + assistant.model_id.clone() + } else { + "default".to_string() + }; + + Ok(RunConfig { + assistant, + personalities, + skills, + mcp_servers, + model_id, + }) +} + +/// The context → model → tools → repeat loop for one turn. +async fn drive_loop( + conv: &mut crate::perspectives::perspective_instance::PerspectiveInstance, + reply_id: &str, + thread_id: &str, + config: &RunConfig, + user_turn: &str, + history: &[Message], +) -> Result<()> { + let ctx = AgentContext::main_agent(); + let ai = AIService::global_instance() + .await + .map_err(|e| anyhow!("AI service unavailable: {e}"))?; + + let toolset = ToolSet::new(vec![ + ToolProvider::Builtin(BuiltinTools::new(conv.uuid.clone())), + ToolProvider::Mcp(McpToolProvider::new(config.mcp_servers.clone())), + ]); + let tool_defs = toolset.tool_defs(); + + // The live prompt, extended with a tool transcript between iterations so + // the continuation stays in the executor's `final_prompt` (tool roles are + // otherwise dropped). + let mut working_prompt = user_turn.to_string(); + // Accumulated tool call + result records, mirrored onto Message.toolCalls. + let mut tool_record: Vec = Vec::new(); + + for iteration in 0..MAX_ITERATIONS { + let messages = assemble_messages( + &config.assistant, + &config.personalities, + &config.skills, + &tool_defs, + history, + &working_prompt, + ); + + // The assistant loop lets the model decide whether to call a tool + // (`ToolChoice::Auto`) — an unconstrained generation whose `` + // blocks are recovered from the text afterwards. `build_tool_call_parser` + // returns `None` for `Auto`; this is the seam where a policy could force + // `Required`/`Named` to hard-constrain decoding on a given iteration. + let constraint: Option> = + build_tool_call_parser(&tool_defs, &ToolChoice::Auto, true); + + // Stream the model output, rewriting Message.content at a cadence. + let full = + stream_completion(conv, reply_id, &ctx, &ai, &config.model_id, messages, constraint) + .await?; + + let calls = extract_tool_calls(&full); + if calls.is_empty() { + // Final answer. + let visible = strip_tool_calls(&full); + store::set_single_target(conv, reply_id, P_CONTENT, &visible, &ctx).await?; + store::set_single_target(conv, reply_id, P_STATUS, "complete", &ctx).await?; + return Ok(()); + } + + let _ = store::upsert_run_state( + conv, + thread_id, + "awaiting_tool", + &iteration.to_string(), + &full, + ) + .await; + + // Execute each requested tool, building the continuation transcript. + let mut transcript = format!("\n{}", full); + for call in &calls { + let result = match toolset.execute(&call.name, &call.arguments).await { + Ok(r) => r, + Err(e) => json_err(&e.to_string()), + }; + tool_record.push(serde_json::json!({ + "name": call.name, + "arguments": call.arguments, + "result": result, + })); + // Persist a role:'tool' message for history + WE display. + let tool_msg = Message { + id: format!("we://message/{}", uuid::Uuid::new_v4()), + thread_id: thread_id.to_string(), + role: "tool".to_string(), + content: result.clone(), + tool_calls: String::new(), + ts: chrono::Utc::now().to_rfc3339(), + status: "complete".to_string(), + }; + let _ = store::write_message(conv, &tool_msg, &ctx).await; + transcript.push_str(&format!( + "\n{}", + call.name, result + )); + } + + // Mirror the accumulated tool calls onto the assistant message. + if let Ok(json) = serde_json::to_string(&tool_record) { + let _ = store::set_single_target(conv, reply_id, P_TOOL_CALLS, &json, &ctx).await; + } + + working_prompt = format!( + "{working_prompt}{transcript}\n\nUse the tool results above to answer the user." + ); + let _ = store::upsert_run_state( + conv, + thread_id, + "running", + &iteration.to_string(), + "", + ) + .await; + } + + // Iteration budget exhausted — finalise with whatever we have. + log::warn!("assistant_runtime: thread {thread_id} hit MAX_ITERATIONS ({MAX_ITERATIONS})"); + store::set_single_target(conv, reply_id, P_STATUS, "complete", &ctx).await?; + Ok(()) +} + +/// Run one streaming completion, rewriting `Message.content` as tokens arrive, +/// and return the full generated text. +async fn stream_completion( + conv: &mut crate::perspectives::perspective_instance::PerspectiveInstance, + reply_id: &str, + ctx: &AgentContext, + ai: &AIService, + model_id: &str, + messages: Vec<(String, String)>, + constraint: Option>, +) -> Result { + let (mut token_rx, done_rx) = ai + .prompt_messages_stream(model_id.to_string(), messages, constraint) + .await + .map_err(|e| anyhow!("model stream failed: {e}"))?; + + let mut buffer = String::new(); + let mut since_update = 0usize; + while let Some(token) = token_rx.recv().await { + buffer.push_str(&token); + since_update += 1; + if since_update >= CONTENT_UPDATE_EVERY_TOKENS { + since_update = 0; + let visible = strip_tool_calls(&buffer); + let _ = store::set_single_target(conv, reply_id, P_CONTENT, &visible, ctx).await; + } + } + // Drain the completion signal (token counts unused here). + let _ = done_rx.await; + + Ok(buffer) +} + +/// Mark the assistant reply + RunState as errored. +async fn fail( + conv: &mut crate::perspectives::perspective_instance::PerspectiveInstance, + reply_id: &str, + thread_id: &str, + message: &str, +) -> Result<()> { + let ctx = AgentContext::main_agent(); + let _ = store::set_single_target(conv, reply_id, P_CONTENT, &format!("Error: {message}"), &ctx) + .await; + store::set_single_target(conv, reply_id, P_STATUS, "error", &ctx).await?; + let _ = store::upsert_run_state(conv, thread_id, "error", "error", message).await; + Ok(()) +} + +/// Remove `…` blocks for display; keep surrounding prose. +fn strip_tool_calls(text: &str) -> String { + const OPEN: &str = ""; + const CLOSE: &str = ""; + let mut out = String::new(); + let mut rest = text; + while let Some(start) = rest.find(OPEN) { + out.push_str(&rest[..start]); + let after = &rest[start + OPEN.len()..]; + match after.find(CLOSE) { + Some(end) => rest = &after[end + CLOSE.len()..], + None => { + rest = ""; + break; + } + } + } + out.push_str(rest); + out.trim().to_string() +} + +fn json_err(message: &str) -> String { + serde_json::json!({ "error": message }).to_string() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn strip_removes_tool_blocks_keeps_prose() { + let text = "Sure!\n\n{\"name\":\"f\",\"arguments\":{}}\n\nDone."; + assert_eq!(strip_tool_calls(text), "Sure!\n\nDone."); + } + + #[test] + fn strip_handles_unterminated_block() { + let text = "Answer\n{\"name\":\"f\""; + assert_eq!(strip_tool_calls(text), "Answer"); + } + + #[test] + fn strip_noop_on_plain_text() { + assert_eq!(strip_tool_calls("just prose"), "just prose"); + } +} diff --git a/rust-executor/src/assistant_runtime/sdna.rs b/rust-executor/src/assistant_runtime/sdna.rs new file mode 100644 index 000000000..9f4ca0923 --- /dev/null +++ b/rust-executor/src/assistant_runtime/sdna.rs @@ -0,0 +1,227 @@ +//! SDNA (Social DNA) subject-class definitions for the WE entity model. +//! +//! Emits Prolog SDNA facts matching the `Ad4mModel`/`buildSDNA` convention so +//! the `we://` classes are first-class subject classes for WE, the MCP +//! `query_subjects` tool, and Prolog `instance/2` queries. The runtime loop +//! itself operates at the link level (see [`super::store`]) and does not depend +//! on these being registered — this is a bootstrap so the classes exist when +//! WE has not yet published its own SHACL/SDNA. +//! +//! Registration is idempotent: classes already present (e.g. published by WE +//! with SHACL) are skipped, and the per-class `uid` is derived deterministically +//! from the class name so a re-run emits byte-identical facts. + +use std::collections::HashSet; + +use crate::agent::AgentContext; +use crate::perspectives::perspective_instance::{PerspectiveInstance, SdnaType}; + +use super::entities::{ + CLASS_ASSISTANT, CLASS_MCP_SERVER, CLASS_MESSAGE, CLASS_PERSONALITY, CLASS_RUN_STATE, + CLASS_SKILL, CLASS_THREAD, FLAG, P_ASSISTANT_ID, P_AUTH, P_BODY, P_COMMAND, P_CONTENT, + P_CREATED_AT, P_CURSOR, P_DESCRIPTION, P_MCP_SERVER_IDS, P_MODEL_ID, P_NAME, + P_PENDING_TOOL_CALL, P_PERSONALITY_IDS, P_ROLE, P_SKILL_IDS, P_STATUS, P_SYSTEM_PROMPT, + P_THREAD_ID, P_TITLE, P_TOOL_CALLS, P_TS, P_UPDATED_AT, REL_MESSAGE, +}; + +struct ClassDef { + /// Registered class name (also the flag target `we://`). + name: &'static str, + /// Scalar property predicates. + properties: &'static [&'static str], + /// HasMany collection relations as `(name, predicate)`. + collections: &'static [(&'static str, &'static str)], +} + +fn class_defs() -> Vec { + vec![ + ClassDef { + name: CLASS_ASSISTANT, + properties: &[ + P_NAME, + P_MODEL_ID, + P_SYSTEM_PROMPT, + P_PERSONALITY_IDS, + P_SKILL_IDS, + P_MCP_SERVER_IDS, + ], + collections: &[], + }, + ClassDef { + name: CLASS_PERSONALITY, + properties: &[P_NAME, P_BODY], + collections: &[], + }, + ClassDef { + name: CLASS_SKILL, + properties: &[P_NAME, P_DESCRIPTION, P_BODY], + collections: &[], + }, + ClassDef { + name: CLASS_MCP_SERVER, + properties: &[P_NAME, "we://transport", P_AUTH, "we://url", P_COMMAND], + collections: &[], + }, + ClassDef { + name: CLASS_THREAD, + properties: &[P_TITLE, P_ASSISTANT_ID, P_MODEL_ID, P_CREATED_AT, P_UPDATED_AT], + collections: &[("messages", REL_MESSAGE)], + }, + ClassDef { + name: CLASS_MESSAGE, + properties: &[P_THREAD_ID, P_ROLE, P_CONTENT, P_TOOL_CALLS, P_TS, P_STATUS], + collections: &[], + }, + ClassDef { + name: CLASS_RUN_STATE, + properties: &[P_THREAD_ID, P_STATUS, P_CURSOR, P_PENDING_TOOL_CALL], + collections: &[], + }, + ] +} + +/// Deterministic positive integer id for a class name (the `uid` slot in the +/// SDNA facts). Stable across runs so re-registration is a no-op. +fn class_uid(name: &str) -> u64 { + // FNV-1a, bounded to 8 digits to stay a compact Prolog integer. + let mut hash: u64 = 0xcbf2_9ce4_8422_2325; + for b in name.as_bytes() { + hash ^= *b as u64; + hash = hash.wrapping_mul(0x0000_0100_0000_01b3); + } + hash % 100_000_000 +} + +/// Short property name used in `property/2` facts — the tail after `we://`. +fn short_name(predicate: &str) -> &str { + predicate.strip_prefix("we://").unwrap_or(predicate) +} + +/// Generate the Prolog SDNA for one class. +fn generate_sdna(def: &ClassDef) -> String { + let uid = class_uid(def.name); + let mut s = String::new(); + + s.push_str(&format!("subject_class(\"{}\", {}).\n", def.name, uid)); + + // An instance is any base carrying the class flag link. + s.push_str(&format!( + "instance({}, Base) :- triple(Base, \"{}\", \"{}\").\n", + uid, FLAG, def.name + )); + + // Constructor writes the flag; destructor removes it. + s.push_str(&format!( + "constructor({}, '[{{\"action\":\"addLink\",\"source\":\"this\",\"predicate\":\"{}\",\"target\":\"{}\"}}]').\n", + uid, FLAG, def.name + )); + s.push_str(&format!( + "destructor({}, '[{{\"action\":\"removeLink\",\"source\":\"this\",\"predicate\":\"{}\",\"target\":\"{}\"}}]').\n", + uid, FLAG, def.name + )); + + for pred in def.properties { + let pname = short_name(pred); + s.push_str(&format!("property({}, \"{}\").\n", uid, pname)); + s.push_str(&format!( + "property_getter({}, Base, \"{}\", Value) :- triple(Base, \"{}\", Value).\n", + uid, pname, pred + )); + s.push_str(&format!( + "property_setter({}, \"{}\", '[{{\"action\":\"setSingleTarget\",\"source\":\"this\",\"predicate\":\"{}\",\"target\":\"value\"}}]').\n", + uid, pname, pred + )); + } + + for (cname, pred) in def.collections { + s.push_str(&format!("collection({}, \"{}\").\n", uid, cname)); + s.push_str(&format!( + "collection_getter({}, Base, \"{}\", List) :- findall(C, triple(Base, \"{}\", C), List).\n", + uid, cname, pred + )); + s.push_str(&format!( + "collection_adder({}, \"{}\", '[{{\"action\":\"addLink\",\"source\":\"this\",\"predicate\":\"{}\",\"target\":\"value\"}}]').\n", + uid, cname, pred + )); + s.push_str(&format!( + "collection_remover({}, \"{}\", '[{{\"action\":\"removeLink\",\"source\":\"this\",\"predicate\":\"{}\",\"target\":\"value\"}}]').\n", + uid, cname, pred + )); + } + + s +} + +/// Register any WE subject classes not already present in the perspective. +/// Best effort — failures are logged and skipped so a bad class never blocks +/// the run loop (which is link-level and independent of SDNA). +pub async fn ensure_subject_classes(p: &mut PerspectiveInstance, ctx: &AgentContext) { + let existing: HashSet = p + .get_subject_classes_from_shacl() + .await + .unwrap_or_default() + .into_iter() + .collect(); + + for def in class_defs() { + if existing.contains(def.name) { + continue; + } + let sdna = generate_sdna(&def); + if let Err(e) = p + .add_sdna(def.name.to_string(), sdna, SdnaType::SubjectClass, None, ctx) + .await + { + log::warn!( + "assistant_runtime: failed to register SDNA class {} in perspective {}: {}", + def.name, + p.uuid, + e + ); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn class_uid_is_deterministic_and_bounded() { + assert_eq!(class_uid(CLASS_MESSAGE), class_uid(CLASS_MESSAGE)); + assert!(class_uid(CLASS_MESSAGE) < 100_000_000); + assert_ne!(class_uid(CLASS_MESSAGE), class_uid(CLASS_THREAD)); + } + + #[test] + fn message_sdna_contains_expected_facts() { + let def = class_defs() + .into_iter() + .find(|d| d.name == CLASS_MESSAGE) + .unwrap(); + let sdna = generate_sdna(&def); + assert!(sdna.contains(&format!("subject_class(\"{}\"", CLASS_MESSAGE))); + // instance condition keys off the flag link + assert!(sdna.contains(&format!( + "instance({}, Base) :- triple(Base, \"{}\", \"{}\").", + class_uid(CLASS_MESSAGE), + FLAG, + CLASS_MESSAGE + ))); + // property getter for content maps to the we://content predicate + assert!(sdna.contains("property_getter")); + assert!(sdna.contains(&format!("triple(Base, \"{}\", Value)", P_CONTENT))); + } + + #[test] + fn thread_sdna_declares_messages_collection() { + let def = class_defs() + .into_iter() + .find(|d| d.name == CLASS_THREAD) + .unwrap(); + let sdna = generate_sdna(&def); + assert!(sdna.contains("collection(")); + assert!(sdna.contains("\"messages\"")); + assert!(sdna.contains(&format!("triple(Base, \"{}\", C)", REL_MESSAGE))); + } +} diff --git a/rust-executor/src/assistant_runtime/store.rs b/rust-executor/src/assistant_runtime/store.rs new file mode 100644 index 000000000..19384ae36 --- /dev/null +++ b/rust-executor/src/assistant_runtime/store.rs @@ -0,0 +1,241 @@ +//! Perspective read/write helpers — the bridge between the pure WE entity +//! model in [`super::entities`] and the executor's in-process +//! [`PerspectiveInstance`] link store. +//! +//! Writes go through `PerspectiveInstance::add_link` / `remove_link` as the +//! main agent (no billing, no capability checks) so the subsystem can persist +//! replies regardless of which user owns the perspective. Reads use +//! `get_links` and decode `we://` property links back into entity structs. + +use std::collections::HashMap; + +use anyhow::Result; + +use crate::agent::AgentContext; +use crate::perspectives::perspective_instance::PerspectiveInstance; +use crate::perspectives::{all_perspectives, get_perspective}; +use crate::types::{Link, LinkExpression, LinkQuery, LinkStatus}; + +use super::entities::{ + decode_literal, encode_literal, Message, RunState, CLASS_MESSAGE, CLASS_RUN_STATE, FLAG, + P_THREAD_ID, REL_MESSAGE, +}; + +/// Load the property map (predicate → decoded value) for a base URI by reading +/// all of its outgoing links. +pub async fn load_props(p: &PerspectiveInstance, base: &str) -> HashMap { + let links = p + .get_links(&LinkQuery { + source: Some(base.to_string()), + ..Default::default() + }) + .await + .unwrap_or_default(); + let mut map = HashMap::new(); + for l in links { + if let Some(pred) = l.data.predicate { + map.insert(pred, decode_literal(&l.data.target)); + } + } + map +} + +/// Return the base URIs of every instance of `class` in the perspective (the +/// `source` of each `--we://flag--> we://` link). +pub async fn find_instances(p: &PerspectiveInstance, class: &str) -> Vec { + p.get_links(&LinkQuery { + predicate: Some(FLAG.to_string()), + target: Some(class.to_string()), + ..Default::default() + }) + .await + .unwrap_or_default() + .into_iter() + .map(|l| l.data.source) + .collect() +} + +/// Find the perspective that holds `base` as an instance of `class`, scanning +/// all perspectives. Used to locate an `Assistant` (and its config) without a +/// separate we-root pointer. +pub async fn find_perspective_with_instance(base: &str, class: &str) -> Option { + for p in all_perspectives() { + let hit = p + .get_links(&LinkQuery { + source: Some(base.to_string()), + predicate: Some(FLAG.to_string()), + target: Some(class.to_string()), + ..Default::default() + }) + .await + .unwrap_or_default(); + if !hit.is_empty() { + return Some(p); + } + } + None +} + +/// Load a single `Message` instance by base URI. +pub async fn load_message(p: &PerspectiveInstance, base: &str) -> Message { + Message::from_props(base.to_string(), &load_props(p, base).await) +} + +/// Load every `Message` in a thread, ordered by `ts`. Uses the canonical +/// `we://thread_id` join, tolerating either literal-encoded (`literal:string:…`) +/// or raw-URI storage of the id by matching on the decoded value. +pub async fn thread_messages(conv: &PerspectiveInstance, thread_id: &str) -> Vec { + let links = conv + .get_links(&LinkQuery { + predicate: Some(P_THREAD_ID.to_string()), + ..Default::default() + }) + .await + .unwrap_or_default(); + + let mut msgs = Vec::with_capacity(links.len()); + for l in links { + if decode_literal(&l.data.target) == thread_id { + msgs.push(load_message(conv, &l.data.source).await); + } + } + msgs.sort_by(|a, b| a.ts.cmp(&b.ts)); + msgs +} + +/// Set a single-valued property link: remove any existing `(base, predicate)` +/// links, then add the new literal-encoded value. Mirrors the `setSingleTarget` +/// SDNA action WE uses for scalar properties. +pub async fn set_single_target( + p: &mut PerspectiveInstance, + base: &str, + predicate: &str, + value: &str, + ctx: &AgentContext, +) -> Result<()> { + let existing = p + .get_links(&LinkQuery { + source: Some(base.to_string()), + predicate: Some(predicate.to_string()), + ..Default::default() + }) + .await + .unwrap_or_default(); + for d in existing { + let le = LinkExpression::from(d); + // Best-effort: a concurrent writer may have already removed it. + let _ = p.remove_link(le, None).await; + } + p.add_link( + Link { + source: base.to_string(), + predicate: Some(predicate.to_string()), + target: encode_literal(value), + }, + LinkStatus::Shared, + None, + ctx, + ) + .await?; + Ok(()) +} + +/// Persist a fresh `Message` instance: its flag + property links plus the +/// Thread→Message collection link so both the canonical (`we://thread_id`) and +/// relational (`Thread.messages`) query paths resolve it. +pub async fn write_message( + conv: &mut PerspectiveInstance, + msg: &Message, + ctx: &AgentContext, +) -> Result<()> { + for link in msg.to_links() { + conv.add_link(link, LinkStatus::Shared, None, ctx).await?; + } + if !msg.thread_id.is_empty() { + conv.add_link( + Link { + source: msg.thread_id.clone(), + predicate: Some(REL_MESSAGE.to_string()), + target: msg.id.clone(), + }, + LinkStatus::Shared, + None, + ctx, + ) + .await?; + } + Ok(()) +} + +/// Create or update the `RunState` for a thread. Reuses the existing instance +/// (found by matching `thread_id`) so a run keeps one durable record. Best +/// effort: failures are surfaced to the caller which logs and continues. +pub async fn upsert_run_state( + conv: &mut PerspectiveInstance, + thread_id: &str, + status: &str, + cursor: &str, + pending_tool_call: &str, +) -> Result<()> { + let ctx = AgentContext::main_agent(); + + // Look for an existing RunState for this thread. + let mut existing_base: Option = None; + for base in find_instances(conv, CLASS_RUN_STATE).await { + let rs = RunState::from_props(base.clone(), &load_props(conv, &base).await); + if rs.thread_id == thread_id { + existing_base = Some(base); + break; + } + } + + if let Some(base) = existing_base { + set_single_target(conv, &base, super::entities::P_STATUS, status, &ctx).await?; + set_single_target(conv, &base, super::entities::P_CURSOR, cursor, &ctx).await?; + set_single_target( + conv, + &base, + super::entities::P_PENDING_TOOL_CALL, + pending_tool_call, + &ctx, + ) + .await?; + } else { + let rs = RunState { + id: format!("we://run_state/{}", uuid::Uuid::new_v4()), + thread_id: thread_id.to_string(), + status: status.to_string(), + cursor: cursor.to_string(), + pending_tool_call: pending_tool_call.to_string(), + }; + for link in rs.to_links() { + conv.add_link(link, LinkStatus::Shared, None, &ctx).await?; + } + } + Ok(()) +} + +/// Load every active `RunState` across all perspectives (for boot resume), +/// returned as `(perspective_uuid, RunState)` pairs. +pub async fn active_run_states() -> Vec<(String, RunState)> { + let mut out = Vec::new(); + for p in all_perspectives() { + for base in find_instances(&p, CLASS_RUN_STATE).await { + let rs = RunState::from_props(base.clone(), &load_props(&p, &base).await); + if rs.is_active() && !rs.thread_id.is_empty() { + out.push((p.uuid.clone(), rs)); + } + } + } + out +} + +/// Fetch a mutable perspective handle by uuid. +pub fn writable(perspective_uuid: &str) -> Option { + get_perspective(perspective_uuid) +} + +/// Return the base URIs of every `Message` instance in a perspective. +pub async fn message_bases(conv: &PerspectiveInstance) -> Vec { + find_instances(conv, CLASS_MESSAGE).await +} diff --git a/rust-executor/src/assistant_runtime/tools.rs b/rust-executor/src/assistant_runtime/tools.rs new file mode 100644 index 000000000..429c78154 --- /dev/null +++ b/rust-executor/src/assistant_runtime/tools.rs @@ -0,0 +1,363 @@ +//! Tool providers for the assistant loop. +//! +//! Tool execution sits behind a small provider abstraction so the model's +//! granted tools can come from more than one source. Two providers ship today: +//! +//! * [`BuiltinTools`] — in-process perspective / neighbourhood graph +//! operations, fully implemented (zero HTTP; direct +//! `PerspectiveInstance`/`neighbourhoods` calls). +//! * [`McpToolProvider`] — external MCP servers configured on an assistant. +//! The live MCP **client** transport is a documented follow-up: `rmcp` is +//! currently built with server features only, so wiring a client means +//! enabling `rmcp`'s `client` + `transport-*-client` features and connecting +//! each `McpServer`. Until then this provider exposes **no** tools and any +//! attempt to call one returns an explicit error (never a silent stub). The +//! provider boundary is the seam that follow-up work slots into without +//! touching the loop. +//! +//! Dispatch is a plain enum (no `async-trait` dependency); [`ToolSet`] +//! aggregates providers and routes a call to whichever owns the tool name. + +use anyhow::{anyhow, Result}; +use serde_json::{json, Value}; + +use crate::agent::AgentContext; +use crate::api::openai_compat::types::{FunctionDef, ToolDef}; +use crate::perspectives::get_perspective; +use crate::types::{Link, LinkQuery, LinkStatus, Perspective}; + +use super::entities::{decode_literal, McpServer}; + +/// One source of tools. +pub enum ToolProvider { + Builtin(BuiltinTools), + Mcp(McpToolProvider), +} + +impl ToolProvider { + pub fn tools(&self) -> Vec { + match self { + ToolProvider::Builtin(b) => b.tools(), + ToolProvider::Mcp(m) => m.tools(), + } + } + + pub fn owns(&self, name: &str) -> bool { + self.tools().iter().any(|t| t.function.name == name) + } + + pub async fn execute(&self, name: &str, arguments: &str) -> Result { + match self { + ToolProvider::Builtin(b) => b.execute(name, arguments).await, + ToolProvider::Mcp(m) => m.execute(name, arguments).await, + } + } +} + +/// The set of tools granted to one assistant, across providers. +pub struct ToolSet { + providers: Vec, +} + +impl ToolSet { + pub fn new(providers: Vec) -> Self { + Self { providers } + } + + /// The OpenAI-shaped tool definitions to render into the system prompt. + pub fn tool_defs(&self) -> Vec { + self.providers.iter().flat_map(|p| p.tools()).collect() + } + + pub fn is_empty(&self) -> bool { + self.providers.iter().all(|p| p.tools().is_empty()) + } + + /// Route a tool call to the provider that owns the name. + pub async fn execute(&self, name: &str, arguments: &str) -> Result { + for p in &self.providers { + if p.owns(name) { + return p.execute(name, arguments).await; + } + } + Err(anyhow!("Unknown tool: {}", name)) + } +} + +// --------------------------------------------------------------------------- +// Built-in perspective / neighbourhood tools +// --------------------------------------------------------------------------- + +/// In-process graph tools scoped to the run's conversation perspective (with +/// an optional `perspective_uuid` override on each call). +pub struct BuiltinTools { + pub perspective_uuid: String, +} + +impl BuiltinTools { + pub fn new(perspective_uuid: String) -> Self { + Self { perspective_uuid } + } + + fn tools(&self) -> Vec { + vec![ + def( + "perspective_add_link", + "Add an RDF-like link (source, predicate, target) to the active perspective's knowledge graph. Use literal:string: targets for scalar values.", + json!({ + "type": "object", + "properties": { + "source": {"type": "string", "description": "Subject URI"}, + "predicate": {"type": "string", "description": "Predicate URI"}, + "target": {"type": "string", "description": "Target URI or literal"}, + "perspective_uuid": {"type": "string", "description": "Optional perspective override; defaults to the active conversation perspective"} + }, + "required": ["source", "predicate", "target"] + }), + ), + def( + "perspective_query_links", + "Query links in the active perspective, filtering by any of source/predicate/target. Returns matching triples with decoded literal targets.", + json!({ + "type": "object", + "properties": { + "source": {"type": "string"}, + "predicate": {"type": "string"}, + "target": {"type": "string"}, + "perspective_uuid": {"type": "string"} + } + }), + ), + def( + "perspective_get_subject", + "Read all properties (outgoing links) of a subject/base URI in the active perspective, returned as a predicate→value map.", + json!({ + "type": "object", + "properties": { + "base": {"type": "string", "description": "The subject/base URI to read"}, + "perspective_uuid": {"type": "string"} + }, + "required": ["base"] + }), + ), + def( + "neighbourhood_publish", + "Publish a perspective as a shared neighbourhood using a link-language template address. Returns the neighbourhood URL that others can join.", + json!({ + "type": "object", + "properties": { + "perspective_uuid": {"type": "string"}, + "link_language": {"type": "string", "description": "Link-language template or cloned language address"}, + "name": {"type": "string", "description": "Optional neighbourhood name"} + }, + "required": ["link_language"] + }), + ), + def( + "neighbourhood_join", + "Join a neighbourhood from its URL (neighbourhood://...), creating a local synced perspective. Returns the new perspective uuid.", + json!({ + "type": "object", + "properties": { + "url": {"type": "string"} + }, + "required": ["url"] + }), + ), + ] + } + + fn perspective_arg(&self, v: &Value) -> String { + v.get("perspective_uuid") + .and_then(Value::as_str) + .map(|s| s.to_string()) + .unwrap_or_else(|| self.perspective_uuid.clone()) + } + + async fn execute(&self, name: &str, arguments: &str) -> Result { + let v: Value = serde_json::from_str(arguments).unwrap_or_else(|_| json!({})); + let ctx = AgentContext::main_agent(); + + match name { + "perspective_add_link" => { + let uuid = self.perspective_arg(&v); + let mut p = + get_perspective(&uuid).ok_or_else(|| anyhow!("Perspective not found: {uuid}"))?; + let link = Link { + source: req_str(&v, "source")?, + predicate: Some(req_str(&v, "predicate")?), + target: req_str(&v, "target")?, + }; + let d = p.add_link(link, LinkStatus::Shared, None, &ctx).await?; + Ok(json!({ + "success": true, + "link": { + "source": d.data.source, + "predicate": d.data.predicate, + "target": d.data.target, + "timestamp": d.timestamp, + } + }) + .to_string()) + } + "perspective_query_links" => { + let uuid = self.perspective_arg(&v); + let p = + get_perspective(&uuid).ok_or_else(|| anyhow!("Perspective not found: {uuid}"))?; + let query = LinkQuery { + source: v.get("source").and_then(Value::as_str).map(str::to_string), + predicate: v.get("predicate").and_then(Value::as_str).map(str::to_string), + target: v.get("target").and_then(Value::as_str).map(str::to_string), + ..Default::default() + }; + let links = p.get_links(&query).await.unwrap_or_default(); + let rows: Vec = links + .into_iter() + .map(|l| { + json!({ + "source": l.data.source, + "predicate": l.data.predicate, + "target": l.data.target, + "value": decode_literal(&l.data.target), + }) + }) + .collect(); + Ok(json!({ "count": rows.len(), "links": rows }).to_string()) + } + "perspective_get_subject" => { + let uuid = self.perspective_arg(&v); + let base = req_str(&v, "base")?; + let p = + get_perspective(&uuid).ok_or_else(|| anyhow!("Perspective not found: {uuid}"))?; + let props = super::store::load_props(&p, &base).await; + Ok(json!({ "base": base, "properties": props }).to_string()) + } + "neighbourhood_publish" => { + let uuid = self.perspective_arg(&v); + let link_language = req_str(&v, "link_language")?; + let meta = Perspective { links: Vec::new() }; + let url = crate::neighbourhoods::neighbourhood_publish_from_perspective_with_context( + &uuid, + link_language, + meta, + &ctx, + ) + .await?; + Ok(json!({ "success": true, "neighbourhood_url": url }).to_string()) + } + "neighbourhood_join" => { + let url = req_str(&v, "url")?; + let handle = + crate::neighbourhoods::install_neighbourhood_with_context(url, &ctx).await?; + Ok(json!({ "success": true, "perspective_uuid": handle.uuid }).to_string()) + } + other => Err(anyhow!("Unknown built-in tool: {other}")), + } + } +} + +// --------------------------------------------------------------------------- +// MCP tool provider (live client transport is a documented follow-up) +// --------------------------------------------------------------------------- + +/// Tools from an assistant's configured external MCP servers. +/// +/// FOLLOW-UP: connecting requires enabling `rmcp`'s client features +/// (`client`, `transport-streamable-http-client`, `transport-sse-client`) in +/// `rust-executor/Cargo.toml` and, per configured [`McpServer`], establishing a +/// session, running `list_tools`, and forwarding `call_tool`. That work slots +/// in behind this provider without touching the run loop or [`ToolSet`]. Until +/// it lands this provider is inert: no tools are advertised and any call fails +/// loudly. +pub struct McpToolProvider { + servers: Vec, +} + +impl McpToolProvider { + pub fn new(servers: Vec) -> Self { + if !servers.is_empty() { + log::warn!( + "assistant_runtime: {} MCP server(s) configured on this assistant but the MCP \ + client transport is not yet wired (follow-up: enable rmcp client features and \ + connect). Their tools are unavailable for this run.", + servers.len() + ); + } + Self { servers } + } + + fn tools(&self) -> Vec { + // No tools until the client transport is wired — deliberately empty so + // the model is never offered a tool the executor cannot fulfil. + Vec::new() + } + + async fn execute(&self, name: &str, _arguments: &str) -> Result { + Err(anyhow!( + "MCP tool '{name}' is unavailable: the MCP client transport is not yet wired \ + (follow-up behind McpToolProvider). {} server(s) configured.", + self.servers.len() + )) + } +} + +// --------------------------------------------------------------------------- +// helpers +// --------------------------------------------------------------------------- + +fn def(name: &str, description: &str, parameters: Value) -> ToolDef { + ToolDef { + kind: "function".to_string(), + function: FunctionDef { + name: name.to_string(), + description: Some(description.to_string()), + parameters: Some(parameters), + }, + } +} + +fn req_str(v: &Value, key: &str) -> Result { + v.get(key) + .and_then(Value::as_str) + .map(|s| s.to_string()) + .ok_or_else(|| anyhow!("Missing required argument: {key}")) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn builtin_advertises_graph_tools() { + let t = BuiltinTools::new("uuid-1".into()); + let names: Vec = t.tools().into_iter().map(|d| d.function.name).collect(); + assert!(names.contains(&"perspective_add_link".to_string())); + assert!(names.contains(&"perspective_query_links".to_string())); + assert!(names.contains(&"neighbourhood_publish".to_string())); + } + + #[test] + fn mcp_provider_is_inert_until_wired() { + let p = McpToolProvider::new(vec![McpServer { + id: "m".into(), + name: "srv".into(), + transport: "http".into(), + url: "http://x".into(), + ..Default::default() + }]); + assert!(p.tools().is_empty()); + } + + #[test] + fn toolset_routes_by_ownership() { + let set = ToolSet::new(vec![ + ToolProvider::Builtin(BuiltinTools::new("u".into())), + ToolProvider::Mcp(McpToolProvider::new(vec![])), + ]); + assert!(!set.is_empty()); + assert!(set + .tool_defs() + .iter() + .any(|d| d.function.name == "perspective_add_link")); + } +} diff --git a/rust-executor/src/config.rs b/rust-executor/src/config.rs index 3b3436990..5853abb6f 100644 --- a/rust-executor/src/config.rs +++ b/rust-executor/src/config.rs @@ -91,6 +91,9 @@ pub struct Ad4mConfig { pub enable_mcp: Option, /// Port for MCP HTTP server (default: 3001) pub mcp_port: Option, + /// Enable the server-side AI-assistant run subsystem (default: enabled). + /// Only `Some(false)` disables it, mirroring how a missing flag defaults on. + pub enable_assistants: Option, /// Path to write PID file (for test harness cleanup) pub pid_file: Option, } @@ -182,6 +185,7 @@ impl Default for Ad4mConfig { smtp_config: None, enable_mcp: None, mcp_port: None, + enable_assistants: None, pid_file: None, }; config.prepare(); diff --git a/rust-executor/src/lib.rs b/rust-executor/src/lib.rs index 34358db1b..060d0a80a 100644 --- a/rust-executor/src/lib.rs +++ b/rust-executor/src/lib.rs @@ -2,6 +2,7 @@ extern crate lazy_static; pub mod api; +pub mod assistant_runtime; pub mod config; pub mod email_service; pub mod entanglement_service; @@ -618,6 +619,23 @@ pub async fn run(mut config: Ad4mConfig) -> JoinHandle<()> { }); } + // Check if the AI-assistant subsystem is enabled — run it alongside the + // MCP + REST servers. Enabled unless explicitly disabled, mirroring the + // default-on convention (a missing flag stays on). This closure captures + // nothing from `config`, so `config` is still moved into the REST server + // below. + if config.enable_assistants != Some(false) { + info!("Starting AI-assistant run subsystem..."); + std::thread::spawn(move || { + let runtime = tokio::runtime::Builder::new_multi_thread() + .thread_name(String::from("assistant_runtime")) + .enable_all() + .build() + .unwrap(); + runtime.block_on(assistant_runtime::start()); + }); + } + info!("Starting REST API server..."); std::thread::spawn(move || { From 35bb6885b1bc9f5bbab7ee5550475490e5816bbc Mon Sep 17 00:00:00 2001 From: Josh Field <10372036+HexaField@users.noreply.github.com> Date: Tue, 28 Jul 2026 20:26:32 +1000 Subject: [PATCH 3/6] feat(assistant): wire the rmcp MCP client (external tool servers) McpToolProvider now connects an assistant's configured McpServers via rmcp (stdio -> child process; http/streamable/sse -> streamable-HTTP; websocket unsupported by the pinned rmcp), discovers tools via list_tools -> ToolDef, and dispatches call_tool. Servers that fail to connect are logged and skipped (one bad server never fails the set). Enables rmcp client + transport-*-client features (server features retained). cargo check clean; 4 tools unit tests pass. Co-Authored-By: Claude Opus 4.8 --- Cargo.lock | 16 ++ rust-executor/Cargo.toml | 10 +- rust-executor/src/assistant_runtime/run.rs | 2 +- rust-executor/src/assistant_runtime/tools.rs | 268 +++++++++++++++---- 4 files changed, 248 insertions(+), 48 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 28ba1aba6..0aa1df5bc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -15153,6 +15153,20 @@ dependencies = [ "yansi", ] +[[package]] +name = "process-wrap" +version = "9.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e842efad9119158434d193c6682e2ebee4b44d6ad801d7b349623b3f57cdf55" +dependencies = [ + "futures", + "indexmap 2.11.1", + "nix 0.31.2", + "tokio", + "tracing", + "windows 0.62.2", +] + [[package]] name = "profiling" version = "1.0.17" @@ -16349,7 +16363,9 @@ dependencies = [ "http-body-util", "pastey", "pin-project-lite", + "process-wrap", "rand 0.9.2", + "reqwest 0.12.28", "rmcp-macros", "schemars 1.2.1", "serde", diff --git a/rust-executor/Cargo.toml b/rust-executor/Cargo.toml index 7a9c28576..58c6a104e 100644 --- a/rust-executor/Cargo.toml +++ b/rust-executor/Cargo.toml @@ -67,7 +67,15 @@ hex = "0.4.3" argon2 = { version = "0.5.0", features = ["simple"] } rand = "0.8.5" base64 = "0.21.0" -rmcp = { version = "0.15.0", features = ["server", "transport-streamable-http-server"] } +rmcp = { version = "0.15.0", features = [ + "server", + "transport-streamable-http-server", + # Client side (assistant_runtime MCP tool provider): + "client", + "transport-streamable-http-client", + "transport-streamable-http-client-reqwest", + "transport-child-process", +] } axum = { version = "0.8", features = ["ws", "multipart"] } axum-server = { version = "0.7", features = ["tls-rustls"] } tower-http = { version = "0.6", features = ["cors", "set-header", "catch-panic"] } diff --git a/rust-executor/src/assistant_runtime/run.rs b/rust-executor/src/assistant_runtime/run.rs index 06630fe20..ef64306c8 100644 --- a/rust-executor/src/assistant_runtime/run.rs +++ b/rust-executor/src/assistant_runtime/run.rs @@ -216,7 +216,7 @@ async fn drive_loop( let toolset = ToolSet::new(vec![ ToolProvider::Builtin(BuiltinTools::new(conv.uuid.clone())), - ToolProvider::Mcp(McpToolProvider::new(config.mcp_servers.clone())), + ToolProvider::Mcp(McpToolProvider::connect(config.mcp_servers.clone()).await), ]); let tool_defs = toolset.tool_defs(); diff --git a/rust-executor/src/assistant_runtime/tools.rs b/rust-executor/src/assistant_runtime/tools.rs index 429c78154..49dcc39f1 100644 --- a/rust-executor/src/assistant_runtime/tools.rs +++ b/rust-executor/src/assistant_runtime/tools.rs @@ -6,14 +6,14 @@ //! * [`BuiltinTools`] — in-process perspective / neighbourhood graph //! operations, fully implemented (zero HTTP; direct //! `PerspectiveInstance`/`neighbourhoods` calls). -//! * [`McpToolProvider`] — external MCP servers configured on an assistant. -//! The live MCP **client** transport is a documented follow-up: `rmcp` is -//! currently built with server features only, so wiring a client means -//! enabling `rmcp`'s `client` + `transport-*-client` features and connecting -//! each `McpServer`. Until then this provider exposes **no** tools and any -//! attempt to call one returns an explicit error (never a silent stub). The -//! provider boundary is the seam that follow-up work slots into without -//! touching the loop. +//! * [`McpToolProvider`] — external MCP servers configured on an assistant, +//! connected via a live `rmcp` client. Each `McpServer` is connected on run +//! start (stdio → child process; http/streamable/sse → streamable-HTTP; +//! websocket unsupported by the pinned rmcp), its tools discovered via +//! `list_tools` and mapped to OpenAI `ToolDef`s, and calls dispatched via +//! `call_tool`. A server that fails to connect or enumerate is logged and +//! skipped, so the model is never offered an unfulfillable tool and one bad +//! server never fails the whole set. //! //! Dispatch is a plain enum (no `async-trait` dependency); [`ToolSet`] //! aggregates providers and routes a call to whichever owns the tool name. @@ -21,6 +21,13 @@ use anyhow::{anyhow, Result}; use serde_json::{json, Value}; +use rmcp::model::{CallToolRequestParams, Tool}; +use rmcp::service::RunningService; +use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig; +use rmcp::transport::{ConfigureCommandExt, StreamableHttpClientTransport, TokioChildProcess}; +use rmcp::{RoleClient, ServiceExt}; +use tokio::process::Command; + use crate::agent::AgentContext; use crate::api::openai_compat::types::{FunctionDef, ToolDef}; use crate::perspectives::get_perspective; @@ -257,47 +264,203 @@ impl BuiltinTools { } // --------------------------------------------------------------------------- -// MCP tool provider (live client transport is a documented follow-up) +// MCP tool provider — live rmcp client // --------------------------------------------------------------------------- -/// Tools from an assistant's configured external MCP servers. +/// One connected external MCP server: the live rmcp client session plus the +/// tool definitions discovered from it. +struct McpConnection { + server_name: String, + client: RunningService, + tools: Vec, +} + +/// Tools from an assistant's configured external MCP servers, connected via +/// rmcp. Servers are connected once when the provider is built (and the +/// sessions are cached for the provider's — i.e. the run's — lifetime). A +/// server that fails to connect or enumerate tools is logged and skipped, so +/// the model is never offered a tool the executor cannot fulfil, and one bad +/// server never fails the whole tool set. /// -/// FOLLOW-UP: connecting requires enabling `rmcp`'s client features -/// (`client`, `transport-streamable-http-client`, `transport-sse-client`) in -/// `rust-executor/Cargo.toml` and, per configured [`McpServer`], establishing a -/// session, running `list_tools`, and forwarding `call_tool`. That work slots -/// in behind this provider without touching the run loop or [`ToolSet`]. Until -/// it lands this provider is inert: no tools are advertised and any call fails -/// loudly. +/// Transports (rmcp 0.15): `stdio` → [`TokioChildProcess`]; +/// `http`/`streamable` → [`StreamableHttpClientTransport`]; `sse` → the same +/// streamable-HTTP client (0.15 folds SSE into `client-side-sse` — there is no +/// standalone SSE client transport); `websocket` is unsupported (this rmcp has +/// no ws client transport) and is skipped with a clear warning. pub struct McpToolProvider { - servers: Vec, + connections: Vec, } impl McpToolProvider { - pub fn new(servers: Vec) -> Self { - if !servers.is_empty() { - log::warn!( - "assistant_runtime: {} MCP server(s) configured on this assistant but the MCP \ - client transport is not yet wired (follow-up: enable rmcp client features and \ - connect). Their tools are unavailable for this run.", - servers.len() - ); + /// Connect to every configured server and discover its tools. Failures are + /// logged and skipped; never panics, never fails the whole set. + pub async fn connect(servers: Vec) -> Self { + let mut connections = Vec::new(); + for server in servers { + match connect_server(&server).await { + Ok(conn) => { + log::info!( + "assistant_runtime: connected MCP server '{}' ({} tools)", + conn.server_name, + conn.tools.len() + ); + connections.push(conn); + } + Err(e) => log::warn!( + "assistant_runtime: MCP server '{}' (transport '{}') unavailable, skipping: {}", + server.name, + server.transport, + e + ), + } + } + Self { connections } + } + + /// A provider with no connections (fallback / tests). + pub fn empty() -> Self { + Self { + connections: Vec::new(), } - Self { servers } } fn tools(&self) -> Vec { - // No tools until the client transport is wired — deliberately empty so - // the model is never offered a tool the executor cannot fulfil. - Vec::new() + self.connections + .iter() + .flat_map(|c| c.tools.clone()) + .collect() } - async fn execute(&self, name: &str, _arguments: &str) -> Result { - Err(anyhow!( - "MCP tool '{name}' is unavailable: the MCP client transport is not yet wired \ - (follow-up behind McpToolProvider). {} server(s) configured.", - self.servers.len() - )) + async fn execute(&self, name: &str, arguments: &str) -> Result { + for conn in &self.connections { + if conn.tools.iter().any(|t| t.function.name == name) { + return call_mcp_tool(&conn.client, name, arguments).await; + } + } + Err(anyhow!("MCP tool '{name}' not found on any connected server")) + } +} + +/// Connect one server (dispatching on its transport) and enumerate its tools. +async fn connect_server(server: &McpServer) -> Result { + let client: RunningService = match server.transport.as_str() { + "stdio" => { + let mut parts = server.command.split_whitespace(); + let program = parts + .next() + .ok_or_else(|| anyhow!("stdio server '{}' has an empty command", server.name))? + .to_string(); + let args: Vec = parts.map(str::to_string).collect(); + let transport = TokioChildProcess::new(Command::new(&program).configure(|cmd| { + for arg in &args { + cmd.arg(arg); + } + })) + .map_err(|e| anyhow!("failed to spawn '{program}': {e}"))?; + () + .serve(transport) + .await + .map_err(|e| anyhow!("stdio MCP handshake failed: {e:?}"))? + } + "http" | "streamable" | "sse" => { + if server.url.is_empty() { + return Err(anyhow!( + "'{}' server '{}' has no url", + server.transport, + server.name + )); + } + if server.transport == "sse" { + log::info!( + "assistant_runtime: MCP server '{}' declares transport 'sse'; the pinned rmcp \ + folds SSE into the streamable-HTTP client (no standalone SSE transport), \ + connecting via streamable-HTTP.", + server.name + ); + } + // `from_config` builds the transport's own reqwest client + // internally, so the executor's reqwest version is never involved. + let mut config = StreamableHttpClientTransportConfig::with_uri(server.url.clone()); + if !server.auth.is_empty() { + // Bearer token value (no `Bearer ` prefix), per the config API. + config.auth_header = Some(server.auth.clone()); + } + let transport = StreamableHttpClientTransport::from_config(config); + () + .serve(transport) + .await + .map_err(|e| anyhow!("streamable-HTTP MCP handshake failed: {e:?}"))? + } + "websocket" | "ws" => { + return Err(anyhow!( + "transport 'websocket' is unsupported by the pinned rmcp (no ws client transport)" + )); + } + other => return Err(anyhow!("unknown transport '{other}'")), + }; + + let raw_tools = client + .list_all_tools() + .await + .map_err(|e| anyhow!("list_tools failed: {e}"))?; + let tools = raw_tools.iter().map(tool_to_def).collect(); + + Ok(McpConnection { + server_name: server.name.clone(), + client, + tools, + }) +} + +/// Call one MCP tool and flatten its result content to a string. +async fn call_mcp_tool( + client: &RunningService, + name: &str, + arguments: &str, +) -> Result { + // The OpenAI `function.arguments` is a JSON *string* of an object. + let parsed: Value = serde_json::from_str(arguments).unwrap_or(Value::Null); + let arguments = match parsed { + Value::Object(map) => Some(map), + _ => None, + }; + + let result = client + .call_tool(CallToolRequestParams { + meta: None, + name: name.to_string().into(), + arguments, + task: None, + }) + .await + .map_err(|e| anyhow!("call_tool '{name}' failed: {e}"))?; + + // Prefer concatenated text blocks; fall back to structured content / raw. + let text = result + .content + .iter() + .filter_map(|c| c.as_text().map(|t| t.text.clone())) + .collect::>() + .join("\n"); + if !text.is_empty() { + Ok(text) + } else if let Some(structured) = result.structured_content { + Ok(structured.to_string()) + } else { + Ok(serde_json::to_string(&result.content).unwrap_or_default()) + } +} + +/// Map an rmcp [`Tool`] to an OpenAI-shaped [`ToolDef`]: name, description, and +/// `inputSchema` → `parameters`. +fn tool_to_def(tool: &Tool) -> ToolDef { + ToolDef { + kind: "function".to_string(), + function: FunctionDef { + name: tool.name.to_string(), + description: tool.description.as_ref().map(|d| d.to_string()), + parameters: Some(Value::Object((*tool.input_schema).clone())), + }, } } @@ -337,22 +500,35 @@ mod tests { } #[test] - fn mcp_provider_is_inert_until_wired() { - let p = McpToolProvider::new(vec![McpServer { - id: "m".into(), - name: "srv".into(), - transport: "http".into(), - url: "http://x".into(), - ..Default::default() - }]); - assert!(p.tools().is_empty()); + fn mcp_tool_maps_to_tooldef() { + use std::sync::Arc; + let mut schema = serde_json::Map::new(); + schema.insert("type".to_string(), json!("object")); + schema.insert( + "properties".to_string(), + json!({ "q": { "type": "string" } }), + ); + let tool = Tool::new("search", "Search the web", Arc::new(schema)); + + let def = tool_to_def(&tool); + assert_eq!(def.kind, "function"); + assert_eq!(def.function.name, "search"); + assert_eq!(def.function.description.as_deref(), Some("Search the web")); + let params = def.function.parameters.expect("parameters mapped"); + assert_eq!(params["type"], json!("object")); + assert_eq!(params["properties"]["q"]["type"], json!("string")); + } + + #[test] + fn empty_mcp_provider_advertises_no_tools() { + assert!(McpToolProvider::empty().tools().is_empty()); } #[test] fn toolset_routes_by_ownership() { let set = ToolSet::new(vec![ ToolProvider::Builtin(BuiltinTools::new("u".into())), - ToolProvider::Mcp(McpToolProvider::new(vec![])), + ToolProvider::Mcp(McpToolProvider::empty()), ]); assert!(!set.is_empty()); assert!(set From 2fe34322713f4ddab99cb45c16208ef4241cd359 Mon Sep 17 00:00:00 2001 From: Josh Field <10372036+HexaField@users.noreply.github.com> Date: Tue, 28 Jul 2026 20:33:46 +1000 Subject: [PATCH 4/6] fix(cli): thread enable_assistants through the executor CLI The assistant-runtime subsystem added Ad4mConfig.enable_assistants (rust-executor); the cli crate constructs Ad4mConfig and must supply it. Adds a --enable-assistants flag on the ad4m-executor bin (mirroring --enable-mcp; default on) and None in the dev/ad4m constructions. Fixes the E0063 that only surfaced building the downstream cli crate. Co-Authored-By: Claude Opus 4.8 --- cli/src/ad4m_executor.rs | 5 +++++ cli/src/dev.rs | 2 ++ cli/src/main.rs | 1 + 3 files changed, 8 insertions(+) diff --git a/cli/src/ad4m_executor.rs b/cli/src/ad4m_executor.rs index f16d131c4..3c410e3ec 100644 --- a/cli/src/ad4m_executor.rs +++ b/cli/src/ad4m_executor.rs @@ -169,6 +169,9 @@ enum Domain { enable_mcp: Option, #[arg(long, action)] mcp_port: Option, + /// Enable the server-side AI-assistant runtime (default: on). + #[arg(long, action)] + enable_assistants: Option, /// Write the executor PID to this file on startup (removed on clean shutdown). /// Useful for test harnesses that need targeted process cleanup. #[arg(long)] @@ -226,6 +229,7 @@ async fn main() -> Result<()> { enable_multi_user, enable_mcp, mcp_port, + enable_assistants, pid_file, } = args.domain { @@ -267,6 +271,7 @@ async fn main() -> Result<()> { smtp_config: None, enable_mcp, mcp_port, + enable_assistants, pid_file, }) .await; diff --git a/cli/src/dev.rs b/cli/src/dev.rs index d708d0897..c3bc6c6e3 100644 --- a/cli/src/dev.rs +++ b/cli/src/dev.rs @@ -59,6 +59,7 @@ pub async fn run(command: DevFunctions) -> Result<()> { enable_multi_user: None, enable_mcp: None, mcp_port: None, + enable_assistants: None, smtp_config: None, pid_file: None, }) @@ -214,6 +215,7 @@ pub async fn run(command: DevFunctions) -> Result<()> { enable_multi_user: None, enable_mcp: None, mcp_port: None, + enable_assistants: None, smtp_config: None, pid_file: None, }) diff --git a/cli/src/main.rs b/cli/src/main.rs index 5cba04569..a6aeefed1 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -269,6 +269,7 @@ async fn main() -> Result<()> { enable_multi_user, enable_mcp, mcp_port, + enable_assistants: None, pid_file, localhost: None, auto_permit_cap_requests: None, From 56a5a18af283522d3419558a2b8704bcb9eb2ba7 Mon Sep 17 00:00:00 2001 From: Josh Field <10372036+HexaField@users.noreply.github.com> Date: Tue, 28 Jul 2026 20:57:24 +1000 Subject: [PATCH 5/6] fix(assistant): resolve model name/id before calling AIService MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The run loop passed Assistant.model_id straight to AIService, whose LLM channel is keyed by the registered model uuid — so a friendly name like 'qwen2.5' never matched (Model not found in LLM channel). Resolve via the /v1 model_selector (id | name | 'default') first, the same as the /v1 endpoint. Caught by the live loop end-to-end test on real hardware. Co-Authored-By: Claude Opus 4.8 --- rust-executor/src/assistant_runtime/run.rs | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/rust-executor/src/assistant_runtime/run.rs b/rust-executor/src/assistant_runtime/run.rs index ef64306c8..75dadbceb 100644 --- a/rust-executor/src/assistant_runtime/run.rs +++ b/rust-executor/src/assistant_runtime/run.rs @@ -183,13 +183,25 @@ async fn resolve_config( } } - let model_id = if !thread.model_id.is_empty() { + let requested_model = if !thread.model_id.is_empty() { thread.model_id.clone() } else if !assistant.model_id.is_empty() { assistant.model_id.clone() } else { "default".to_string() }; + // Resolve the model reference (id | name | "default") to the registered + // model id that AIService keys its LLM channel by — the same resolution + // the /v1 endpoint performs. Without this a friendly name like "qwen2.5" + // never matches the channel (which is keyed by the model uuid). + let model_id = crate::api::openai_compat::model_selector::resolve_model( + &requested_model, + crate::types::ModelType::Llm, + ) + .await + .map_err(|_| { + anyhow!("assistant model '{requested_model}' is not registered on this executor") + })?; Ok(RunConfig { assistant, From edff8dd271f0fa18d811b9703d3476c47d626a94 Mon Sep 17 00:00:00 2001 From: Josh Field <10372036+HexaField@users.noreply.github.com> Date: Tue, 28 Jul 2026 22:13:53 +1000 Subject: [PATCH 6/6] =?UTF-8?q?test(assistant):=20genuine=20e2e=20tests=20?= =?UTF-8?q?=E2=80=94=20real=20subsystem,=20mocked=20LLM=20with=20captured?= =?UTF-8?q?=20Qwen2.5=20responses?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a ModelBackend seam (AiServiceBackend prod pass-through + FixtureModelBackend test) so the run loop can be driven by recorded responses. Two e2e tests run the FULL subsystem — real in-process PerspectiveInstance, real subject-class read/write, the real loop, real built-in tool execution, real persistence — against raw model text derived from Qwen2.5-7B /v1 responses captured live on Apple Silicon (tests/fixtures/): - plain turn -> assistant message with the captured content, status complete, RunState done - tool loop -> real perspective_add_link executed (sky_color_blue--hasColor-->blue added), post-tool final answer, toolCalls recorded, role:tool message persisted, RunState done 24 assistant_runtime tests pass (22 unit + 2 e2e). Production path unchanged (real backend is a faithful pass-through). Co-Authored-By: Claude Opus 4.8 --- .../src/assistant_runtime/e2e_tests.rs | 357 ++++++++++++++++++ rust-executor/src/assistant_runtime/mod.rs | 9 +- .../src/assistant_runtime/model_backend.rs | 64 ++++ .../src/assistant_runtime/registry.rs | 15 +- rust-executor/src/assistant_runtime/run.rs | 52 ++- rust-executor/src/perspectives/mod.rs | 13 + rust-executor/tests/fixtures/chat_plain.json | 1 + .../tests/fixtures/final_after_tool.json | 1 + .../tests/fixtures/toolcall_builtin.json | 1 + .../tests/fixtures/toolcall_weather.json | 1 + 10 files changed, 487 insertions(+), 27 deletions(-) create mode 100644 rust-executor/src/assistant_runtime/e2e_tests.rs create mode 100644 rust-executor/src/assistant_runtime/model_backend.rs create mode 100644 rust-executor/tests/fixtures/chat_plain.json create mode 100644 rust-executor/tests/fixtures/final_after_tool.json create mode 100644 rust-executor/tests/fixtures/toolcall_builtin.json create mode 100644 rust-executor/tests/fixtures/toolcall_weather.json diff --git a/rust-executor/src/assistant_runtime/e2e_tests.rs b/rust-executor/src/assistant_runtime/e2e_tests.rs new file mode 100644 index 000000000..07183edcc --- /dev/null +++ b/rust-executor/src/assistant_runtime/e2e_tests.rs @@ -0,0 +1,357 @@ +//! End-to-end tests for the assistant run subsystem. +//! +//! These drive the FULL subsystem — real in-process perspectives, real +//! subject-class link read/write, the real `run_thread` loop, real built-in +//! tool execution, and real streaming/persistence — with only the LLM replaced +//! by [`FixtureModelBackend`], which plays back RAW model text derived from +//! Qwen2.5-7B `/v1` responses captured live (`tests/fixtures/*.json`). +//! +//! The perspective is built with the same in-memory primitives the executor's +//! own perspective tests use (`Ad4mDb(:memory:)` + `AgentService` test +//! instance) and registered in the global registry via +//! [`crate::perspectives::insert_perspective_for_test`] so `get_perspective` / +//! `all_perspectives` resolve it exactly as in production. +//! +//! The tests share process-global state (`Ad4mDb`, `AgentService`), so they +//! serialize on a suite lock; run them with e.g. +//! `cargo test assistant_runtime::e2e_tests`. + +use std::collections::VecDeque; +use std::future::Future; +use std::pin::Pin; +use std::sync::{Arc, Mutex, OnceLock}; + +use kalosm::language::ArcParser; +use serde_json::Value; +use tokio::sync::mpsc; +use tokio::sync::Mutex as AsyncMutex; + +use crate::agent::{AgentContext, AgentService}; +use crate::db::Ad4mDb; +use crate::perspectives::insert_perspective_for_test; +use crate::perspectives::perspective_instance::PerspectiveInstance; +use crate::types::{ + Link, LinkQuery, LinkStatus, LocalModelInput, ModelInput, ModelType, PerspectiveHandle, + PerspectiveState, +}; + +use super::entities; +use super::model_backend::{ModelBackend, TokenStream}; +use super::run::{run_thread, RunInput}; +use super::store; + +// --------------------------------------------------------------------------- +// Fixtures (captured live from Qwen2.5-7B via /v1/chat/completions) +// --------------------------------------------------------------------------- + +const CHAT_PLAIN: &str = + include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/chat_plain.json")); +const TOOLCALL_BUILTIN: &str = include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/toolcall_builtin.json" +)); +const FINAL_AFTER_TOOL: &str = include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/final_after_tool.json" +)); + +/// Convert a captured `/v1` chat-completion response into the RAW model text +/// the local backend would have produced: plain `content` verbatim, or each +/// tool call rendered as a Hermes `` block with `arguments` parsed +/// from the OpenAI JSON *string* back into an object — the exact form the +/// subsystem's real `extract_tool_calls` + fold-back expects. +fn fixture_to_raw(fixture_json: &str) -> String { + let v: Value = serde_json::from_str(fixture_json).expect("fixture is valid JSON"); + let message = &v["choices"][0]["message"]; + + if let Some(content) = message["content"].as_str() { + return content.to_string(); + } + + if let Some(tool_calls) = message["tool_calls"].as_array() { + let mut out = String::new(); + for call in tool_calls { + let function = &call["function"]; + let name = function["name"].as_str().expect("tool call name"); + let args_str = function["arguments"] + .as_str() + .expect("tool call arguments are a JSON string"); + let args_obj: Value = + serde_json::from_str(args_str).expect("arguments parse to an object"); + let block = serde_json::json!({ "name": name, "arguments": args_obj }); + out.push_str(&format!("\n{}\n\n", block)); + } + return out.trim().to_string(); + } + + String::new() +} + +// --------------------------------------------------------------------------- +// Fixture model backend +// --------------------------------------------------------------------------- + +/// A [`ModelBackend`] that replays scripted raw responses in call order. Each +/// `stream` call pops the next response and streams it as space-delimited +/// chunks (so the real per-token content-rewrite/cadence path runs), then +/// closes the channel to signal completion. +struct FixtureModelBackend { + responses: Mutex>, +} + +impl FixtureModelBackend { + fn new(raw_texts: Vec) -> Self { + Self { + responses: Mutex::new(raw_texts.into_iter().collect()), + } + } +} + +impl ModelBackend for FixtureModelBackend { + fn stream( + &self, + _model_id: String, + _messages: Vec<(String, String)>, + _constraint: Option>, + ) -> Pin> + Send + '_>> { + // Pop synchronously (no lock held across the await point). + let next = self + .responses + .lock() + .unwrap() + .pop_front() + .unwrap_or_default(); + Box::pin(async move { + let (tx, rx) = mpsc::unbounded_channel(); + for token in next.split_inclusive(' ') { + let _ = tx.send(token.to_string()); + } + // tx drops here → channel closes → the loop sees completion. + Ok(rx) + }) + } +} + +// --------------------------------------------------------------------------- +// Harness +// --------------------------------------------------------------------------- + +/// Serialize e2e tests — they share the process-global `Ad4mDb`/`AgentService`. +fn suite_lock() -> &'static AsyncMutex<()> { + static LOCK: OnceLock> = OnceLock::new(); + LOCK.get_or_init(|| AsyncMutex::new(())) +} + +/// Build a fresh in-memory perspective and register it globally. +async fn setup_perspective() -> PerspectiveInstance { + crate::test_utils::setup_wallet(); + let _ = Ad4mDb::init_global_instance(":memory:"); + AgentService::init_global_test_instance(); + + let handle = PerspectiveHandle { + uuid: uuid::Uuid::new_v4().to_string(), + name: Some("assistant-e2e".to_string()), + shared_url: None, + neighbourhood: None, + state: PerspectiveState::Private, + owners: None, + }; + let instance = PerspectiveInstance::new(handle, None); + insert_perspective_for_test(instance.clone()); + instance +} + +/// Register a minimal LLM model row so the REAL model resolution in +/// `resolve_config` succeeds (the fixture backend ignores which model it is). +fn register_model(name: &str) { + Ad4mDb::with_global_instance(|db| { + db.add_model(&ModelInput { + name: name.to_string(), + api: None, + local: Some(LocalModelInput { + file_name: name.to_string(), + tokenizer_source: None, + huggingface_repo: None, + revision: None, + }), + model_type: ModelType::Llm, + }) + }) + .expect("register model row"); +} + +async fn add_links(p: &mut PerspectiveInstance, links: Vec) { + let ctx = AgentContext::main_agent(); + for link in links { + p.add_link(link, LinkStatus::Shared, None, &ctx) + .await + .expect("add_link"); + } +} + +/// Seed an `Assistant` + `Thread`, returning `(assistant_id, thread_id)`. +async fn seed_assistant_and_thread( + p: &mut PerspectiveInstance, + model_name: &str, +) -> (String, String) { + let assistant_id = format!("we://assistant/{}", uuid::Uuid::new_v4()); + let thread_id = format!("we://thread/{}", uuid::Uuid::new_v4()); + add_links( + p, + vec![ + entities::flag_link(&assistant_id, entities::CLASS_ASSISTANT), + entities::property_link(&assistant_id, entities::P_NAME, "Ada"), + entities::property_link(&assistant_id, entities::P_MODEL_ID, model_name), + entities::property_link( + &assistant_id, + entities::P_SYSTEM_PROMPT, + "You are Ada, a concise assistant.", + ), + entities::flag_link(&thread_id, entities::CLASS_THREAD), + entities::property_link(&thread_id, entities::P_ASSISTANT_ID, &assistant_id), + ], + ) + .await; + (assistant_id, thread_id) +} + +/// Write a completed user `Message` (timestamped before any reply). +async fn write_user_message(p: &mut PerspectiveInstance, thread_id: &str, content: &str) { + let ctx = AgentContext::main_agent(); + let msg = entities::Message { + id: format!("we://message/{}", uuid::Uuid::new_v4()), + thread_id: thread_id.to_string(), + role: "user".to_string(), + content: content.to_string(), + tool_calls: String::new(), + ts: "2020-01-01T00:00:00Z".to_string(), + status: "complete".to_string(), + }; + store::write_message(p, &msg, &ctx) + .await + .expect("write user message"); +} + +fn assistant_reply(msgs: &[entities::Message]) -> &entities::Message { + msgs.iter() + .find(|m| m.role == "assistant") + .expect("an assistant reply exists") +} + +async fn run_state_status(p: &PerspectiveInstance, thread_id: &str) -> Option { + for base in store::find_instances(p, entities::CLASS_RUN_STATE).await { + let rs = entities::RunState::from_props(base.clone(), &store::load_props(p, &base).await); + if rs.thread_id == thread_id { + return Some(rs.status); + } + } + None +} + +// --------------------------------------------------------------------------- +// Scenarios +// --------------------------------------------------------------------------- + +/// Plain turn: a completed user message → the loop → the model returns plain +/// content → an assistant message with that content, `status=complete`, and a +/// `RunState` of `done`. +#[tokio::test] +async fn e2e_plain_turn_persists_complete_assistant_message() { + let _guard = suite_lock().lock().await; + + let mut p = setup_perspective().await; + register_model("qwen2.5"); + let (_assistant_id, thread_id) = seed_assistant_and_thread(&mut p, "qwen2.5").await; + write_user_message(&mut p, &thread_id, "Hello, who are you?").await; + + let backend = Arc::new(FixtureModelBackend::new(vec![fixture_to_raw(CHAT_PLAIN)])); + run_thread( + RunInput { + perspective_uuid: p.uuid.clone(), + thread_id: thread_id.clone(), + }, + backend, + ) + .await + .expect("run_thread"); + + let msgs = store::thread_messages(&p, &thread_id).await; + let reply = assistant_reply(&msgs); + let expected = "Hi there! I'm Qwen, an artificial intelligence developed by Alibaba Cloud. \ + My main function is to assist with various tasks and provide information on a \ + wide range of topics. How can I help you today?"; + assert_eq!(reply.content, expected); + assert_eq!(reply.status, "complete"); + assert_eq!( + run_state_status(&p, &thread_id).await.as_deref(), + Some("done") + ); +} + +/// Tool loop: the model returns a `perspective_add_link` tool call → the +/// subsystem EXECUTES the real built-in tool (a real link is added to the +/// perspective) → the model returns the final answer. Proves real tool +/// execution + fold-back + persistence. +#[tokio::test] +async fn e2e_tool_loop_executes_builtin_tool_and_folds_result() { + let _guard = suite_lock().lock().await; + + let mut p = setup_perspective().await; + register_model("qwen2.5"); + let (_assistant_id, thread_id) = seed_assistant_and_thread(&mut p, "qwen2.5").await; + write_user_message(&mut p, &thread_id, "Record that the sky is blue.").await; + + // Call 1 → perspective_add_link tool call; Call 2 → final answer. + let backend = Arc::new(FixtureModelBackend::new(vec![ + fixture_to_raw(TOOLCALL_BUILTIN), + fixture_to_raw(FINAL_AFTER_TOOL), + ])); + run_thread( + RunInput { + perspective_uuid: p.uuid.clone(), + thread_id: thread_id.clone(), + }, + backend, + ) + .await + .expect("run_thread"); + + // (a) The tool really added the link to the perspective. + let links = p + .get_links(&LinkQuery { + source: Some("sky_color_blue".to_string()), + predicate: Some("hasColor".to_string()), + ..Default::default() + }) + .await + .expect("get_links"); + assert_eq!(links.len(), 1, "expected exactly the tool-added link"); + assert_eq!(links[0].data.target, "blue"); + + // (b) The assistant reply carries the post-tool final answer, complete. + let msgs = store::thread_messages(&p, &thread_id).await; + let reply = assistant_reply(&msgs); + let expected_final = "The current temperature in San Francisco is 14 degrees Celsius \ + and the weather condition is foggy."; + assert_eq!(reply.content, expected_final); + assert_eq!(reply.status, "complete"); + + // (c) Message.toolCalls records the executed call. + assert!( + reply.tool_calls.contains("perspective_add_link"), + "toolCalls JSON should record the call, got: {}", + reply.tool_calls + ); + + // (d) A role:'tool' message was persisted for history/WE display. + assert!( + msgs.iter().any(|m| m.role == "tool"), + "a tool-result message should be persisted" + ); + + // (e) Durable RunState finished. + assert_eq!( + run_state_status(&p, &thread_id).await.as_deref(), + Some("done") + ); +} diff --git a/rust-executor/src/assistant_runtime/mod.rs b/rust-executor/src/assistant_runtime/mod.rs index c55542fe7..5785b5af5 100644 --- a/rust-executor/src/assistant_runtime/mod.rs +++ b/rust-executor/src/assistant_runtime/mod.rs @@ -22,14 +22,18 @@ pub mod context; pub mod entities; +pub mod model_backend; pub mod registry; pub mod run; pub mod sdna; pub mod store; pub mod tools; +#[cfg(test)] +mod e2e_tests; + use std::collections::{HashMap, HashSet}; -use std::sync::{Mutex, OnceLock}; +use std::sync::{Arc, Mutex, OnceLock}; use tokio::sync::broadcast::error::RecvError; @@ -39,6 +43,7 @@ use crate::pubsub::{get_global_pubsub, PERSPECTIVE_LINK_ADDED_TOPIC}; use crate::types::PerspectiveLinkWithOwner; use entities::{Message, CLASS_ASSISTANT, CLASS_MESSAGE, CLASS_THREAD}; +use model_backend::AiServiceBackend; use registry::RunRegistry; use run::RunInput; @@ -53,7 +58,7 @@ fn ensured_set() -> &'static Mutex> { /// Entry point — never returns. Bootstraps, then watches for new turns. pub async fn start() { log::info!("assistant_runtime: starting"); - let registry = RunRegistry::new(); + let registry = RunRegistry::new(Arc::new(AiServiceBackend)); // Give the executor a moment to finish loading perspectives before the // boot-time resume scan; the live watcher below catches everything after. diff --git a/rust-executor/src/assistant_runtime/model_backend.rs b/rust-executor/src/assistant_runtime/model_backend.rs new file mode 100644 index 000000000..a454106f5 --- /dev/null +++ b/rust-executor/src/assistant_runtime/model_backend.rs @@ -0,0 +1,64 @@ +//! The model seam. +//! +//! The run loop streams model output through [`ModelBackend`] rather than +//! calling [`AIService`] directly, so it can be driven by the real local/remote +//! model in production and by recorded fixtures in tests. The production impl +//! ([`AiServiceBackend`]) is a thin pass-through: it behaves exactly as a +//! direct `AIService::prompt_messages_stream` call did — same model, same +//! constrained decoding, same token stream — the only difference being that the +//! completion oneshot is drained on a detached task (the loop already treats +//! the token channel closing as completion, so this is behaviourally identical). + +use std::future::Future; +use std::pin::Pin; + +use anyhow::{anyhow, Result}; +use kalosm::language::ArcParser; +use tokio::sync::mpsc; + +use crate::ai_service::AIService; + +/// A receiver of output *tokens*. The stream is complete when the channel +/// closes — the same signal the loop already keys off. +pub type TokenStream = mpsc::UnboundedReceiver; + +/// Abstraction over the LLM used by the run loop. +pub trait ModelBackend: Send + Sync { + /// Start a streaming completion. Mirrors the token half of + /// `AIService::prompt_messages_stream`. + fn stream( + &self, + model_id: String, + messages: Vec<(String, String)>, + constraint: Option>, + ) -> Pin> + Send + '_>>; +} + +/// Production backend: the real local/remote model via the global [`AIService`]. +pub struct AiServiceBackend; + +impl ModelBackend for AiServiceBackend { + fn stream( + &self, + model_id: String, + messages: Vec<(String, String)>, + constraint: Option>, + ) -> Pin> + Send + '_>> { + Box::pin(async move { + let ai = AIService::global_instance() + .await + .map_err(|e| anyhow!("AI service unavailable: {e}"))?; + let (token_rx, done_rx) = ai + .prompt_messages_stream(model_id, messages, constraint) + .await + .map_err(|e| anyhow!("model stream failed: {e}"))?; + // The token channel closing already signals completion to the loop; + // drain the completion oneshot (token counts unused) on a detached + // task so lifecycle/cleanup runs without blocking the caller. + tokio::spawn(async move { + let _ = done_rx.await; + }); + Ok(token_rx) + }) + } +} diff --git a/rust-executor/src/assistant_runtime/registry.rs b/rust-executor/src/assistant_runtime/registry.rs index ce4d5ee9a..8255ded54 100644 --- a/rust-executor/src/assistant_runtime/registry.rs +++ b/rust-executor/src/assistant_runtime/registry.rs @@ -11,6 +11,7 @@ use std::sync::Arc; use tokio::sync::Mutex; use tokio::task::JoinHandle; +use super::model_backend::ModelBackend; use super::run::{self, RunInput}; struct RunHandle { @@ -18,15 +19,18 @@ struct RunHandle { } /// Cheaply-cloneable handle to the set of live runs, keyed by thread id. +/// Holds the production [`ModelBackend`] handed to each spawned run. #[derive(Clone)] pub struct RunRegistry { inner: Arc>>, + backend: Arc, } impl RunRegistry { - pub fn new() -> Self { + pub fn new(backend: Arc) -> Self { Self { inner: Arc::new(Mutex::new(HashMap::new())), + backend, } } @@ -45,9 +49,10 @@ impl RunRegistry { } let registry = self.clone(); + let backend = self.backend.clone(); let tid = thread_id.clone(); let handle = tokio::spawn(async move { - if let Err(e) = run::run_thread(input).await { + if let Err(e) = run::run_thread(input, backend).await { log::error!("assistant_runtime: run_thread error for {tid}: {e}"); } registry.remove(&tid).await; @@ -66,9 +71,3 @@ impl RunRegistry { self.inner.lock().await.len() } } - -impl Default for RunRegistry { - fn default() -> Self { - Self::new() - } -} diff --git a/rust-executor/src/assistant_runtime/run.rs b/rust-executor/src/assistant_runtime/run.rs index 75dadbceb..a1bf9c840 100644 --- a/rust-executor/src/assistant_runtime/run.rs +++ b/rust-executor/src/assistant_runtime/run.rs @@ -17,12 +17,13 @@ //! auto-publishes on `PERSPECTIVE_LINK_ADDED_TOPIC` so WE's live query //! re-renders — no separate token channel, matching WE. +use std::sync::Arc; + use anyhow::{anyhow, Result}; use kalosm::language::ArcParser; use crate::agent::AgentContext; -use crate::ai_service::AIService; use crate::api::openai_compat::tool_grammar::{build_tool_call_parser, extract_tool_calls, ToolChoice}; use super::context::assemble_messages; @@ -30,6 +31,7 @@ use super::entities::{ Assistant, McpServer, Message, Personality, Skill, Thread, CLASS_ASSISTANT, CLASS_MCP_SERVER, CLASS_THREAD, P_CONTENT, P_STATUS, P_TOOL_CALLS, }; +use super::model_backend::ModelBackend; use super::store; use super::tools::{BuiltinTools, McpToolProvider, ToolProvider, ToolSet}; @@ -45,10 +47,12 @@ pub struct RunInput { pub thread_id: String, } -/// Run one user turn for `input.thread_id`. Errors are handled internally +/// Run one user turn for `input.thread_id`, streaming model output through +/// `backend` (the real [`AiServiceBackend`](super::model_backend::AiServiceBackend) +/// in production; a fixture backend in tests). Errors are handled internally /// (assistant message + RunState marked `error`); the returned `Result` is for /// the registry's logging only. -pub async fn run_thread(input: RunInput) -> Result<()> { +pub async fn run_thread(input: RunInput, backend: Arc) -> Result<()> { let RunInput { perspective_uuid, thread_id, @@ -97,7 +101,17 @@ pub async fn run_thread(input: RunInput) -> Result<()> { let _ = store::upsert_run_state(&mut conv, &thread_id, "running", "0", "").await; - match drive_loop(&mut conv, &reply_id, &thread_id, &config, &user_turn, &history).await { + match drive_loop( + &mut conv, + &reply_id, + &thread_id, + &config, + &user_turn, + &history, + backend.as_ref(), + ) + .await + { Ok(()) => { let _ = store::upsert_run_state(&mut conv, &thread_id, "done", "final", "").await; Ok(()) @@ -213,6 +227,7 @@ async fn resolve_config( } /// The context → model → tools → repeat loop for one turn. +#[allow(clippy::too_many_arguments)] async fn drive_loop( conv: &mut crate::perspectives::perspective_instance::PerspectiveInstance, reply_id: &str, @@ -220,11 +235,9 @@ async fn drive_loop( config: &RunConfig, user_turn: &str, history: &[Message], + backend: &dyn ModelBackend, ) -> Result<()> { let ctx = AgentContext::main_agent(); - let ai = AIService::global_instance() - .await - .map_err(|e| anyhow!("AI service unavailable: {e}"))?; let toolset = ToolSet::new(vec![ ToolProvider::Builtin(BuiltinTools::new(conv.uuid.clone())), @@ -258,9 +271,16 @@ async fn drive_loop( build_tool_call_parser(&tool_defs, &ToolChoice::Auto, true); // Stream the model output, rewriting Message.content at a cadence. - let full = - stream_completion(conv, reply_id, &ctx, &ai, &config.model_id, messages, constraint) - .await?; + let full = stream_completion( + conv, + reply_id, + &ctx, + backend, + &config.model_id, + messages, + constraint, + ) + .await?; let calls = extract_tool_calls(&full); if calls.is_empty() { @@ -335,19 +355,19 @@ async fn drive_loop( /// Run one streaming completion, rewriting `Message.content` as tokens arrive, /// and return the full generated text. +#[allow(clippy::too_many_arguments)] async fn stream_completion( conv: &mut crate::perspectives::perspective_instance::PerspectiveInstance, reply_id: &str, ctx: &AgentContext, - ai: &AIService, + backend: &dyn ModelBackend, model_id: &str, messages: Vec<(String, String)>, constraint: Option>, ) -> Result { - let (mut token_rx, done_rx) = ai - .prompt_messages_stream(model_id.to_string(), messages, constraint) - .await - .map_err(|e| anyhow!("model stream failed: {e}"))?; + let mut token_rx = backend + .stream(model_id.to_string(), messages, constraint) + .await?; let mut buffer = String::new(); let mut since_update = 0usize; @@ -360,8 +380,6 @@ async fn stream_completion( let _ = store::set_single_target(conv, reply_id, P_CONTENT, &visible, ctx).await; } } - // Drain the completion signal (token counts unused here). - let _ = done_rx.await; Ok(buffer) } diff --git a/rust-executor/src/perspectives/mod.rs b/rust-executor/src/perspectives/mod.rs index 6efd0e2e4..93d9c2f1f 100644 --- a/rust-executor/src/perspectives/mod.rs +++ b/rust-executor/src/perspectives/mod.rs @@ -264,6 +264,19 @@ pub fn get_perspective(uuid: &str) -> Option { }) } +/// Register a pre-built [`PerspectiveInstance`] directly in the global registry +/// so `get_perspective`/`all_perspectives` resolve it — without the DB write, +/// background tasks, or pubsub of [`add_perspective`]. Test-only: used by the +/// assistant-runtime end-to-end tests to drive the real subsystem against an +/// in-memory perspective. +#[cfg(test)] +pub fn insert_perspective_for_test(instance: PerspectiveInstance) { + PERSPECTIVES + .write() + .expect("Couldn't get write lock on PERSPECTIVES") + .insert(instance.uuid.clone(), RwLock::new(instance)); +} + pub async fn update_perspective(handle: &PerspectiveHandle) -> Result<(), String> { { if PERSPECTIVES.read().unwrap().get(&handle.uuid).is_none() { diff --git a/rust-executor/tests/fixtures/chat_plain.json b/rust-executor/tests/fixtures/chat_plain.json new file mode 100644 index 000000000..f656d26ac --- /dev/null +++ b/rust-executor/tests/fixtures/chat_plain.json @@ -0,0 +1 @@ +{"id":"chatcmpl-86e8bad4-3b21-496d-8426-60f93b48febe","object":"chat.completion","created":1785239489,"model":"qwen2.5","choices":[{"index":0,"message":{"role":"assistant","content":"Hi there! I'm Qwen, an artificial intelligence developed by Alibaba Cloud. My main function is to assist with various tasks and provide information on a wide range of topics. How can I help you today?"},"finish_reason":"stop"}],"usage":{"prompt_tokens":5,"completion_tokens":50,"total_tokens":55}} \ No newline at end of file diff --git a/rust-executor/tests/fixtures/final_after_tool.json b/rust-executor/tests/fixtures/final_after_tool.json new file mode 100644 index 000000000..12a0748fb --- /dev/null +++ b/rust-executor/tests/fixtures/final_after_tool.json @@ -0,0 +1 @@ +{"id":"chatcmpl-f01933c5-d944-48b0-9755-1d5df5a36003","object":"chat.completion","created":1785239529,"model":"qwen2.5","choices":[{"index":0,"message":{"role":"assistant","content":"The current temperature in San Francisco is 14 degrees Celsius and the weather condition is foggy."},"finish_reason":"stop"}],"usage":{"prompt_tokens":21,"completion_tokens":25,"total_tokens":46}} \ No newline at end of file diff --git a/rust-executor/tests/fixtures/toolcall_builtin.json b/rust-executor/tests/fixtures/toolcall_builtin.json new file mode 100644 index 000000000..bff0f667b --- /dev/null +++ b/rust-executor/tests/fixtures/toolcall_builtin.json @@ -0,0 +1 @@ +{"id":"chatcmpl-059b020c-ee59-44f4-a6e2-5a1a15c79b42","object":"chat.completion","created":1785239521,"model":"qwen2.5","choices":[{"index":0,"message":{"role":"assistant","tool_calls":[{"id":"call_b2730c94-43db-414d-8256-ad21d8b559f4","type":"function","function":{"name":"perspective_add_link","arguments":"{\"source\":\"sky_color_blue\",\"predicate\":\"hasColor\",\"target\":\"blue\"}"}}]},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":18,"completion_tokens":51,"total_tokens":69}} \ No newline at end of file diff --git a/rust-executor/tests/fixtures/toolcall_weather.json b/rust-executor/tests/fixtures/toolcall_weather.json new file mode 100644 index 000000000..ee8e2f136 --- /dev/null +++ b/rust-executor/tests/fixtures/toolcall_weather.json @@ -0,0 +1 @@ +{"id":"chatcmpl-8e211116-20d6-407b-bb76-7f13be9891be","object":"chat.completion","created":1785239505,"model":"qwen2.5","choices":[{"index":0,"message":{"role":"assistant","tool_calls":[{"id":"call_59844f59-a71f-4ddb-a79a-1703efd32ff6","type":"function","function":{"name":"get_weather","arguments":"{\"location\":\"San Francisco\",\"unit\":\"celsius\"}"}}]},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":12,"completion_tokens":40,"total_tokens":52}} \ No newline at end of file