First Version

This commit is contained in:
2026-07-10 15:02:09 +01:00
commit 38494a85a9
562 changed files with 196313 additions and 0 deletions
+17
View File
@@ -0,0 +1,17 @@
[package]
name = "core-api"
version = "0.1.0"
edition = "2024"
[dependencies]
serde = { version = "1", features = ["derive"] }
serde_json = "1"
tokio = { version = "1", features = ["sync", "macros"] }
tokio-util = { version = "0.7" }
chrono = { version = "0.4", default-features = false, features = ["clock", "std"] }
async-trait = "0.1"
anyhow = "1"
axum = { version = "0.8", default-features = false }
# SqlitePool exposed to plugins via PluginContext (plugin.md §12.1). Version
# pinned to match the rest of the workspace.
sqlx = { version = "0.9.0", features = ["runtime-tokio", "sqlite"] }
+22
View File
@@ -0,0 +1,22 @@
use async_trait::async_trait;
/// Minimal approval API exposed to plugins.
///
/// Plugins receive `Arc<dyn ApprovalApi>` via `PluginContext` and use it to
/// resolve pending tool-call approvals without depending on the main crate's
/// `ApprovalManager` directly.
#[async_trait]
pub trait ApprovalApi: Send + Sync {
/// Approve a pending tool-call request.
async fn approve(&self, request_id: i64);
/// Reject a pending tool-call request with an optional note.
async fn reject(&self, request_id: i64, note: String);
/// Approve + register a session bypass so future tool calls of the same
/// category/MCP-server are skipped automatically.
///
/// - `bypass_secs = Some(n)`: bypass lasts `n` seconds (e.g. 900 = 15 min)
/// - `bypass_secs = None`: bypass lasts until the session ends
async fn approve_with_bypass(&self, request_id: i64, bypass_secs: Option<u64>);
}
+160
View File
@@ -0,0 +1,160 @@
//! In-process event bus for chat turns and system events.
//!
//! Every completed chat turn (user message + assistant response) is published
//! here **after** both messages have been persisted to the DB successfully.
//! Context compaction events are published when the compactor completes a
//! summarisation cycle.
//!
//! Subscribers receive a clone of each [`BusEvent`] via a
//! `tokio::sync::broadcast` channel.
//!
//! # Current consumers
//! - `honcho` plugin: processes `UserMessage` and `AssistantResponse` variants.
//!
//! # Usage
//! ```rust,ignore
//! // Producer (ChatSessionHandler):
//! bus.user_message(ChatEvent { ... });
//! bus.assistant_response(ChatEvent { ... });
//!
//! // Producer (ContextCompactor):
//! bus.compaction_done(CompactionEvent { ... });
//!
//! // Consumer (spawn once, keep running):
//! let mut rx = bus.subscribe();
//! tokio::spawn(async move {
//! loop {
//! match rx.recv().await {
//! Ok(BusEvent::UserMessage(e)) => { /* ... */ }
//! Ok(BusEvent::AssistantResponse(e)) => { /* ... */ }
//! Ok(BusEvent::CompactionDone(e)) => { /* ... */ }
//! Err(RecvError::Lagged(n)) => warn!("lagged by {n}"),
//! Err(RecvError::Closed) => break,
//! }
//! }
//! });
//! ```
use chrono::{DateTime, Utc};
use tokio::sync::broadcast;
pub use tokio::sync::broadcast::error::RecvError;
/// Default channel capacity. At 256 events the bus can absorb a burst of 128
/// turns before a slow consumer starts lagging.
const DEFAULT_CAPACITY: usize = 256;
// ── Sub-event types ───────────────────────────────────────────────────────────
#[derive(Debug, Clone)]
pub enum ChatEventRole {
User,
Assistant,
/// Sub-agent invocation (role = "agent" in DB).
Agent,
}
/// Per-tool-call detail attached to an assistant `ChatEvent`.
#[derive(Debug, Clone)]
pub struct ToolCallEvent {
pub name: String,
pub arguments: Option<String>,
pub result: Option<String>,
/// "done" | "failed"
pub status: String,
}
/// A single message in a completed chat turn.
#[derive(Debug, Clone)]
pub struct ChatEvent {
pub session_id: i64,
pub stack_id: i64,
/// `chat_history.id` for this message.
pub message_id: i64,
pub role: ChatEventRole,
pub content: String,
/// True for system-generated messages that look like user turns
/// (TicManager ticks, notification briefings).
pub is_synthetic: bool,
/// True when a real user is actively participating in the session
/// (web, telegram). False for automated sessions (cron, tic).
pub is_interactive: bool,
/// True for short-lived task sessions (cron, tic) that have no
/// long-term conversational value (e.g. skip Honcho memory sink).
pub is_ephemeral: bool,
/// Non-empty only for assistant messages that triggered tool calls.
pub tool_calls: Vec<ToolCallEvent>,
pub created_at: DateTime<Utc>,
}
/// Emitted after the compactor successfully persists a new summary.
#[derive(Debug, Clone)]
pub struct CompactionEvent {
pub session_id: i64,
pub stack_id: i64,
/// `chat_summaries.id` of the newly created summary row.
pub summary_id: i64,
/// All `chat_history` rows with `id <= covers_up_to_message_id` are now
/// covered by the summary and will no longer be sent to the LLM raw.
pub covers_up_to_message_id: i64,
/// Input token count of the turn that triggered this compaction.
pub triggered_by_tokens: u32,
}
// ── Top-level bus event ───────────────────────────────────────────────────────
/// All events that flow through the [`ChatEventBus`].
#[derive(Debug, Clone)]
pub enum BusEvent {
/// A user (or synthetic) message was saved after a completed turn.
UserMessage(ChatEvent),
/// The assistant's final response was saved after a completed turn.
AssistantResponse(ChatEvent),
/// The context compactor created a new summary for a stack.
CompactionDone(CompactionEvent),
}
// ── Bus ───────────────────────────────────────────────────────────────────────
pub struct ChatEventBus {
tx: broadcast::Sender<BusEvent>,
}
impl ChatEventBus {
pub fn new() -> Self {
Self::with_capacity(DEFAULT_CAPACITY)
}
pub fn with_capacity(capacity: usize) -> Self {
let (tx, _) = broadcast::channel(capacity);
Self { tx }
}
/// Publish a user message event. No-op if there are no active subscribers.
pub fn user_message(&self, event: ChatEvent) {
let _ = self.tx.send(BusEvent::UserMessage(event));
}
/// Publish an assistant response event. No-op if there are no active subscribers.
pub fn assistant_response(&self, event: ChatEvent) {
let _ = self.tx.send(BusEvent::AssistantResponse(event));
}
/// Publish a compaction-done event. No-op if there are no active subscribers.
pub fn compaction_done(&self, event: CompactionEvent) {
let _ = self.tx.send(BusEvent::CompactionDone(event));
}
/// Returns a new receiver. Each subscriber gets every future event
/// independently. If the subscriber falls behind by more than the channel
/// capacity it will receive `RecvError::Lagged(n)` — handle gracefully.
pub fn subscribe(&self) -> broadcast::Receiver<BusEvent> {
self.tx.subscribe()
}
}
impl Default for ChatEventBus {
fn default() -> Self {
Self::new()
}
}
+202
View File
@@ -0,0 +1,202 @@
use std::collections::HashMap;
use async_trait::async_trait;
use tokio::sync::broadcast;
use crate::events::GlobalEvent;
use crate::interface_tool::InterfaceTool;
use crate::message_meta::MessageMetadata;
// ── SendMessageOptions ────────────────────────────────────────────────────────
/// Optional parameters for a [`ChatHubApi::send_message`] call.
#[derive(Default)]
pub struct SendMessageOptions {
/// Agent to use for this source's session. Defaults to `"main"` if not set.
/// Only takes effect when a new session is created — ignored for existing sessions.
pub agent_id: Option<String>,
/// Named substitutions applied to the agent's system prompt.
/// Each entry replaces the sentinel `__KEY__` in the loaded prompt text.
/// Matches `<!-- KEY -->` placeholders that `agents::resolve_includes` converts to sentinels.
pub system_substitutions: HashMap<String, String>,
pub client_name: Option<String>,
/// Extra text prepended to the agent's system prompt for this turn only.
/// STATIC: safe to cache — use for interface-specific formatting rules that
/// never change turn-to-turn (e.g. Telegram HTML mode).
pub extra_system_context: Option<String>,
/// Extra system message injected AFTER the conversation history for this turn only.
/// DYNAMIC: not cached — use for per-turn context (e.g. notification framing).
pub extra_system_dynamic: Option<String>,
/// Short reminder injected near the tail of the message list to prevent drift.
pub tail_reminder: Option<String>,
pub interface_tools: Vec<InterfaceTool>,
/// True for system-generated messages injected as user turns (notification briefings).
pub is_synthetic: bool,
/// Opaque structured metadata persisted on the user turn (e.g. file attachments).
/// ChatHub forwards it verbatim; the MessageBuilder/UI derive their own views.
pub metadata: Option<MessageMetadata>,
}
// ── ChatHubApi ────────────────────────────────────────────────────────────────
/// Abstraction over [`ChatHub`](crate) that plugins and external crates depend on.
///
/// Implementing this trait for `ChatHub` in the main crate is the only coupling
/// point needed: plugins can accept `Arc<dyn ChatHubApi>` and stay independent.
#[async_trait]
pub trait ChatHubApi: Send + Sync {
/// Register a source. No-op for duplicate registrations.
async fn register(&self, source_id: &str);
/// Send a user message for a source, running a full LLM turn.
/// Creates a session lazily if none exists yet.
async fn send_message(
&self,
source_id: &str,
prompt: &str,
opts: SendMessageOptions,
) -> anyhow::Result<()>;
/// Create a new session for the source, discarding the previous one.
async fn clear(&self, source_id: &str) -> anyhow::Result<i64>;
/// Subscribe to the global event bus.
/// Filtering by source is the caller's responsibility.
fn events(&self, source_id: &str) -> broadcast::Receiver<GlobalEvent>;
/// Set which source is the "home" for background agent notifications.
async fn set_home(&self, source_id: &str) -> anyhow::Result<()>;
/// Returns token usage `(input, output)` for the last message in the source's session.
async fn context_info(
&self,
source_id: &str,
) -> anyhow::Result<(Option<i64>, Option<i64>)>;
/// Total spend (USD) of the source's active session, including synchronous
/// sub-agent frames and excluding asynchronous tasks (which run in their own
/// session). `None` when no provider reported a cost.
async fn cost_info(&self, source_id: &str) -> anyhow::Result<Option<f64>>;
/// Force compaction of the source's active session history.
/// Returns `true` if compaction occurred.
async fn force_compact(&self, source_id: &str) -> anyhow::Result<bool>;
/// Resume any interrupted turn for a source's active session.
async fn resume(&self, source_id: &str) -> anyhow::Result<()>;
/// Approve a pending tool-call approval request.
async fn approve(&self, request_id: i64);
/// Reject a pending tool-call approval request.
async fn reject(&self, request_id: i64, note: String);
/// Resolve a pending `ask_user_clarification` question.
/// Collapses `session_handler(source_id).resolve_question(...)` into a single
/// hub-level call so callers never need to know about `ChatSessionHandler`.
async fn resolve_question(&self, source_id: &str, request_id: i64, answer: String);
/// Cancel the active LLM turn for a source, clearing any pending approvals
/// and clarification questions. No-op if no session is active.
async fn cancel(&self, source_id: &str);
/// Revoke all session-scoped MCP grants for a source's active session.
/// The next LLM turn will start with no MCP tools activated.
async fn reset_mcp(&self, source_id: &str) -> anyhow::Result<()>;
/// Returns `(models, default)` where `models` is the ordered list of usable
/// LLM client names (first entry is always `"auto"`) and `default` is the
/// configured default client name. Used by `/models` and the web selector.
async fn list_clients(&self) -> (Vec<String>, String);
/// Returns the client name pinned for the source, or `None` when unset
/// (the caller should fall back to AUTO resolution).
async fn get_selected_client(&self, source_id: &str) -> Option<String>;
/// Pin a client name for the source and broadcast `ClientSelected` to every
/// client of the source. `client` must be a `list_clients()` entry (e.g.
/// `"auto"` or a model name); the caller is responsible for validation.
async fn set_selected_client(&self, source_id: &str, client: String);
/// Clear any pinned client for the source (revert to AUTO) and broadcast
/// `ClientSelected { client: "auto" }`.
async fn clear_selected_client(&self, source_id: &str);
/// Snapshot of the model list with the per-source current selection marked.
/// Returns `(index, name, is_current)` tuples — call sites format them as
/// HTML or Markdown without re-querying the LLM manager.
async fn list_clients_marked(
&self,
source_id: &str,
) -> Vec<(usize, String, bool)>;
/// Apply a `/model {arg}` command: resolve the argument, mutate the
/// per-source pinned client (broadcasting `ClientSelected`), return a
/// structured outcome the caller can format for its medium (HTML/Markdown).
async fn apply_model_command(
&self,
source_id: &str,
arg: &str,
) -> ModelCommandOutcome;
}
// ── Model command helpers (shared business logic) ────────────────────────────
/// Outcome of [`ChatHubApi::apply_model_command`]. The caller formats each
/// variant for its medium (Telegram HTML, web Markdown, …).
#[derive(Debug, Clone)]
pub enum ModelCommandOutcome {
/// A model was pinned. The backend has already broadcast `ClientSelected`.
Set(String),
/// The pin was cleared (back to AUTO). The backend has already broadcast
/// `ClientSelected { client: "auto" }`.
Cleared,
/// The argument was empty, out of range, or ambiguous. Carries a
/// user-facing message (no formatting — the caller wraps it as needed).
Error(String),
}
/// Resolve a `/model` (or analogous — e.g. a future `/reasoning`) argument
/// against an ordered list.
///
/// Returns:
/// - `Ok(Some(client))` for a unique match (caller pins it)
/// - `Ok(None)` for `auto` / index 0 (caller clears the pin)
/// - `Err(user_facing_message)` when the input is empty / out of range /
/// ambiguous
///
/// Accepts (in order):
/// 1. `auto` (case-insensitive) — or numeric `0` which is conventionally the
/// "auto" slot in `client_names()`
/// 2. Numeric index `N` → exact lookup
/// 3. Exact case-insensitive name match
/// 4. Substring match (case-insensitive) — must be unique
pub fn resolve_list_arg(models: &[String], arg: &str) -> Result<Option<String>, String> {
let arg = arg.trim();
if arg.is_empty() {
return Err("Usage: /model N or /model name or /model auto".to_string());
}
if arg.eq_ignore_ascii_case("auto") {
return Ok(None);
}
if let Ok(n) = arg.parse::<usize>() {
return match models.get(n) {
Some(m) if m == "auto" => Ok(None),
Some(m) => Ok(Some(m.clone())),
None => Err(format!("Index {n} out of range. Use /models to see the list.")),
};
}
if let Some(m) = models.iter().find(|m| m.eq_ignore_ascii_case(arg)) {
return Ok(if m == "auto" { None } else { Some(m.clone()) });
}
let lower = arg.to_ascii_lowercase();
let hits: Vec<&String> = models.iter().filter(|m| m.to_ascii_lowercase().contains(&lower)).collect();
match hits.len() {
1 => Ok(if hits[0] == "auto" { None } else { Some(hits[0].clone()) }),
0 => Err(format!("No model matches '{arg}'. Use /models to see the list.")),
_ => Err(format!(
"Multiple models match '{arg}': {}. Be more specific.",
hits.iter().map(|h| h.as_str()).collect::<Vec<_>>().join(", ")
)),
}
}
+156
View File
@@ -0,0 +1,156 @@
use async_trait::async_trait;
use serde_json::Value;
/// A single message in a conversation.
#[derive(Debug, Clone)]
pub struct Message {
pub role: Role,
pub content: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Role {
System,
User,
Assistant,
}
impl Message {
pub fn system(content: impl Into<String>) -> Self {
Self { role: Role::System, content: content.into() }
}
pub fn user(content: impl Into<String>) -> Self {
Self { role: Role::User, content: content.into() }
}
pub fn assistant(content: impl Into<String>) -> Self {
Self { role: Role::Assistant, content: content.into() }
}
}
/// Options for a single chat completion request.
#[derive(Debug, Clone)]
pub struct ChatOptions {
pub model: String,
pub max_tokens: Option<u32>,
pub temperature: Option<f32>,
/// Session/stack IDs for request logging. Set by the LLM loop; ignored by
/// providers — only the logging wrapper reads them.
pub session_id: Option<i64>,
pub stack_id: Option<i64>,
}
/// Raw HTTP metadata captured during a provider call.
/// Sensitive header values (api_key) are redacted before storage.
#[derive(Debug, Default)]
pub struct LlmRawMeta {
pub request_headers: Option<Value>,
pub request_body: Option<Value>,
pub response_headers: Option<Value>,
pub response_body: Option<Value>,
}
/// The response from a chat completion (text only).
#[derive(Debug, Clone)]
pub struct ChatResponse {
pub content: String,
pub input_tokens: Option<u32>,
pub output_tokens: Option<u32>,
/// True when the model stopped due to hitting the token limit.
pub truncated: bool,
/// Chain-of-thought produced by reasoning models (e.g. DeepSeek thinking mode).
/// Must be echoed back in the assistant message on subsequent turns.
pub reasoning_content: Option<String>,
/// Tokens served from the provider's prompt cache (Anthropic: cache_read_input_tokens,
/// OpenAI: prompt_tokens_details.cached_tokens). None when the provider does not
/// report cache metrics.
pub cache_read_tokens: Option<u32>,
/// Tokens written into the provider's prompt cache (Anthropic only:
/// cache_creation_input_tokens). None for providers that do not expose this.
pub cache_creation_tokens: Option<u32>,
/// Cost of the request in USD, when the provider reports it (OpenRouter
/// returns it under `usage.cost`). None for providers that do not bill
/// per-request or do not expose the figure.
pub cost: Option<f64>,
}
/// A single tool call requested by the LLM.
#[derive(Debug, Clone)]
pub struct ToolCall {
pub id: String,
pub name: String,
pub arguments: Value,
}
/// Result of one LLM turn when tools are available.
#[derive(Debug)]
pub enum LlmTurn {
Message(ChatResponse),
ToolCalls {
content: String,
calls: Vec<ToolCall>,
input_tokens: Option<u32>,
output_tokens: Option<u32>,
reasoning_content: Option<String>,
cache_read_tokens: Option<u32>,
cache_creation_tokens: Option<u32>,
cost: Option<f64>,
},
}
/// Stateless LLM client. Implementations hold only connection config (base URL,
/// API key). No memory, no database, no session state.
#[async_trait]
pub trait ChatbotClient: Send + Sync {
async fn chat(
&self,
messages: &[Message],
options: &ChatOptions,
) -> anyhow::Result<ChatResponse>;
/// Extracts the request cost in USD from a provider's raw JSON response,
/// when the provider reports it. OpenRouter (and other OpenAI-compatible
/// gateways) return it under `usage.cost`; the default reads that path and
/// yields None when absent. Providers with a different shape override this.
fn extract_cost(&self, response: &Value) -> Option<f64> {
response["usage"]["cost"].as_f64()
}
/// Chat with tool support. Default implementation ignores tools and falls
/// back to `chat()`.
async fn chat_with_tools(
&self,
messages: &[Value],
tools: &[Value],
options: &ChatOptions,
) -> anyhow::Result<LlmTurn> {
let simple: Vec<Message> = messages
.iter()
.filter_map(|m| {
let role = m["role"].as_str()?;
let content = m["content"].as_str().unwrap_or("").to_string();
match role {
"system" => Some(Message::system(content)),
"user" => Some(Message::user(content)),
"assistant" => Some(Message::assistant(content)),
_ => None,
}
})
.collect();
let _ = tools;
let resp = self.chat(&simple, options).await?;
Ok(LlmTurn::Message(resp))
}
/// Like `chat_with_tools` but also returns raw HTTP metadata for logging.
/// Providers that make real HTTP calls should override this.
async fn chat_with_tools_raw(
&self,
messages: &[Value],
tools: &[Value],
options: &ChatOptions,
) -> anyhow::Result<(LlmTurn, Option<LlmRawMeta>)> {
self.chat_with_tools(messages, tools, options).await.map(|t| (t, None))
}
}
+58
View File
@@ -0,0 +1,58 @@
//! Custom slash-command abstraction shared between the main crate and plugins.
//!
//! `LlmCommandManager` (main crate) implements [`CommandApi`]; plugins (e.g. the
//! Telegram bot) accept `Arc<dyn CommandApi>` and stay decoupled from the manager's
//! concrete type and from the file-system discovery logic. Commands live in
//! `commands/<name>/` and are read at request time (see `src/core/command/mod.rs`).
//!
//! This module is intentionally synchronous: command discovery is trivial file I/O
//! over tiny files with no DB/network dependency, so it does not need the
//! `async_trait` ceremony of the other plugin API traits.
use serde::{Deserialize, Serialize};
/// One enabled custom command — the listing DTO (no template body). Returned by
/// [`CommandApi::list_enabled`] for `/help` and the composer autocomplete.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct CommandInfo {
/// Canonical command name, without the leading `/` (e.g. `review`).
pub name: String,
/// One-line description shown in `/help` and the autocomplete dropdown.
pub description: String,
}
/// A resolved command ready for template expansion. Returned by
/// [`CommandApi::resolve`]. `name` is canonical (lowercase, no `/`); `template` is
/// the raw `COMMAND.md` body.
#[derive(Debug, Clone)]
pub struct ResolvedCommand {
pub name: String,
pub template: String,
}
/// Expand a command template by substituting the user's arguments. Replaces
/// `{{args}}` / `{{prompt}}` with `args`; if neither placeholder is present, the
/// arguments are appended after a `---` separator (opencode-like default). An empty
/// `args` with no placeholder leaves the template verbatim.
pub fn expand_template(template: &str, args: &str) -> String {
if template.contains("{{args}}") || template.contains("{{prompt}}") {
template.replace("{{args}}", args).replace("{{prompt}}", args)
} else if args.is_empty() {
template.to_string()
} else {
format!("{template}\n\n---\n\n{args}")
}
}
/// Abstraction over the command manager that plugins depend on. The main crate's
/// `LlmCommandManager` is the only implementation.
pub trait CommandApi: Send + Sync {
/// Every enabled, non-reserved command (metadata only — no template body),
/// sorted by name. Tolerant of a missing `commands/` directory (returns empty).
fn list_enabled(&self) -> Vec<CommandInfo>;
/// Resolve a command by name (case-insensitive), loading its template body.
/// Returns `None` when the command does not exist, is disabled, or is reserved
/// (collides with a hard-coded system command).
fn resolve(&self, name: &str) -> Option<ResolvedCommand>;
}
+29
View File
@@ -0,0 +1,29 @@
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum PropertyType {
String,
Int,
Bool,
SecurityGroup,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConfigProperty {
pub key: String,
pub name: String,
pub description: String,
pub property_type: PropertyType,
/// Value used when the key is absent from the DB config table.
pub default_value: Option<String>,
}
/// A named group of related [`ConfigProperty`] items, shown as a distinct
/// section in the Config UI.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConfigSet {
pub name: String,
pub description: String,
pub properties: Vec<ConfigProperty>,
}
+280
View File
@@ -0,0 +1,280 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::message_meta::Attachment;
// ── Client → Server ───────────────────────────────────────────────────────────
#[derive(Deserialize)]
pub struct ClientMessage {
pub content: String,
/// Files attached to this message (uploaded beforehand via `POST /api/{source}/uploads`).
#[serde(default)]
pub attachments: Vec<Attachment>,
}
/// Typed data push from remote clients (iOS app, etc.).
/// Sent over the existing WebSocket as `{"type":"data","stream":"...","payload":{...}}`.
#[derive(Deserialize)]
pub struct InboundDataMessage {
pub stream: String,
pub payload: Value,
}
// ── Global event envelope ─────────────────────────────────────────────────────
/// Envelope that wraps every event on the global broadcast bus.
/// `source` is `None` for system/background events (cron, tic, plugins).
#[derive(Clone)]
pub struct GlobalEvent {
pub source: Option<String>,
pub session_id: Option<i64>,
pub event: ServerEvent,
}
// ── Server → Client ───────────────────────────────────────────────────────────
#[derive(Clone, Serialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ServerEvent {
/// A tool call was started. DB status: running.
ToolStart {
tool_call_id: i64,
message_id: i64,
name: String,
arguments: Value,
/// Concise human-readable label (≤60 chars): tool + primary argument.
label_short: String,
/// Verbose human-readable label (≤120 chars): tool + all meaningful arguments.
label_full: String,
/// Path to a single viewable file this call targets, if any. The
/// frontend renders it as a clickable link to the file viewer.
#[serde(skip_serializing_if = "Option::is_none")]
path: Option<String>,
},
/// A tool call completed successfully. DB status: done.
ToolDone {
tool_call_id: i64,
result: String,
/// Result type tag: `"string"` (plain text) or `"json"` (structured
/// payload, e.g. MCP `structuredContent`). The frontend uses it to render
/// typed results instead of a raw text blob. Always populated by the
/// server; the frontend treats an absent/unknown value as plain text, so
/// older clients degrade gracefully.
result_type: String,
},
/// A tool call failed. DB status: error.
ToolError {
tool_call_id: i64,
error: String,
},
/// A tool call was stopped by the user via `/stop`. DB status: cancelled.
/// Distinct from `ToolError`: a cancellation is deliberate, not a failure.
ToolCancelled {
tool_call_id: i64,
},
/// A tool call was denied by an approval policy or a human. DB status: rejected.
/// Distinct from `ToolError`: a denial is a policy decision, not a failure.
ToolRejected {
tool_call_id: i64,
reason: String,
},
/// A sub-agent stack frame was opened.
AgentStart {
stack_id: i64,
parent_tool_call_id: i64,
agent_id: String,
parent_agent_id: String,
depth: i64,
/// The prompt sent to the sub-agent (truncated to 500 chars by the sender).
prompt_preview: String,
},
/// A sub-agent stack frame was closed.
AgentDone {
stack_id: i64,
agent_id: String,
parent_agent_id: String,
/// The sub-agent's final response (truncated to 500 chars by the sender).
result_preview: String,
},
/// The assistant response is complete.
Done {
message_id: i64,
stack_id: i64,
content: String,
input_tokens: Option<u32>,
output_tokens: Option<u32>,
},
/// A fatal error occurred processing the request.
Error {
message: String,
},
/// The LLM was cut off by the token limit (finish_reason="length").
Truncated {
output_tokens: Option<u32>,
},
/// The LLM produced text alongside tool calls (reasoning before acting).
Thinking {
message_id: i64,
content: String,
input_tokens: Option<u32>,
output_tokens: Option<u32>,
},
/// A write operation requires user approval before executing (shows a diff).
PendingWrite {
request_id: i64,
tool_call_id: i64,
path: String,
old_content: Option<String>,
new_content: String,
},
/// A non-file tool call requires user approval before executing.
/// Used for MCP tools, execute_cmd, restart, and any other tool
/// that the ApprovalManager flags as `Require`.
ApprovalRequired {
request_id: i64,
tool_call_id: i64,
tool_name: String,
arguments: Value,
},
/// A sub-agent needs clarification from the user before continuing.
AgentQuestion {
request_id: i64,
tool_call_id: i64,
title: String,
question: String,
suggested_answers: Vec<String>,
},
/// A book file was written by a tool; the frontend should reload if it has it open.
FileChanged {
path: String,
},
/// Ask the frontend to open a file for the user. Behaves like
/// `window.openFile(path)`: navigates to the file viewer page for markdown /
/// text / images, or opens an HTML file in a new browser tab. Emitted by
/// the future `show_file_to_user` interface tool (not wired yet).
OpenFile {
path: String,
},
/// The active LLM model failed and the system switched to a fallback automatically.
ModelFallback {
from: String,
to: String,
reason: String,
},
/// All LLM fallback attempts were exhausted; the turn could not complete.
LlmFailed {
tried: Vec<String>,
last_error: String,
},
/// A new approval entered the Inbox. Emitted on the global bus when the
/// `ApprovalManager` registers a pending request, so bus subscribers (e.g.
/// the mobile-connector plugin) can re-snapshot the Inbox. Distinct from
/// `ApprovalRequired`, which is the per-session WS event carrying full args
/// for the active client.
ApprovalRequested {
request_id: i64,
tool_call_id: i64,
tool_name: String,
},
/// A pending approval or pending-write was resolved (approved or rejected).
/// Emitted on the global bus so all clients (e.g. Telegram) can update their UI.
ApprovalResolved {
request_id: i64,
tool_call_id: i64,
approved: bool,
},
/// A new clarification entered the Inbox. Emitted on the global bus when the
/// `ClarificationManager` registers a pending question. Distinct from
/// `AgentQuestion`, which is the per-session WS event for the active client.
ClarificationRequested {
request_id: i64,
title: String,
},
/// A pending clarification was resolved (answered). Emitted on the global bus
/// so all clients can update their Inbox view.
ClarificationResolved {
request_id: i64,
},
/// A server-initiated MCP elicitation entered the Inbox (e.g. an MCP server
/// asking for a sudo password mid tool-call). Carries only `request_id` +
/// `title` — never the requested value.
ElicitationRequested {
request_id: i64,
title: String,
},
/// A pending elicitation was resolved (accepted / declined / cancelled).
/// Emitted on the global bus so all clients can update their Inbox view.
ElicitationResolved {
request_id: i64,
},
/// The active session for a source was replaced (e.g. /new, /clear).
NewSession {
session_id: i64,
},
/// A user message was persisted to history; broadcast so every client (the
/// sender included) renders the bubble. Emitted at save time — when the row
/// is appended, either at turn start or at a round boundary for messages
/// injected mid-turn — so the bubble appears exactly where the agent saw it.
/// Clients render purely from this echo (no optimistic local rendering).
UserMessage {
/// Id of the `chat_history` row just written.
message_id: i64,
content: String,
/// Files attached to the message; lets secondary clients render chips live.
#[serde(default, skip_serializing_if = "Vec::is_empty")]
attachments: Vec<Attachment>,
},
/// Sent to a client right after it (re)connects, reporting whether a turn is
/// currently in flight for its session. Lets a reloaded page restore the
/// SEND→STOP button state instead of assuming idle.
TurnRunning {
running: bool,
},
/// The selected LLM client (model) for a source changed. Broadcast to every
/// client of the source so dropdowns/selects stay in sync. `client` is a
/// `client_names()` entry (typically `"auto"` or a model name). Driven by
/// `ChatHub::set_selected_client` — the backend is the single source of truth.
ClientSelected {
client: String,
},
}
impl ServerEvent {
pub fn to_json(&self) -> String {
serde_json::to_string(self).expect("ServerEvent serialization failed")
}
pub fn type_name(&self) -> &'static str {
match self {
Self::ToolStart { .. } => "tool_start",
Self::ToolDone { .. } => "tool_done",
Self::ToolError { .. } => "tool_error",
Self::ToolCancelled { .. } => "tool_cancelled",
Self::ToolRejected { .. } => "tool_rejected",
Self::AgentStart { .. } => "agent_start",
Self::AgentDone { .. } => "agent_done",
Self::Done { .. } => "done",
Self::Error { .. } => "error",
Self::Thinking { .. } => "thinking",
Self::PendingWrite { .. } => "pending_write",
Self::ApprovalRequired { .. } => "approval_required",
Self::AgentQuestion { .. } => "agent_question",
Self::FileChanged { .. } => "file_changed",
Self::OpenFile { .. } => "open_file",
Self::Truncated { .. } => "truncated",
Self::ModelFallback { .. } => "model_fallback",
Self::LlmFailed { .. } => "llm_failed",
Self::ApprovalRequested { .. } => "approval_requested",
Self::ApprovalResolved { .. } => "approval_resolved",
Self::ClarificationRequested { .. } => "clarification_requested",
Self::ClarificationResolved { .. } => "clarification_resolved",
Self::ElicitationRequested { .. } => "elicitation_requested",
Self::ElicitationResolved { .. } => "elicitation_resolved",
Self::NewSession { .. } => "new_session",
Self::UserMessage { .. } => "user_message",
Self::TurnRunning { .. } => "turn_running",
Self::ClientSelected { .. } => "client_selected",
}
}
}
+49
View File
@@ -0,0 +1,49 @@
use std::sync::Arc;
use anyhow::Result;
use async_trait::async_trait;
use serde_json::Value;
// ── Record types (DB ↔ manager) ───────────────────────────────────────────────
/// Full model record, mirroring one row in `image_generate_models`.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ImageGenerateModelRecord {
pub id: i64,
pub provider_id: i64,
pub model_id: String,
/// Display alias (also used as the generator `id()`).
pub name: String,
/// Lower number = tried first by `get()`.
pub priority: i32,
}
/// Implemented by any provider that can generate an image from a text prompt.
/// Returns raw image bytes (PNG expected).
#[async_trait]
pub trait ImageGenerate: Send + Sync {
/// A stable, unique identifier for this provider (e.g. `"comfyui-portrait"`).
fn id(&self) -> &str;
/// Human-readable display name.
fn name(&self) -> &str;
/// Human-readable description shown to the LLM in `image_generate_providers_list`.
/// Should mention format, style, default dimensions, and ideal use cases.
fn description(&self) -> Option<&str> { None }
/// JSON Schema for the `extra_params` argument accepted by this provider.
/// Returned as-is in `image_generate_providers_list` so the LLM knows what to pass.
fn extra_params_schema(&self) -> Option<Value> { None }
/// Generate an image. `extra_params` is the provider-specific JSON object
/// passed by the LLM (validated against `extra_params_schema`).
async fn generate(&self, prompt: &str, extra_params: Option<&Value>) -> Result<Vec<u8>>;
}
/// Write-side of the image generator manager: register and remove ephemeral providers.
///
/// Implemented by `ImageGeneratorManager` in the main crate. Plugins that provide
/// their own image generation (e.g. a ComfyUI plugin) use this to register providers
/// at start and unregister them at stop.
#[async_trait]
pub trait ImageGenerateRegistry: Send + Sync {
async fn register(&self, provider: Arc<dyn ImageGenerate>);
async fn unregister(&self, id: &str);
}
+92
View File
@@ -0,0 +1,92 @@
//! Inbox API exposed to plugins.
//!
//! The Skald `Inbox` façade (src/core/inbox.rs) wraps the `ApprovalManager` and
//! `ClarificationManager`. Plugins receive `Arc<dyn InboxApi>` via `PluginContext`
//! and use it to read pending items and resolve them, without depending on the
//! main crate. `request_id` is the integer rowid used both for resolution and
//! for idempotency (plugin.md §2): `approve`/`reject`/`answer` on an already
//! resolved id are no-ops.
use async_trait::async_trait;
use serde::Serialize;
use serde_json::Value;
/// One pending approval, surfaced to plugins (mirrors the main crate's
/// `PendingApprovalInfo`, trimmed to the fields a plugin needs for the
/// `inbox_update` payload — see payloads.md §3.1).
#[derive(Debug, Clone, Serialize)]
pub struct InboxApprovalItem {
pub request_id: i64,
pub tool_name: String,
/// Human-readable one-line label for the tool call, from `Tool::describe(Short)`.
/// Used for the card / push notification text.
pub summary: String,
/// Raw tool arguments (untruncated). Source of truth for the detail dialog so
/// the user sees exactly what is being approved — critical for `execute_cmd`.
pub arguments: Value,
pub agent_id: String,
pub source: String,
pub context_label: Option<String>,
/// ISO-8601 timestamp string (UTC).
pub created_at: String,
}
/// One pending clarification, surfaced to plugins (mirrors the main crate's
/// `PendingClarificationInfo`).
#[derive(Debug, Clone, Serialize)]
pub struct InboxClarificationItem {
pub request_id: i64,
pub agent_id: String,
pub source: String,
pub context_label: Option<String>,
pub title: String,
pub question: String,
pub suggested_answers: Vec<String>,
/// ISO-8601 timestamp string (UTC).
pub created_at: String,
}
/// One pending elicitation, surfaced to plugins (mirrors the main crate's
/// `PendingElicitationInfo`). Holds only prompt metadata — never the value.
#[derive(Debug, Clone, Serialize)]
pub struct InboxElicitationItem {
pub request_id: i64,
pub server_name: String,
pub message: String,
pub field_name: Option<String>,
pub sensitive: bool,
pub is_confirmation: bool,
/// ISO-8601 timestamp string (UTC).
pub created_at: String,
}
/// A snapshot of all pending Inbox items.
#[derive(Debug, Clone, Serialize)]
pub struct InboxSnapshot {
pub total: usize,
pub approvals: Vec<InboxApprovalItem>,
pub clarifications: Vec<InboxClarificationItem>,
pub elicitations: Vec<InboxElicitationItem>,
}
/// Inbox operations available to plugins.
#[async_trait]
pub trait InboxApi: Send + Sync {
/// Snapshot of all currently pending approvals + clarifications.
async fn list_pending(&self) -> InboxSnapshot;
/// Approve a pending tool-call request. No-op if already resolved.
async fn approve(&self, request_id: i64);
/// Reject a pending tool-call request with a reason. No-op if already resolved.
async fn reject(&self, request_id: i64, reason: String);
/// Answer a pending clarification. Returns `true` if a pending entry was
/// found and resolved, `false` otherwise (idempotent).
async fn answer(&self, request_id: i64, answer: String) -> bool;
/// Resolve a pending MCP elicitation. `action` is `"accept"`/`"decline"`/
/// `"cancel"`; `content` carries the field values for `accept` (it may hold
/// a secret — do not log it). Returns `true` if a pending entry was found.
async fn resolve_elicitation(&self, request_id: i64, action: String, content: Option<Value>) -> bool;
}
+24
View File
@@ -0,0 +1,24 @@
use std::pin::Pin;
use std::sync::Arc;
use serde_json::Value;
/// Future returned by an [`InterfaceTool`] handler.
pub type ToolFuture = Pin<Box<dyn std::future::Future<Output = anyhow::Result<String>> + Send>>;
/// A single LLM-callable tool injected by a specific interface (Telegram, Web, Cron, …).
///
/// The handler closure captures interface-specific state (e.g. `Arc<Bot>` + `ChatId`).
pub struct InterfaceTool {
/// OpenAI-format tool definition sent to the LLM in the tools array.
pub definition: Value,
/// Async handler invoked when the LLM calls this tool.
pub handler: Arc<dyn Fn(Value) -> ToolFuture + Send + Sync>,
}
impl InterfaceTool {
/// Returns the tool's name as declared in the definition.
pub fn name(&self) -> &str {
self.definition["function"]["name"].as_str().unwrap_or("")
}
}
+25
View File
@@ -0,0 +1,25 @@
/// Application name, sent as `X-Title` HTTP header to LLM/image/audio providers.
pub const APP_NAME: &str = "Skald";
pub mod approval;
pub mod bus;
pub mod system_bus;
pub mod chatbot;
pub mod chat_hub;
pub mod command;
pub mod events;
pub mod image_generate;
pub mod inbox;
pub mod interface_tool;
pub mod location;
pub mod memory;
pub mod message_meta;
pub mod plugin;
pub mod provider;
pub mod remote;
pub mod tool;
pub mod secrets;
pub mod transcribe;
pub mod tts;
pub mod config_property;
pub use config_property::{ConfigProperty, ConfigSet, PropertyType};
+69
View File
@@ -0,0 +1,69 @@
use std::collections::HashMap;
use std::sync::RwLock;
use chrono::{DateTime, Utc};
#[derive(Clone, Debug)]
pub struct GpsCoord {
pub latitude: f64,
pub longitude: f64,
}
#[derive(Clone, Debug)]
pub struct LocationEntry {
pub coord: GpsCoord,
pub accuracy: Option<f64>,
pub updated_at: DateTime<Utc>,
pub is_live: bool,
}
pub struct LocationManager {
locations: RwLock<HashMap<String, LocationEntry>>,
}
impl LocationManager {
pub fn new() -> Self {
Self { locations: RwLock::new(HashMap::new()) }
}
pub fn update(&self, source: &str, coord: GpsCoord, accuracy: Option<f64>, is_live: bool) {
let entry = LocationEntry { coord, accuracy, updated_at: Utc::now(), is_live };
self.locations.write().unwrap().insert(source.to_string(), entry);
}
pub fn get(&self, source: &str) -> Option<LocationEntry> {
self.locations.read().unwrap().get(source).cloned()
}
pub fn latest(&self) -> Option<(String, LocationEntry)> {
self.locations.read().unwrap()
.iter()
.max_by_key(|(_, e)| e.updated_at)
.map(|(k, v)| (k.clone(), v.clone()))
}
pub fn all(&self) -> Vec<(String, LocationEntry)> {
self.locations.read().unwrap()
.iter()
.map(|(k, v)| (k.clone(), v.clone()))
.collect()
}
}
impl Default for LocationManager {
fn default() -> Self { Self::new() }
}
/// Write-only view of a location store.
///
/// Plugins store `Arc<dyn LocationUpdater>` so they can report GPS fixes
/// without depending on the concrete `LocationManager` struct.
pub trait LocationUpdater: Send + Sync {
fn update(&self, source: &str, coord: GpsCoord, accuracy: Option<f64>, is_live: bool);
}
impl LocationUpdater for LocationManager {
fn update(&self, source: &str, coord: GpsCoord, accuracy: Option<f64>, is_live: bool) {
LocationManager::update(self, source, coord, accuracy, is_live);
}
}
+29
View File
@@ -0,0 +1,29 @@
use std::sync::Arc;
use async_trait::async_trait;
use crate::tool::Tool;
/// Pluggable long-term memory backend.
///
/// Implementations are registered with `MemoryManager` in the main crate.
/// At most one backend is active at a time (singleton rule enforced by the manager).
#[async_trait]
pub trait Memory: Send + Sync {
/// Unique identifier for this backend (e.g. `"honcho"`).
fn id(&self) -> &str;
/// Returns `true` when the backend is reachable and ready.
fn is_available(&self) -> bool;
/// Retrieves context for the upcoming turn to inject into the system prompt.
/// Returns `None` on cold start, backend down, or nothing useful available.
async fn query_context(&self, session_id: i64, user_message: &str) -> Option<String>;
/// Optional LLM-callable tools exposed by this backend (e.g. `memory_query`).
/// Called per turn — added to the live tool list and dispatched before the
/// global tool registry.
fn tools(&self) -> Vec<Arc<dyn Tool>> {
vec![]
}
}
+79
View File
@@ -0,0 +1,79 @@
//! Structured, reusable metadata attached to a `chat_history` row.
//!
//! Persisted as a single JSON column (`chat_history.metadata`) and intentionally
//! generic: today it carries user file **attachments**, but new keys can be added
//! later without a schema change. Two independent readers derive different views
//! from the same source:
//! - the **LLM context** builder appends [`attachments_block`] to the user turn,
//! - the **history UI** renders the structured attachments as chips.
//!
//! The raw `[SYSTEM INFO]` text block is therefore never persisted — it is
//! generated on the fly from this metadata.
use serde::{Deserialize, Serialize};
/// One file attached by the user to a message. `path` is relative to the project
/// root (e.g. `data/uploads/123/file.pdf`) so it is both servable under `/data/…`
/// and resolvable by the filesystem tools.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Attachment {
pub path: String,
pub name: String,
/// Best-effort MIME type (e.g. `application/pdf`); `None` if unknown.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub mimetype: Option<String>,
/// Size in bytes, when known.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub filesize: Option<u64>,
}
/// Generic metadata bag for a chat message. Extra keys may be added over time;
/// `#[serde(default)]` keeps deserialization tolerant of older/newer shapes.
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct MessageMetadata {
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub attachments: Vec<Attachment>,
/// Present when this user turn was produced by a custom slash command.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub command: Option<CommandRef>,
}
impl MessageMetadata {
/// True when there is nothing worth persisting.
pub fn is_empty(&self) -> bool {
self.attachments.is_empty() && self.command.is_none()
}
}
/// Identifies a user message produced by a custom slash command. The history row's
/// `content` holds the **expanded template** (replayed to the LLM verbatim); the UI
/// renders `display` — the original `/command …` the user typed — instead.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct CommandRef {
/// Canonical command name, without the leading `/` (e.g. `review`).
pub name: String,
/// The original text the user typed (e.g. `/review revisiona Cat.java`).
pub display: String,
}
/// Renders the human-readable block appended to a user turn so the LLM learns
/// which files were attached. Returns an empty string when there are none, so
/// callers can unconditionally concatenate it.
///
/// Shared by the web/mobile path and the Telegram plugin so every surface emits
/// an identical format.
pub fn attachments_block(attachments: &[Attachment]) -> String {
if attachments.is_empty() {
return String::new();
}
let noun = if attachments.len() == 1 { "file" } else { "files" };
let mut block = format!(
"\n\n[SYSTEM INFO]\n{} attached {}:",
attachments.len(),
noun
);
for a in attachments {
block.push_str(&format!("\n* {}", a.path));
}
block
}
+95
View File
@@ -0,0 +1,95 @@
use std::sync::Arc;
use anyhow::Result;
use async_trait::async_trait;
use serde_json::Value;
use tokio::sync::RwLock;
use crate::approval::ApprovalApi;
use crate::bus::ChatEventBus;
use crate::command::CommandApi;
use crate::system_bus::SystemEventBus;
use crate::chat_hub::ChatHubApi;
use crate::image_generate::ImageGenerateRegistry;
use crate::inbox::InboxApi;
use crate::location::LocationUpdater;
use crate::memory::Memory;
use crate::provider::ApiProviderRegistry;
use crate::remote::RemoteAccess;
use crate::secrets::SecretsApi;
use crate::transcribe::{TranscribeProvider, TranscribeRegistry};
use crate::tts::{TtsProvider, TtsRegistry};
/// Closure that builds a fresh Axum router (e.g. for the mesh-facing server).
pub type RouterFactory = Arc<dyn Fn() -> axum::Router + Send + Sync>;
/// All deps a plugin may need — passed to [`Plugin::start`] and [`Plugin::reload`].
///
/// Fields are `Arc<dyn Trait>` sourced from `core-api`. Plugins use only the
/// fields relevant to them; unused fields are ignored.
/// `router_factory` and `remote_slot` are networking-specific — used only by
/// `RemotePlugin`.
#[derive(Clone)]
pub struct PluginContext {
pub chat_hub: Arc<dyn ChatHubApi>,
/// Custom file-based slash commands (`commands/<name>/`). Read-only from the
/// plugin side — lets the Telegram bot resolve `/command` expansions.
pub command: Arc<dyn CommandApi>,
pub approval: Arc<dyn ApprovalApi>,
/// Unified Inbox façade (approvals + clarifications). See plugin.md §12.2.
pub inbox: Arc<dyn InboxApi>,
/// Skald's shared SQLite pool — lets plugins create/use their own tables
/// (e.g. `relay_*`) in the main DB. See plugin.md §12.1.
pub db: Arc<sqlx::SqlitePool>,
pub secrets: Arc<dyn SecretsApi>,
pub transcribe: Arc<dyn TranscribeProvider>,
pub transcribe_registry: Arc<dyn TranscribeRegistry>,
pub image_generate_registry: Arc<dyn ImageGenerateRegistry>,
pub tts_registry: Arc<dyn TtsRegistry>,
pub tts_provider: Arc<dyn TtsProvider>,
pub api_provider_registry: Arc<dyn ApiProviderRegistry>,
pub location: Arc<dyn LocationUpdater>,
pub event_bus: Arc<ChatEventBus>,
pub system_bus: Arc<SystemEventBus>,
pub web_port: u16,
pub remote_slot: Arc<RwLock<Option<Arc<dyn RemoteAccess>>>>,
pub router_factory: RouterFactory,
}
/// Plugin lifecycle contract.
///
/// Each plugin implements this trait. The `PluginManager` in the main crate
/// manages their lifecycle and passes a `PluginContext` on every start/reload.
#[async_trait]
pub trait Plugin: Send + Sync {
fn id(&self) -> &str;
fn name(&self) -> &str;
fn description(&self) -> &str;
fn is_running(&self) -> bool;
/// JSON Schema describing the plugin's config fields.
fn config_schema(&self) -> Value { serde_json::json!({}) }
/// Called whenever the enabled flag or config changes — including at startup.
/// The plugin is responsible for diffing state and restarting only what changed.
async fn reload(&self, enabled: bool, config: Value, ctx: PluginContext) -> Result<()>;
async fn start(&self, ctx: PluginContext) -> Result<()>;
async fn stop(&self) -> Result<()>;
/// Runtime state surfaced to the UI and to agents (e.g. mesh IP).
fn runtime_status(&self) -> Option<Value> { None }
/// Optional Axum router contributed by the plugin. When `Some`, the main
/// `WebFrontend` nests it under `/api/plugin/<id>/` behind Skald's normal
/// auth (plugin.md §12.3). The router must close over the plugin's own state
/// (it receives no `State`). Default: no routes — existing plugins are
/// unaffected.
fn http_router(&self) -> Option<axum::Router> { None }
/// Returns a [`Memory`] backend if this plugin provides one.
fn memory(&self) -> Option<Arc<dyn Memory>> { None }
fn as_any(&self) -> &dyn std::any::Any;
fn as_arc_any(self: Arc<Self>) -> Arc<dyn std::any::Any + Send + Sync>;
}
+251
View File
@@ -0,0 +1,251 @@
use std::sync::Arc;
use anyhow::Result;
use async_trait::async_trait;
use crate::chatbot::ChatbotClient;
use crate::image_generate::{ImageGenerate, ImageGenerateModelRecord};
use crate::tts::{TextToSpeech, TtsModelRecord, RemoteTtsModelInfo};
use crate::transcribe::{Transcribe, TranscribeModelRecord, RemoteTranscribeModelInfo};
// ── LlmStrength ───────────────────────────────────────────────────────────────
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Deserialize, serde::Serialize)]
#[serde(rename_all = "snake_case")]
pub enum LlmStrength {
VeryLow,
Low,
Average,
High,
VeryHigh,
}
// ── Provider record types (DB ↔ manager) ──────────────────────────────────────
/// Full provider record (includes secrets — only expose over trusted local connections).
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct LlmProviderRecord {
pub id: i64,
pub name: String,
/// Provider type_id string as stored in DB (e.g. "open_ai", "anthropic").
#[serde(rename = "type")]
pub provider: String,
pub api_key: Option<String>,
/// Only used by ollama and lm_studio.
pub base_url: Option<String>,
pub description: Option<String>,
}
/// Full model record.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct LlmModelRecord {
pub id: i64,
pub provider_id: i64,
pub model_id: String,
pub name: String,
pub strength: Option<LlmStrength>,
pub scope: Vec<String>,
pub is_default: bool,
pub priority: i32,
pub extra_params: Option<serde_json::Value>,
pub context_length: Option<i64>,
pub max_output_tokens: Option<i64>,
pub knowledge_cutoff: Option<String>,
pub capabilities: Vec<String>,
/// Selected reasoning value (interpreted per provider via `ReasoningMode`):
/// a JSON string for a `ValueSet` (e.g. `"high"`, `"enabled"`) or a JSON
/// number for a `Range` (e.g. `8000`). `None` = reasoning off / unset.
pub reasoning: Option<serde_json::Value>,
}
/// Remote model info returned by a provider's `list_llm_models()`.
#[derive(Debug, Clone, serde::Serialize)]
pub struct RemoteLlmModelInfo {
pub id: String,
pub name: String,
pub context_length: Option<u64>,
pub max_completion_tokens: Option<u64>,
pub knowledge_cutoff: Option<String>,
pub capabilities: Vec<String>,
pub vision: Option<bool>,
pub price_input_per_million: Option<f64>,
pub price_output_per_million: Option<f64>,
/// Reasoning control descriptor for this model, if it supports reasoning.
/// Filled by the API layer from `ApiProvider::reasoning_mode`; providers'
/// `list_llm_models` may leave it `None`.
#[serde(default)]
pub reasoning: Option<ReasoningMode>,
}
// ── ReasoningMode ─────────────────────────────────────────────────────────────
/// Describes how a model exposes its "reasoning"/"thinking" knob, so the UI can
/// render the right control and the caller knows the space of valid values.
/// Provider-owned (each `ApiProvider` returns it per model); the actual request
/// translation is done by `ApiProvider::reasoning_request`.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ReasoningMode {
/// A fixed set of discrete string choices — e.g. `["low","medium","high"]`
/// (effort levels) or `["disabled","enabled"]` (on/off toggle).
ValueSet {
values: Vec<String>,
default: Option<String>,
},
/// A continuous numeric range — e.g. Anthropic thinking `budget_tokens`.
Range {
min: i64,
max: i64,
step: Option<i64>,
default: Option<i64>,
unit: Option<String>,
},
}
// ── ServiceType ───────────────────────────────────────────────────────────────
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize)]
#[serde(rename_all = "snake_case")]
pub enum ServiceType {
Llm,
Transcribe,
ImageGenerate,
Tts,
}
// ── UI metadata ───────────────────────────────────────────────────────────────
#[derive(Debug, Clone, serde::Serialize)]
pub struct ProviderUiMeta {
pub type_id: &'static str,
pub display_name: &'static str,
pub description: Option<&'static str>,
pub color: &'static str,
pub icon: &'static str,
pub fields: &'static [ProviderField],
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct ProviderField {
pub key: &'static str,
pub label: &'static str,
pub required: bool,
pub secret: bool,
}
// ── BuiltLlmClient ────────────────────────────────────────────────────────────
pub struct BuiltLlmClient {
pub client: Arc<dyn ChatbotClient>,
pub prompt_cache: bool,
}
// ── ApiProvider trait ─────────────────────────────────────────────────────────
#[async_trait]
pub trait ApiProvider: Send + Sync {
/// Short stable identifier stored in the DB (e.g. "open_ai", "anthropic").
fn type_id(&self) -> &'static str;
fn display_name(&self) -> &'static str;
fn supported_types(&self) -> &'static [ServiceType];
async fn list_llm_models(
&self,
_record: &LlmProviderRecord,
) -> Result<Option<Vec<RemoteLlmModelInfo>>> {
Ok(None)
}
/// Reasoning control descriptor for a given model, or `None` if the model
/// does not support reasoning. Drives the UI control (a `ValueSet` dropdown
/// or a numeric `Range`). Provider-owned so each provider decides which of
/// its models support reasoning and in what form.
fn reasoning_mode(
&self,
_model_id: &str,
_capabilities: &[String],
) -> Option<ReasoningMode> {
None
}
/// Translates a selected reasoning `value` (see `LlmModelRecord::reasoning`)
/// into the provider-specific JSON fragment to merge into the request body
/// (e.g. `{"reasoning":{"effort":"high"}}` or `{"thinking":{"type":"enabled"}}`).
/// `None` = nothing to inject.
fn reasoning_request(
&self,
_value: &serde_json::Value,
) -> Option<serde_json::Value> {
None
}
async fn llm_model_info(
&self,
_record: &LlmProviderRecord,
_model_id: &str,
) -> Result<Option<RemoteLlmModelInfo>> {
Ok(None)
}
async fn list_tts_models(
&self,
_record: &LlmProviderRecord,
) -> Result<Option<Vec<RemoteTtsModelInfo>>> {
Ok(None)
}
async fn list_transcribe_models(
&self,
_record: &LlmProviderRecord,
) -> Result<Option<Vec<RemoteTranscribeModelInfo>>> {
Ok(None)
}
fn build_llm(
&self,
_record: &LlmProviderRecord,
_model: &LlmModelRecord,
) -> Option<Result<BuiltLlmClient>> {
None
}
fn build_tts(
&self,
_record: &LlmProviderRecord,
_model: &TtsModelRecord,
) -> Option<Result<Arc<dyn TextToSpeech>>> {
None
}
fn build_transcriber(
&self,
_record: &LlmProviderRecord,
_model: &TranscribeModelRecord,
) -> Option<Result<Arc<dyn Transcribe>>> {
None
}
fn build_image_generator(
&self,
_record: &LlmProviderRecord,
_model: &ImageGenerateModelRecord,
) -> Option<Result<Arc<dyn ImageGenerate>>> {
None
}
fn ui_meta(&self) -> ProviderUiMeta;
}
// ── ApiProviderRegistry trait ─────────────────────────────────────────────────
/// Write-side of the provider registry: register and remove plugin-provided
/// `ApiProvider` implementations at runtime.
///
/// Implemented by `ProviderRegistry` in the main crate. Plugins that supply
/// their own API provider (e.g. `plugin-elevenlabs`) use this to register at
/// start and unregister at stop.
#[async_trait]
pub trait ApiProviderRegistry: Send + Sync {
fn register_plugin(&self, provider: Arc<dyn ApiProvider>);
fn unregister_plugin(&self, type_id: &str);
}
+22
View File
@@ -0,0 +1,22 @@
use std::net::Ipv4Addr;
use anyhow::Result;
use async_trait::async_trait;
/// Abstraction over a mesh/remote-connectivity provider.
/// The core (AppState, plugins) talks only to this trait —
/// no Tailscale-specific types leak outside the implementation.
#[async_trait]
pub trait RemoteAccess: Send + Sync {
/// Short provider name used in logs and status fields (e.g. `"tailscale"`).
fn provider_name(&self) -> &str;
/// IP address of this node on the mesh network.
async fn device_ip(&self) -> Result<Ipv4Addr>;
/// True once the provider has joined the mesh and obtained an IP.
fn is_connected(&self) -> bool;
/// Graceful shutdown — disconnect from the mesh and release resources.
async fn shutdown(&self);
}
+39
View File
@@ -0,0 +1,39 @@
use std::sync::Arc;
use anyhow::Result;
use async_trait::async_trait;
/// Full CRUD interface to the secrets store.
///
/// Implemented by `SecretsStore` in the main crate. Both plugins (via
/// `PluginContext`) and agent tools (via `AppState`) receive an
/// `Arc<dyn SecretsApi>` so neither needs to depend on the main crate.
///
/// Security notes:
/// - Values should never appear in log output.
/// - Keys are safe to log/list.
/// - Storage is SQLite — same protection level as the rest of the DB.
#[async_trait]
pub trait SecretsApi: Send + Sync {
/// Returns the value for `key`, or `None` if not set.
async fn get(&self, key: &str) -> Option<String>;
/// Inserts or replaces the secret for `key`.
async fn set(&self, key: &str, value: &str) -> Result<()>;
/// Removes the secret for `key`. No-op if not present.
async fn delete(&self, key: &str) -> Result<()>;
/// Returns all stored keys (never values).
async fn list_keys(&self) -> Vec<String>;
}
/// Resolves a required secret, returning a descriptive error if absent.
pub async fn require(secrets: &Arc<dyn SecretsApi>, key: &str) -> Result<String> {
secrets.get(key).await.ok_or_else(|| {
anyhow::anyhow!(
"secret '{}' is not set — tell the agent: \"set the secret {} to <value>\"",
key, key,
)
})
}
+92
View File
@@ -0,0 +1,92 @@
//! In-process event bus for system-level lifecycle events.
//!
//! Distinct from [`crate::bus::ChatEventBus`] which carries chat-turn events.
//! This bus carries infrastructure events: provider registration, plugin state
//! changes, etc. Any component can subscribe and react without direct coupling.
//!
//! # Usage
//! ```rust,ignore
//! // Producer (e.g. a plugin):
//! bus.send(SystemEvent::ApiProviderRegistered { type_id: "elevenlabs".into() });
//!
//! // Consumer (spawn once at startup):
//! let mut rx = bus.subscribe();
//! tokio::spawn(async move {
//! loop {
//! match rx.recv().await {
//! Ok(SystemEvent::ApiProviderRegistered { type_id }) => { /* reload */ }
//! Ok(SystemEvent::ApiProviderUnregistered { type_id }) => { /* reload */ }
//! Err(RecvError::Lagged(n)) => warn!("system_bus lagged by {n}"),
//! Err(RecvError::Closed) => break,
//! }
//! }
//! });
//! ```
use tokio::sync::broadcast;
pub use tokio::sync::broadcast::error::RecvError;
const DEFAULT_CAPACITY: usize = 64;
// ── Events ────────────────────────────────────────────────────────────────────
/// All system-level events that flow through the [`SystemEventBus`].
#[derive(Debug, Clone)]
pub enum SystemEvent {
/// A plugin registered a new `ApiProvider` (e.g. ElevenLabs on plugin start).
ApiProviderRegistered { type_id: String },
/// A plugin unregistered an `ApiProvider` (e.g. on plugin stop/disable).
ApiProviderUnregistered { type_id: String },
/// A config key was changed via the API (only fires when the value actually changes).
ConfigKeyUpdated { key: String, old_value: Option<String>, new_value: String },
/// A scheduled job finished (success or failure). `origin_ref` is the opaque
/// string stored in `scheduled_jobs.origin_ref` (e.g. `"PROJECT_TASK:42"`).
JobCompleted {
job_id: i64,
origin_ref: Option<String>,
result: Option<String>,
error: Option<String>,
},
/// A session was forcibly cancelled (e.g. via the kill-task API).
/// Subscribers should cancel any in-flight LLM turn, pending approvals,
/// and pending clarifications for that session.
SessionCancelled {
session_id: i64,
},
}
// ── Bus ───────────────────────────────────────────────────────────────────────
pub struct SystemEventBus {
tx: broadcast::Sender<SystemEvent>,
}
impl SystemEventBus {
pub fn new() -> Self {
Self::with_capacity(DEFAULT_CAPACITY)
}
pub fn with_capacity(capacity: usize) -> Self {
let (tx, _) = broadcast::channel(capacity);
Self { tx }
}
/// Publish an event. No-op if there are no active subscribers.
pub fn send(&self, event: SystemEvent) {
let _ = self.tx.send(event);
}
/// Returns a new independent receiver. Each subscriber gets every future
/// event independently. If the subscriber falls behind by more than the
/// channel capacity it receives `RecvError::Lagged(n)`.
pub fn subscribe(&self) -> broadcast::Receiver<SystemEvent> {
self.tx.subscribe()
}
}
impl Default for SystemEventBus {
fn default() -> Self {
Self::new()
}
}
+355
View File
@@ -0,0 +1,355 @@
use std::future::Future;
use std::pin::Pin;
use std::sync::Mutex;
use anyhow::Result;
use serde_json::Value;
use tokio_util::sync::CancellationToken;
// ── ToolDescriptionLength ─────────────────────────────────────────────────────
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ToolDescriptionLength {
Short,
Full,
}
// ── ToolCategory ──────────────────────────────────────────────────────────────
/// Logical grouping for a tool.
///
/// Used for access-control filtering and for display/audit purposes.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ToolCategory {
/// Read or write files on disk.
Filesystem,
/// Run shell commands or restart the process.
Shell,
/// Invoke sub-agents via call_agent.
Subagent,
/// Read-only discovery of system state.
Introspection,
/// Mutate system configuration.
Config,
}
// ── Tool trait ────────────────────────────────────────────────────────────────
/// A single LLM-callable tool.
pub trait Tool: Send + Sync {
fn name(&self) -> &str;
fn description(&self) -> &str;
/// Human-readable label for this tool invocation shown in UI / notifications.
fn describe(&self, _args: &Value, _length: ToolDescriptionLength) -> String {
self.name().to_string()
}
/// If this invocation targets a single file the user can open in the file
/// viewer, return its path (relative to the project root, or absolute).
/// Tools that target a directory (list/grep) or no file at all return
/// `None`. The frontend renders the returned path as a clickable link.
fn target_path(&self, _args: &Value) -> Option<String> {
None
}
/// JSON Schema for the `parameters` field in the OpenAI function definition.
fn parameters_schema(&self) -> Value;
/// Execute the tool synchronously and return a plain-text result (or error string).
/// Tools that require async I/O should override `execute_async` instead.
fn execute(&self, _args: Value) -> Result<String> {
Err(anyhow::anyhow!("tool '{}': sync execute not implemented — use execute_async", self.name()))
}
/// Execute the tool asynchronously. The default wraps `execute`; async tools
/// (e.g. image generation) override this directly to avoid `block_in_place`.
fn execute_async<'a>(&'a self, args: Value) -> Pin<Box<dyn Future<Output = Result<String>> + Send + 'a>> {
Box::pin(async move { self.execute(args) })
}
/// Execute and produce a typed [`ToolResult`]. The default bridges
/// [`execute_async`](Self::execute_async) (plain string) into [`ToolResult::Text`],
/// so existing tools need no changes. Override to return [`ToolResult::Json`]
/// for tools with structured output (e.g. MCP tools exposing `structuredContent`).
fn execute_typed<'a>(&'a self, args: Value)
-> Pin<Box<dyn Future<Output = Result<ToolResult>> + Send + 'a>> {
Box::pin(async move { Ok(ToolResult::Text(self.execute_async(args).await?)) })
}
/// Start a single execution of this tool and return a live [`ToolExecution`]
/// handle. This is the entry point the session driver uses: it lets the tool
/// own its in-flight state and implement its own `stop()` (e.g. ComfyUI sends
/// an `/interrupt`; `execute_cmd` relies on `kill_on_drop`).
///
/// The default wraps `execute_typed` in a [`SimpleExecution`], whose `stop()`
/// drops the work future — already enough to make `/stop` responsive for any
/// I/O-bound tool. Tools needing remote/child teardown override this and
/// return their own `ToolExecution` with a bespoke `stop()`.
fn run<'a>(&'a self, args: Value) -> Box<dyn ToolExecution + 'a> {
Box::new(SimpleExecution::new(self.execute_typed(args)))
}
/// Logical category of this tool.
fn category(&self) -> ToolCategory;
/// If true, this tool is only included in the tool list for sub-agents (depth > 0).
fn sub_agents_only(&self) -> bool { false }
/// If true, this tool is only included in the tool list for the root agent (depth == 0).
fn root_agent_only(&self) -> bool { false }
/// If true, this tool is only available to interactive sessions (web, telegram, mobile, voice).
/// Non-interactive background sessions (cron, tic) will not receive this tool definition.
fn interactive_only(&self) -> bool { false }
/// Full OpenAI-format tool definition ready to be sent to the LLM.
fn openai_definition(&self) -> Value {
serde_json::json!({
"type": "function",
"function": {
"name": self.name(),
"description": self.description(),
"parameters": self.parameters_schema(),
}
})
}
}
// ── ToolExecutionState ────────────────────────────────────────────────────────
/// Lifecycle state of a single tool execution.
///
/// Richer than the persisted `chat_llm_tools.status` string: it distinguishes a
/// user `/stop` (`Cancelled`) and a policy/human denial (`Rejected`) from a real
/// tool error (`Failed`). The session driver owns the approval-phase states
/// (`Pending`, `AwaitingApproval`, `Rejected`); a [`ToolExecution`] itself only
/// ever reports `Running → Completed | Failed | Cancelled`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ToolExecutionState {
/// Intent recorded, not yet started (transient; not persisted on its own).
Pending,
/// Blocked waiting for a human approval / clarification answer.
AwaitingApproval,
/// Actively executing.
Running,
/// Finished successfully.
Completed,
/// Finished with a tool/runtime error.
Failed,
/// Stopped by the user via `/stop` — not an error.
Cancelled,
/// Denied by an approval policy or a human — not an error.
Rejected,
}
impl ToolExecutionState {
/// String persisted in `chat_llm_tools.status`. `AwaitingApproval` maps to the
/// legacy `pending` value so existing resume logic keeps working; the brand-new
/// `Pending` state is never persisted (the row is created on the first real
/// transition) and defaults to `running` defensively.
pub fn as_db_str(self) -> &'static str {
match self {
Self::Pending => "running",
Self::AwaitingApproval => "pending",
Self::Running => "running",
Self::Completed => "done",
Self::Failed => "failed",
Self::Cancelled => "cancelled",
Self::Rejected => "rejected",
}
}
}
// ── ToolResult ────────────────────────────────────────────────────────────────
/// The typed result of a successful tool execution.
///
/// Most tools produce plain text ([`ToolResult::Text`], persisted as
/// `result_type = "string"`). Tools with structured output — notably MCP servers
/// returning `structuredContent` — produce [`ToolResult::Json`] (persisted as
/// `result_type = "json"`), so the host/frontend can render the typed payload
/// instead of a raw text blob. At the LLM wire both variants become the same
/// `{"role":"tool","content": <string>}` bytes (see [`ToolResult::to_wire`]); the
/// type tag only matters to the host.
#[derive(Debug, Clone)]
pub enum ToolResult {
/// Plain-text result. The default for every built-in tool.
Text(String),
/// Structured JSON result (e.g. MCP `structuredContent`).
Json(serde_json::Value),
}
impl ToolResult {
/// Tag persisted in `chat_llm_tools.result_type` and sent over the WS as
/// `ServerEvent::ToolDone.result_type`. Either `"string"` or `"json"`.
pub fn kind(&self) -> &'static str {
match self {
Self::Text(_) => "string",
Self::Json(_) => "json",
}
}
/// Wire content for the LLM tool message: text as-is, Json serialized to a
/// compact JSON string. Both OpenAI and Anthropic encode tool results as
/// text/JSON, so this is the canonical string form persisted in
/// `chat_llm_tools.result` and replayed by the message builder.
pub fn to_wire(&self) -> String {
match self {
Self::Text(s) => s.clone(),
Self::Json(v) => serde_json::to_string(v).unwrap_or_else(|_| "null".to_string()),
}
}
}
impl From<String> for ToolResult {
fn from(s: String) -> Self { Self::Text(s) }
}
impl From<&str> for ToolResult {
fn from(s: &str) -> Self { Self::Text(s.to_string()) }
}
// ── ExecutionOutcome ──────────────────────────────────────────────────────────
/// Terminal outcome produced by [`ToolExecution::wait`]. A running execution can
/// only end in one of these three ways; `Rejected`/`AwaitingApproval` are decided
/// by the approval gate *before* the work runs and never appear here.
#[derive(Debug, Clone)]
pub enum ExecutionOutcome {
Completed(ToolResult),
Failed(String),
Cancelled,
}
impl ExecutionOutcome {
pub fn state(&self) -> ToolExecutionState {
match self {
Self::Completed(_) => ToolExecutionState::Completed,
Self::Failed(_) => ToolExecutionState::Failed,
Self::Cancelled => ToolExecutionState::Cancelled,
}
}
}
/// A boxed, owned unit of asynchronous tool work producing a typed [`ToolResult`].
/// This is what [`Tool::execute_typed`] returns; [`SimpleExecution`] wraps one.
pub type ToolWork<'a> = Pin<Box<dyn Future<Output = Result<ToolResult>> + Send + 'a>>;
// ── ToolExecution ─────────────────────────────────────────────────────────────
/// A single, live execution of a [`Tool`]. Owns its in-memory state and decides
/// how it stops. Pure: it never touches the DB or the WebSocket — the session
/// driver mirrors its state transitions to persistence and transport.
///
/// `Send + Sync` is required because the driver shares `&self` across the two
/// concurrent branches of the cancellation race (`wait` vs `stop`).
pub trait ToolExecution: Send + Sync {
/// Current in-memory lifecycle state.
fn state(&self) -> ToolExecutionState;
/// Drive the work to its terminal outcome. Called exactly once by the driver.
fn wait<'a>(&'a self) -> Pin<Box<dyn Future<Output = ExecutionOutcome> + Send + 'a>>;
/// Tool-specific cancellation: signal the work to stop and tear down any
/// remote/child resources. The default relies on the driver dropping the
/// `wait` future; [`SimpleExecution`] overrides it to cancel its stop-token.
fn stop<'a>(&'a self) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>> {
Box::pin(async {})
}
}
// ── SimpleExecution ───────────────────────────────────────────────────────────
/// Default [`ToolExecution`] for any tool that is a single async unit of work.
///
/// Holds the work future plus a stop-token; `wait` races the two, so a `stop()`
/// (or the driver dropping `wait`) drops the work future — aborting the in-flight
/// I/O (a `reqwest` connection, a `kill_on_drop` child, …). Enough to make
/// `/stop` responsive for every I/O-bound tool with zero per-tool code.
pub struct SimpleExecution<'a> {
state: Mutex<ToolExecutionState>,
stop: CancellationToken,
work: tokio::sync::Mutex<Option<ToolWork<'a>>>,
}
impl<'a> SimpleExecution<'a> {
pub fn new(work: ToolWork<'a>) -> Self {
Self {
state: Mutex::new(ToolExecutionState::Running),
stop: CancellationToken::new(),
work: tokio::sync::Mutex::new(Some(work)),
}
}
}
impl<'a> ToolExecution for SimpleExecution<'a> {
fn state(&self) -> ToolExecutionState {
*self.state.lock().unwrap()
}
fn wait<'b>(&'b self) -> Pin<Box<dyn Future<Output = ExecutionOutcome> + Send + 'b>> {
Box::pin(async move {
let work = self.work.lock().await.take();
let Some(work) = work else {
// Already consumed (e.g. a second wait after cancellation).
return ExecutionOutcome::Cancelled;
};
let outcome = tokio::select! {
biased;
_ = self.stop.cancelled() => ExecutionOutcome::Cancelled,
r = work => match r {
Ok(s) => ExecutionOutcome::Completed(s),
Err(e) => ExecutionOutcome::Failed(e.to_string()),
},
};
*self.state.lock().unwrap() = outcome.state();
outcome
})
}
fn stop<'b>(&'b self) -> Pin<Box<dyn Future<Output = ()> + Send + 'b>> {
Box::pin(async move { self.stop.cancel(); })
}
}
// ── drive_execution ───────────────────────────────────────────────────────────
/// Run a [`ToolExecution`] to completion while honouring a cancellation token.
///
/// When `cancel` fires we call `exec.stop()` once (tool-specific teardown); the
/// execution then resolves `wait` to `Cancelled`. Both methods take `&self`, so
/// the two concurrent borrows are shared and the borrow checker is happy.
pub async fn drive_execution(
exec: &dyn ToolExecution,
cancel: &CancellationToken,
) -> ExecutionOutcome {
let work = exec.wait();
tokio::pin!(work);
let mut stopped = false;
loop {
tokio::select! {
biased;
outcome = &mut work => return outcome,
_ = cancel.cancelled(), if !stopped => {
exec.stop().await;
stopped = true;
}
}
}
}
// ── Helpers ───────────────────────────────────────────────────────────────────
/// Truncates a label to `max` chars, appending `…` if cut.
pub fn truncate_label(s: &str, max: usize) -> String {
if s.chars().count() <= max {
return s.to_string();
}
let cut = max.saturating_sub(1);
let mut end = cut;
while !s.is_char_boundary(end) { end -= 1; }
format!("{}", &s[..end])
}
+59
View File
@@ -0,0 +1,59 @@
use std::sync::Arc;
use async_trait::async_trait;
// ── Record types (DB ↔ manager) ───────────────────────────────────────────────
/// Full model record, mirroring one row in `transcribe_models`.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct TranscribeModelRecord {
pub id: i64,
pub provider_id: i64,
pub model_id: String,
/// Display alias (also used as the transcriber `id()`).
pub name: String,
/// BCP-47 language hint, e.g. `"it"`. `None` → auto-detect.
pub language: Option<String>,
/// Lower number = tried first by `get()`.
pub priority: i32,
}
/// Remote model info returned by a provider's `list_transcribe_models()`.
#[derive(Debug, Clone, serde::Serialize)]
pub struct RemoteTranscribeModelInfo {
pub id: String,
pub name: String,
pub description: Option<String>,
/// BCP-47 language codes supported by this model (empty = unknown).
pub languages: Vec<String>,
}
/// Implemented by any provider that can convert audio bytes to text.
/// The `format` hint (e.g. `"ogg"`, `"mp3"`) is advisory.
#[async_trait]
pub trait Transcribe: Send + Sync {
/// A stable, unique identifier for this provider (e.g. `"whisper_local"`).
fn id(&self) -> &str;
async fn transcribe(&self, audio: Vec<u8>, format: &str) -> anyhow::Result<String>;
}
/// Resolves the currently active [`Transcribe`] provider.
///
/// Implemented by `TranscribeManager` in the main crate. Plugins store
/// `Arc<dyn TranscribeProvider>` so they can resolve the active transcriber
/// per-call without holding a reference to `AppState`.
#[async_trait]
pub trait TranscribeProvider: Send + Sync {
async fn get(&self) -> Option<Arc<dyn Transcribe>>;
}
/// Write-side of the transcribe manager: register and remove ephemeral providers.
///
/// Implemented by `TranscribeManager`. Plugins that provide their own STT
/// (e.g. `WhisperLocalPlugin`) use this to register themselves at start and
/// unregister at stop.
#[async_trait]
pub trait TranscribeRegistry: Send + Sync {
async fn register(&self, provider: Arc<dyn Transcribe>);
async fn unregister(&self, id: &str);
}
+88
View File
@@ -0,0 +1,88 @@
use std::sync::Arc;
use anyhow::Result;
use async_trait::async_trait;
// ── Record types (DB ↔ manager) ───────────────────────────────────────────────
/// Full model record, mirroring one row in `tts_models`.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct TtsModelRecord {
pub id: i64,
pub provider_id: i64,
pub model_id: String,
/// Voice/speaker identifier, if the provider requires it separately from the model.
pub voice_id: Option<String>,
/// Display alias (also used as the synthesiser `id()`).
pub name: String,
pub description: Option<String>,
/// Default voice instructions (tone, speed, style).
pub instructions: Option<String>,
/// Audio container/codec requested from the provider (`response_format`):
/// `mp3`, `opus`, `aac`, `flac`, `wav`, `pcm`. `None` ⇒ provider default (`mp3`).
/// Some models only accept a specific value (e.g. Gemini TTS requires `pcm`).
pub response_format: Option<String>,
/// Lower number = tried first by `get()`.
pub priority: i32,
}
/// Remote model info returned by a provider's `list_tts_models()`.
#[derive(Debug, Clone, serde::Serialize)]
pub struct RemoteTtsModelInfo {
pub id: String,
pub name: String,
pub description: Option<String>,
/// BCP-47 language codes supported by this model (empty = unknown).
pub languages: Vec<String>,
/// Cost multiplier relative to the provider's base rate (1.0 = standard).
pub cost_factor: Option<f64>,
/// Usage instructions: supported tags, markup, etc. Shown in UI and passed
/// to the LLM when generating text destined for this synthesiser.
pub instructions: Option<String>,
}
/// Implemented by any provider that can convert text to audio bytes.
/// Returns raw audio bytes (MP3 expected unless the provider states otherwise).
#[async_trait]
pub trait TextToSpeech: Send + Sync {
/// A stable, unique identifier for this provider (e.g. `"openai_tts_alloy"`).
fn id(&self) -> &str;
/// Human-readable display name.
fn name(&self) -> &str;
/// Human-readable description (voice style, language, ideal use cases).
fn description(&self) -> Option<&str> { None }
/// Default synthesis instructions: voice style, tone, speed, and any
/// provider-specific text markup syntax (e.g. emotion tags).
/// Surfaced to the LLM via `TtsModelInfo` so it knows how to format input text.
/// Individual call-time instructions passed to `synthesize` take precedence.
fn instructions(&self) -> Option<&str> { None }
/// Audio format (container/codec) of the bytes returned by `synthesize`,
/// e.g. `mp3`, `opus`, `wav`, `pcm`. Consumers that require a specific
/// container (e.g. Telegram voice messages need Ogg/Opus) use this to decide
/// whether and how to transcode. Default `"mp3"`.
fn output_format(&self) -> &str { "mp3" }
/// Synthesise `text` to audio bytes.
/// `instructions` overrides the provider's default instructions for this call only.
async fn synthesize(&self, text: &str, instructions: Option<&str>) -> Result<Vec<u8>>;
}
/// Resolves the currently active [`TextToSpeech`] provider.
///
/// Implemented by `TtsManager` in the main crate. Plugins store
/// `Arc<dyn TtsProvider>` to resolve the active synthesiser per-call
/// without holding a reference to `AppState`.
#[async_trait]
pub trait TtsProvider: Send + Sync {
async fn get(&self) -> Option<Arc<dyn TextToSpeech>>;
}
/// Write-side of the TTS manager: register and remove ephemeral providers.
///
/// Implemented by `TtsManager`. Plugins that supply their own TTS engine
/// (e.g. a local Kokoro or Piper plugin) use this to register at start
/// and unregister at stop.
#[async_trait]
pub trait TtsRegistry: Send + Sync {
async fn register(&self, provider: Arc<dyn TextToSpeech>);
async fn unregister(&self, id: &str);
}
+10
View File
@@ -0,0 +1,10 @@
[package]
name = "honcho-client"
version = "0.1.0"
edition = "2024"
[dependencies]
reqwest = { version = "0.13", default-features = false, features = ["rustls-no-provider", "charset", "http2", "system-proxy", "json"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
tracing = "0.1"
+173
View File
@@ -0,0 +1,173 @@
use reqwest::{Client, RequestBuilder, Response};
use serde::de::DeserializeOwned;
use serde::Serialize;
use tracing::debug;
use crate::error::{HonchoError, Result};
const DEFAULT_BASE_URL: &str = "https://api.honcho.dev";
/// Honcho REST API client.
///
/// Cheaply clone-able — the inner `reqwest::Client` holds an `Arc`.
#[derive(Debug, Clone)]
pub struct HonchoClient {
pub(crate) http: Client,
pub(crate) base_url: String,
pub(crate) token: String,
}
impl HonchoClient {
/// Create a client pointing at the production SaaS endpoint.
pub fn new(token: impl Into<String>) -> Self {
Self::with_base_url(DEFAULT_BASE_URL, token)
}
/// Create a client pointing at a custom base URL (e.g. `http://localhost:8000`).
pub fn with_base_url(base_url: impl Into<String>, token: impl Into<String>) -> Self {
Self {
http: Client::new(),
base_url: base_url.into().trim_end_matches('/').to_owned(),
token: token.into(),
}
}
// ── internal helpers ──────────────────────────────────────────────────
pub(crate) fn url(&self, path: &str) -> String {
format!("{}{}", self.base_url, path)
}
/// Build a URL with optional query params (key-value pairs).
pub(crate) fn url_with_query(&self, path: &str, params: &[(&str, String)]) -> String {
if params.is_empty() {
return self.url(path);
}
let qs: String = params
.iter()
.map(|(k, v)| format!("{}={}", urlencoding(k), urlencoding(v)))
.collect::<Vec<_>>()
.join("&");
format!("{}{}?{}", self.base_url, path, qs)
}
fn auth(&self, rb: RequestBuilder) -> RequestBuilder {
rb.bearer_auth(&self.token)
}
pub(crate) async fn get<T: DeserializeOwned>(&self, path: &str) -> Result<T> {
let url = self.url(path);
debug!("GET {url}");
let resp = self.auth(self.http.get(&url)).send().await?;
parse_response(resp).await
}
pub(crate) async fn get_with_query<T: DeserializeOwned>(
&self,
path: &str,
query: &[(&str, String)],
) -> Result<T> {
let url = self.url_with_query(path, query);
debug!("GET {url}");
let resp = self.auth(self.http.get(&url)).send().await?;
parse_response(resp).await
}
pub(crate) async fn post<B: Serialize, T: DeserializeOwned>(
&self,
path: &str,
body: &B,
) -> Result<T> {
let url = self.url(path);
debug!("POST {url}");
let resp = self.auth(self.http.post(&url)).json(body).send().await?;
parse_response(resp).await
}
pub(crate) async fn post_with_query<B: Serialize, T: DeserializeOwned>(
&self,
path: &str,
query: &[(&str, String)],
body: &B,
) -> Result<T> {
let url = self.url_with_query(path, query);
debug!("POST {url}");
let resp = self.auth(self.http.post(&url)).json(body).send().await?;
parse_response(resp).await
}
pub(crate) async fn put<B: Serialize, T: DeserializeOwned>(
&self,
path: &str,
body: &B,
) -> Result<T> {
let url = self.url(path);
debug!("PUT {url}");
let resp = self.auth(self.http.put(&url)).json(body).send().await?;
parse_response(resp).await
}
pub(crate) async fn put_with_query<B: Serialize, T: DeserializeOwned>(
&self,
path: &str,
query: &[(&str, String)],
body: &B,
) -> Result<T> {
let url = self.url_with_query(path, query);
debug!("PUT {url}");
let resp = self.auth(self.http.put(&url)).json(body).send().await?;
parse_response(resp).await
}
pub(crate) async fn delete_ok(&self, path: &str) -> Result<()> {
let url = self.url(path);
debug!("DELETE {url}");
let resp = self.auth(self.http.delete(&url)).send().await?;
parse_no_content(resp).await
}
pub(crate) async fn delete_json<B: Serialize>(&self, path: &str, body: &B) -> Result<()> {
let url = self.url(path);
debug!("DELETE {url} (with body)");
let resp = self
.auth(self.http.delete(&url))
.json(body)
.send()
.await?;
parse_no_content(resp).await
}
}
async fn parse_response<T: DeserializeOwned>(resp: Response) -> Result<T> {
let status = resp.status();
if status.is_success() {
Ok(resp.json::<T>().await?)
} else {
let code = status.as_u16();
let body = resp.text().await.unwrap_or_default();
Err(HonchoError::Http { status: code, body })
}
}
pub(crate) async fn parse_no_content(resp: Response) -> Result<()> {
if resp.status().is_success() {
Ok(())
} else {
let status = resp.status().as_u16();
let body = resp.text().await.unwrap_or_default();
Err(HonchoError::Http { status, body })
}
}
/// Minimal percent-encoding for query param keys and values.
fn urlencoding(s: &str) -> String {
let mut out = String::with_capacity(s.len());
for b in s.bytes() {
match b {
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9'
| b'-' | b'_' | b'.' | b'~' => out.push(b as char),
_ => out.push_str(&format!("%{b:02X}")),
}
}
out
}
+80
View File
@@ -0,0 +1,80 @@
use crate::{
client::{parse_no_content, HonchoClient},
error::Result,
models::*,
workspaces::page_query,
};
impl HonchoClient {
/// Create up to 100 conclusions in a single batch.
pub async fn add_conclusions(
&self,
workspace_id: &str,
batch: &ConclusionBatchCreate,
) -> Result<Vec<Conclusion>> {
self.post(
&format!("/v3/workspaces/{workspace_id}/conclusions"),
batch,
)
.await
}
/// Convenience: add a single conclusion.
pub async fn add_conclusion(
&self,
workspace_id: &str,
c: ConclusionCreate,
) -> Result<Vec<Conclusion>> {
self.add_conclusions(
workspace_id,
&ConclusionBatchCreate {
conclusions: vec![c],
},
)
.await
}
pub async fn list_conclusions(
&self,
workspace_id: &str,
params: &PageParams,
filter: &ConclusionGet,
) -> Result<Page<Conclusion>> {
self.post_with_query(
&format!("/v3/workspaces/{workspace_id}/conclusions/list"),
&page_query(params),
filter,
)
.await
}
/// Semantic search over conclusions.
pub async fn query_conclusions(
&self,
workspace_id: &str,
query: &ConclusionQuery,
) -> Result<Vec<Conclusion>> {
self.post(
&format!("/v3/workspaces/{workspace_id}/conclusions/query"),
query,
)
.await
}
pub async fn delete_conclusion(
&self,
workspace_id: &str,
conclusion_id: &str,
) -> Result<()> {
let url = self.url(&format!(
"/v3/workspaces/{workspace_id}/conclusions/{conclusion_id}"
));
let resp = self
.http
.delete(&url)
.bearer_auth(&self.token)
.send()
.await?;
parse_no_content(resp).await
}
}
+56
View File
@@ -0,0 +1,56 @@
use std::error::Error as StdError;
use std::fmt;
#[derive(Debug)]
pub enum HonchoError {
Http { status: u16, body: String },
Request(reqwest::Error),
Json(serde_json::Error),
}
impl fmt::Display for HonchoError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Http { status, body } => write!(f, "HTTP error {status}: {body}"),
Self::Request(e) => {
// reqwest's own Display stops at "error sending request for url
// (...)" and hides the transport cause. Walk the source chain so
// the real reason (e.g. "Connection reset by peer", "operation
// timed out") shows up in logs instead of a useless top line.
write!(f, "Request failed: {e}")?;
let mut src = e.source();
while let Some(cause) = src {
write!(f, ": {cause}")?;
src = cause.source();
}
Ok(())
}
Self::Json(e) => write!(f, "JSON error: {e}"),
}
}
}
impl std::error::Error for HonchoError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Request(e) => Some(e),
Self::Json(e) => Some(e),
_ => None,
}
}
}
impl From<reqwest::Error> for HonchoError {
fn from(e: reqwest::Error) -> Self {
Self::Request(e)
}
}
impl From<serde_json::Error> for HonchoError {
fn from(e: serde_json::Error) -> Self {
Self::Json(e)
}
}
pub type Result<T> = std::result::Result<T, HonchoError>;
+63
View File
@@ -0,0 +1,63 @@
//! Honcho v3 REST API client.
//!
//! # Quick start
//!
//! ```rust,no_run
//! use honcho_client::{HonchoClient, models::*};
//!
//! #[tokio::main]
//! async fn main() -> anyhow::Result<()> {
//! // Point at your local Docker instance
//! let client = HonchoClient::with_base_url("http://localhost:8000", "your-api-key");
//!
//! // Create or retrieve a workspace
//! let ws = client.create_workspace(&WorkspaceCreate {
//! id: "my-agent".into(),
//! ..Default::default()
//! }).await?;
//!
//! // Create a peer (= one user / agent identity)
//! let peer = client.create_peer(&ws.id, &PeerCreate {
//! id: "daniele".into(),
//! metadata: None,
//! configuration: None,
//! }).await?;
//!
//! // Open a session
//! let session = client.create_session(&ws.id, &SessionCreate::default()).await?;
//!
//! // Add messages
//! client.add_message(&ws.id, &session.id, MessageCreate {
//! content: "Hello!".into(),
//! peer_id: peer.id.clone(),
//! metadata: None,
//! configuration: None,
//! created_at: None,
//! }).await?;
//!
//! // Query memory (Dialectic API)
//! let answer = client.peer_chat(&ws.id, &peer.id, &DialecticOptions {
//! query: "What did the user say?".into(),
//! session_id: Some(session.id.clone()),
//! target: None,
//! stream: Some(false),
//! reasoning_level: None,
//! }).await?;
//!
//! println!("{answer:#?}");
//! Ok(())
//! }
//! ```
pub mod client;
pub mod conclusions;
pub mod error;
pub mod messages;
pub mod models;
pub mod peers;
pub mod sessions;
pub mod workspaces;
// Flat re-exports for convenience
pub use client::HonchoClient;
pub use error::{HonchoError, Result};
+80
View File
@@ -0,0 +1,80 @@
use crate::{
client::HonchoClient,
error::Result,
models::*,
workspaces::page_query,
};
impl HonchoClient {
/// Add one or more messages to a session in a single request.
pub async fn add_messages(
&self,
workspace_id: &str,
session_id: &str,
batch: &MessageBatchCreate,
) -> Result<Vec<Message>> {
self.post(
&format!("/v3/workspaces/{workspace_id}/sessions/{session_id}/messages"),
batch,
)
.await
}
/// Convenience: add a single message.
pub async fn add_message(
&self,
workspace_id: &str,
session_id: &str,
msg: MessageCreate,
) -> Result<Vec<Message>> {
self.add_messages(
workspace_id,
session_id,
&MessageBatchCreate { messages: vec![msg] },
)
.await
}
pub async fn list_messages(
&self,
workspace_id: &str,
session_id: &str,
params: &PageParams,
filter: &MessageGet,
) -> Result<Page<Message>> {
self.post_with_query(
&format!("/v3/workspaces/{workspace_id}/sessions/{session_id}/messages/list"),
&page_query(params),
filter,
)
.await
}
pub async fn get_message(
&self,
workspace_id: &str,
session_id: &str,
message_id: &str,
) -> Result<Message> {
self.get(&format!(
"/v3/workspaces/{workspace_id}/sessions/{session_id}/messages/{message_id}"
))
.await
}
pub async fn update_message(
&self,
workspace_id: &str,
session_id: &str,
message_id: &str,
body: &MessageUpdate,
) -> Result<Message> {
self.put(
&format!(
"/v3/workspaces/{workspace_id}/sessions/{session_id}/messages/{message_id}"
),
body,
)
.await
}
}
+308
View File
@@ -0,0 +1,308 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
// ──────────────────────────────────────────
// Pagination
// ──────────────────────────────────────────
#[derive(Debug, Clone, Deserialize)]
pub struct Page<T> {
pub items: Vec<T>,
pub total: u64,
pub page: u64,
pub size: u64,
pub pages: u64,
}
// ──────────────────────────────────────────
// Workspace
// ──────────────────────────────────────────
#[derive(Debug, Clone, Deserialize)]
pub struct Workspace {
pub id: String,
pub metadata: Option<Value>,
pub configuration: Option<Value>,
pub created_at: String,
}
#[derive(Debug, Clone, Serialize, Default)]
pub struct WorkspaceCreate {
pub id: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub configuration: Option<Value>,
}
#[derive(Debug, Clone, Serialize, Default)]
pub struct WorkspaceUpdate {
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub configuration: Option<Value>,
}
/// Filter body for `POST /workspaces/list`
#[derive(Debug, Clone, Serialize, Default)]
pub struct WorkspaceGet {
#[serde(skip_serializing_if = "Option::is_none")]
pub filters: Option<Value>,
}
// ──────────────────────────────────────────
// Peer
// ──────────────────────────────────────────
#[derive(Debug, Clone, Deserialize)]
pub struct Peer {
pub id: String,
pub workspace_id: String,
pub created_at: String,
pub metadata: Option<Value>,
pub configuration: Option<Value>,
}
#[derive(Debug, Clone, Serialize)]
pub struct PeerCreate {
pub id: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub configuration: Option<Value>,
}
#[derive(Debug, Clone, Serialize, Default)]
pub struct PeerUpdate {
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub configuration: Option<Value>,
}
#[derive(Debug, Clone, Serialize, Default)]
pub struct PeerGet {
#[serde(skip_serializing_if = "Option::is_none")]
pub filters: Option<Value>,
}
// ──────────────────────────────────────────
// Session
// ──────────────────────────────────────────
#[derive(Debug, Clone, Deserialize)]
pub struct Session {
pub id: String,
pub is_active: bool,
pub workspace_id: String,
pub metadata: Option<Value>,
pub configuration: Option<Value>,
pub created_at: String,
}
#[derive(Debug, Clone, Serialize, Default)]
pub struct SessionCreate {
#[serde(skip_serializing_if = "Option::is_none")]
pub id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<Value>,
/// Map of peer_id → SessionPeerConfig
#[serde(skip_serializing_if = "Option::is_none")]
pub peers: Option<std::collections::HashMap<String, SessionPeerConfig>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub configuration: Option<Value>,
}
#[derive(Debug, Clone, Serialize, Default)]
pub struct SessionUpdate {
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub configuration: Option<Value>,
}
#[derive(Debug, Clone, Serialize, Default)]
pub struct SessionGet {
#[serde(skip_serializing_if = "Option::is_none")]
pub filters: Option<Value>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct SessionPeerConfig {
#[serde(skip_serializing_if = "Option::is_none")]
pub observe_me: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub observe_others: Option<bool>,
}
// ──────────────────────────────────────────
// Message
// ──────────────────────────────────────────
#[derive(Debug, Clone, Deserialize)]
pub struct Message {
pub id: String,
pub content: String,
pub peer_id: String,
pub session_id: String,
pub workspace_id: String,
pub metadata: Option<Value>,
pub created_at: String,
pub token_count: i64,
}
#[derive(Debug, Clone, Serialize)]
pub struct MessageCreate {
pub content: String,
pub peer_id: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub configuration: Option<Value>,
/// RFC3339 datetime; if None the server assigns now
#[serde(skip_serializing_if = "Option::is_none")]
pub created_at: Option<String>,
}
#[derive(Debug, Clone, Serialize)]
pub struct MessageBatchCreate {
pub messages: Vec<MessageCreate>,
}
#[derive(Debug, Clone, Serialize, Default)]
pub struct MessageGet {
#[serde(skip_serializing_if = "Option::is_none")]
pub filters: Option<Value>,
}
#[derive(Debug, Clone, Serialize, Default)]
pub struct MessageUpdate {
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<Value>,
}
// ──────────────────────────────────────────
// Conclusion
// ──────────────────────────────────────────
#[derive(Debug, Clone, Deserialize)]
pub struct Conclusion {
pub id: String,
pub content: String,
pub observer_id: String,
pub observed_id: String,
pub session_id: Option<String>,
pub created_at: String,
}
#[derive(Debug, Clone, Serialize)]
pub struct ConclusionCreate {
/// 165535 characters
pub content: String,
pub observer_id: String,
pub observed_id: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub session_id: Option<String>,
}
#[derive(Debug, Clone, Serialize)]
pub struct ConclusionBatchCreate {
/// 1100 conclusions
pub conclusions: Vec<ConclusionCreate>,
}
#[derive(Debug, Clone, Serialize, Default)]
pub struct ConclusionGet {
#[serde(skip_serializing_if = "Option::is_none")]
pub filters: Option<Value>,
}
#[derive(Debug, Clone, Serialize)]
pub struct ConclusionQuery {
pub query: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub top_k: Option<u32>,
/// 0.01.0
#[serde(skip_serializing_if = "Option::is_none")]
pub distance: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub filters: Option<Value>,
}
// ──────────────────────────────────────────
// Dialectic / Peer context
// ──────────────────────────────────────────
#[derive(Debug, Clone, Serialize)]
pub struct DialecticOptions {
/// Natural language query
pub query: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub session_id: Option<String>,
/// peer_id to get the representation for (defaults to the caller)
#[serde(skip_serializing_if = "Option::is_none")]
pub target: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub stream: Option<bool>,
/// minimal | low | medium | high | max
#[serde(skip_serializing_if = "Option::is_none")]
pub reasoning_level: Option<String>,
}
#[derive(Debug, Clone, Serialize, Default)]
pub struct PeerRepresentationGet {
#[serde(skip_serializing_if = "Option::is_none")]
pub session_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub target: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub search_query: Option<String>,
/// 1100
#[serde(skip_serializing_if = "Option::is_none")]
pub search_top_k: Option<u32>,
/// 0.01.0
#[serde(skip_serializing_if = "Option::is_none")]
pub search_max_distance: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub include_most_frequent: Option<bool>,
/// 1100
#[serde(skip_serializing_if = "Option::is_none")]
pub max_conclusions: Option<u32>,
}
// ──────────────────────────────────────────
// Search
// ──────────────────────────────────────────
#[derive(Debug, Clone, Serialize)]
pub struct MessageSearchOptions {
pub query: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub filters: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub limit: Option<u32>,
}
// ──────────────────────────────────────────
// Queue
// ──────────────────────────────────────────
#[derive(Debug, Clone, Deserialize)]
pub struct QueueStatus {
pub total_work_units: u64,
pub completed_work_units: u64,
pub in_progress_work_units: u64,
pub pending_work_units: u64,
pub sessions: Option<Value>,
}
// ──────────────────────────────────────────
// List query params (pagination)
// ──────────────────────────────────────────
#[derive(Debug, Clone, Default)]
pub struct PageParams {
pub page: Option<u64>,
pub size: Option<u64>,
pub reverse: Option<bool>,
}
+180
View File
@@ -0,0 +1,180 @@
use crate::{
client::HonchoClient,
error::Result,
models::*,
workspaces::page_query,
};
impl HonchoClient {
// ── CRUD ──────────────────────────────────────────────────────────────
pub async fn create_peer(&self, workspace_id: &str, body: &PeerCreate) -> Result<Peer> {
self.post(&format!("/v3/workspaces/{workspace_id}/peers"), body)
.await
}
pub async fn list_peers(
&self,
workspace_id: &str,
params: &PageParams,
filter: &PeerGet,
) -> Result<Page<Peer>> {
self.post_with_query(
&format!("/v3/workspaces/{workspace_id}/peers/list"),
&page_query(params),
filter,
)
.await
}
pub async fn update_peer(
&self,
workspace_id: &str,
peer_id: &str,
body: &PeerUpdate,
) -> Result<Peer> {
self.put(
&format!("/v3/workspaces/{workspace_id}/peers/{peer_id}"),
body,
)
.await
}
// ── Sessions for a peer ───────────────────────────────────────────────
pub async fn list_peer_sessions(
&self,
workspace_id: &str,
peer_id: &str,
params: &PageParams,
filter: &SessionGet,
) -> Result<Page<Session>> {
self.post_with_query(
&format!("/v3/workspaces/{workspace_id}/peers/{peer_id}/sessions"),
&page_query(params),
filter,
)
.await
}
// ── Dialectic (memory query via LLM) ──────────────────────────────────
/// Query a peer's memory using natural language. Returns the LLM answer.
/// Note: `stream: true` is not supported by this client (use raw reqwest if needed).
pub async fn peer_chat(
&self,
workspace_id: &str,
peer_id: &str,
opts: &DialecticOptions,
) -> Result<serde_json::Value> {
self.post(
&format!("/v3/workspaces/{workspace_id}/peers/{peer_id}/chat"),
opts,
)
.await
}
// ── Representation ────────────────────────────────────────────────────
/// Get a structured representation of a peer's memory.
pub async fn peer_representation(
&self,
workspace_id: &str,
peer_id: &str,
opts: &PeerRepresentationGet,
) -> Result<serde_json::Value> {
self.post(
&format!("/v3/workspaces/{workspace_id}/peers/{peer_id}/representation"),
opts,
)
.await
}
// ── Context ───────────────────────────────────────────────────────────
/// Get context for a peer (conclusions + card, ready for injection into prompts).
pub async fn peer_context(
&self,
workspace_id: &str,
peer_id: &str,
opts: &PeerRepresentationGet,
) -> Result<serde_json::Value> {
let mut q: Vec<(&str, String)> = vec![];
if let Some(ref v) = opts.target {
q.push(("target", v.clone()));
}
if let Some(ref v) = opts.search_query {
q.push(("search_query", v.clone()));
}
if let Some(v) = opts.search_top_k {
q.push(("search_top_k", v.to_string()));
}
if let Some(v) = opts.search_max_distance {
q.push(("search_max_distance", v.to_string()));
}
if let Some(v) = opts.include_most_frequent {
q.push(("include_most_frequent", v.to_string()));
}
if let Some(v) = opts.max_conclusions {
q.push(("max_conclusions", v.to_string()));
}
self.get_with_query(
&format!("/v3/workspaces/{workspace_id}/peers/{peer_id}/context"),
&q,
)
.await
}
// ── Card ──────────────────────────────────────────────────────────────
pub async fn get_peer_card(
&self,
workspace_id: &str,
peer_id: &str,
target: Option<&str>,
) -> Result<serde_json::Value> {
let mut q: Vec<(&str, String)> = vec![];
if let Some(t) = target {
q.push(("target", t.to_owned()));
}
self.get_with_query(
&format!("/v3/workspaces/{workspace_id}/peers/{peer_id}/card"),
&q,
)
.await
}
pub async fn set_peer_card(
&self,
workspace_id: &str,
peer_id: &str,
target: Option<&str>,
card: serde_json::Value,
) -> Result<serde_json::Value> {
let mut q: Vec<(&str, String)> = vec![];
if let Some(t) = target {
q.push(("target", t.to_owned()));
}
self.put_with_query(
&format!("/v3/workspaces/{workspace_id}/peers/{peer_id}/card"),
&q,
&card,
)
.await
}
// ── Search ────────────────────────────────────────────────────────────
pub async fn search_peer(
&self,
workspace_id: &str,
peer_id: &str,
opts: &MessageSearchOptions,
) -> Result<Vec<Message>> {
self.post(
&format!("/v3/workspaces/{workspace_id}/peers/{peer_id}/search"),
opts,
)
.await
}
}
+211
View File
@@ -0,0 +1,211 @@
use std::collections::HashMap;
use crate::{
client::HonchoClient,
error::Result,
models::*,
workspaces::page_query,
};
impl HonchoClient {
// ── CRUD ──────────────────────────────────────────────────────────────
pub async fn create_session(
&self,
workspace_id: &str,
body: &SessionCreate,
) -> Result<Session> {
self.post(&format!("/v3/workspaces/{workspace_id}/sessions"), body)
.await
}
pub async fn list_sessions(
&self,
workspace_id: &str,
params: &PageParams,
filter: &SessionGet,
) -> Result<Page<Session>> {
self.post_with_query(
&format!("/v3/workspaces/{workspace_id}/sessions/list"),
&page_query(params),
filter,
)
.await
}
pub async fn update_session(
&self,
workspace_id: &str,
session_id: &str,
body: &SessionUpdate,
) -> Result<Session> {
self.put(
&format!("/v3/workspaces/{workspace_id}/sessions/{session_id}"),
body,
)
.await
}
pub async fn delete_session(&self, workspace_id: &str, session_id: &str) -> Result<()> {
self.delete_ok(&format!(
"/v3/workspaces/{workspace_id}/sessions/{session_id}"
))
.await
}
/// Clone a session, optionally up to a specific message.
pub async fn clone_session(
&self,
workspace_id: &str,
session_id: &str,
up_to_message_id: Option<&str>,
) -> Result<Session> {
let q: Vec<(&str, String)> = up_to_message_id
.map(|mid| vec![("message_id", mid.to_owned())])
.unwrap_or_default();
self.post_with_query(
&format!("/v3/workspaces/{workspace_id}/sessions/{session_id}/clone"),
&q,
&serde_json::Value::Null, // no body
)
.await
}
// ── Peers in session ──────────────────────────────────────────────────
/// Add peers to a session.
pub async fn add_session_peers(
&self,
workspace_id: &str,
session_id: &str,
peers: &HashMap<String, SessionPeerConfig>,
) -> Result<serde_json::Value> {
self.post(
&format!("/v3/workspaces/{workspace_id}/sessions/{session_id}/peers"),
peers,
)
.await
}
/// Update peer configs in a session.
pub async fn update_session_peers(
&self,
workspace_id: &str,
session_id: &str,
peers: &HashMap<String, SessionPeerConfig>,
) -> Result<serde_json::Value> {
self.put(
&format!("/v3/workspaces/{workspace_id}/sessions/{session_id}/peers"),
peers,
)
.await
}
/// Remove peers from a session by their ids.
pub async fn remove_session_peers(
&self,
workspace_id: &str,
session_id: &str,
peer_ids: &[&str],
) -> Result<()> {
self.delete_json(
&format!("/v3/workspaces/{workspace_id}/sessions/{session_id}/peers"),
&peer_ids,
)
.await
}
/// List peers in a session.
pub async fn list_session_peers(
&self,
workspace_id: &str,
session_id: &str,
params: &PageParams,
) -> Result<Page<Peer>> {
self.get_with_query(
&format!("/v3/workspaces/{workspace_id}/sessions/{session_id}/peers"),
&page_query(params),
)
.await
}
/// Get config for a specific peer in a session.
pub async fn get_session_peer_config(
&self,
workspace_id: &str,
session_id: &str,
peer_id: &str,
) -> Result<SessionPeerConfig> {
self.get(&format!(
"/v3/workspaces/{workspace_id}/sessions/{session_id}/peers/{peer_id}/config"
))
.await
}
/// Update config for a specific peer in a session.
pub async fn update_session_peer_config(
&self,
workspace_id: &str,
session_id: &str,
peer_id: &str,
config: &SessionPeerConfig,
) -> Result<SessionPeerConfig> {
self.put(
&format!(
"/v3/workspaces/{workspace_id}/sessions/{session_id}/peers/{peer_id}/config"
),
config,
)
.await
}
// ── Context / Summaries ───────────────────────────────────────────────
/// Retrieve context for a session (messages + peer conclusions, token-budgeted).
pub async fn session_context(
&self,
workspace_id: &str,
session_id: &str,
tokens: Option<u32>,
search_query: Option<&str>,
) -> Result<serde_json::Value> {
let mut q: Vec<(&str, String)> = vec![];
if let Some(t) = tokens {
q.push(("tokens", t.to_string()));
}
if let Some(sq) = search_query {
q.push(("search_query", sq.to_owned()));
}
self.get_with_query(
&format!("/v3/workspaces/{workspace_id}/sessions/{session_id}/context"),
&q,
)
.await
}
pub async fn session_summaries(
&self,
workspace_id: &str,
session_id: &str,
) -> Result<serde_json::Value> {
self.get(&format!(
"/v3/workspaces/{workspace_id}/sessions/{session_id}/summaries"
))
.await
}
// ── Search ────────────────────────────────────────────────────────────
pub async fn search_session(
&self,
workspace_id: &str,
session_id: &str,
opts: &MessageSearchOptions,
) -> Result<Vec<Message>> {
self.post(
&format!("/v3/workspaces/{workspace_id}/sessions/{session_id}/search"),
opts,
)
.await
}
}
+97
View File
@@ -0,0 +1,97 @@
use crate::{
client::{parse_no_content, HonchoClient},
error::Result,
models::*,
};
impl HonchoClient {
// ── CRUD ──────────────────────────────────────────────────────────────
/// Create (or retrieve if already exists) a workspace.
pub async fn create_workspace(&self, body: &WorkspaceCreate) -> Result<Workspace> {
self.post("/v3/workspaces", body).await
}
/// List workspaces (admin only).
pub async fn list_workspaces(
&self,
params: &PageParams,
filter: &WorkspaceGet,
) -> Result<Page<Workspace>> {
self.post_with_query("/v3/workspaces/list", &page_query(params), filter)
.await
}
/// Update a workspace.
pub async fn update_workspace(
&self,
workspace_id: &str,
body: &WorkspaceUpdate,
) -> Result<Workspace> {
self.put(&format!("/v3/workspaces/{workspace_id}"), body)
.await
}
/// Delete a workspace (queued, async on server side).
pub async fn delete_workspace(&self, workspace_id: &str) -> Result<()> {
let url = self.url(&format!("/v3/workspaces/{workspace_id}"));
let resp = self
.http
.delete(&url)
.bearer_auth(&self.token)
.send()
.await?;
parse_no_content(resp).await
}
// ── Search ────────────────────────────────────────────────────────────
pub async fn search_workspace(
&self,
workspace_id: &str,
opts: &MessageSearchOptions,
) -> Result<Vec<Message>> {
self.post(&format!("/v3/workspaces/{workspace_id}/search"), opts)
.await
}
// ── Queue ─────────────────────────────────────────────────────────────
pub async fn queue_status(
&self,
workspace_id: &str,
observer_id: Option<&str>,
sender_id: Option<&str>,
session_id: Option<&str>,
) -> Result<QueueStatus> {
let mut q: Vec<(&str, String)> = vec![];
if let Some(v) = observer_id {
q.push(("observer_id", v.to_owned()));
}
if let Some(v) = sender_id {
q.push(("sender_id", v.to_owned()));
}
if let Some(v) = session_id {
q.push(("session_id", v.to_owned()));
}
self.get_with_query(&format!("/v3/workspaces/{workspace_id}/queue/status"), &q)
.await
}
}
// ── helpers ───────────────────────────────────────────────────────────────
pub(crate) fn page_query(p: &PageParams) -> Vec<(&'static str, String)> {
let mut q = vec![];
if let Some(v) = p.page {
q.push(("page", v.to_string()));
}
if let Some(v) = p.size {
q.push(("size", v.to_string()));
}
if let Some(v) = p.reverse {
q.push(("reverse", v.to_string()));
}
q
}
+13
View File
@@ -0,0 +1,13 @@
[package]
name = "llm-client"
version = "0.1.0"
edition = "2024"
[dependencies]
core-api = { path = "../core-api" }
reqwest = { version = "0.13", default-features = false, features = ["rustls-no-provider", "charset", "http2", "system-proxy", "json"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
async-trait = "0.1"
anyhow = "1"
tracing = "0.1"
+366
View File
@@ -0,0 +1,366 @@
use async_trait::async_trait;
use serde_json::{Value, json};
use tracing::{debug, info, trace, warn};
use crate::{ChatOptions, ChatResponse, ChatbotClient, LlmRawMeta, LlmTurn, Message, Role, ToolCall, headers_to_json, redact_key};
const DEFAULT_BASE_URL: &str = "https://api.anthropic.com";
const ANTHROPIC_VERSION: &str = "2023-06-01";
pub struct AnthropicClient {
base_url: String,
api_key: String,
/// Extra top-level request-body keys merged into every request (e.g. the
/// `thinking` config for extended reasoning). See `apply_extra`.
extra_body: Option<Value>,
http: reqwest::Client,
}
impl AnthropicClient {
pub fn new(api_key: impl Into<String>) -> Self {
Self::with_base_url(DEFAULT_BASE_URL, api_key)
}
pub fn with_base_url(base_url: impl Into<String>, api_key: impl Into<String>) -> Self {
Self {
base_url: base_url.into(),
api_key: api_key.into(),
extra_body: None,
http: reqwest::Client::new(),
}
}
/// Like `new` but with extra request-body keys (e.g. `{"thinking": {...}}`).
pub fn with_extra_body(api_key: impl Into<String>, extra_body: Option<Value>) -> Self {
Self {
base_url: DEFAULT_BASE_URL.to_string(),
api_key: api_key.into(),
extra_body,
http: reqwest::Client::new(),
}
}
/// Merges `extra_body` into `body` and enforces Anthropic's extended-thinking
/// constraints: when `thinking` is enabled, `temperature` is not allowed and
/// `max_tokens` must be strictly greater than `budget_tokens`.
fn apply_extra(&self, body: &mut Value) {
let Some(extra) = self.extra_body.as_ref().and_then(|v| v.as_object()) else { return };
let Some(obj) = body.as_object_mut() else { return };
for (k, v) in extra {
obj.insert(k.clone(), v.clone());
}
if obj.get("thinking").map(|t| t["type"] == json!("enabled")).unwrap_or(false) {
obj.remove("temperature");
let budget = obj["thinking"]["budget_tokens"].as_i64().unwrap_or(0);
let cur_max = obj.get("max_tokens").and_then(|v| v.as_i64()).unwrap_or(4096);
if budget > 0 && cur_max <= budget {
obj.insert("max_tokens".to_string(), json!(budget + 4096));
}
}
}
/// Converts OpenAI-format tool definitions to Anthropic format.
/// OpenAI: { "type": "function", "function": { "name", "description", "parameters" } }
/// Anthropic: { "name", "description", "input_schema" }
fn convert_tools(tools: &[Value]) -> Vec<Value> {
tools
.iter()
.filter_map(|t| {
let func = &t["function"];
let name = func["name"].as_str()?;
Some(json!({
"name": name,
"description": func["description"].as_str().unwrap_or(""),
"input_schema": func["parameters"],
}))
})
.collect()
}
/// Converts OpenAI-format message array to Anthropic format.
///
/// Key differences:
/// - System messages are skipped (extracted separately).
/// - Assistant messages with `tool_calls` become content arrays with `tool_use` blocks.
/// - `tool` role messages are grouped into `user` messages with `tool_result` blocks.
fn convert_messages(messages: &[Value]) -> Vec<Value> {
let mut out: Vec<Value> = Vec::new();
let mut i = 0;
while i < messages.len() {
let msg = &messages[i];
let role = msg["role"].as_str().unwrap_or("");
match role {
"system" => { i += 1; }
"user" => {
out.push(json!({
"role": "user",
"content": msg["content"].as_str().unwrap_or(""),
}));
i += 1;
}
"assistant" => {
if let Some(tool_calls) = msg["tool_calls"].as_array() {
let mut content: Vec<Value> = Vec::new();
let text = msg["content"].as_str().unwrap_or("");
if !text.is_empty() {
content.push(json!({ "type": "text", "text": text }));
}
for tc in tool_calls {
let id = tc["id"].as_str().unwrap_or("");
let name = tc["function"]["name"].as_str().unwrap_or("");
let args_str = tc["function"]["arguments"].as_str().unwrap_or("{}");
let input: Value = serde_json::from_str(args_str)
.unwrap_or(Value::Object(Default::default()));
content.push(json!({
"type": "tool_use",
"id": id,
"name": name,
"input": input,
}));
}
out.push(json!({ "role": "assistant", "content": content }));
} else {
out.push(json!({
"role": "assistant",
"content": msg["content"].as_str().unwrap_or(""),
}));
}
i += 1;
}
"tool" => {
// Group all consecutive tool-result messages into a single user message.
let mut results: Vec<Value> = Vec::new();
while i < messages.len() && messages[i]["role"].as_str() == Some("tool") {
let tm = &messages[i];
results.push(json!({
"type": "tool_result",
"tool_use_id": tm["tool_call_id"].as_str().unwrap_or(""),
"content": tm["content"].as_str().unwrap_or(""),
}));
i += 1;
}
out.push(json!({ "role": "user", "content": results }));
}
_ => { i += 1; }
}
}
out
}
}
#[async_trait]
impl ChatbotClient for AnthropicClient {
async fn chat(
&self,
messages: &[Message],
options: &ChatOptions,
) -> anyhow::Result<ChatResponse> {
// Merge all system-role messages into a single `system:` parameter.
let system: Option<String> = {
let parts: Vec<&str> = messages
.iter()
.filter(|m| m.role == Role::System)
.map(|m| m.content.as_str())
.collect();
if parts.is_empty() { None } else { Some(parts.join("\n\n---\n\n")) }
};
let msgs: Vec<Value> = messages
.iter()
.filter(|m| m.role != Role::System)
.map(|m| {
let role = match m.role {
Role::User => "user",
Role::Assistant => "assistant",
Role::System => unreachable!(),
};
json!({ "role": role, "content": m.content })
})
.collect();
let max_tokens = options.max_tokens.unwrap_or(4096);
let mut body = json!({
"model": options.model,
"max_tokens": max_tokens,
"messages": msgs,
});
if let Some(sys) = system { body["system"] = sys.into(); }
if let Some(t) = options.temperature { body["temperature"] = t.into(); }
self.apply_extra(&mut body);
let url = format!("{}/v1/messages", self.base_url.trim_end_matches('/'));
debug!(model = %options.model, "anthropic: sending chat request");
trace!(body = %body, "anthropic: chat request body");
let resp: Value = self
.http
.post(&url)
.header("x-api-key", &self.api_key)
.header("anthropic-version", ANTHROPIC_VERSION)
.json(&body)
.send()
.await?
.error_for_status()?
.json()
.await?;
let content = resp["content"]
.as_array()
.and_then(|arr| arr.first())
.and_then(|block| block["text"].as_str())
.ok_or_else(|| anyhow::anyhow!("Missing content in Anthropic response"))?
.to_string();
let input_tokens = resp["usage"]["input_tokens"].as_u64().map(|n| n as u32);
let output_tokens = resp["usage"]["output_tokens"].as_u64().map(|n| n as u32);
let cache_read_tokens = resp["usage"]["cache_read_input_tokens"].as_u64().map(|n| n as u32);
let cache_creation_tokens = resp["usage"]["cache_creation_input_tokens"].as_u64().map(|n| n as u32);
info!(model = %options.model, ?input_tokens, ?output_tokens, "anthropic: chat response received");
let cost = self.extract_cost(&resp);
Ok(ChatResponse { content, input_tokens, output_tokens, truncated: false, reasoning_content: None, cache_read_tokens, cache_creation_tokens, cost })
}
async fn chat_with_tools(
&self,
messages: &[Value],
tools: &[Value],
options: &ChatOptions,
) -> anyhow::Result<LlmTurn> {
self.chat_with_tools_raw(messages, tools, options).await.map(|(t, _)| t)
}
async fn chat_with_tools_raw(
&self,
messages: &[Value],
tools: &[Value],
options: &ChatOptions,
) -> anyhow::Result<(LlmTurn, Option<LlmRawMeta>)> {
// Collect ALL system-role messages (main prompt, mid-conversation
// summary, tail_reminder) and merge them into a single `system:`
// string. The Anthropic API only accepts a single system parameter;
// mid-conversation system messages generated by build_openai_messages
// are intentionally used for injecting compaction summaries and tail
// reminders — they must not be silently dropped.
let system: Option<String> = {
let parts: Vec<&str> = messages
.iter()
.filter(|m| m["role"].as_str() == Some("system"))
.filter_map(|m| m["content"].as_str())
.collect();
if parts.is_empty() { None } else { Some(parts.join("\n\n---\n\n")) }
};
let anthropic_messages = Self::convert_messages(messages);
let anthropic_tools = Self::convert_tools(tools);
let max_tokens = options.max_tokens.unwrap_or(4096);
let mut body = json!({
"model": options.model,
"max_tokens": max_tokens,
"messages": anthropic_messages,
"tools": anthropic_tools,
});
if let Some(sys) = system { body["system"] = sys.into(); }
if let Some(t) = options.temperature { body["temperature"] = t.into(); }
self.apply_extra(&mut body);
let url = format!("{}/v1/messages", self.base_url.trim_end_matches('/'));
debug!(model = %options.model, tools = tools.len(), "anthropic: sending chat_with_tools request");
trace!(body = %body, "anthropic: chat_with_tools request body");
// Capture request metadata for logging.
let request_body = body.clone();
let request_headers = json!({
"x-api-key": redact_key(&self.api_key),
"anthropic-version": ANTHROPIC_VERSION,
"content-type": "application/json",
});
let http_resp = self
.http
.post(&url)
.header("x-api-key", &self.api_key)
.header("anthropic-version", ANTHROPIC_VERSION)
.header("X-Title", core_api::APP_NAME)
.json(&body)
.send()
.await?
.error_for_status()?;
let response_headers = headers_to_json(http_resp.headers());
let resp_text = http_resp.text().await?;
let resp: Value = serde_json::from_str(&resp_text)
.map_err(|e| anyhow::anyhow!("anthropic: failed to parse response JSON: {e}\nbody: {resp_text}"))?;
let response_body: Value = serde_json::from_str(&resp_text).unwrap_or(Value::Null);
let raw_meta = LlmRawMeta {
request_headers: Some(request_headers),
request_body: Some(request_body),
response_headers: Some(response_headers),
response_body: Some(response_body),
};
let stop_reason = resp["stop_reason"].as_str().unwrap_or("");
let input_tokens = resp["usage"]["input_tokens"].as_u64().map(|n| n as u32);
let output_tokens = resp["usage"]["output_tokens"].as_u64().map(|n| n as u32);
let cache_read_tokens = resp["usage"]["cache_read_input_tokens"].as_u64().map(|n| n as u32);
let cache_creation_tokens = resp["usage"]["cache_creation_input_tokens"].as_u64().map(|n| n as u32);
let content_blocks = resp["content"].as_array().cloned().unwrap_or_default();
let cost = self.extract_cost(&resp);
info!(model = %options.model, ?input_tokens, ?output_tokens, stop_reason, "anthropic: chat_with_tools response received");
if stop_reason == "max_tokens" {
warn!(model = %options.model, ?output_tokens, "anthropic: response truncated (max_tokens reached)");
}
let has_tool_use = content_blocks.iter().any(|b| b["type"].as_str() == Some("tool_use"));
// Check content blocks directly: Anthropic sometimes returns stop_reason "end_turn"
// even when tool_use blocks are present, so stop_reason alone is not reliable.
let turn = if stop_reason == "tool_use" || has_tool_use {
let text: String = content_blocks
.iter()
.filter(|b| b["type"].as_str() == Some("text"))
.filter_map(|b| b["text"].as_str())
.collect::<Vec<_>>()
.join("\n");
let calls: Vec<ToolCall> = content_blocks
.iter()
.filter(|b| b["type"].as_str() == Some("tool_use"))
.map(|b| ToolCall {
id: b["id"].as_str().unwrap_or("").to_string(),
name: b["name"].as_str().unwrap_or("").to_string(),
arguments: b["input"].clone(),
})
.collect();
LlmTurn::ToolCalls { content: text, calls, input_tokens, output_tokens, reasoning_content: None, cache_read_tokens, cache_creation_tokens, cost }
} else {
let content = content_blocks
.iter()
.find(|b| b["type"].as_str() == Some("text"))
.and_then(|b| b["text"].as_str())
.unwrap_or("")
.to_string();
let truncated = stop_reason == "max_tokens";
LlmTurn::Message(ChatResponse { content, input_tokens, output_tokens, truncated, reasoning_content: None, cache_read_tokens, cache_creation_tokens, cost })
};
Ok((turn, Some(raw_meta)))
}
}
+33
View File
@@ -0,0 +1,33 @@
pub mod anthropic;
pub mod lm_studio;
pub mod ollama;
pub mod openai;
// Re-export the trait and all associated types from core-api so existing
// callers that import from `llm_client` continue to work unchanged.
pub use core_api::chatbot::{
ChatOptions, ChatResponse, ChatbotClient, LlmRawMeta, LlmTurn, Message, Role, ToolCall,
};
use serde_json::Value;
/// Converts a reqwest `HeaderMap` into a `serde_json::Value` object.
pub fn headers_to_json(headers: &reqwest::header::HeaderMap) -> Value {
let map: serde_json::Map<String, Value> = headers
.iter()
.map(|(k, v)| (
k.as_str().to_string(),
v.to_str().unwrap_or("<binary>").into(),
))
.collect();
Value::Object(map)
}
/// Returns a redacted preview of an API key: first 7 chars + "***".
pub fn redact_key(key: &str) -> String {
if key.len() > 7 {
format!("{}***", &key[..7])
} else {
"***".to_string()
}
}
+51
View File
@@ -0,0 +1,51 @@
use async_trait::async_trait;
use serde_json::Value;
use crate::{ChatOptions, ChatResponse, ChatbotClient, LlmRawMeta, LlmTurn, Message, openai::OpenAiClient};
/// LM Studio client.
///
/// LM Studio exposes an OpenAI-compatible `/v1` endpoint, so this is a thin
/// wrapper that defaults to `http://localhost:1234/v1` and requires no API key.
pub struct LmStudioClient {
inner: OpenAiClient,
}
impl LmStudioClient {
/// `base_url` defaults to `http://localhost:1234/v1` if `None`.
pub fn new(base_url: Option<impl Into<String>>) -> Self {
let url = base_url
.map(|u| u.into())
.unwrap_or_else(|| "http://localhost:1234/v1".to_string());
Self { inner: OpenAiClient::new(url, "", None, false) }
}
}
#[async_trait]
impl ChatbotClient for LmStudioClient {
async fn chat(
&self,
messages: &[Message],
options: &ChatOptions,
) -> anyhow::Result<ChatResponse> {
self.inner.chat(messages, options).await
}
async fn chat_with_tools(
&self,
messages: &[Value],
tools: &[Value],
options: &ChatOptions,
) -> anyhow::Result<LlmTurn> {
self.inner.chat_with_tools(messages, tools, options).await
}
async fn chat_with_tools_raw(
&self,
messages: &[Value],
tools: &[Value],
options: &ChatOptions,
) -> anyhow::Result<(LlmTurn, Option<LlmRawMeta>)> {
self.inner.chat_with_tools_raw(messages, tools, options).await
}
}
+76
View File
@@ -0,0 +1,76 @@
use async_trait::async_trait;
use serde_json::{Value, json};
use crate::{ChatOptions, ChatResponse, ChatbotClient, Message, Role};
/// Ollama client using the native `/api/chat` endpoint.
///
/// Defaults to `http://localhost:11434`. No API key required.
pub struct OllamaClient {
base_url: String,
http: reqwest::Client,
}
impl OllamaClient {
/// `base_url` defaults to `http://localhost:11434` if `None`.
pub fn new(base_url: Option<impl Into<String>>) -> Self {
let url = base_url
.map(|u| u.into())
.unwrap_or_else(|| "http://localhost:11434".to_string());
Self { base_url: url, http: reqwest::Client::new() }
}
}
#[async_trait]
impl ChatbotClient for OllamaClient {
async fn chat(
&self,
messages: &[Message],
options: &ChatOptions,
) -> anyhow::Result<ChatResponse> {
let msgs: Vec<Value> = messages
.iter()
.map(|m| {
let role = match m.role {
Role::System => "system",
Role::User => "user",
Role::Assistant => "assistant",
};
json!({ "role": role, "content": m.content })
})
.collect();
let mut options_obj = json!({});
if let Some(t) = options.temperature { options_obj["temperature"] = t.into(); }
if let Some(n) = options.max_tokens { options_obj["num_predict"] = n.into(); }
let body = json!({
"model": options.model,
"messages": msgs,
"stream": false,
"options": options_obj,
});
let url = format!("{}/api/chat", self.base_url.trim_end_matches('/'));
let resp: Value = self
.http
.post(&url)
.json(&body)
.send()
.await?
.error_for_status()?
.json()
.await?;
let content = resp["message"]["content"]
.as_str()
.ok_or_else(|| anyhow::anyhow!("Missing content in Ollama response"))?
.to_string();
let input_tokens = resp["prompt_eval_count"].as_u64().map(|n| n as u32);
let output_tokens = resp["eval_count"].as_u64().map(|n| n as u32);
Ok(ChatResponse { content, input_tokens, output_tokens, truncated: false, reasoning_content: None, cache_read_tokens: None, cache_creation_tokens: None, cost: None })
}
}
+262
View File
@@ -0,0 +1,262 @@
use async_trait::async_trait;
use serde_json::{Value, json};
use tracing::{debug, info, trace, warn};
use crate::{ChatOptions, ChatResponse, ChatbotClient, LlmRawMeta, LlmTurn, Message, Role, ToolCall, headers_to_json, redact_key};
use core_api::APP_NAME;
/// OpenAI ChatGPT client (also compatible with any OpenAI-spec endpoint).
pub struct OpenAiClient {
base_url: String,
api_key: String,
extra_params: Option<serde_json::Value>,
/// When true, Anthropic-compatible prompt-caching hints are injected:
/// - `anthropic-beta: prompt-caching-2024-07-31` header is sent.
/// - The last tool definition is tagged with `cache_control: {"type":"ephemeral"}`.
/// - System message content is expected to already be a content array with
/// `cache_control` on the static block (set by `build_openai_messages`).
/// Used for OpenRouter when routing to Anthropic models.
enable_prompt_cache: bool,
http: reqwest::Client,
}
impl OpenAiClient {
pub fn new(base_url: impl Into<String>, api_key: impl Into<String>, extra_params: Option<serde_json::Value>, enable_prompt_cache: bool) -> Self {
Self {
base_url: base_url.into(),
api_key: api_key.into(),
extra_params,
enable_prompt_cache,
http: reqwest::Client::new(),
}
}
/// Merges `extra_params` (if any) into `body`. Only top-level object keys are merged.
fn apply_extra(&self, body: &mut serde_json::Value) {
if let Some(serde_json::Value::Object(extra)) = &self.extra_params {
if let Some(b) = body.as_object_mut() {
for (k, v) in extra {
b.insert(k.clone(), v.clone());
}
}
}
}
fn url(&self) -> String {
format!("{}/chat/completions", self.base_url.trim_end_matches('/'))
}
}
#[async_trait]
impl ChatbotClient for OpenAiClient {
async fn chat(
&self,
messages: &[Message],
options: &ChatOptions,
) -> anyhow::Result<ChatResponse> {
let msgs: Vec<Value> = messages
.iter()
.map(|m| {
let role = match m.role {
Role::System => "system",
Role::User => "user",
Role::Assistant => "assistant",
};
json!({ "role": role, "content": m.content })
})
.collect();
let mut body = json!({
"model": options.model,
"messages": msgs,
});
if let Some(t) = options.max_tokens { body["max_tokens"] = t.into(); }
if let Some(t) = options.temperature { body["temperature"] = t.into(); }
self.apply_extra(&mut body);
debug!(model = %options.model, "openai: sending chat request");
trace!(body = %body, "openai: chat request body");
let resp: Value = self
.http
.post(self.url())
.bearer_auth(&self.api_key)
.header("X-Title", APP_NAME)
.json(&body)
.send()
.await?
.error_for_status()?
.json()
.await?;
let content = match resp["choices"][0]["message"]["content"].as_str() {
Some(s) => s.to_string(),
None => {
warn!(raw_response = %resp, "openai: chat() response has null content");
String::new()
}
};
let input_tokens = resp["usage"]["prompt_tokens"].as_u64().map(|n| n as u32);
let output_tokens = resp["usage"]["completion_tokens"].as_u64().map(|n| n as u32);
let cache_read_tokens = resp["usage"]["prompt_tokens_details"]["cached_tokens"].as_u64().map(|n| n as u32);
let truncated = resp["choices"][0]["finish_reason"].as_str() == Some("length");
let cost = self.extract_cost(&resp);
info!(model = %options.model, ?input_tokens, ?output_tokens, ?cost, truncated, "openai: chat response received");
Ok(ChatResponse { content, input_tokens, output_tokens, truncated, reasoning_content: None, cache_read_tokens, cache_creation_tokens: None, cost })
}
async fn chat_with_tools(
&self,
messages: &[Value],
tools: &[Value],
options: &ChatOptions,
) -> anyhow::Result<LlmTurn> {
self.chat_with_tools_raw(messages, tools, options).await.map(|(t, _)| t)
}
async fn chat_with_tools_raw(
&self,
messages: &[Value],
tools: &[Value],
options: &ChatOptions,
) -> anyhow::Result<(LlmTurn, Option<LlmRawMeta>)> {
let mut body = json!({
"model": options.model,
"messages": messages,
});
if !tools.is_empty() {
// When prompt caching is enabled, tag the last tool with cache_control
// so the entire tools array is included in the Anthropic KV cache prefix.
let tools_value: Value = if self.enable_prompt_cache {
let mut tagged = tools.to_vec();
if let Some(last) = tagged.last_mut() {
last["cache_control"] = json!({"type": "ephemeral"});
}
tagged.into()
} else {
tools.into()
};
body["tools"] = tools_value;
body["tool_choice"] = "auto".into();
}
if let Some(t) = options.max_tokens { body["max_tokens"] = t.into(); }
if let Some(t) = options.temperature { body["temperature"] = t.into(); }
self.apply_extra(&mut body);
debug!(model = %options.model, tools = tools.len(), prompt_cache = self.enable_prompt_cache, "openai: sending chat_with_tools request");
trace!(body = %body, "openai: chat_with_tools request body");
// Capture request metadata for logging.
let mut logged_headers = json!({
"authorization": format!("Bearer {}", redact_key(&self.api_key)),
"content-type": "application/json",
});
if self.enable_prompt_cache {
logged_headers["anthropic-beta"] = "prompt-caching-2024-07-31".into();
}
let request_body = body.clone();
let request_headers = logged_headers;
let mut req = self.http.post(self.url()).bearer_auth(&self.api_key).header("X-Title", APP_NAME);
if self.enable_prompt_cache {
req = req.header("anthropic-beta", "prompt-caching-2024-07-31");
}
let http_resp = req
.json(&body)
.send()
.await?;
let response_headers = headers_to_json(http_resp.headers());
let status = http_resp.status();
let resp_text = http_resp.text().await?;
if !status.is_success() {
return Err(anyhow::anyhow!(
"openai: HTTP {status} from {url}\nbody: {resp_text}",
url = self.url(),
));
}
let resp: Value = serde_json::from_str(&resp_text)
.map_err(|e| anyhow::anyhow!("openai: failed to parse response JSON: {e}\nbody: {resp_text}"))?;
let response_body: Value = serde_json::from_str(&resp_text).unwrap_or(Value::Null);
let raw_meta = LlmRawMeta {
request_headers: Some(request_headers),
request_body: Some(request_body),
response_headers: Some(response_headers),
response_body: Some(response_body),
};
let input_tokens = resp["usage"]["prompt_tokens"].as_u64().map(|n| n as u32);
let output_tokens = resp["usage"]["completion_tokens"].as_u64().map(|n| n as u32);
let cache_read_tokens = resp["usage"]["prompt_tokens_details"]["cached_tokens"].as_u64().map(|n| n as u32);
let cost = self.extract_cost(&resp);
let choice = &resp["choices"][0];
let message = &choice["message"];
let finish = choice["finish_reason"].as_str().unwrap_or("stop");
info!(model = %options.model, ?input_tokens, ?output_tokens, finish_reason = finish, "openai: chat_with_tools response received");
if finish == "length" {
warn!(model = %options.model, ?output_tokens, "openai: response truncated (max_tokens reached)");
}
// Thinking/reasoning content varies by provider:
// - DeepSeek: "reasoning_content" (must be echoed back on subsequent turns, even as "")
// - MiniMax M3 and others: "reasoning"
// We normalize to a single field and echo under both names in message_builder.
let reasoning_content = message["reasoning_content"].as_str()
.or_else(|| message["reasoning"].as_str())
.map(str::to_string);
let tool_calls_array = message["tool_calls"].as_array().filter(|a| !a.is_empty());
// Some models (e.g. Qwen via OpenRouter) return finish_reason "stop" even when
// tool_calls are present, so check the array directly rather than relying on finish_reason.
let turn = if finish == "tool_calls" || tool_calls_array.is_some() {
let content = message["content"].as_str().unwrap_or("").to_string();
let calls = tool_calls_array
.ok_or_else(|| anyhow::anyhow!("finish_reason=tool_calls but tool_calls array missing or empty"))?
.iter()
.map(|tc| {
let id = tc["id"].as_str().unwrap_or("").to_string();
let name = tc["function"]["name"].as_str().unwrap_or("").to_string();
let args: Value = tc["function"]["arguments"]
.as_str()
.and_then(|s| serde_json::from_str(s).ok())
.unwrap_or(Value::Object(Default::default()));
ToolCall { id, name, arguments: args }
})
.collect();
LlmTurn::ToolCalls { content, calls, input_tokens, output_tokens, reasoning_content, cache_read_tokens, cache_creation_tokens: None, cost }
} else {
// content can be null for thinking/reasoning models or when finish_reason="length".
// Fall back to empty string rather than erroring — the partial response is still
// useful and a hard error breaks the session.
let content = match message["content"].as_str() {
Some(s) => s.to_string(),
None => {
tracing::warn!(
finish_reason = finish,
?input_tokens,
?output_tokens,
raw_message = %message,
"OpenAI response has null content",
);
String::new()
}
};
let truncated = finish == "length";
LlmTurn::Message(ChatResponse { content, input_tokens, output_tokens, truncated, reasoning_content, cache_read_tokens, cache_creation_tokens: None, cost })
};
Ok((turn, Some(raw_meta)))
}
}
+14
View File
@@ -0,0 +1,14 @@
[package]
name = "mcp-client"
version = "0.1.0"
edition = "2024"
[dependencies]
tokio = { version = "1", features = ["sync", "process", "io-util", "time", "rt", "macros"] }
reqwest = { version = "0.13", default-features = false, features = ["rustls-no-provider", "charset", "http2", "system-proxy", "json"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
anyhow = "1"
tracing = "0.1"
async-trait = "0.1"
base64 = "0.22"
+29
View File
@@ -0,0 +1,29 @@
use std::collections::HashMap;
use serde::Deserialize;
#[derive(Debug, Deserialize, Clone)]
pub struct McpServerConfig {
pub name: String,
pub transport: McpTransport,
/// stdio only: the executable to spawn.
pub command: Option<String>,
/// stdio only: arguments passed to the command.
pub args: Option<Vec<String>>,
/// stdio only: extra environment variables (values support `${VAR}` interpolation).
pub env: Option<HashMap<String, String>>,
/// http only: base URL of the MCP server.
pub url: Option<String>,
/// http only: API key sent as `Authorization: Bearer <key>` (supports `${VAR}` interpolation).
pub api_key: Option<String>,
}
#[derive(Debug, Deserialize, Clone)]
#[serde(rename_all = "snake_case")]
pub enum McpTransport {
Stdio,
/// Streamable HTTP (remote MCP servers — previously called SSE).
Http,
/// Alias kept for backwards compatibility.
Sse,
}
+423
View File
@@ -0,0 +1,423 @@
use std::sync::Mutex;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use anyhow::Result;
use async_trait::async_trait;
use reqwest::header::{ACCEPT, AUTHORIZATION, CONTENT_TYPE, HeaderMap, HeaderValue};
use serde_json::{Value, json};
use tracing::{debug, warn};
use crate::{McpCallResult, McpServerClient, McpTool, extract_text, interpolate_env};
use crate::config::McpServerConfig;
const CALL_TIMEOUT_SECS: u64 = 120;
/// Best-effort cancellation for an in-flight HTTP `tools/call`: if dropped while
/// armed (a `/stop` drops the request future, or the request timed out), it POSTs
/// `notifications/cancelled` so the server can stop. Correlation is weaker than on
/// stdio — the server must map `requestId` to the abandoned POST, which not every
/// server does — hence best-effort. Disarmed once the server responds (or when a
/// non-timeout send error proves the server never received the request).
struct HttpCancelOnDrop {
id: u64,
client: reqwest::Client,
url: String,
headers: HeaderMap,
name: String,
reason: &'static str,
armed: bool,
}
impl HttpCancelOnDrop {
fn disarm(&mut self) {
self.armed = false;
}
}
impl Drop for HttpCancelOnDrop {
fn drop(&mut self) {
if !self.armed {
return;
}
let (id, client, url, headers, name, reason) =
(self.id, self.client.clone(), self.url.clone(), self.headers.clone(), self.name.clone(), self.reason);
tokio::spawn(async move {
debug!("MCP http '{name}': notifications/cancelled for request {id} ({reason})");
let _ = client.post(&url).headers(headers)
.json(&crate::cancelled_notification(id, reason))
.send().await;
});
}
}
/// Cooperative `tasks/cancel` for a block-and-poll `poll_task` over HTTP: if the
/// poll future is dropped while still polling (a `/stop`) or hits its deadline,
/// POST `tasks/cancel` best-effort. Disarmed once the task reaches a terminal state.
struct HttpTaskCancelOnDrop {
request_id: u64,
task_id: String,
client: reqwest::Client,
url: String,
headers: HeaderMap,
name: String,
armed: bool,
}
impl HttpTaskCancelOnDrop {
fn disarm(&mut self) {
self.armed = false;
}
}
impl Drop for HttpTaskCancelOnDrop {
fn drop(&mut self) {
if !self.armed {
return;
}
let (request_id, task_id, client, url, headers, name) =
(self.request_id, self.task_id.clone(), self.client.clone(), self.url.clone(), self.headers.clone(), self.name.clone());
tokio::spawn(async move {
debug!("MCP http '{name}': tasks/cancel for task {task_id}");
let _ = client.post(&url).headers(headers)
.json(&crate::tasks_cancel_request(request_id, &task_id))
.send().await;
});
}
}
pub struct McpHttpServer {
name: String,
url: String,
client: reqwest::Client,
headers: HeaderMap,
/// Set after the `initialize` response — required by stateful servers like Tavily.
session_id: Mutex<Option<String>>,
/// Protocol version negotiated in the `initialize` response (falls back to
/// [`crate::PROTOCOL_VERSION`]). Once set, echoed in the `MCP-Protocol-Version`
/// header on every post-initialize request, per the Streamable HTTP spec.
protocol_version: Mutex<Option<String>>,
next_id: AtomicU64,
tools: Vec<McpTool>,
/// Capabilities the server advertised in its `InitializeResult`. Captured so a
/// future Tasks polling loop can gate on `tasks` support; unused for now.
server_capabilities: Value,
}
impl McpHttpServer {
pub async fn start(cfg: &McpServerConfig) -> Result<Self> {
let url = cfg.url.as_deref()
.ok_or_else(|| anyhow::anyhow!("http server '{}' requires 'url'", cfg.name))?
.trim_end_matches('/')
.to_string();
let mut headers = HeaderMap::new();
headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
headers.insert(ACCEPT, HeaderValue::from_static("application/json, text/event-stream"));
if let Some(key) = &cfg.api_key {
let val = interpolate_env(key);
let bearer = format!("Bearer {val}");
headers.insert(AUTHORIZATION, bearer.parse()
.map_err(|_| anyhow::anyhow!("invalid api_key for '{}'", cfg.name))?);
}
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(CALL_TIMEOUT_SECS))
.build()?;
let server = McpHttpServer {
name: cfg.name.clone(),
url,
client,
headers,
session_id: Mutex::new(None),
protocol_version: Mutex::new(None),
next_id: AtomicU64::new(1),
tools: Vec::new(),
server_capabilities: json!({}),
};
let init = server.request("initialize", json!({
// The HTTP transport doesn't service the ElicitationHandler (stdio-only),
// so it must NOT advertise the elicitation capability.
"protocolVersion": crate::PROTOCOL_VERSION,
// Experimental Tasks marker only (recognise-but-don't-poll); see the
// stdio transport for the rationale behind keeping it under `experimental`.
"capabilities": { "experimental": { "tasks": {} } },
"clientInfo": { "name": "skald", "version": env!("CARGO_PKG_VERSION") },
})).await?;
// Capture the negotiated version (fall back to our own) so post-initialize
// requests can echo it in the MCP-Protocol-Version header; tolerate a
// downgrade with a warning rather than disconnecting.
let negotiated = init["protocolVersion"].as_str().unwrap_or(crate::PROTOCOL_VERSION);
if negotiated != crate::PROTOCOL_VERSION {
warn!("MCP http '{}': server negotiated protocol {negotiated} (we requested {}); proceeding",
server.name, crate::PROTOCOL_VERSION);
}
*server.protocol_version.lock().unwrap() = Some(negotiated.to_string());
// Capture the server's advertised capabilities for a future Tasks poller.
let server_capabilities = init.get("capabilities").cloned().unwrap_or_else(|| json!({}));
if let Err(e) = server.notify("notifications/initialized", json!({})).await {
warn!("MCP http '{}': initialized notification failed (ignoring): {e}", server.name);
}
// Follow `nextCursor` across pages so large tool lists aren't silently
// truncated; capped at `MAX_TOOL_PAGES` against a stuck cursor.
let mut tools: Vec<McpTool> = Vec::new();
let mut cursor: Option<String> = None;
for page_n in 0..crate::MAX_TOOL_PAGES {
let params = match &cursor {
Some(c) => json!({ "cursor": c }),
None => json!({}),
};
let page = server.request("tools/list", params).await?;
if let Some(arr) = page["tools"].as_array() {
tools.extend(arr.iter().map(|t| McpTool::from_json(&cfg.name, t)));
}
cursor = page["nextCursor"].as_str().filter(|s| !s.is_empty()).map(str::to_string);
if cursor.is_none() {
break;
}
if page_n + 1 == crate::MAX_TOOL_PAGES {
warn!("MCP http '{}': tools/list hit {}-page cap; some tools may be omitted",
server.name, crate::MAX_TOOL_PAGES);
}
}
Ok(McpHttpServer { tools, server_capabilities, ..server })
}
pub fn tools(&self) -> &[McpTool] {
&self.tools
}
/// Capabilities the server advertised at `initialize`. Exposed for a future
/// Tasks polling loop to gate on `tasks` support.
pub fn server_capabilities(&self) -> &Value {
&self.server_capabilities
}
pub async fn call_tool(&self, name: &str, args: Value) -> Result<McpCallResult> {
let mut params = json!({ "name": name, "arguments": args });
if self.wants_task(name) {
// Opt into deferred execution for a task-capable tool (experimental Tasks).
params["task"] = json!({});
}
let result = self.request("tools/call", params).await?;
if result["isError"].as_bool().unwrap_or(false) {
anyhow::bail!("MCP tool error: {}", extract_text(&result));
}
match crate::extract_call_result(&result) {
McpCallResult::Task(task) => self.poll_task(task).await,
other => Ok(other),
}
}
/// True when tool `name` advertises `execution.taskSupport` as `required`/
/// `optional`, so we opt into deferred (Task) execution.
fn wants_task(&self, name: &str) -> bool {
self.tools.iter()
.find(|t| t.name == name)
.and_then(|t| t.task_support.as_deref())
.is_some_and(|s| s == "required" || s == "optional")
}
/// Drives a deferred Task to completion (experimental Tasks, block-and-poll):
/// polls `tasks/get` until a terminal status, then fetches the real result via
/// `tasks/result`. A [`HttpTaskCancelOnDrop`] guard POSTs `tasks/cancel` if this
/// future is dropped (a `/stop`) or the deadline is hit. The overall wait is
/// bounded only by the task's `ttl`, so long tasks no longer hit the 120s wall.
async fn poll_task(&self, task: crate::CreateTaskResult) -> Result<McpCallResult> {
let task_id = task.task_id.as_str();
let cancel_id = self.next_id.fetch_add(1, Ordering::SeqCst);
let mut guard = HttpTaskCancelOnDrop {
request_id: cancel_id,
task_id: task.task_id.clone(),
client: self.client.clone(),
url: self.url.clone(),
headers: self.request_headers(),
name: self.name.clone(),
armed: true,
};
let deadline = crate::poll_deadline(task.ttl_ms);
let mut interval = task.poll_interval_ms;
loop {
tokio::time::sleep(crate::clamp_poll_interval(interval)).await;
if std::time::Instant::now() >= deadline {
anyhow::bail!("MCP http '{}' task '{task_id}' exceeded max wait", self.name);
}
let get = self.request("tasks/get", json!({ "taskId": task_id })).await?;
let Some(state) = crate::CreateTaskResult::parse(&get) else {
anyhow::bail!("MCP http '{}' task '{task_id}': malformed tasks/get response", self.name);
};
interval = state.poll_interval_ms.or(interval);
match state.status {
crate::TaskStatus::Working => continue,
crate::TaskStatus::Completed => break,
crate::TaskStatus::Failed => {
guard.disarm();
anyhow::bail!("MCP http '{}' task '{task_id}' failed: {}", self.name, extract_text(&get));
}
crate::TaskStatus::Cancelled => {
guard.disarm();
anyhow::bail!("MCP http '{}' task '{task_id}' was cancelled by the server", self.name);
}
crate::TaskStatus::InputRequired =>
anyhow::bail!("MCP http '{}' task '{task_id}' requires input mid-task, which isn't supported yet", self.name),
}
}
// Task is terminal (completed) — nothing left to cancel.
guard.disarm();
let result = self.request("tasks/result", json!({ "taskId": task_id })).await?;
if result["isError"].as_bool().unwrap_or(false) {
anyhow::bail!("MCP tool error: {}", extract_text(&result));
}
Ok(crate::extract_call_result(&result))
}
/// Builds per-request headers: the static base plus the captured
/// `Mcp-Session-Id` and `MCP-Protocol-Version`. Both are set only after the
/// `initialize` response, so they're naturally absent on the initialize call
/// itself (the spec scopes the version header to post-initialize requests).
fn request_headers(&self) -> HeaderMap {
let mut headers = self.headers.clone();
if let Some(sid) = self.session_id.lock().unwrap().as_deref() {
if let Ok(val) = HeaderValue::from_str(sid) {
headers.insert("mcp-session-id", val);
}
}
if let Some(ver) = self.protocol_version.lock().unwrap().as_deref() {
if let Ok(val) = HeaderValue::from_str(ver) {
headers.insert("mcp-protocol-version", val);
}
}
headers
}
async fn request(&self, method: &str, params: Value) -> Result<Value> {
let id = self.next_id.fetch_add(1, Ordering::SeqCst);
let body = json!({
"jsonrpc": "2.0",
"id": id,
"method": method,
"params": params,
});
let req_headers = self.request_headers();
// Arm a best-effort cancellation guard for cancellable operations only
// (`tools/call`): a `/stop` that drops this future, or a request timeout,
// then POSTs `notifications/cancelled`. Disarmed once the server responds.
let mut cancel_guard = (method == "tools/call").then(|| HttpCancelOnDrop {
id,
client: self.client.clone(),
url: self.url.clone(),
headers: req_headers.clone(),
name: self.name.clone(),
reason: "cancelled by client",
armed: true,
});
let resp = match self.client
.post(&self.url)
.headers(req_headers)
.json(&body)
.send()
.await
{
Ok(r) => r,
Err(e) => {
// A timeout may leave the server still working → cancel it. Any
// other send failure means the server never got the request, so
// there is nothing to cancel.
if let Some(g) = cancel_guard.as_mut() {
if e.is_timeout() { g.reason = "timeout"; } else { g.disarm(); }
}
anyhow::bail!("MCP http '{}' request failed: {e}", self.name);
}
};
// The server responded — the request completed on its side; disarm.
if let Some(g) = cancel_guard.as_mut() { g.disarm(); }
if let Some(sid) = resp.headers().get("mcp-session-id") {
if let Ok(sid_str) = sid.to_str() {
debug!("MCP http '{}': captured session id", self.name);
*self.session_id.lock().unwrap() = Some(sid_str.to_string());
}
}
let status = resp.status();
let content_type = resp.headers()
.get(CONTENT_TYPE)
.and_then(|v| v.to_str().ok())
.unwrap_or("")
.to_string();
let msg: Value = if content_type.contains("text/event-stream") {
parse_sse_response(resp).await
.map_err(|e| anyhow::anyhow!("MCP http '{}' SSE parse error: {e}", self.name))?
} else {
if !status.is_success() {
let body = resp.text().await.unwrap_or_default();
anyhow::bail!("MCP http '{}' HTTP {status}: {body}", self.name);
}
resp.json::<Value>().await
.map_err(|e| anyhow::anyhow!("MCP http '{}' JSON decode error: {e}", self.name))?
};
if let Some(error) = msg.get("error") {
anyhow::bail!("MCP http '{}' protocol error: {error}", self.name);
}
Ok(msg["result"].clone())
}
async fn notify(&self, method: &str, params: Value) -> Result<()> {
let body = json!({
"jsonrpc": "2.0",
"method": method,
"params": params,
});
let req_headers = self.request_headers();
self.client
.post(&self.url)
.headers(req_headers)
.json(&body)
.send()
.await
.map_err(|e| anyhow::anyhow!("MCP http notify failed: {e}"))?;
Ok(())
}
}
#[async_trait]
impl McpServerClient for McpHttpServer {
fn tools(&self) -> &[McpTool] { self.tools() }
async fn call_tool(&self, name: &str, args: Value) -> Result<McpCallResult> { self.call_tool(name, args).await }
}
async fn parse_sse_response(resp: reqwest::Response) -> Result<Value> {
let text = resp.text().await?;
for line in text.lines() {
let data = match line.strip_prefix("data:") {
Some(d) => d.trim(),
None => continue,
};
if data == "[DONE]" { break; }
if let Ok(msg) = serde_json::from_str::<Value>(data) {
if msg.get("result").is_some() || msg.get("error").is_some() {
return Ok(msg);
}
}
}
anyhow::bail!("no JSON-RPC result found in SSE response")
}
+565
View File
@@ -0,0 +1,565 @@
pub mod config;
pub mod http_server;
pub mod log;
pub mod server;
use async_trait::async_trait;
use serde_json::{Value, json};
pub use config::{McpServerConfig, McpTransport};
pub use log::{McpLogLine, McpLogTx};
pub use server::{
ElicitationAction, ElicitationHandler, ElicitationReply, ElicitationRequest, McpNotification,
};
/// MCP protocol version this client advertises in `initialize` and (over HTTP)
/// echoes in the `MCP-Protocol-Version` header on every post-initialize request.
/// Shared by both transports so they can never drift apart. The host tolerates a
/// server negotiating a different (older) version — see the HTTP transport, which
/// captures the server's reply and warns rather than disconnecting.
pub const PROTOCOL_VERSION: &str = "2025-11-25";
/// Safety cap on `tools/list` pagination: stop following `nextCursor` after this
/// many pages so a buggy or hostile server that never clears the cursor can't
/// loop the client forever.
pub(crate) const MAX_TOOL_PAGES: usize = 50;
/// Builds a `notifications/cancelled` message for an in-flight request. Shared by
/// both transports (like [`PROTOCOL_VERSION`]) so the wire shape can't drift. Per
/// the MCP spec the client MUST NOT cancel the `initialize` request; callers only
/// arm this for cancellable operations (`tools/call`).
pub(crate) fn cancelled_notification(request_id: u64, reason: &str) -> Value {
json!({
"jsonrpc": "2.0",
"method": "notifications/cancelled",
"params": { "requestId": request_id, "reason": reason },
})
}
/// Builds a `tasks/cancel` request for a task the client is abandoning
/// (experimental Tasks). Sent fire-and-forget from the poll cancel-guard — the
/// response is ignored — so it carries a pre-allocated `request_id`. Shared by both
/// transports.
pub(crate) fn tasks_cancel_request(request_id: u64, task_id: &str) -> Value {
json!({
"jsonrpc": "2.0",
"id": request_id,
"method": "tasks/cancel",
"params": { "taskId": task_id },
})
}
// ── Public types ──────────────────────────────────────────────────────────────
#[derive(Debug, Clone)]
pub struct McpTool {
pub server_name: String,
pub name: String,
pub description: String,
pub input_schema: Value,
/// Optional human-readable title (MCP 2025-06-18+).
pub title: Option<String>,
/// Optional JSON Schema describing the tool's structured output
/// (`outputSchema`, MCP 2025-06-18+). Captured for future validation and to
/// know a tool can return `structuredContent`.
pub output_schema: Option<Value>,
/// Optional tool annotations (`readOnlyHint`, `destructiveHint`, …), treated
/// as untrusted hints per the spec.
pub annotations: Option<Value>,
/// Optional per-tool Tasks negotiation (`execution.taskSupport`:
/// `required`/`optional`/`forbidden`, MCP 2025-11-25 experimental). Captured
/// as an untrusted hint so a future polling implementation knows which tools
/// may return a [`CreateTaskResult`]. See [`McpCallResult::Task`].
pub task_support: Option<String>,
}
impl McpTool {
/// Builds an [`McpTool`] from one entry of a `tools/list` `tools[]` array.
/// Shared by both transports so the field mapping (incl. the 2025-06-18+
/// `title`/`outputSchema`/`annotations`) stays in one place.
pub fn from_json(server_name: &str, t: &Value) -> McpTool {
McpTool {
server_name: server_name.to_string(),
name: t["name"].as_str().unwrap_or("").to_string(),
description: t["description"].as_str().unwrap_or("").to_string(),
input_schema: t.get("inputSchema").cloned().unwrap_or_else(|| json!({
"type": "object", "properties": {},
})),
title: t.get("title").and_then(Value::as_str).map(str::to_string),
output_schema: t.get("outputSchema").cloned(),
annotations: t.get("annotations").cloned(),
task_support: t.get("execution")
.and_then(|e| e.get("taskSupport"))
.and_then(Value::as_str)
.map(str::to_string),
}
}
pub fn tool_id(&self) -> String {
format!("mcp__{}__{}", self.server_name, self.name)
}
pub fn to_openai_definition(&self) -> Value {
let params = if self.input_schema.is_object() {
self.input_schema.clone()
} else {
json!({ "type": "object", "properties": {} })
};
json!({
"type": "function",
"function": {
"name": self.tool_id(),
"description": format!("[{}] {}", self.server_name, self.description),
"parameters": params,
}
})
}
}
/// Status of a configured MCP server.
#[derive(Debug, Clone, serde::Serialize)]
#[serde(tag = "status", rename_all = "snake_case")]
pub enum McpServerStatus {
Running { tools: Vec<String> },
Error { message: String },
Disabled,
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct McpServerInfo {
pub name: String,
pub transport: String,
pub description: Option<String>,
pub friendly_name: Option<String>,
#[serde(flatten)]
pub status: McpServerStatus,
}
// ── Tool-result media ──────────────────────────────────────────────────────────
/// Kind of a non-text content block carried in an MCP `tools/call` result.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum McpMediaKind {
Image,
Audio,
/// Embedded or linked resource (any other binary/file payload).
Resource,
}
/// Payload of one media content block: either decoded inline bytes (from an
/// `image`/`audio`/embedded-`resource` base64 field) or a link to a remote
/// resource (`resource_link`), which the crate does not download.
#[derive(Debug, Clone)]
pub enum McpMediaData {
/// Decoded bytes plus the server-declared MIME type.
Inline { bytes: Vec<u8>, mime: String },
/// A `resource_link` URI, passed through untouched (no fetch here).
Link { uri: String, mime: Option<String> },
}
/// One non-text content block extracted from a `tools/call` result. The crate
/// stays a generic transport: it decodes base64 to bytes but never touches the
/// disk — persistence is the host's job (`McpManager`).
#[derive(Debug, Clone)]
pub struct McpMedia {
pub kind: McpMediaKind,
pub data: McpMediaData,
}
// ── Tasks (MCP 2025-11-25, experimental) ────────────────────────────────────────
/// Lifecycle state of a Task (`CreateTaskResult.status`). `Working` transitions to
/// `Completed`/`Failed`/`Cancelled`; `InputRequired` means the receiver needs more
/// input (e.g. an elicitation) before it can continue.
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum TaskStatus {
Working,
Completed,
Failed,
Cancelled,
InputRequired,
}
impl TaskStatus {
/// A terminal status: the task will not change further, so polling stops.
/// `InputRequired` is **not** terminal (the task resumes once input is given),
/// but the current block-and-poll client can't fulfil input mid-task, so it
/// treats it as an error rather than looping forever — see `poll_task`.
pub fn is_terminal(self) -> bool {
matches!(self, TaskStatus::Completed | TaskStatus::Failed | TaskStatus::Cancelled)
}
}
/// Clamps a server-suggested `pollInterval` (ms) into a sane range: default 1s when
/// absent, floored at 500ms (don't hammer the server) and capped at 30s (stay
/// responsive). Shared by both transports so their poll cadence can't drift.
pub(crate) fn clamp_poll_interval(ms: Option<u64>) -> std::time::Duration {
std::time::Duration::from_millis(ms.unwrap_or(1000).clamp(500, 30_000))
}
/// Deadline after which the block-and-poll loop gives up on a task and cancels it:
/// the server's `ttl` when given, otherwise a generous 1-hour cap so a stuck task
/// can't pin a session forever.
pub(crate) fn poll_deadline(ttl_ms: Option<u64>) -> std::time::Instant {
let ttl = std::time::Duration::from_millis(ttl_ms.unwrap_or(3_600_000));
std::time::Instant::now() + ttl
}
/// A durable task handle returned by a receiver when a request is executed as a
/// deferred, pollable operation instead of blocking (MCP 2025-11-25 *experimental*
/// Tasks). The client polls `tasks/get` until a terminal status, then fetches the
/// real result via `tasks/result` (see `poll_task` on each transport). Also used to
/// parse each `tasks/get` response (same shape).
#[derive(Debug, Clone)]
pub struct CreateTaskResult {
pub task_id: String,
pub status: TaskStatus,
/// Suggested delay between `tasks/get` polls, in milliseconds.
pub poll_interval_ms: Option<u64>,
/// Time-to-live of the task handle, in milliseconds.
pub ttl_ms: Option<u64>,
}
impl CreateTaskResult {
/// Parses a `tools/call` result that carries a task handle. Returns `None` if
/// `v` is not a task-shaped object (no `taskId`). The handle may live at the
/// top level or under a `task` field, depending on the server.
pub(crate) fn parse(v: &Value) -> Option<CreateTaskResult> {
let obj = v.get("task").filter(|t| t.is_object()).unwrap_or(v);
let task_id = obj.get("taskId").and_then(Value::as_str)?.to_string();
let status = obj.get("status")
.and_then(|s| serde_json::from_value::<TaskStatus>(s.clone()).ok())
.unwrap_or(TaskStatus::Working);
Some(CreateTaskResult {
task_id,
status,
poll_interval_ms: obj.get("pollInterval").and_then(Value::as_u64),
ttl_ms: obj.get("ttl").and_then(Value::as_u64),
})
}
}
// ── Transport trait ───────────────────────────────────────────────────────────
/// Typed result of an MCP `tools/call`. Mirrors the host's tool-result shape
/// without coupling this crate to it; the host (`McpManager`) maps it. Per the
/// MCP spec, `structuredContent` is canonical when present (servers SHOULD also
/// mirror it in a `TextContent` block for backwards compatibility); we prefer it
/// and fall back to the joined `text` items, which fixes the silent empty-result
/// case for servers that omit the text mirror. Non-text content blocks
/// (`image`/`audio`/`resource`/`resource_link`) are surfaced via [`McpCallResult::Media`]
/// so the host can persist them instead of dropping them.
#[derive(Debug, Clone)]
pub enum McpCallResult {
/// Joined `text` content items.
Text(String),
/// Canonical structured payload from `structuredContent`.
Json(Value),
/// At least one non-text media block was present. `text`/`structured` carry
/// any accompanying textual/structured payload from the same result.
Media {
text: Option<String>,
structured: Option<Value>,
items: Vec<McpMedia>,
},
/// The server deferred the call as a Task (experimental Tasks, 2025-11-25) and
/// returned a durable handle instead of a result. Skald surfaces the handle;
/// polling for the real result is a follow-up.
Task(CreateTaskResult),
}
#[async_trait]
pub trait McpServerClient: Send + Sync {
fn tools(&self) -> &[McpTool];
async fn call_tool(&self, name: &str, args: Value) -> anyhow::Result<McpCallResult>;
}
// ── Helpers ───────────────────────────────────────────────────────────────────
/// Parses `mcp__<server>__<tool>` → `(server, tool)`.
pub fn parse_mcp_tool_name(name: &str) -> Option<(&str, &str)> {
let rest = name.strip_prefix("mcp__")?;
let sep = rest.find("__")?;
Some((&rest[..sep], &rest[sep + 2..]))
}
/// Extracts text content from an MCP tool result value (the `text` items of
/// `content[]`, plus the inline `text` of embedded text resources). Used for the
/// `isError` path and as the text component of a successful result.
pub(crate) fn extract_text(result: &Value) -> String {
result["content"]
.as_array()
.map(|arr| classify_content(arr).0)
.unwrap_or_default()
}
/// Decodes a standard-base64 string to bytes, returning `None` on malformed input.
fn decode_b64(s: &str) -> Option<Vec<u8>> {
use base64::Engine;
base64::engine::general_purpose::STANDARD.decode(s).ok()
}
/// Walks `content[]`, returning the joined text and any non-text media blocks
/// (`image`/`audio`/embedded-`resource`/`resource_link`). Base64 payloads are
/// decoded to bytes here; persistence is left to the host.
fn classify_content(content: &[Value]) -> (String, Vec<McpMedia>) {
let mut texts: Vec<String> = Vec::new();
let mut media: Vec<McpMedia> = Vec::new();
for item in content {
match item.get("type").and_then(Value::as_str) {
Some("text") => {
if let Some(t) = item.get("text").and_then(Value::as_str) {
texts.push(t.to_string());
}
}
Some(kind @ ("image" | "audio")) => {
let mime = item.get("mimeType").and_then(Value::as_str)
.unwrap_or("application/octet-stream").to_string();
if let Some(bytes) = item.get("data").and_then(Value::as_str).and_then(decode_b64) {
let kind = if kind == "image" { McpMediaKind::Image } else { McpMediaKind::Audio };
media.push(McpMedia { kind, data: McpMediaData::Inline { bytes, mime } });
}
}
Some("resource") => {
// Embedded resource: a base64 `blob` (binary) or inline `text`.
let res = item.get("resource");
let mime = res.and_then(|r| r.get("mimeType")).and_then(Value::as_str)
.unwrap_or("application/octet-stream").to_string();
if let Some(bytes) = res.and_then(|r| r.get("blob")).and_then(Value::as_str).and_then(decode_b64) {
media.push(McpMedia { kind: McpMediaKind::Resource, data: McpMediaData::Inline { bytes, mime } });
} else if let Some(t) = res.and_then(|r| r.get("text")).and_then(Value::as_str) {
texts.push(t.to_string());
}
}
Some("resource_link") => {
if let Some(uri) = item.get("uri").and_then(Value::as_str) {
let mime = item.get("mimeType").and_then(Value::as_str).map(str::to_string);
media.push(McpMedia {
kind: McpMediaKind::Resource,
data: McpMediaData::Link { uri: uri.to_string(), mime },
});
}
}
// Unknown/typeless block: capture a bare `text` field if present.
_ => {
if let Some(t) = item.get("text").and_then(Value::as_str) {
texts.push(t.to_string());
}
}
}
}
(texts.join("\n"), media)
}
/// Builds the typed result of an MCP `tools/call`. When any non-text media block
/// is present it returns [`McpCallResult::Media`] (so the host can persist the
/// bytes instead of dropping them). Otherwise it preserves the original
/// precedence: `structuredContent` is canonical when present, else the joined
/// `text` items — which also fixes the silent empty-result case for servers that
/// return only `structuredContent` without the recommended text mirror.
pub(crate) fn extract_call_result(result: &Value) -> McpCallResult {
// A Task handle (deferred execution) has no `content[]`; recognise it first so
// it isn't mistaken for an empty result.
if result.get("content").is_none() {
if let Some(task) = CreateTaskResult::parse(result) {
return McpCallResult::Task(task);
}
}
let content = result["content"].as_array().cloned().unwrap_or_default();
let (text, media) = classify_content(&content);
let structured = result.get("structuredContent").filter(|v| !v.is_null()).cloned();
if !media.is_empty() {
return McpCallResult::Media {
text: (!text.is_empty()).then_some(text),
structured,
items: media,
};
}
if let Some(sc) = structured {
return McpCallResult::Json(sc);
}
McpCallResult::Text(text)
}
/// Interpolates `${VAR}` references in a string from the process environment.
pub(crate) fn interpolate_env(s: &str) -> String {
let mut result = s.to_string();
loop {
let Some(start) = result.find("${") else { break };
let Some(rel_end) = result[start..].find('}') else { break };
let var_name = result[start + 2..start + rel_end].to_string();
let value = std::env::var(&var_name).unwrap_or_else(|_| {
tracing::warn!("MCP env var ${{{var_name}}} not set");
String::new()
});
result = format!("{}{}{}", &result[..start], value, &result[start + rel_end + 1..]);
}
result
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn b64(bytes: &[u8]) -> String {
use base64::Engine;
base64::engine::general_purpose::STANDARD.encode(bytes)
}
#[test]
fn image_block_becomes_media_with_decoded_bytes() {
let png = [0x89u8, b'P', b'N', b'G'];
let result = json!({
"content": [
{ "type": "text", "text": "here is your image" },
{ "type": "image", "data": b64(&png), "mimeType": "image/png" }
]
});
match extract_call_result(&result) {
McpCallResult::Media { text, items, structured } => {
assert_eq!(text.as_deref(), Some("here is your image"));
assert!(structured.is_none());
assert_eq!(items.len(), 1);
assert_eq!(items[0].kind, McpMediaKind::Image);
match &items[0].data {
McpMediaData::Inline { bytes, mime } => {
assert_eq!(bytes.as_slice(), &png);
assert_eq!(mime, "image/png");
}
_ => panic!("expected inline media"),
}
}
other => panic!("expected Media, got {other:?}"),
}
}
#[test]
fn text_only_stays_text() {
let result = json!({ "content": [ { "type": "text", "text": "plain" } ] });
match extract_call_result(&result) {
McpCallResult::Text(t) => assert_eq!(t, "plain"),
other => panic!("expected Text, got {other:?}"),
}
}
#[test]
fn structured_without_media_stays_json() {
let result = json!({
"content": [ { "type": "text", "text": "mirror" } ],
"structuredContent": { "ok": true }
});
match extract_call_result(&result) {
McpCallResult::Json(v) => assert_eq!(v, json!({ "ok": true })),
other => panic!("expected Json, got {other:?}"),
}
}
#[test]
fn cancelled_notification_has_request_id_and_reason() {
let msg = cancelled_notification(7, "timeout");
assert_eq!(msg["jsonrpc"], "2.0");
assert_eq!(msg["method"], "notifications/cancelled");
assert!(msg.get("id").is_none(), "notifications MUST NOT carry an id");
assert_eq!(msg["params"]["requestId"], 7);
assert_eq!(msg["params"]["reason"], "timeout");
}
#[test]
fn from_json_captures_task_support() {
let t = json!({
"name": "gen_video",
"execution": { "taskSupport": "required" }
});
let tool = McpTool::from_json("srv", &t);
assert_eq!(tool.task_support.as_deref(), Some("required"));
// Absent execution → None.
let plain = McpTool::from_json("srv", &json!({ "name": "echo" }));
assert!(plain.task_support.is_none());
}
#[test]
fn create_task_result_parses_top_level_and_nested() {
// Top-level handle.
let top = json!({
"taskId": "abc-123", "status": "working",
"pollInterval": 500, "ttl": 60000
});
let r = CreateTaskResult::parse(&top).expect("top-level task");
assert_eq!(r.task_id, "abc-123");
assert_eq!(r.status, TaskStatus::Working);
assert_eq!(r.poll_interval_ms, Some(500));
assert_eq!(r.ttl_ms, Some(60000));
// Nested under `task`, unknown status → defaults to Working.
let nested = json!({ "task": { "taskId": "x", "status": "input_required" } });
let r = CreateTaskResult::parse(&nested).expect("nested task");
assert_eq!(r.task_id, "x");
assert_eq!(r.status, TaskStatus::InputRequired);
// Not task-shaped → None.
assert!(CreateTaskResult::parse(&json!({ "content": [] })).is_none());
}
#[test]
fn task_status_is_terminal() {
assert!(TaskStatus::Completed.is_terminal());
assert!(TaskStatus::Failed.is_terminal());
assert!(TaskStatus::Cancelled.is_terminal());
assert!(!TaskStatus::Working.is_terminal());
assert!(!TaskStatus::InputRequired.is_terminal());
}
#[test]
fn clamp_poll_interval_defaults_and_bounds() {
use std::time::Duration;
assert_eq!(clamp_poll_interval(None), Duration::from_millis(1000)); // default
assert_eq!(clamp_poll_interval(Some(10)), Duration::from_millis(500)); // floor
assert_eq!(clamp_poll_interval(Some(99_999)), Duration::from_millis(30_000)); // cap
assert_eq!(clamp_poll_interval(Some(2000)), Duration::from_millis(2000)); // passthrough
}
#[test]
fn tasks_cancel_request_shape() {
let msg = tasks_cancel_request(42, "job-1");
assert_eq!(msg["jsonrpc"], "2.0");
assert_eq!(msg["id"], 42);
assert_eq!(msg["method"], "tasks/cancel");
assert_eq!(msg["params"]["taskId"], "job-1");
}
#[test]
fn extract_call_result_recognises_task_handle() {
let result = json!({ "taskId": "job-9", "status": "working", "ttl": 120000 });
match extract_call_result(&result) {
McpCallResult::Task(t) => {
assert_eq!(t.task_id, "job-9");
assert_eq!(t.status, TaskStatus::Working);
assert_eq!(t.ttl_ms, Some(120000));
}
other => panic!("expected Task, got {other:?}"),
}
}
#[test]
fn resource_link_passes_through_without_fetch() {
let result = json!({
"content": [ { "type": "resource_link", "uri": "https://x/y.mp4", "mimeType": "video/mp4" } ]
});
match extract_call_result(&result) {
McpCallResult::Media { items, .. } => match &items[0].data {
McpMediaData::Link { uri, mime } => {
assert_eq!(uri, "https://x/y.mp4");
assert_eq!(mime.as_deref(), Some("video/mp4"));
}
_ => panic!("expected link"),
},
other => panic!("expected Media, got {other:?}"),
}
}
}
+116
View File
@@ -0,0 +1,116 @@
//! Per-server diagnostic log lines.
//!
//! An MCP server produces diagnostics from three places: its process `stderr`
//! (stdio transport), MCP `notifications/message` log records, and connection
//! lifecycle events (start failure / disconnect). This crate stays a generic
//! transport — it *emits* [`McpLogLine`]s over a channel but never touches the
//! disk. The host (`McpManager`) owns the channel and writes each line to a
//! per-server file `logs/mcp/<name>.log`.
//!
//! Note: the MCP `logging` utility (`notifications/message` + `logging/setLevel`)
//! is deprecated from the 2026-07-28 draft (SEP-2577) in favour of `stderr`, so
//! `stderr` is the primary, future-proof source; `notifications/message` is
//! captured for interop with servers that still emit it.
use std::borrow::Cow;
use serde_json::Value;
use tokio::sync::mpsc;
/// A single diagnostic line from an MCP server, routed by the host to that
/// server's log file.
#[derive(Debug, Clone)]
pub struct McpLogLine {
/// The MCP server that produced the line (selects the target file).
pub server: String,
/// Severity/kind tag, written inside `[...]`: `"stderr"` for raw child
/// stderr, an MCP log level (`"debug"`..`"emergency"`) for
/// `notifications/message`, or `"lifecycle"` for start/disconnect events.
pub level: Cow<'static, str>,
/// Human-readable text, already flattened to a single line (no trailing `\n`).
pub text: String,
}
/// Channel the host installs to receive [`McpLogLine`]s from a running server.
pub type McpLogTx = mpsc::UnboundedSender<McpLogLine>;
impl McpLogLine {
/// A raw line drained from the child process's `stderr` (stdio transport).
/// The level is unknown, so it's tagged `stderr`; the text itself usually
/// carries the server's own `ERROR`/`WARNING` marker.
pub fn stderr(server: impl Into<String>, text: impl Into<String>) -> Self {
Self { server: server.into(), level: Cow::Borrowed("stderr"), text: text.into() }
}
/// A connection lifecycle event (start failure, disconnect). Emitted for every
/// transport so even HTTP servers — which have no `stderr` — get a connection
/// history in their file.
pub fn lifecycle(server: impl Into<String>, text: impl Into<String>) -> Self {
Self { server: server.into(), level: Cow::Borrowed("lifecycle"), text: text.into() }
}
/// Builds a line from the `params` of an MCP `notifications/message`
/// (`{ level, logger?, data }`). Missing `level` falls back to `info`; `data`
/// strings pass through, other JSON is compacted to one line so nothing is
/// lost and the file stays greppable.
pub fn from_message(server: impl Into<String>, params: &Value) -> Self {
let level = params.get("level").and_then(Value::as_str).unwrap_or("info").to_string();
let logger = params.get("logger").and_then(Value::as_str);
let text = format_message_data(logger, params.get("data"));
Self { server: server.into(), level: Cow::Owned(level), text }
}
}
/// Flattens the `data` of a `notifications/message` into one line, prefixing the
/// optional `logger` name. Strings pass through; any other JSON is compacted.
fn format_message_data(logger: Option<&str>, data: Option<&Value>) -> String {
let body = match data {
Some(Value::String(s)) => s.clone(),
Some(v) => serde_json::to_string(v).unwrap_or_else(|_| v.to_string()),
None => String::new(),
};
match logger {
Some(l) if !l.is_empty() => format!("{l}: {body}"),
_ => body,
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn from_message_uses_level_and_string_data() {
let line = McpLogLine::from_message("srv", &json!({ "level": "error", "data": "boom" }));
assert_eq!(line.level, "error");
assert_eq!(line.text, "boom");
}
#[test]
fn from_message_compacts_object_data_and_prefixes_logger() {
// firecrawl-shaped payload.
let params = json!({
"level": "info",
"logger": "scraper",
"data": { "message": "Scraping URL", "context": { "url": "https://x/y" } }
});
let line = McpLogLine::from_message("firecrawl", &params);
assert_eq!(line.level, "info");
assert!(line.text.starts_with("scraper: "));
assert!(line.text.contains("Scraping URL"));
assert!(line.text.contains("https://x/y"));
}
#[test]
fn from_message_defaults_level_to_info() {
let line = McpLogLine::from_message("srv", &json!({ "data": "hi" }));
assert_eq!(line.level, "info");
}
#[test]
fn stderr_and_lifecycle_tags() {
assert_eq!(McpLogLine::stderr("s", "x").level, "stderr");
assert_eq!(McpLogLine::lifecycle("s", "x").level, "lifecycle");
}
}
+608
View File
@@ -0,0 +1,608 @@
use std::collections::HashMap;
use std::process::Stdio;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::time::Duration;
use anyhow::Result;
use async_trait::async_trait;
use serde_json::{Value, json};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::process::{ChildStdin, Command};
use tokio::sync::{Mutex, mpsc, oneshot};
use tracing::{debug, warn};
use crate::{McpCallResult, McpLogLine, McpLogTx, McpServerClient, McpTool, extract_text, interpolate_env};
use crate::config::McpServerConfig;
/// A server-initiated notification from an MCP server: `(server_name, full JSON-RPC message)`.
pub type McpNotification = (String, Value);
const CALL_TIMEOUT_SECS: u64 = 120;
// ── Elicitation (MCP spec 2025-06-18) ──────────────────────────────────────────
/// A server-initiated elicitation request: the server asks the user for input
/// *during* a tool call (`elicitation/create`). The secret/value never reaches
/// the LLM and is never persisted.
#[derive(Debug, Clone)]
pub struct ElicitationRequest {
/// Human-readable message to show the user.
pub message: String,
/// JSON Schema (flat object of typed fields) describing the requested input.
pub requested_schema: Value,
}
/// The user's decision on an elicitation request.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ElicitationAction {
Accept,
Decline,
Cancel,
}
impl ElicitationAction {
fn as_str(self) -> &'static str {
match self {
ElicitationAction::Accept => "accept",
ElicitationAction::Decline => "decline",
ElicitationAction::Cancel => "cancel",
}
}
}
/// The reply sent back to the server for an `elicitation/create` request.
#[derive(Debug, Clone)]
pub struct ElicitationReply {
pub action: ElicitationAction,
/// Field values for `accept`; `None` for `decline`/`cancel`.
pub content: Option<Value>,
}
/// Bridges a server-initiated elicitation to whatever surfaces it to the user
/// (in Skald: the Agent Inbox via `ElicitationManager`). The crate writes the
/// JSON-RPC reply itself — the handler only produces the decision. The returned
/// value (and any secret it carries) flows straight to the server's stdin and is
/// never logged here.
#[async_trait]
pub trait ElicitationHandler: Send + Sync {
async fn handle(&self, server_name: &str, request: ElicitationRequest) -> ElicitationReply;
}
/// Serialises `msg` as a single JSON-RPC line and writes it to the child's stdin.
async fn write_json_line(stdin: &Arc<Mutex<ChildStdin>>, msg: &Value) {
if let Ok(mut line) = serde_json::to_string(msg) {
line.push('\n');
let _ = stdin.lock().await.write_all(line.as_bytes()).await;
}
}
/// Handles a server→client JSON-RPC request (e.g. `elicitation/create`) without
/// blocking the read-loop: the user reply may take minutes, so the work is
/// spawned and the JSON-RPC response is written back when it resolves.
fn handle_server_request(
server_name: &str,
msg: Value,
stdin: &Arc<Mutex<ChildStdin>>,
handler: &Option<Arc<dyn ElicitationHandler>>,
pending_elicitations: &Arc<AtomicUsize>,
) {
let id = msg.get("id").cloned().unwrap_or(Value::Null);
let method = msg.get("method").and_then(Value::as_str).unwrap_or("");
if method == "elicitation/create" {
let params = msg.get("params").cloned().unwrap_or_else(|| json!({}));
let request = ElicitationRequest {
message: params.get("message").and_then(Value::as_str).unwrap_or("").to_string(),
requested_schema: params.get("requestedSchema").cloned().unwrap_or_else(|| json!({})),
};
let stdin = Arc::clone(stdin);
let Some(handler) = handler.clone() else {
// Capability declared but no handler wired: cancel so the server
// doesn't hang waiting for input we can't collect.
tokio::spawn(async move {
write_json_line(&stdin, &json!({
"jsonrpc": "2.0", "id": id, "result": { "action": "cancel" },
})).await;
});
return;
};
pending_elicitations.fetch_add(1, Ordering::SeqCst);
let counter = Arc::clone(pending_elicitations);
let server = server_name.to_string();
tokio::spawn(async move {
let reply = handler.handle(&server, request).await;
let mut result = json!({ "action": reply.action.as_str() });
if reply.action == ElicitationAction::Accept {
if let Some(content) = reply.content {
result["content"] = content;
}
}
write_json_line(&stdin, &json!({
"jsonrpc": "2.0", "id": id, "result": result,
})).await;
counter.fetch_sub(1, Ordering::SeqCst);
});
} else {
// Unknown server→client request: reply method-not-found so the server
// isn't left hanging.
let stdin = Arc::clone(stdin);
let method = method.to_string();
tokio::spawn(async move {
write_json_line(&stdin, &json!({
"jsonrpc": "2.0", "id": id,
"error": { "code": -32601, "message": format!("method not found: {method}") },
})).await;
});
}
}
/// Guards an in-flight `tools/call`: if this is dropped while still armed — the
/// caller aborted the tool (a `/stop` drops the work future) or the call timed out
/// — it tells the server to stop via `notifications/cancelled` and drops the now
/// orphaned `pending` entry (which would otherwise leak). Disarmed once the server
/// replies or disconnects, where cancelling is pointless. The spec forbids
/// cancelling `initialize`, so only `tools/call` arms a guard.
struct CancelOnDrop {
id: u64,
stdin: Arc<Mutex<ChildStdin>>,
pending: Arc<Mutex<HashMap<u64, oneshot::Sender<Value>>>>,
name: String,
reason: &'static str,
armed: bool,
}
impl CancelOnDrop {
fn disarm(&mut self) {
self.armed = false;
}
}
impl Drop for CancelOnDrop {
fn drop(&mut self) {
if !self.armed {
return;
}
let (id, stdin, pending, name, reason) =
(self.id, Arc::clone(&self.stdin), Arc::clone(&self.pending), self.name.clone(), self.reason);
tokio::spawn(async move {
pending.lock().await.remove(&id);
debug!(target: "mcp_client", "[{name}] notifications/cancelled for request {id} ({reason})");
write_json_line(&stdin, &crate::cancelled_notification(id, reason)).await;
});
}
}
/// Cooperative `tasks/cancel` for a block-and-poll `poll_task`: if the poll future
/// is dropped while still polling (a `/stop`) or hits its deadline, tell the server
/// to abandon the task. Fire-and-forget — the response carries no `pending` entry,
/// so the read-loop simply discards it. Disarmed once the task reaches a terminal
/// state (or fails), where cancelling is pointless.
struct TaskCancelOnDrop {
request_id: u64,
task_id: String,
stdin: Arc<Mutex<ChildStdin>>,
name: String,
armed: bool,
}
impl TaskCancelOnDrop {
fn disarm(&mut self) {
self.armed = false;
}
}
impl Drop for TaskCancelOnDrop {
fn drop(&mut self) {
if !self.armed {
return;
}
let (request_id, task_id, stdin, name) =
(self.request_id, self.task_id.clone(), Arc::clone(&self.stdin), self.name.clone());
tokio::spawn(async move {
debug!(target: "mcp_client", "[{name}] tasks/cancel for task {task_id}");
write_json_line(&stdin, &crate::tasks_cancel_request(request_id, &task_id)).await;
});
}
}
pub struct McpServer {
name: String,
stdin: Arc<Mutex<ChildStdin>>,
pending: Arc<Mutex<HashMap<u64, oneshot::Sender<Value>>>>,
next_id: AtomicU64,
tools: Vec<McpTool>,
/// Number of in-flight server→client elicitations awaiting a user reply.
/// While > 0, an in-flight `tools/call` on this server must not time out
/// (the user may still be typing a password into the Inbox).
pending_elicitations: Arc<AtomicUsize>,
/// Capabilities the server advertised in its `InitializeResult`. Captured so a
/// future Tasks polling loop can gate on `tasks` support; unused for now.
server_capabilities: Value,
}
impl McpServer {
pub async fn start(
cfg: &McpServerConfig,
notification_tx: Option<mpsc::UnboundedSender<McpNotification>>,
log_tx: Option<McpLogTx>,
elicitation_handler: Option<Arc<dyn ElicitationHandler>>,
) -> Result<Self> {
let command = cfg.command.as_deref()
.ok_or_else(|| anyhow::anyhow!("stdio server '{}' requires 'command'", cfg.name))?;
let mut cmd = Command::new(command);
if let Some(args) = &cfg.args {
cmd.args(args);
}
if let Some(env_map) = &cfg.env {
for (k, v) in env_map {
cmd.env(k, interpolate_env(v));
}
}
cmd.stdin(Stdio::piped())
.stdout(Stdio::piped())
// Capture the child's stderr instead of inheriting it: many MCP
// servers (e.g. FastMCP) print a startup banner, deprecation
// warnings and INFO logs there, which would otherwise spill raw
// onto our console. We drain it into tracing at `debug` level so it
// stays quiet by default but is still available for diagnostics.
.stderr(Stdio::piped())
.kill_on_drop(true);
// Detach the child into its own process group so that a terminal
// Ctrl+C (SIGINT delivered to the whole foreground process group)
// does not reach it directly. Otherwise Python-based MCP servers
// catch the SIGINT and dump a KeyboardInterrupt traceback to stderr.
// Cleanup still happens via `kill_on_drop`: when the app shuts down
// and the reader task is dropped, the child receives SIGKILL silently.
#[cfg(unix)]
cmd.process_group(0);
let mut child = cmd.spawn()
.map_err(|e| anyhow::anyhow!("failed to spawn '{}': {e}", cfg.name))?;
let stdin = child.stdin.take()
.ok_or_else(|| anyhow::anyhow!("could not capture stdin for '{}'", cfg.name))?;
// Wrap stdin early so both the struct and the read-loop (which writes
// elicitation replies back to the server) can share the same handle.
let stdin = Arc::new(Mutex::new(stdin));
let stdout = child.stdout.take()
.ok_or_else(|| anyhow::anyhow!("could not capture stdout for '{}'", cfg.name))?;
// Drain the child's stderr into tracing at `debug` (so banners/warnings
// don't pollute our console at the default level) *and* forward each line
// to the host's per-server log file via `log_tx`. Per the MCP 2026 draft,
// `stderr` is the primary place servers should log, so this is the main
// diagnostic source for stdio transports.
if let Some(stderr) = child.stderr.take() {
let server_name_err = cfg.name.clone();
let log_tx_err = log_tx.clone();
tokio::spawn(async move {
let mut lines = BufReader::new(stderr).lines();
while let Ok(Some(line)) = lines.next_line().await {
if !line.trim().is_empty() {
debug!(target: "mcp_client", "[{server_name_err}] {line}");
if let Some(tx) = &log_tx_err {
let _ = tx.send(McpLogLine::stderr(server_name_err.clone(), &line));
}
}
}
});
}
let pending: Arc<Mutex<HashMap<u64, oneshot::Sender<Value>>>> =
Arc::new(Mutex::new(HashMap::new()));
let pending_elicitations = Arc::new(AtomicUsize::new(0));
let pending_bg = pending.clone();
let server_name_bg = cfg.name.clone();
let notification_tx_bg = notification_tx;
let log_tx_bg = log_tx;
let stdin_bg = Arc::clone(&stdin);
let elicitation_handler_bg = elicitation_handler;
let pending_elicitations_bg = Arc::clone(&pending_elicitations);
tokio::spawn(async move {
let mut child = child;
let mut lines = BufReader::new(stdout).lines();
loop {
match lines.next_line().await {
Ok(Some(line)) if !line.trim().is_empty() => {
if let Ok(msg) = serde_json::from_str::<Value>(&line) {
let has_method = msg.get("method").is_some();
let has_id = msg.get("id").map(|v| !v.is_null()).unwrap_or(false);
if has_method && has_id {
// Server → client request (e.g. elicitation/create):
// has *both* `method` and `id`, so it must be checked
// before the response branch below.
handle_server_request(
&server_name_bg, msg, &stdin_bg,
&elicitation_handler_bg, &pending_elicitations_bg,
);
} else if let Some(id) = msg["id"].as_u64() {
if let Some(tx) = pending_bg.lock().await.remove(&id) {
let _ = tx.send(msg);
}
} else if has_method {
// `notifications/message` is the MCP logging utility
// (deprecated 2026-07-28): route it to the per-server
// log file, not to the notification queue that feeds
// TIC — otherwise log records masquerade as business
// events. Every other notification (e.g. the custom
// `event/*` methods) flows on to `notification_tx`.
if msg.get("method").and_then(Value::as_str) == Some("notifications/message") {
if let Some(tx) = &log_tx_bg {
let params = &msg["params"];
let _ = tx.send(McpLogLine::from_message(server_name_bg.clone(), params));
}
} else if let Some(tx) = &notification_tx_bg {
let _ = tx.send((server_name_bg.clone(), msg));
}
}
}
}
Ok(Some(_)) => {}
_ => break,
}
}
let exit_info = match child.wait().await {
Ok(status) if !status.success() => format!(
"process exited with {}",
status.code().map(|c| c.to_string()).unwrap_or_else(|| "signal".into())
),
_ => "process exited unexpectedly".into(),
};
let error_msg = format!("MCP '{}' disconnected: {exit_info}", server_name_bg);
if let Some(tx) = &log_tx_bg {
let _ = tx.send(McpLogLine::lifecycle(server_name_bg.clone(), format!("disconnected: {exit_info}")));
}
for (_, tx) in pending_bg.lock().await.drain() {
let _ = tx.send(json!({
"jsonrpc": "2.0",
"error": { "code": -32000, "message": error_msg }
}));
}
});
let server = McpServer {
name: cfg.name.clone(),
stdin,
pending,
next_id: AtomicU64::new(1),
tools: Vec::new(),
pending_elicitations,
server_capabilities: json!({}),
};
let init = server.request("initialize", json!({
// Declare the elicitation capability (form mode) so servers know they
// may request input mid-call.
"protocolVersion": crate::PROTOCOL_VERSION,
"capabilities": {
"elicitation": {},
// Experimental Tasks marker: we *recognise* a CreateTaskResult
// (see McpCallResult::Task) but don't poll yet, so we deliberately
// avoid claiming the full `tasks` capability — that could make a
// server defer results we can't retrieve. Kept under `experimental`
// until the polling loop lands.
"experimental": { "tasks": {} },
},
"clientInfo": { "name": "skald", "version": env!("CARGO_PKG_VERSION") },
})).await?;
// Tolerate a server that negotiates a different (older) version — warn but
// keep going rather than disconnecting.
if let Some(v) = init["protocolVersion"].as_str() {
if v != crate::PROTOCOL_VERSION {
warn!("MCP '{}': server negotiated protocol {v} (we requested {}); proceeding",
server.name, crate::PROTOCOL_VERSION);
}
}
// Capture the server's advertised capabilities for a future Tasks poller.
let server_capabilities = init.get("capabilities").cloned().unwrap_or_else(|| json!({}));
server.notify("notifications/initialized", json!({})).await?;
// Follow `nextCursor` across pages so servers with large tool lists aren't
// silently truncated; capped at `MAX_TOOL_PAGES` against a stuck cursor.
let mut tools: Vec<McpTool> = Vec::new();
let mut cursor: Option<String> = None;
for page_n in 0..crate::MAX_TOOL_PAGES {
let params = match &cursor {
Some(c) => json!({ "cursor": c }),
None => json!({}),
};
let page = server.request("tools/list", params).await?;
if let Some(arr) = page["tools"].as_array() {
tools.extend(arr.iter().map(|t| McpTool::from_json(&cfg.name, t)));
}
cursor = page["nextCursor"].as_str().filter(|s| !s.is_empty()).map(str::to_string);
if cursor.is_none() {
break;
}
if page_n + 1 == crate::MAX_TOOL_PAGES {
warn!("MCP '{}': tools/list hit {}-page cap; some tools may be omitted",
server.name, crate::MAX_TOOL_PAGES);
}
}
Ok(McpServer { tools, server_capabilities, ..server })
}
pub fn tools(&self) -> &[McpTool] {
&self.tools
}
/// Capabilities the server advertised at `initialize`. Exposed for a future
/// Tasks polling loop to gate on `tasks` support.
pub fn server_capabilities(&self) -> &Value {
&self.server_capabilities
}
pub async fn call_tool(&self, name: &str, args: Value) -> Result<McpCallResult> {
let mut params = json!({ "name": name, "arguments": args });
if self.wants_task(name) {
// Opt into deferred execution for a task-capable tool (experimental
// Tasks). Per spec, adding the `task` field to the request is the
// opt-in for `tools/call` (the client is the requestor).
params["task"] = json!({});
}
let result = self.request("tools/call", params).await?;
if result["isError"].as_bool().unwrap_or(false) {
anyhow::bail!("MCP tool error: {}", extract_text(&result));
}
match crate::extract_call_result(&result) {
McpCallResult::Task(task) => self.poll_task(task).await,
other => Ok(other),
}
}
/// True when tool `name` advertises `execution.taskSupport` as `required`/
/// `optional`, so we opt into deferred (Task) execution.
fn wants_task(&self, name: &str) -> bool {
self.tools.iter()
.find(|t| t.name == name)
.and_then(|t| t.task_support.as_deref())
.is_some_and(|s| s == "required" || s == "optional")
}
/// Drives a deferred Task to completion (experimental Tasks, block-and-poll):
/// polls `tasks/get` until a terminal status, then fetches the real result via
/// `tasks/result`. A [`TaskCancelOnDrop`] guard sends `tasks/cancel` if this
/// future is dropped (a `/stop`) or the deadline is hit. Each poll request is a
/// normal short `request()` (subject to the 120s timeout); the *overall* wait is
/// bounded only by the task's `ttl` — so long tasks no longer hit that wall.
async fn poll_task(&self, task: crate::CreateTaskResult) -> Result<McpCallResult> {
let task_id = task.task_id.as_str();
let cancel_id = self.next_id.fetch_add(1, Ordering::SeqCst);
let mut guard = TaskCancelOnDrop {
request_id: cancel_id,
task_id: task.task_id.clone(),
stdin: Arc::clone(&self.stdin),
name: self.name.clone(),
armed: true,
};
let deadline = crate::poll_deadline(task.ttl_ms);
let mut interval = task.poll_interval_ms;
loop {
tokio::time::sleep(crate::clamp_poll_interval(interval)).await;
if std::time::Instant::now() >= deadline {
anyhow::bail!("MCP '{}' task '{task_id}' exceeded max wait", self.name);
}
let get = self.request("tasks/get", json!({ "taskId": task_id })).await?;
let Some(state) = crate::CreateTaskResult::parse(&get) else {
anyhow::bail!("MCP '{}' task '{task_id}': malformed tasks/get response", self.name);
};
interval = state.poll_interval_ms.or(interval);
match state.status {
crate::TaskStatus::Working => continue,
crate::TaskStatus::Completed => break,
crate::TaskStatus::Failed => {
guard.disarm();
anyhow::bail!("MCP '{}' task '{task_id}' failed: {}", self.name, extract_text(&get));
}
crate::TaskStatus::Cancelled => {
guard.disarm();
anyhow::bail!("MCP '{}' task '{task_id}' was cancelled by the server", self.name);
}
// Still alive but blocked on input we can't supply mid-task: abandon
// it (guard stays armed → tasks/cancel). See follow-up.
crate::TaskStatus::InputRequired =>
anyhow::bail!("MCP '{}' task '{task_id}' requires input mid-task, which isn't supported yet", self.name),
}
}
// Task is terminal (completed) — nothing left to cancel.
guard.disarm();
let result = self.request("tasks/result", json!({ "taskId": task_id })).await?;
if result["isError"].as_bool().unwrap_or(false) {
anyhow::bail!("MCP tool error: {}", extract_text(&result));
}
Ok(crate::extract_call_result(&result))
}
async fn request(&self, method: &str, params: Value) -> Result<Value> {
let id = self.next_id.fetch_add(1, Ordering::SeqCst);
let (tx, rx) = oneshot::channel();
self.pending.lock().await.insert(id, tx);
// Arm a cancellation guard for cancellable operations only (`tools/call`):
// if this future is dropped by a `/stop` or times out before the server
// replies, the guard notifies the server. Disarmed on a real reply.
let mut cancel_guard = (method == "tools/call").then(|| CancelOnDrop {
id,
stdin: Arc::clone(&self.stdin),
pending: Arc::clone(&self.pending),
name: self.name.clone(),
reason: "cancelled by client",
armed: true,
});
let msg = json!({
"jsonrpc": "2.0",
"id": id,
"method": method,
"params": params,
});
let mut line = serde_json::to_string(&msg)?;
line.push('\n');
self.stdin.lock().await.write_all(line.as_bytes()).await?;
// Wait for the response, but re-arm the timeout while an elicitation is
// in flight on this server: the user may still be typing a secret into
// the Inbox, and the server won't answer `tools/call` until then.
tokio::pin!(rx);
let response = loop {
tokio::select! {
res = &mut rx => match res {
Ok(v) => break v,
Err(_) => {
// Server disconnected: nothing to cancel.
if let Some(g) = cancel_guard.as_mut() { g.disarm(); }
anyhow::bail!("MCP '{}' disconnected", self.name);
}
},
_ = tokio::time::sleep(Duration::from_secs(CALL_TIMEOUT_SECS)) => {
if self.pending_elicitations.load(Ordering::SeqCst) == 0 {
// Leave the guard armed so it fires `notifications/cancelled`;
// tag the reason as a timeout.
if let Some(g) = cancel_guard.as_mut() { g.reason = "timeout"; }
anyhow::bail!("MCP '{}' timed out on '{method}'", self.name);
}
// Elicitation pending: keep waiting for the user.
}
}
};
// Server replied (even with an error): the request completed — disarm.
if let Some(g) = cancel_guard.as_mut() { g.disarm(); }
if let Some(error) = response.get("error") {
anyhow::bail!("MCP '{}' protocol error: {error}", self.name);
}
Ok(response["result"].clone())
}
async fn notify(&self, method: &str, params: Value) -> Result<()> {
let msg = json!({
"jsonrpc": "2.0",
"method": method,
"params": params,
});
let mut line = serde_json::to_string(&msg)?;
line.push('\n');
self.stdin.lock().await.write_all(line.as_bytes()).await?;
Ok(())
}
}
#[async_trait]
impl McpServerClient for McpServer {
fn tools(&self) -> &[McpTool] { self.tools() }
async fn call_tool(&self, name: &str, args: Value) -> Result<McpCallResult> { self.call_tool(name, args).await }
}
+130
View File
@@ -0,0 +1,130 @@
//! End-to-end test of the elicitation path against a real stdio MCP subprocess.
//!
//! A tiny Python "server" exposes one tool that, when called, issues a
//! server→client `elicitation/create` and echoes back whatever value it receives.
//! This exercises the read-loop request branch, the spawned reply writer, the
//! capability handshake, and `content` passthrough. Skipped if `python3` is absent.
use std::io::Write;
use std::process::Command;
use std::sync::Arc;
use async_trait::async_trait;
use serde_json::json;
use mcp_client::config::{McpServerConfig, McpTransport};
use mcp_client::server::{
ElicitationAction, ElicitationHandler, ElicitationReply, ElicitationRequest, McpServer,
};
use mcp_client::McpCallResult;
const FAKE_SERVER: &str = r#"
import sys, json
def send(obj):
sys.stdout.write(json.dumps(obj) + "\n")
sys.stdout.flush()
def readline():
line = sys.stdin.readline()
if not line:
return None
line = line.strip()
return line if line else readline()
while True:
raw = readline()
if raw is None:
break
msg = json.loads(raw)
mid = msg.get("id")
method = msg.get("method")
if method == "initialize":
send({"jsonrpc": "2.0", "id": mid, "result": {
"protocolVersion": "2025-06-18", "capabilities": {},
"serverInfo": {"name": "fake", "version": "0"}}})
elif method == "notifications/initialized":
pass
elif method == "tools/list":
send({"jsonrpc": "2.0", "id": mid, "result": {"tools": [
{"name": "need_secret", "description": "asks for a secret",
"inputSchema": {"type": "object", "properties": {}}}]}})
elif method == "tools/call":
eid = "elicit-1"
send({"jsonrpc": "2.0", "id": eid, "method": "elicitation/create", "params": {
"message": "Enter password",
"requestedSchema": {"type": "object", "properties": {
"password": {"type": "string", "format": "password"}}}}})
reply = None
while reply is None:
r = readline()
if r is None:
break
rr = json.loads(r)
if rr.get("id") == eid:
reply = rr
val = (reply or {}).get("result", {}).get("content", {}).get("password", "<none>")
send({"jsonrpc": "2.0", "id": mid, "result": {
"content": [{"type": "text", "text": "got:" + val}], "isError": False}})
elif mid is not None:
send({"jsonrpc": "2.0", "id": mid, "error": {"code": -32601, "message": "unknown"}})
"#;
struct AcceptHandler;
#[async_trait]
impl ElicitationHandler for AcceptHandler {
async fn handle(&self, _server: &str, req: ElicitationRequest) -> ElicitationReply {
assert_eq!(req.message, "Enter password");
assert!(req.requested_schema.get("properties").is_some());
ElicitationReply {
action: ElicitationAction::Accept,
content: Some(json!({ "password": "hunter2" })),
}
}
}
fn python3_available() -> bool {
Command::new("python3").arg("--version").output().is_ok()
}
#[tokio::test]
async fn elicitation_roundtrip_returns_secret_to_server() {
if !python3_available() {
eprintln!("python3 not found — skipping elicitation integration test");
return;
}
let script_path = std::env::temp_dir().join(format!("skald_elicit_{}.py", std::process::id()));
std::fs::File::create(&script_path)
.unwrap()
.write_all(FAKE_SERVER.as_bytes())
.unwrap();
let cfg = McpServerConfig {
name: "fake".to_string(),
transport: McpTransport::Stdio,
command: Some("python3".to_string()),
args: Some(vec![script_path.to_string_lossy().to_string()]),
env: None,
url: None,
api_key: None,
};
let server = McpServer::start(&cfg, None, None, Some(Arc::new(AcceptHandler)))
.await
.expect("server should start");
let result = server
.call_tool("need_secret", json!({}))
.await
.expect("tool call should succeed");
// The fake server returns text content (no structuredContent) → Text variant.
match result {
McpCallResult::Text(s) => assert_eq!(s, "got:hunter2"),
other => panic!("expected Text, got {other:?}"),
}
let _ = std::fs::remove_file(&script_path);
}
+131
View File
@@ -0,0 +1,131 @@
//! End-to-end test of per-server diagnostic capture against a real stdio MCP
//! subprocess.
//!
//! A tiny Python "server" prints a banner to **stderr** and, right after
//! `notifications/initialized`, emits both a `notifications/message` log record and
//! a business `event/ping` notification to stdout. The test asserts that:
//! - the stderr banner arrives on `log_tx` tagged `stderr`,
//! - the `notifications/message` arrives on `log_tx` with its MCP level, and is
//! **not** delivered to `notification_tx` (it's diverted away from TIC),
//! - the business `event/ping` still arrives on `notification_tx`.
//! Skipped if `python3` is absent.
use std::io::Write;
use std::process::Command;
use std::time::Duration;
use mcp_client::config::{McpServerConfig, McpTransport};
use mcp_client::server::{McpNotification, McpServer};
use mcp_client::McpLogLine;
use serde_json::Value;
use tokio::sync::mpsc;
const FAKE_SERVER: &str = r#"
import sys, json
def send(obj):
sys.stdout.write(json.dumps(obj) + "\n")
sys.stdout.flush()
def readline():
line = sys.stdin.readline()
if not line:
return None
line = line.strip()
return line if line else readline()
while True:
raw = readline()
if raw is None:
break
msg = json.loads(raw)
mid = msg.get("id")
method = msg.get("method")
if method == "initialize":
send({"jsonrpc": "2.0", "id": mid, "result": {
"protocolVersion": "2025-11-25", "capabilities": {},
"serverInfo": {"name": "fake", "version": "0"}}})
elif method == "notifications/initialized":
# A diagnostic banner on stderr (the primary, future-proof log source).
print("startup banner on stderr", file=sys.stderr, flush=True)
# An MCP logging record (should be diverted to the log file, NOT TIC).
send({"jsonrpc": "2.0", "method": "notifications/message",
"params": {"level": "warning", "logger": "test", "data": "disk almost full"}})
# A business event (should still reach the notification queue / TIC).
send({"jsonrpc": "2.0", "method": "event/ping", "params": {"n": 1}})
elif method == "tools/list":
send({"jsonrpc": "2.0", "id": mid, "result": {"tools": []}})
elif mid is not None:
send({"jsonrpc": "2.0", "id": mid, "error": {"code": -32601, "message": "unknown"}})
"#;
fn python3_available() -> bool {
Command::new("python3").arg("--version").output().is_ok()
}
fn write_script() -> std::path::PathBuf {
let path = std::env::temp_dir().join(format!("skald_logging_{}.py", std::process::id()));
std::fs::File::create(&path).unwrap().write_all(FAKE_SERVER.as_bytes()).unwrap();
path
}
/// Drains a channel for up to `budget`, returning everything received.
async fn drain<T>(rx: &mut mpsc::UnboundedReceiver<T>, budget: Duration) -> Vec<T> {
let mut out = Vec::new();
let deadline = tokio::time::Instant::now() + budget;
while let Ok(Some(item)) = tokio::time::timeout_at(deadline, rx.recv()).await {
out.push(item);
}
out
}
#[tokio::test]
async fn stderr_and_log_records_are_captured_and_diverted() {
if !python3_available() {
eprintln!("python3 not found — skipping per-server logging integration test");
return;
}
let script = write_script();
let cfg = McpServerConfig {
name: "fake".to_string(),
transport: McpTransport::Stdio,
command: Some("python3".to_string()),
args: Some(vec![script.to_string_lossy().to_string()]),
env: None,
url: None,
api_key: None,
};
let (notif_tx, mut notif_rx) = mpsc::unbounded_channel::<McpNotification>();
let (log_tx, mut log_rx) = mpsc::unbounded_channel::<McpLogLine>();
let _server = McpServer::start(&cfg, Some(notif_tx), Some(log_tx), None)
.await
.expect("server should start");
let logs: Vec<McpLogLine> = drain(&mut log_rx, Duration::from_millis(1500)).await;
let notes: Vec<McpNotification> = drain(&mut notif_rx, Duration::from_millis(500)).await;
// stderr banner captured on the log channel.
assert!(
logs.iter().any(|l| l.level == "stderr" && l.text.contains("startup banner on stderr")),
"expected a [stderr] log line, got: {logs:?}",
);
// notifications/message captured with its MCP level + logger prefix.
assert!(
logs.iter().any(|l| l.level == "warning" && l.text.contains("disk almost full") && l.text.contains("test")),
"expected a [warning] log line from notifications/message, got: {logs:?}",
);
// The log record was diverted away from the notification queue…
assert!(
!notes.iter().any(|(_, msg)| msg.get("method").and_then(Value::as_str) == Some("notifications/message")),
"notifications/message must NOT reach the notification queue, got: {notes:?}",
);
// …while the business event still flowed through it.
assert!(
notes.iter().any(|(_, msg)| msg.get("method").and_then(Value::as_str) == Some("event/ping")),
"expected event/ping on the notification queue, got: {notes:?}",
);
let _ = std::fs::remove_file(&script);
}
+103
View File
@@ -0,0 +1,103 @@
//! End-to-end test of `tools/list` cursor pagination against a real stdio MCP
//! subprocess.
//!
//! A tiny Python "server" returns its tools across two pages: the first
//! `tools/list` (no `cursor`) yields one tool plus a `nextCursor`; the follow-up
//! (with that `cursor`) yields the rest and no `nextCursor`. This exercises the
//! cursor loop in `McpServer::start`. Skipped if `python3` is absent.
use std::io::Write;
use std::process::Command;
use mcp_client::config::{McpServerConfig, McpTransport};
use mcp_client::server::McpServer;
const FAKE_SERVER: &str = r#"
import sys, json
def send(obj):
sys.stdout.write(json.dumps(obj) + "\n")
sys.stdout.flush()
def readline():
line = sys.stdin.readline()
if not line:
return None
line = line.strip()
return line if line else readline()
while True:
raw = readline()
if raw is None:
break
msg = json.loads(raw)
mid = msg.get("id")
method = msg.get("method")
if method == "initialize":
send({"jsonrpc": "2.0", "id": mid, "result": {
"protocolVersion": "2025-11-25", "capabilities": {},
"serverInfo": {"name": "fake", "version": "0"}}})
elif method == "notifications/initialized":
pass
elif method == "tools/list":
cursor = (msg.get("params") or {}).get("cursor")
if cursor is None:
# Page 1: one tool + a cursor pointing at the next page.
send({"jsonrpc": "2.0", "id": mid, "result": {
"tools": [{"name": "alpha", "description": "first",
"inputSchema": {"type": "object", "properties": {}}}],
"nextCursor": "page-2"}})
elif cursor == "page-2":
# Page 2: the rest, no further cursor → loop terminates.
send({"jsonrpc": "2.0", "id": mid, "result": {
"tools": [{"name": "beta", "description": "second",
"inputSchema": {"type": "object", "properties": {}}},
{"name": "gamma", "description": "third",
"inputSchema": {"type": "object", "properties": {}}}]}})
else:
send({"jsonrpc": "2.0", "id": mid, "result": {"tools": []}})
elif mid is not None:
send({"jsonrpc": "2.0", "id": mid, "error": {"code": -32601, "message": "unknown"}})
"#;
fn python3_available() -> bool {
Command::new("python3").arg("--version").output().is_ok()
}
#[tokio::test]
async fn tools_list_follows_next_cursor_across_pages() {
if !python3_available() {
eprintln!("python3 not found — skipping pagination integration test");
return;
}
let script_path =
std::env::temp_dir().join(format!("skald_paginate_{}.py", std::process::id()));
std::fs::File::create(&script_path)
.unwrap()
.write_all(FAKE_SERVER.as_bytes())
.unwrap();
let cfg = McpServerConfig {
name: "fake".to_string(),
transport: McpTransport::Stdio,
command: Some("python3".to_string()),
args: Some(vec![script_path.to_string_lossy().to_string()]),
env: None,
url: None,
api_key: None,
};
let server = McpServer::start(&cfg, None, None, None)
.await
.expect("server should start");
let names: Vec<&str> = server.tools().iter().map(|t| t.name.as_str()).collect();
assert_eq!(
names,
vec!["alpha", "beta", "gamma"],
"all pages should be collected"
);
let _ = std::fs::remove_file(&script_path);
}
+160
View File
@@ -0,0 +1,160 @@
//! End-to-end test of the block-and-poll Tasks client against a real stdio MCP
//! subprocess.
//!
//! A tiny Python "server" exposes one tool with `execution.taskSupport: "required"`.
//! On `tools/call` it returns a `CreateTaskResult` (a durable handle, no `content`)
//! instead of a result; the client then polls `tasks/get` until `completed` and
//! fetches the real answer via `tasks/result`. A second mode never completes, so the
//! test can drop the call mid-poll and assert the client sends `tasks/cancel`.
//! Skipped if `python3` is absent.
use std::io::Write;
use std::process::Command;
use std::time::Duration;
use mcp_client::config::{McpServerConfig, McpTransport};
use mcp_client::server::McpServer;
use mcp_client::McpCallResult;
/// argv[1] = mode ("complete" | "cancel"); argv[2] = marker file (cancel mode).
const FAKE_SERVER: &str = r#"
import sys, json
mode = sys.argv[1] if len(sys.argv) > 1 else "complete"
marker = sys.argv[2] if len(sys.argv) > 2 else None
get_count = 0
def send(obj):
sys.stdout.write(json.dumps(obj) + "\n")
sys.stdout.flush()
def readline():
line = sys.stdin.readline()
if not line:
return None
line = line.strip()
return line if line else readline()
while True:
raw = readline()
if raw is None:
break
msg = json.loads(raw)
mid = msg.get("id")
method = msg.get("method")
if method == "initialize":
send({"jsonrpc": "2.0", "id": mid, "result": {
"protocolVersion": "2025-11-25", "capabilities": {},
"serverInfo": {"name": "fake", "version": "0"}}})
elif method == "notifications/initialized":
pass
elif method == "tools/list":
send({"jsonrpc": "2.0", "id": mid, "result": {"tools": [
{"name": "gen", "description": "deferred generator",
"inputSchema": {"type": "object", "properties": {}},
"execution": {"taskSupport": "required"}}]}})
elif method == "tools/call":
# Defer: return a task handle (no `content`).
send({"jsonrpc": "2.0", "id": mid, "result": {
"taskId": "job-1", "status": "working", "pollInterval": 500}})
elif method == "tasks/get":
get_count += 1
if mode == "complete" and get_count >= 2:
send({"jsonrpc": "2.0", "id": mid, "result": {
"taskId": "job-1", "status": "completed"}})
else:
send({"jsonrpc": "2.0", "id": mid, "result": {
"taskId": "job-1", "status": "working", "pollInterval": 500}})
elif method == "tasks/result":
send({"jsonrpc": "2.0", "id": mid, "result": {
"content": [{"type": "text", "text": "the real answer"}]}})
elif method == "tasks/cancel":
if marker:
with open(marker, "w") as f:
f.write((msg.get("params") or {}).get("taskId", ""))
if mid is not None:
send({"jsonrpc": "2.0", "id": mid, "result": {}})
elif mid is not None:
send({"jsonrpc": "2.0", "id": mid, "error": {"code": -32601, "message": "unknown"}})
"#;
fn python3_available() -> bool {
Command::new("python3").arg("--version").output().is_ok()
}
fn write_script() -> std::path::PathBuf {
let path = std::env::temp_dir().join(format!("skald_tasks_{}.py", std::process::id()));
std::fs::File::create(&path).unwrap().write_all(FAKE_SERVER.as_bytes()).unwrap();
path
}
fn cfg(script: &std::path::Path, mode: &str, marker: Option<&std::path::Path>) -> McpServerConfig {
let mut args = vec![script.to_string_lossy().to_string(), mode.to_string()];
if let Some(m) = marker {
args.push(m.to_string_lossy().to_string());
}
McpServerConfig {
name: "fake".to_string(),
transport: McpTransport::Stdio,
command: Some("python3".to_string()),
args: Some(args),
env: None,
url: None,
api_key: None,
}
}
#[tokio::test]
async fn task_is_polled_to_completion_and_returns_real_result() {
if !python3_available() {
eprintln!("python3 not found — skipping tasks integration test");
return;
}
let script = write_script();
let server = McpServer::start(&cfg(&script, "complete", None), None, None, None)
.await
.expect("server should start");
// The tool opts into tasks (taskSupport: required); call_tool must add the
// `task` field, poll to completion, and return the real result — not the handle.
let result = server.call_tool("gen", serde_json::json!({}))
.await
.expect("task should complete");
match result {
McpCallResult::Text(t) => assert_eq!(t, "the real answer"),
other => panic!("expected the polled Text result, got {other:?}"),
}
let _ = std::fs::remove_file(&script);
}
#[tokio::test]
async fn dropping_a_polling_call_sends_tasks_cancel() {
if !python3_available() {
eprintln!("python3 not found — skipping tasks cancel integration test");
return;
}
let script = write_script();
let marker = std::env::temp_dir().join(format!("skald_tasks_marker_{}", std::process::id()));
let _ = std::fs::remove_file(&marker);
let server = McpServer::start(&cfg(&script, "cancel", Some(&marker)), None, None, None)
.await
.expect("server should start");
// In "cancel" mode tasks/get never completes, so the call polls forever; time out
// to drop the future mid-poll, which must fire tasks/cancel via the drop-guard.
let timed = tokio::time::timeout(
Duration::from_millis(900),
server.call_tool("gen", serde_json::json!({})),
).await;
assert!(timed.is_err(), "call should still be polling (timed out), not resolved");
// Give the guard's spawned tasks/cancel time to reach the server.
tokio::time::sleep(Duration::from_millis(700)).await;
let recorded = std::fs::read_to_string(&marker).unwrap_or_default();
assert_eq!(recorded, "job-1", "server should have received tasks/cancel for job-1");
let _ = std::fs::remove_file(&script);
let _ = std::fs::remove_file(&marker);
}
+14
View File
@@ -0,0 +1,14 @@
[package]
name = "plugin-comfyui"
version = "0.1.0"
edition = "2024"
[dependencies]
core-api = { path = "../core-api" }
anyhow = "1"
async-trait = "0.1"
serde = { version = "1", features = ["derive"] }
serde_json = "1"
tokio = { version = "1", features = ["full"] }
tracing = "0.1"
reqwest = { version = "0.13", default-features = false, features = ["rustls-no-provider", "charset", "http2", "system-proxy", "json", "multipart"] }
+601
View File
@@ -0,0 +1,601 @@
use std::collections::{HashMap, HashSet};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use anyhow::{anyhow, Result};
use async_trait::async_trait;
use core_api::image_generate::{ImageGenerate, ImageGenerateRegistry};
use core_api::plugin::{Plugin, PluginContext};
use serde::Deserialize;
use serde_json::{Value, json};
use tokio::sync::Mutex;
use tokio::task::JoinHandle;
use tokio::time::timeout;
use tracing::{info, warn};
// ── Config ────────────────────────────────────────────────────────────────────
#[derive(Debug, Clone, Deserialize)]
struct PluginConfig {
#[serde(default = "default_base_url")]
base_url: String,
#[serde(default = "default_workflows_dir")]
workflows_dir: String,
#[serde(default)]
default_negative: String,
}
fn default_base_url() -> String { "http://localhost:8188".into() }
fn default_workflows_dir() -> String { "data/comfyui/workflows".into() }
impl Default for PluginConfig {
fn default() -> Self {
Self {
base_url: default_base_url(),
workflows_dir: default_workflows_dir(),
default_negative: String::new(),
}
}
}
// ── _personal_agent metadata ──────────────────────────────────────────────────
#[derive(Debug, Deserialize, Default)]
struct PersonalAgentMeta {
name: Option<String>,
description: Option<String>,
prompt_node: Option<String>,
negative_prompt_node: Option<String>,
#[serde(default)]
prompt_field: Option<String>,
#[serde(default)]
prompt_field_extra: Option<Vec<String>>,
#[serde(default)]
negative_prompt_field: Option<String>,
#[serde(default)]
negative_prompt_field_extra: Option<Vec<String>>,
#[serde(default)]
extra_params: ExtraParamNodeMap,
#[serde(default)]
input_image_node: Option<String>,
}
#[derive(Debug, Deserialize, Default)]
struct ExtraParamNodeMap {
width_node: Option<String>,
height_node: Option<String>,
steps_node: Option<String>,
}
// ── Runtime status ────────────────────────────────────────────────────────────
#[derive(Default)]
struct RuntimeStatus {
comfyui_online: bool,
registered_providers: usize,
}
// ── ComfyUiWorkflowGenerator ──────────────────────────────────────────────────
pub struct ComfyUiWorkflowGenerator {
id: String,
name: String,
description: Option<String>,
extra_params_schema: Option<Value>,
workflow_path: PathBuf,
prompt_node: Option<String>,
negative_prompt_node: Option<String>,
prompt_field: Option<String>,
prompt_field_extra: Option<Vec<String>>,
negative_prompt_field: Option<String>,
negative_prompt_field_extra: Option<Vec<String>>,
extra_param_nodes: ExtraParamNodeMap,
input_image_node: Option<String>,
default_negative: String,
base_url: String,
http: Arc<reqwest::Client>,
}
impl ComfyUiWorkflowGenerator {
fn from_file(
path: &Path,
config: &PluginConfig,
http: Arc<reqwest::Client>,
) -> Result<Self> {
let content = std::fs::read_to_string(path)
.map_err(|e| anyhow!("failed to read {:?}: {e}", path))?;
let workflow: Value = serde_json::from_str(&content)
.map_err(|e| anyhow!("invalid JSON in {:?}: {e}", path))?;
let stem = path.file_stem().and_then(|s| s.to_str()).unwrap_or("unknown");
let id = format!("comfyui-{stem}");
let meta: PersonalAgentMeta = workflow.get("_personal_agent")
.and_then(|v| serde_json::from_value(v.clone()).ok())
.unwrap_or_default();
let name = meta.name.unwrap_or_else(|| stem.to_string());
// Require either a prompt node or an input image node (or both)
if meta.prompt_node.is_none()
&& find_first_node(&workflow, "CLIPTextEncode").is_none()
&& meta.input_image_node.is_none()
{
anyhow::bail!("no CLIPTextEncode or input_image_node in {:?}", path);
}
let prompt_node = meta.prompt_node
.or_else(|| find_first_node(&workflow, "CLIPTextEncode"));
let extra_params_schema = build_extra_params_schema(&workflow, &meta.extra_params);
Ok(Self {
id,
name,
description: meta.description,
extra_params_schema,
workflow_path: path.to_path_buf(),
prompt_node,
negative_prompt_node: meta.negative_prompt_node,
prompt_field: meta.prompt_field,
prompt_field_extra: meta.prompt_field_extra,
negative_prompt_field: meta.negative_prompt_field,
negative_prompt_field_extra: meta.negative_prompt_field_extra,
extra_param_nodes: meta.extra_params,
input_image_node: meta.input_image_node,
default_negative: config.default_negative.clone(),
base_url: config.base_url.clone(),
http,
})
}
async fn poll_until_done(&self, prompt_id: &str) -> Result<ImageOutputInfo> {
let url = format!("{}/history/{prompt_id}", self.base_url.trim_end_matches('/'));
let deadline = tokio::time::Instant::now() + Duration::from_secs(300);
loop {
if tokio::time::Instant::now() > deadline {
anyhow::bail!("ComfyUI generation timed out after 300s");
}
tokio::time::sleep(Duration::from_secs(2)).await;
let json: Value = self.http.get(&url).send().await
.map_err(|e| anyhow!("ComfyUI /history request failed: {e}"))?
.json().await
.map_err(|e| anyhow!("ComfyUI /history parse failed: {e}"))?;
// Empty object {} = still in queue
let Some(entry) = json.get(prompt_id) else { continue };
let status = &entry["status"];
// Execution error — extract the message from status.messages.
// Must be checked BEFORE the "still running" guard: a failed prompt
// reports `completed: false`, so checking completion first would
// mask the error and poll until the deadline.
if status["status_str"].as_str() == Some("error") {
let error_msg = status["messages"].as_array()
.and_then(|msgs| {
msgs.iter().find(|m| m[0].as_str() == Some("execution_error"))
})
.and_then(|m| m.get(1))
.map(|details| {
let node = details["node_type"].as_str().unwrap_or("unknown");
let exc = details["exception_message"].as_str().unwrap_or("unknown error");
format!("{node}: {exc}")
})
.unwrap_or_else(|| "unknown execution error".to_string());
anyhow::bail!("ComfyUI execution error: {error_msg}");
}
// Still running
if status["completed"].as_bool() == Some(false) {
continue;
}
if let Some(outputs) = entry["outputs"].as_object() {
for node_output in outputs.values() {
if let Some(img) = node_output["images"].as_array().and_then(|a| a.first()) {
return Ok(ImageOutputInfo {
filename: img["filename"].as_str().unwrap_or("").to_string(),
subfolder: img["subfolder"].as_str().unwrap_or("").to_string(),
});
}
}
}
anyhow::bail!("ComfyUI: generation completed but no image in outputs");
}
}
}
struct ImageOutputInfo {
filename: String,
subfolder: String,
}
#[async_trait]
impl ImageGenerate for ComfyUiWorkflowGenerator {
fn id(&self) -> &str { &self.id }
fn name(&self) -> &str { &self.name }
fn description(&self) -> Option<&str> { self.description.as_deref() }
fn extra_params_schema(&self) -> Option<Value> { self.extra_params_schema.clone() }
async fn generate(&self, prompt: &str, extra_params: Option<&Value>) -> Result<Vec<u8>> {
// Re-read from disk so modified workflows are picked up immediately
let content = tokio::fs::read_to_string(&self.workflow_path).await
.map_err(|e| anyhow!("failed to read workflow: {e}"))?;
let mut workflow: Value = serde_json::from_str(&content)
.map_err(|e| anyhow!("invalid workflow JSON: {e}"))?;
if let Some(obj) = workflow.as_object_mut() {
obj.remove("_personal_agent");
}
// Inject prompt — skip if the workflow has no prompt node (e.g. upscale-only)
if let Some(ref node) = self.prompt_node {
let prompt_field = self.prompt_field.as_deref().unwrap_or("text");
workflow[node]["inputs"][prompt_field] = json!(prompt);
if let Some(ref extras) = self.prompt_field_extra {
for field in extras {
workflow[node]["inputs"][field] = json!(prompt);
}
}
}
// Inject negative prompt
if let Some(neg_node) = &self.negative_prompt_node {
if !self.default_negative.is_empty() {
let neg_field = self.negative_prompt_field.as_deref().unwrap_or("text");
workflow[neg_node]["inputs"][neg_field] = json!(self.default_negative);
if let Some(ref extras) = self.negative_prompt_field_extra {
for field in extras {
workflow[neg_node]["inputs"][field] = json!(self.default_negative);
}
}
}
}
// Inject extra_params into declared nodes
if let Some(params) = extra_params {
if let (Some(node), Some(v)) = (&self.extra_param_nodes.width_node, params["width"].as_i64()) { workflow[node]["inputs"]["width"] = json!(v); }
if let (Some(node), Some(v)) = (&self.extra_param_nodes.height_node, params["height"].as_i64()) { workflow[node]["inputs"]["height"] = json!(v); }
if let (Some(node), Some(v)) = (&self.extra_param_nodes.steps_node, params["steps"].as_i64()) { workflow[node]["inputs"]["steps"] = json!(v); }
}
// ── Input image (img2img) ──────────────────────────────────────────────────
if let (Some(node), Some(image_path)) = (&self.input_image_node, extra_params.and_then(|p| p["input_image"].as_str())) {
// 1. Read the image file
let image_bytes = tokio::fs::read(image_path).await
.map_err(|e| anyhow!("failed to read input image '{image_path}': {e}"))?;
// 2. Get just the filename (strip path)
let filename = Path::new(image_path)
.file_name()
.and_then(|s| s.to_str())
.ok_or_else(|| anyhow!("invalid image path: {image_path}"))?;
// 3. Upload to ComfyUI
let upload_url = format!("{}/upload/image", self.base_url.trim_end_matches('/'));
let form = reqwest::multipart::Form::new()
.part("image", reqwest::multipart::Part::bytes(image_bytes)
.file_name(filename.to_string())
.mime_str("image/png")
.unwrap_or_else(|_| reqwest::multipart::Part::bytes(vec![])))
.text("overwrite", "true");
let upload_resp = self.http.post(&upload_url)
.multipart(form)
.send().await
.map_err(|e| anyhow!("ComfyUI /upload/image failed: {e}"))?;
if !upload_resp.status().is_success() {
let status = upload_resp.status();
let body = upload_resp.text().await.unwrap_or_default();
anyhow::bail!("ComfyUI /upload/image error {status}: {body}");
}
// 4. Set the filename in the LoadImage node
// Use just the name (ComfyUI prepends its input/ directory)
workflow[node]["inputs"]["image"] = json!(filename);
}
// POST /prompt
let url = format!("{}/prompt", self.base_url.trim_end_matches('/'));
let resp = self.http.post(&url)
.json(&json!({ "prompt": workflow }))
.send().await
.map_err(|e| anyhow!("ComfyUI /prompt failed: {e}"))?;
if !resp.status().is_success() {
let status = resp.status();
let body = resp.text().await.unwrap_or_default();
anyhow::bail!("ComfyUI /prompt error {status}: {body}");
}
let resp_json: Value = resp.json().await
.map_err(|e| anyhow!("ComfyUI /prompt response parse failed: {e}"))?;
let prompt_id = resp_json["prompt_id"].as_str()
.ok_or_else(|| anyhow!("ComfyUI: no prompt_id in response"))?
.to_string();
info!(provider = %self.id, %prompt_id, "ComfyUI: queued");
let image_info = self.poll_until_done(&prompt_id).await?;
// Download image
let view_url = format!(
"{}/view?filename={}&subfolder={}&type=output",
self.base_url.trim_end_matches('/'),
image_info.filename,
image_info.subfolder,
);
let img_resp = self.http.get(&view_url).send().await
.map_err(|e| anyhow!("ComfyUI /view failed: {e}"))?;
let bytes = img_resp.bytes().await
.map_err(|e| anyhow!("ComfyUI /view download failed: {e}"))?;
info!(provider = %self.id, bytes = bytes.len(), "ComfyUI: generation complete");
Ok(bytes.to_vec())
}
}
// ── ComfyUIPlugin ─────────────────────────────────────────────────────────────
pub struct ComfyUIPlugin {
http: Arc<reqwest::Client>,
watcher: Mutex<Option<JoinHandle<()>>>,
status: Arc<Mutex<RuntimeStatus>>,
/// IDs currently registered in ImageGeneratorManager — shared with watcher task.
registered_ids: Arc<Mutex<HashSet<String>>>,
/// Last registry received in reload() — used to unregister on stop/disable.
last_registry: Mutex<Option<Arc<dyn ImageGenerateRegistry>>>,
}
impl ComfyUIPlugin {
pub fn new() -> Self {
Self {
http: Arc::new(reqwest::Client::builder()
.timeout(Duration::from_secs(30))
.build()
.expect("reqwest client")),
watcher: Mutex::new(None),
status: Arc::new(Mutex::new(RuntimeStatus::default())),
registered_ids: Arc::new(Mutex::new(HashSet::new())),
last_registry: Mutex::new(None),
}
}
async fn stop_watcher_and_cleanup(&self) {
if let Some(handle) = self.watcher.lock().await.take() {
handle.abort();
}
let ids: Vec<String> = self.registered_ids.lock().await.drain().collect();
if let Some(reg) = self.last_registry.lock().await.as_ref() {
for id in &ids {
reg.unregister(id).await;
}
}
*self.status.lock().await = RuntimeStatus::default();
}
}
#[async_trait]
impl Plugin for ComfyUIPlugin {
fn id(&self) -> &str { "comfyui" }
fn name(&self) -> &str { "ComfyUI" }
fn description(&self) -> &str { "Local image generation via ComfyUI. Each JSON file in data/comfyui/workflows/ becomes a separate provider." }
fn is_running(&self) -> bool { self.watcher.try_lock().map_or(false, |g| g.is_some()) }
fn config_schema(&self) -> Value {
json!({
"type": "object",
"properties": {
"base_url": { "type": "string", "default": "http://localhost:8188", "description": "ComfyUI server URL" },
"workflows_dir": { "type": "string", "default": "data/comfyui/workflows", "description": "Directory with workflow JSON files" },
"default_negative": { "type": "string", "default": "", "description": "Default negative prompt for all workflows" }
}
})
}
async fn reload(&self, enabled: bool, config: Value, ctx: PluginContext) -> Result<()> {
// Abort previous watcher and unregister its providers
self.stop_watcher_and_cleanup().await;
// Store registry for future cleanup
*self.last_registry.lock().await = Some(Arc::clone(&ctx.image_generate_registry));
if !enabled {
info!("ComfyUI plugin disabled");
return Ok(());
}
let plugin_config: PluginConfig = serde_json::from_value(config).unwrap_or_default();
// Ensure workflows directory exists
tokio::fs::create_dir_all(&plugin_config.workflows_dir).await.ok();
let registry = Arc::clone(&ctx.image_generate_registry);
let http = Arc::clone(&self.http);
let status = Arc::clone(&self.status);
let registered_ids = Arc::clone(&self.registered_ids);
let handle = tokio::spawn(watcher_loop(
registry, plugin_config, http, status, registered_ids,
));
*self.watcher.lock().await = Some(handle);
info!("ComfyUI plugin started");
Ok(())
}
async fn start(&self, ctx: PluginContext) -> Result<()> {
self.reload(true, json!({}), ctx).await
}
async fn stop(&self) -> Result<()> {
self.stop_watcher_and_cleanup().await;
info!("ComfyUI plugin stopped");
Ok(())
}
fn runtime_status(&self) -> Option<Value> {
self.status.try_lock().ok().map(|s| json!({
"comfyui_online": s.comfyui_online,
"registered_providers": s.registered_providers,
}))
}
fn as_any(&self) -> &dyn std::any::Any { self }
fn as_arc_any(self: Arc<Self>) -> Arc<dyn std::any::Any + Send + Sync> { self }
}
// ── Watcher loop ──────────────────────────────────────────────────────────────
async fn watcher_loop(
registry: Arc<dyn ImageGenerateRegistry>,
config: PluginConfig,
http: Arc<reqwest::Client>,
status: Arc<Mutex<RuntimeStatus>>,
registered_ids: Arc<Mutex<HashSet<String>>>,
) {
let mut interval = tokio::time::interval(Duration::from_secs(5));
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
let mut was_online = false;
// path → mtime of the registered version
let mut known: HashMap<PathBuf, SystemTime> = HashMap::new();
loop {
interval.tick().await;
// ── 1. Health check ───────────────────────────────────────────────────
let is_online = health_check(&http, &config.base_url).await;
if !is_online && was_online {
for path in known.keys() {
let id = path_to_id(path);
registry.unregister(&id).await;
registered_ids.lock().await.remove(&id);
}
known.clear();
warn!("ComfyUI unreachable — all providers unregistered");
}
was_online = is_online;
{
let mut s = status.lock().await;
s.comfyui_online = is_online;
s.registered_providers = known.len();
}
if !is_online { continue; }
// ── 2. File scan ──────────────────────────────────────────────────────
let current = match scan_workflows(&config.workflows_dir).await {
Ok(m) => m,
Err(e) => { warn!(error = %e, "ComfyUI: workflow scan failed"); continue; }
};
// Removed files
for path in known.keys() {
if !current.contains_key(path) {
let id = path_to_id(path);
registry.unregister(&id).await;
registered_ids.lock().await.remove(&id);
info!(%id, "ComfyUI: workflow removed");
}
}
// New or modified files
for (path, mtime) in &current {
let changed = known.get(path).map_or(true, |old| old != mtime);
if !changed { continue; }
if known.contains_key(path) {
// modified — unregister old version first
let id = path_to_id(path);
registry.unregister(&id).await;
registered_ids.lock().await.remove(&id);
}
match ComfyUiWorkflowGenerator::from_file(path, &config, Arc::clone(&http)) {
Ok(provider) => {
info!(id = %provider.id, "ComfyUI: workflow registered");
registered_ids.lock().await.insert(provider.id.clone());
registry.register(Arc::new(provider)).await;
}
Err(e) => warn!(path = %path.display(), error = %e, "ComfyUI: skipping workflow"),
}
}
known = current;
status.lock().await.registered_providers = known.len();
}
}
// ── Helpers ───────────────────────────────────────────────────────────────────
async fn health_check(http: &reqwest::Client, base_url: &str) -> bool {
let url = format!("{}/system_stats", base_url.trim_end_matches('/'));
timeout(Duration::from_secs(2), http.get(&url).send())
.await
.map(|r| r.map(|r| r.status().is_success()).unwrap_or(false))
.unwrap_or(false)
}
async fn scan_workflows(dir: &str) -> Result<HashMap<PathBuf, SystemTime>> {
let mut map = HashMap::new();
let mut entries = tokio::fs::read_dir(dir).await
.map_err(|e| anyhow!("cannot read workflows dir '{dir}': {e}"))?;
while let Some(entry) = entries.next_entry().await? {
let path = entry.path();
if path.extension().and_then(|e| e.to_str()) != Some("json") { continue; }
if let Ok(meta) = entry.metadata().await {
if let Ok(mtime) = meta.modified() {
map.insert(path, mtime);
}
}
}
Ok(map)
}
fn path_to_id(path: &Path) -> String {
let stem = path.file_stem().and_then(|s| s.to_str()).unwrap_or("unknown");
format!("comfyui-{stem}")
}
fn find_first_node(workflow: &Value, class_type: &str) -> Option<String> {
let obj = workflow.as_object()?;
let mut keys: Vec<u64> = obj.keys()
.filter_map(|k| k.parse().ok())
.collect();
keys.sort_unstable();
keys.into_iter()
.map(|n| n.to_string())
.find(|k| workflow[k]["class_type"].as_str() == Some(class_type))
}
fn build_extra_params_schema(workflow: &Value, nodes: &ExtraParamNodeMap) -> Option<Value> {
let mut props = serde_json::Map::new();
if let Some(id) = &nodes.width_node {
let default = workflow[id]["inputs"]["width"].as_i64().unwrap_or(512);
props.insert("width".into(), json!({ "type": "integer", "default": default }));
}
if let Some(id) = &nodes.height_node {
let default = workflow[id]["inputs"]["height"].as_i64().unwrap_or(512);
props.insert("height".into(), json!({ "type": "integer", "default": default }));
}
if let Some(id) = &nodes.steps_node {
let default = workflow[id]["inputs"]["steps"].as_i64().unwrap_or(20);
props.insert("steps".into(), json!({ "type": "integer", "default": default }));
}
if props.is_empty() { None } else { Some(json!({ "type": "object", "properties": props })) }
}
+14
View File
@@ -0,0 +1,14 @@
[package]
name = "plugin-elevenlabs"
version = "0.1.0"
edition = "2024"
[dependencies]
core-api = { path = "../core-api" }
anyhow = "1"
async-trait = "0.1"
serde_json = "1"
tokio = { version = "1", features = ["full"] }
tracing = "0.1"
reqwest = { version = "0.13", default-features = false, features = ["rustls-no-provider", "charset", "http2", "system-proxy", "json", "multipart"] }
parking_lot = "0.12"
+378
View File
@@ -0,0 +1,378 @@
use std::sync::Arc;
use anyhow::{Context, Result, anyhow};
use async_trait::async_trait;
use parking_lot::Mutex;
use serde_json::Value;
use tracing::{debug, info};
use core_api::plugin::{Plugin, PluginContext};
use core_api::provider::{
ApiProvider, LlmProviderRecord, ProviderField, ProviderUiMeta, RemoteLlmModelInfo,
ServiceType,
};
use core_api::transcribe::{RemoteTranscribeModelInfo, Transcribe, TranscribeModelRecord};
use core_api::tts::{RemoteTtsModelInfo, TextToSpeech, TtsModelRecord};
// ── Constants ─────────────────────────────────────────────────────────────────
const EL_BASE_URL: &str = "https://api.elevenlabs.io/v1";
const EL_DEFAULT_MODEL: &str = "eleven_multilingual_v2";
// ── TTS Synthesiser ───────────────────────────────────────────────────────────
/// ElevenLabsTtsSynthesiser — cloud TTS via the ElevenLabs v1 API.
///
/// Endpoint: `POST https://api.elevenlabs.io/v1/text-to-speech/{voice_id}`
/// Auth: `xi-api-key` header (not Bearer).
pub struct ElevenLabsTtsSynthesiser {
id: String,
api_key: String,
model_id: String,
voice_id: String,
instructions: Option<String>,
http: reqwest::Client,
}
impl ElevenLabsTtsSynthesiser {
pub fn new(
id: impl Into<String>,
api_key: impl Into<String>,
model_id: impl Into<String>,
voice_id: Option<String>,
instructions: Option<String>,
) -> Self {
let model_id = model_id.into();
// Legacy fallback: if no voice_id, treat model_id as voice_id and use default model.
let (resolved_model, resolved_voice) = match voice_id {
Some(v) => (model_id, v),
None => (EL_DEFAULT_MODEL.to_string(), model_id),
};
Self {
id: id.into(),
api_key: api_key.into(),
model_id: resolved_model,
voice_id: resolved_voice,
instructions,
http: reqwest::Client::new(),
}
}
}
#[async_trait]
impl TextToSpeech for ElevenLabsTtsSynthesiser {
fn id(&self) -> &str { &self.id }
fn name(&self) -> &str { &self.id }
fn instructions(&self) -> Option<&str> { self.instructions.as_deref() }
async fn synthesize(&self, text: &str, _instructions: Option<&str>) -> Result<Vec<u8>> {
debug!(
chars = text.len(),
voice_id = %self.voice_id,
"elevenlabs_tts: synthesising",
);
let url = format!("{EL_BASE_URL}/text-to-speech/{}", self.voice_id);
let body = serde_json::json!({
"text": text,
"model_id": self.model_id,
});
let resp = self.http
.post(&url)
.header("xi-api-key", &self.api_key)
.json(&body)
.send()
.await
.map_err(|e| anyhow!("elevenlabs_tts: request failed: {e}"))?;
let status = resp.status();
if !status.is_success() {
let err: serde_json::Value = resp.json().await.unwrap_or_default();
let msg = err["detail"]["message"].as_str()
.or_else(|| err["detail"].as_str())
.unwrap_or("unknown error");
anyhow::bail!("elevenlabs_tts: API error {status}: {msg}");
}
let audio = resp
.bytes()
.await
.map_err(|e| anyhow!("elevenlabs_tts: failed to read audio bytes: {e}"))?
.to_vec();
info!(bytes = audio.len(), voice_id = %self.voice_id, "elevenlabs_tts: synthesis complete");
Ok(audio)
}
}
// ── Transcriber ───────────────────────────────────────────────────────────────
/// ElevenLabsTranscriber — cloud Speech-to-Text via the ElevenLabs Scribe API.
///
/// Endpoint: `POST https://api.elevenlabs.io/v1/speech-to-text`
/// Auth: `xi-api-key` header (not Bearer).
pub struct ElevenLabsTranscriber {
id: String,
api_key: String,
model_id: String,
http: reqwest::Client,
}
impl ElevenLabsTranscriber {
pub fn new(
id: impl Into<String>,
api_key: impl Into<String>,
model_id: impl Into<String>,
) -> Self {
Self {
id: id.into(),
api_key: api_key.into(),
model_id: model_id.into(),
http: reqwest::Client::new(),
}
}
}
#[async_trait]
impl Transcribe for ElevenLabsTranscriber {
fn id(&self) -> &str { &self.id }
async fn transcribe(&self, audio: Vec<u8>, format: &str) -> Result<String> {
debug!(
bytes = audio.len(),
format = %format,
model_id = %self.model_id,
"elevenlabs_transcribe: transcribing",
);
let url = format!("{EL_BASE_URL}/speech-to-text");
let filename = format!("audio.{format}");
let part = reqwest::multipart::Part::bytes(audio)
.file_name(filename)
.mime_str("audio/wav")?;
let form = reqwest::multipart::Form::new()
.text("model_id", self.model_id.clone())
.part("file", part);
let resp = self.http
.post(&url)
.header("xi-api-key", &self.api_key)
.multipart(form)
.send()
.await
.map_err(|e| anyhow!("elevenlabs_transcribe: request failed: {e}"))?;
let status = resp.status();
if !status.is_success() {
let err: serde_json::Value = resp.json().await.unwrap_or_default();
let msg = err["detail"]["message"].as_str()
.or_else(|| err["detail"].as_str())
.unwrap_or("unknown error");
anyhow::bail!("elevenlabs_transcribe: API error {status}: {msg}");
}
let body: serde_json::Value = resp
.json()
.await
.map_err(|e| anyhow!("elevenlabs_transcribe: failed to parse response: {e}"))?;
let text = body["text"]
.as_str()
.ok_or_else(|| anyhow!("elevenlabs_transcribe: missing 'text' in response"))?
.to_string();
info!(chars = text.len(), "elevenlabs_transcribe: done");
Ok(text)
}
}
// ── ApiProvider ───────────────────────────────────────────────────────────────
/// ElevenLabs supports TTS and Transcription only — no LLM chat/completion.
pub struct ElevenLabsProvider {
http: reqwest::Client,
}
impl ElevenLabsProvider {
pub fn new() -> Self {
Self { http: reqwest::Client::new() }
}
async fn fetch_models(&self, api_key: &str) -> Result<serde_json::Value> {
self.http
.get("https://api.elevenlabs.io/v1/models")
.header("xi-api-key", api_key)
.send()
.await
.map_err(|e| anyhow!("ElevenLabs request failed: {e}"))?
.error_for_status()
.map_err(|e| anyhow!("ElevenLabs error response: {e}"))?
.json()
.await
.map_err(|e| anyhow!("ElevenLabs response parse failed: {e}"))
}
}
#[async_trait]
impl ApiProvider for ElevenLabsProvider {
fn type_id(&self) -> &'static str { "elevenlabs" }
fn display_name(&self) -> &'static str { "ElevenLabs" }
fn supported_types(&self) -> &'static [ServiceType] {
&[ServiceType::Tts, ServiceType::Transcribe]
}
async fn list_llm_models(&self, _record: &LlmProviderRecord) -> Result<Option<Vec<RemoteLlmModelInfo>>> {
Ok(None)
}
async fn list_tts_models(&self, record: &LlmProviderRecord) -> Result<Option<Vec<RemoteTtsModelInfo>>> {
let api_key = record.api_key.as_deref()
.ok_or_else(|| anyhow!("provider '{}': api_key required for elevenlabs model listing", record.name))?;
let resp = self.fetch_models(api_key).await?;
let models = resp.as_array()
.ok_or_else(|| anyhow!("unexpected ElevenLabs response shape"))?
.iter()
.filter(|m| m["can_do_text_to_speech"].as_bool().unwrap_or(false))
.map(|m| {
let id = m["model_id"].as_str().unwrap_or("").to_string();
let name = m["name"].as_str().unwrap_or(&id).to_string();
let description = m["description"].as_str().map(str::to_string);
let languages = m["languages"].as_array()
.map(|langs| langs.iter()
.filter_map(|l| l["language_id"].as_str().map(str::to_string))
.collect())
.unwrap_or_default();
let cost_factor = m["token_cost_factor"].as_f64();
let instructions = elevenlabs_tts_instructions(&id);
RemoteTtsModelInfo { id, name, description, languages, cost_factor, instructions }
})
.collect();
Ok(Some(models))
}
async fn list_transcribe_models(&self, record: &LlmProviderRecord) -> Result<Option<Vec<RemoteTranscribeModelInfo>>> {
let api_key = record.api_key.as_deref()
.ok_or_else(|| anyhow!("provider '{}': api_key required for elevenlabs model listing", record.name))?;
let resp = self.fetch_models(api_key).await?;
let models = resp.as_array()
.ok_or_else(|| anyhow!("unexpected ElevenLabs response shape"))?
.iter()
.filter(|m| m["can_do_voice_conversion"].as_bool().unwrap_or(false)
|| m["model_id"].as_str().map(|id| id.starts_with("scribe")).unwrap_or(false))
.map(|m| {
let id = m["model_id"].as_str().unwrap_or("").to_string();
let name = m["name"].as_str().unwrap_or(&id).to_string();
let description = m["description"].as_str().map(str::to_string);
let languages = m["languages"].as_array()
.map(|langs| langs.iter()
.filter_map(|l| l["language_id"].as_str().map(str::to_string))
.collect())
.unwrap_or_default();
RemoteTranscribeModelInfo { id, name, description, languages }
})
.collect();
Ok(Some(models))
}
fn build_tts(&self, record: &LlmProviderRecord, model: &TtsModelRecord) -> Option<Result<Arc<dyn TextToSpeech>>> {
Some((|| {
let api_key = record.api_key.clone()
.with_context(|| format!("provider '{}': api_key required for elevenlabs", record.name))?;
Ok(Arc::new(ElevenLabsTtsSynthesiser::new(
&model.name, api_key, &model.model_id, model.voice_id.clone(), model.instructions.clone(),
)) as Arc<dyn TextToSpeech>)
})())
}
fn build_transcriber(&self, record: &LlmProviderRecord, model: &TranscribeModelRecord) -> Option<Result<Arc<dyn Transcribe>>> {
Some((|| {
let api_key = record.api_key.clone()
.with_context(|| format!("provider '{}': api_key required for elevenlabs", record.name))?;
Ok(Arc::new(ElevenLabsTranscriber::new(
&model.name, api_key, &model.model_id,
)) as Arc<dyn Transcribe>)
})())
}
fn ui_meta(&self) -> ProviderUiMeta {
ProviderUiMeta {
type_id: "elevenlabs",
display_name: "ElevenLabs",
description: Some("Text-to-speech and transcription"),
color: "#f59e0b",
icon: "bi-waveform",
fields: &[
ProviderField { key: "api_key", label: "API Key", required: true, secret: true },
],
}
}
}
fn elevenlabs_tts_instructions(model_id: &str) -> Option<String> {
match model_id {
"eleven_multilingual_v2" | "eleven_turbo_v2_5" | "eleven_turbo_v2" | "eleven_flash_v2_5" | "eleven_flash_v2" => Some(
"You can use the following markers in text to add expressiveness:\n\
- <break time=\"0.5s\" /> — pause of given duration\n\
- <phoneme alphabet=\"ipa\" ph=\"...\">word</phoneme> — explicit pronunciation\n\
Laugh, cough, or sigh naturally by writing them as actions in parentheses, e.g. (laughs), (sighs), (coughs).\n\
Emphasise a word with ALL CAPS or by repeating letters (e.g. sooo goood).\n\
Keep sentences short for best pacing.".to_string()
),
"eleven_monolingual_v1" => Some(
"English-only model. Supports (laughs), (sighs), (coughs) for non-verbal sounds.\n\
Use ALL CAPS for emphasis. Avoid non-English characters.".to_string()
),
_ => None,
}
}
// ── Plugin ────────────────────────────────────────────────────────────────────
pub struct ElevenLabsPlugin {
running: Mutex<bool>,
}
impl ElevenLabsPlugin {
pub fn new() -> Self {
Self { running: Mutex::new(false) }
}
}
#[async_trait]
impl Plugin for ElevenLabsPlugin {
fn id(&self) -> &str { "elevenlabs" }
fn name(&self) -> &str { "ElevenLabs" }
fn description(&self) -> &str { "ElevenLabs TTS and transcription provider" }
fn is_running(&self) -> bool { *self.running.lock() }
async fn start(&self, ctx: PluginContext) -> Result<()> {
ctx.api_provider_registry.register_plugin(Arc::new(ElevenLabsProvider::new()));
*self.running.lock() = true;
Ok(())
}
async fn reload(&self, enabled: bool, _config: Value, ctx: PluginContext) -> Result<()> {
if enabled {
ctx.api_provider_registry.register_plugin(Arc::new(ElevenLabsProvider::new()));
*self.running.lock() = true;
} else {
ctx.api_provider_registry.unregister_plugin("elevenlabs");
*self.running.lock() = false;
}
Ok(())
}
async fn stop(&self) -> Result<()> {
*self.running.lock() = false;
Ok(())
}
fn as_any(&self) -> &dyn std::any::Any { self }
fn as_arc_any(self: Arc<Self>) -> Arc<dyn std::any::Any + Send + Sync> { self }
}
+14
View File
@@ -0,0 +1,14 @@
[package]
name = "plugin-honcho"
version = "0.1.0"
edition = "2024"
[dependencies]
core-api = { path = "../core-api" }
honcho-client = { path = "../honcho-client" }
anyhow = "1"
async-trait = "0.1"
serde_json = "1"
tokio = { version = "1", features = ["full"] }
tokio-util = { version = "0.7", features = ["rt"] }
tracing = "0.1"
+997
View File
@@ -0,0 +1,997 @@
//! Honcho memory plugin — streams completed chat turns to a Honcho server
//! and exposes a [`Memory`] read path via [`HonchoMemory`].
//!
//! # Write path
//! Subscribes to the [`ChatEventBus`] and forwards every user/assistant message
//! from **interactive, non-ephemeral** sessions to Honcho so that the server can
//! build long-term memory (conclusions) about the user.
//!
//! # Read path
//! [`HonchoMemory`] implements the [`Memory`] trait. Before each LLM turn,
//! `query_context` calls Honcho's `session_context` API to retrieve a
//! token-budgeted summary of what is known so far and injects it into the
//! system prompt.
//!
//! # Filtering (write path)
//! An event is forwarded only when **all** of the following hold:
//! - `is_interactive = true` — a real user is in the conversation
//! - `is_ephemeral = false` — not a short-lived automated session (cron, tic)
//! - `is_synthetic = false` — message content was typed by a user, not
//! injected by the system
//!
//! # Honcho object model
//! ```
//! workspace (one per agent instance, from config)
//! ├── peer "user" (observe_others = true)
//! ├── peer "assistant" (observe_me = true)
//! └── session (one per local chat_sessions.id, created lazily)
//! ├── message peer_id="user"
//! └── message peer_id="assistant"
//! ```
//!
//! The `session_map` (local session_id → Honcho session UUID) is shared between
//! the write-path listener task and `HonchoMemory` so both sides see the same
//! mapping without duplication.
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use anyhow::Result;
use async_trait::async_trait;
use serde_json::{Value, json};
use tokio::sync::{Mutex, RwLock};
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;
use tracing::{debug, info, trace, warn};
use core_api::bus::{BusEvent, ChatEvent, ChatEventRole, RecvError};
use core_api::memory::Memory;
use core_api::plugin::PluginContext;
use core_api::tool::{Tool, ToolCategory};
use honcho_client::HonchoClient;
use honcho_client::models::{
ConclusionCreate, MessageCreate, PeerCreate, PeerRepresentationGet,
SessionCreate, SessionPeerConfig, WorkspaceCreate,
};
const PLUGIN_ID: &str = "honcho";
const PEER_USER: &str = "user";
const PEER_ASSISTANT: &str = "assistant";
/// Token budget for session_context queries.
const CONTEXT_TOKENS: u32 = 2000;
// ── Config ────────────────────────────────────────────────────────────────────
#[derive(Clone, PartialEq)]
struct HonchoConfig {
base_url: String,
api_key: String,
workspace_id: String,
}
// ── HonchoMemory ──────────────────────────────────────────────────────────────
/// Implements the [`Memory`] trait for Honcho.
///
/// Created once in [`HonchoPlugin::new`] and shared for the plugin's lifetime.
/// The plugin calls [`HonchoMemory::activate`] on start and
/// [`HonchoMemory::deactivate`] on stop to swap the live client in/out without
/// replacing the `Arc`.
pub struct HonchoMemory {
/// Mirrors `HonchoPlugin::running`; false when the plugin is stopped.
running: Arc<AtomicBool>,
/// Active client + workspace_id; None when the plugin is not running.
inner: std::sync::RwLock<Option<HonchoInner>>,
/// Shared with the write-path listener task.
session_map: Arc<RwLock<HashMap<i64, String>>>,
}
#[derive(Clone)]
struct HonchoInner {
client: Arc<HonchoClient>,
workspace_id: String,
}
impl HonchoMemory {
fn new(running: Arc<AtomicBool>) -> Self {
Self {
running,
inner: std::sync::RwLock::new(None),
session_map: Arc::new(RwLock::new(HashMap::new())),
}
}
fn activate(&self, client: Arc<HonchoClient>, workspace_id: String) {
*self.inner.write().unwrap() = Some(HonchoInner { client, workspace_id });
}
fn deactivate(&self) {
*self.inner.write().unwrap() = None;
// Clear the session map so a fresh start builds a clean mapping.
// (The Honcho sessions themselves are not deleted — they keep accumulating.)
// Use try_write: if somehow a query is in flight we skip the clear and
// it will be corrected on the next restart anyway.
if let Ok(mut map) = self.session_map.try_write() {
map.clear();
}
}
fn inner(&self) -> Option<HonchoInner> {
self.inner.read().unwrap().clone()
}
}
#[async_trait]
impl Memory for HonchoMemory {
fn id(&self) -> &str { PLUGIN_ID }
fn is_available(&self) -> bool {
self.running.load(Ordering::Relaxed)
&& self.inner.read().unwrap().is_some()
}
async fn query_context(&self, session_id: i64, user_message: &str) -> Option<String> {
// Truncate to at most 120 *characters* (not bytes) to avoid a panic on
// multi-byte UTF-8 codepoints (e.g. 'è' spans two bytes, so a fixed
// byte-index like 120 can land in the middle of it).
let preview_end = user_message
.char_indices()
.nth(120) // byte offset of the 121st char = end of first 120 chars
.map(|(i, _)| i)
.unwrap_or(user_message.len());
trace!(
session_id,
msg_preview = &user_message[..preview_end],
"honcho: query_context invoked"
);
let HonchoInner { client, workspace_id } = self.inner()?;
// ── Strategy: peer_context (global) + session_context (current session) ──
//
// peer_context with search_query searches conclusions derived from ALL past
// sessions — this is the only way cross-session references ("remember when
// we talked about X last week?") can be resolved automatically.
//
// session_context is kept as a secondary call for the current session only,
// to surface conclusions/summaries specific to the ongoing conversation that
// may not yet be reflected in the peer-level representation.
//
// Two embeddings per turn is the cost; the benefit is that the LLM always
// has both global long-term memory AND current-session context.
//
// NOTE: session_context is skipped on the first turn (404 — session not yet
// created in Honcho by the write path) to avoid a wasted HTTP round-trip.
// ── 1. Global peer context (cross-session, semantic search) ──────────────
trace!(session_id, "honcho: querying peer_context (global, with search_query)");
let peer_ctx = match client.peer_context(
&workspace_id,
PEER_USER,
&PeerRepresentationGet {
search_query: Some(user_message.to_string()),
..Default::default()
},
).await {
Ok(ctx) => {
trace!(session_id, raw_json = %ctx, "honcho: peer_context raw response");
let f = format_context(ctx);
debug!(
"honcho: peer_context (global) for session {session_id} ({} chars)",
f.as_deref().map_or(0, |s| s.len())
);
f
}
Err(e) => {
warn!("honcho: peer_context failed: {e}");
None
}
};
// ── 2. Current-session context (session-scoped, no extra embedding) ──────
//
// session_context is a GET with search_query but Honcho re-uses the same
// embedding vector already computed for the peer_context call above
// (server-side caching). No additional LM Studio call in practice.
let deterministic_id = format!("{workspace_id}-{session_id}");
trace!(session_id, honcho_session_id = %deterministic_id, "honcho: querying session_context");
let session_ctx = match client.session_context(
&workspace_id,
&deterministic_id,
Some(CONTEXT_TOKENS),
Some(user_message),
).await {
Ok(ctx) => {
trace!(session_id, raw_json = %ctx, "honcho: session_context raw response");
let f = format_context(ctx);
debug!(
"honcho: session_context for session {session_id} ({} chars)",
f.as_deref().map_or(0, |s| s.len())
);
f
}
Err(honcho_client::error::HonchoError::Http { status: 404, .. }) => {
debug!("honcho: session {deterministic_id} not yet in Honcho (first turn) — skipping session_context");
None
}
Err(e) => {
warn!("honcho: session_context failed for session {session_id}: {e}");
None
}
};
// ── 3. Merge: peer (global) first, then session-specific ─────────────────
let merged = match (peer_ctx, session_ctx) {
(Some(p), Some(s)) if p != s => {
trace!(session_id, "honcho: merging peer + session context");
Some(format!("{p}\n\n{s}"))
}
(Some(p), _) => Some(p),
(_, Some(s)) => Some(s),
(None, None) => None,
};
if let Some(ref text) = merged {
trace!(session_id, injected = %text, "honcho: context injected into system prompt");
} else {
trace!(session_id, "honcho: no context to inject");
}
merged
}
fn tools(&self) -> Vec<Arc<dyn Tool>> {
match self.inner() {
Some(HonchoInner { client, workspace_id }) => vec![
Arc::new(MemoryQueryTool {
client: Arc::clone(&client),
workspace_id: workspace_id.clone(),
}),
Arc::new(HonchoProfileTool {
client: Arc::clone(&client),
workspace_id: workspace_id.clone(),
}),
Arc::new(HonchoSearchTool {
client: Arc::clone(&client),
workspace_id: workspace_id.clone(),
}),
Arc::new(HonchoContextTool {
client: Arc::clone(&client),
workspace_id: workspace_id.clone(),
}),
Arc::new(HonchoConcludeTool { client, workspace_id }),
],
None => vec![],
}
}
}
// ── MemoryQueryTool ───────────────────────────────────────────────────────────
/// LLM-callable tool that queries Honcho's Dialectic API.
///
/// The official Honcho documentation explicitly recommends exposing `peer.chat()`
/// as a tool for agents: the LLM decides on its own when extra memory context
/// is needed and calls this tool with a natural-language question.
///
/// Uses `tokio::task::block_in_place` to bridge the sync `Tool::execute` interface
/// with the async HTTP call, safely running inside the existing Tokio runtime.
struct MemoryQueryTool {
client: Arc<HonchoClient>,
workspace_id: String,
}
impl Tool for MemoryQueryTool {
fn name(&self) -> &str { "memory_query" }
fn description(&self) -> &str {
"Query long-term memory about the user using natural language. \
Ask anything about the user's preferences, past conversations, \
or known facts. Returns a synthesized answer from Honcho's memory. \
Use when you need specific information about the user that is not \
already present in the current conversation."
}
fn parameters_schema(&self) -> Value {
json!({
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "Natural language question about the user. \
E.g. 'What programming languages does the user prefer?'"
}
},
"required": ["query"]
})
}
fn category(&self) -> ToolCategory {
ToolCategory::Introspection
}
fn execute(&self, args: Value) -> anyhow::Result<String> {
let query = args["query"]
.as_str()
.ok_or_else(|| anyhow::anyhow!("memory_query: missing 'query' argument"))?
.to_string();
let client = Arc::clone(&self.client);
let workspace_id = self.workspace_id.clone();
// Bridge sync Tool::execute → async HTTP call.
// block_in_place yields the thread to the Tokio scheduler while the
// nested block_on drives the future to completion — safe inside an
// existing multi-thread Tokio runtime.
tokio::task::block_in_place(|| {
tokio::runtime::Handle::current().block_on(async move {
let opts = honcho_client::models::DialecticOptions {
query,
session_id: None,
target: None,
stream: Some(false),
reasoning_level: Some("low".to_string()),
};
let response = client
.peer_chat(&workspace_id, PEER_USER, &opts)
.await
.map_err(|e| anyhow::anyhow!("memory_query: {e}"))?;
// The Dialectic endpoint returns a JSON object.
// Try known content fields; fall back to pretty-printed JSON.
let text = response.get("content")
.or_else(|| response.get("response"))
.or_else(|| response.get("message"))
.and_then(|v| v.as_str())
.map(str::to_string)
.unwrap_or_else(|| {
serde_json::to_string_pretty(&response)
.unwrap_or_else(|_| response.to_string())
});
Ok(text)
})
})
}
}
/// Bridge a synchronous `Tool::execute` to an async Honcho call.
///
/// `block_in_place` yields the worker thread back to the Tokio scheduler while
/// the nested `block_on` drives the future to completion — safe inside the
/// existing multi-thread runtime without spawning a new thread. Shared by all
/// Honcho tools.
fn run_blocking<F, T>(fut: F) -> T
where
F: std::future::Future<Output = T>,
{
tokio::task::block_in_place(|| tokio::runtime::Handle::current().block_on(fut))
}
// ── HonchoProfileTool ─────────────────────────────────────────────────────────
/// Reads or overwrites the user's *peer card* — a curated list of key facts
/// (name, role, preferences, communication style) maintained by Honcho.
struct HonchoProfileTool {
client: Arc<HonchoClient>,
workspace_id: String,
}
impl Tool for HonchoProfileTool {
fn name(&self) -> &str { "honcho_profile" }
fn description(&self) -> &str {
"Read or update the peer card for the user in Honcho — a curated list of \
key facts (name, role, preferences, communication style). Omit `card` to \
read the current card; pass `card` as a list of fact strings to overwrite it."
}
fn parameters_schema(&self) -> Value {
json!({
"type": "object",
"properties": {
"card": {
"type": "array",
"items": { "type": "string" },
"description": "New peer card as a list of fact strings. \
Omit to read the current card."
}
}
})
}
fn category(&self) -> ToolCategory { ToolCategory::Introspection }
fn execute(&self, args: Value) -> anyhow::Result<String> {
let client = Arc::clone(&self.client);
let workspace_id = self.workspace_id.clone();
let card_update = args.get("card").and_then(|v| v.as_array()).cloned();
run_blocking(async move {
match card_update {
Some(facts) => {
client
.set_peer_card(&workspace_id, PEER_USER, None, json!(facts))
.await
.map_err(|e| anyhow::anyhow!("honcho_profile: {e}"))?;
Ok(format!("Peer card updated ({} facts).", facts.len()))
}
None => {
let card = client
.get_peer_card(&workspace_id, PEER_USER, None)
.await
.map_err(|e| anyhow::anyhow!("honcho_profile: {e}"))?;
Ok(serde_json::to_string_pretty(&card)
.unwrap_or_else(|_| card.to_string()))
}
}
})
}
}
// ── HonchoSearchTool ──────────────────────────────────────────────────────────
/// Semantic search over the conclusions Honcho has derived about the user.
/// Returns raw ranked excerpts — no LLM synthesis — including their IDs so the
/// model can later delete a specific one via `honcho_conclude`.
struct HonchoSearchTool {
client: Arc<HonchoClient>,
workspace_id: String,
}
impl Tool for HonchoSearchTool {
fn name(&self) -> &str { "honcho_search" }
fn description(&self) -> &str {
"Semantic search over facts Honcho has derived about the user. Returns raw \
excerpts ranked by relevance to `query` — no LLM synthesis. Faster and \
cheaper than memory_query. Each fact includes its id (usable with \
honcho_conclude) when available."
}
fn parameters_schema(&self) -> Value {
json!({
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "What to search for in Honcho's memory about the user."
}
},
"required": ["query"]
})
}
fn category(&self) -> ToolCategory { ToolCategory::Introspection }
fn execute(&self, args: Value) -> anyhow::Result<String> {
let query = args["query"]
.as_str()
.ok_or_else(|| anyhow::anyhow!("honcho_search: missing 'query' argument"))?
.to_string();
let client = Arc::clone(&self.client);
let workspace_id = self.workspace_id.clone();
// Honcho's `conclusions/query` endpoint requires observer/observed
// filters; the proven path (shared with the read-path) is `peer_context`
// with a `search_query`, which ranks the user's conclusions by relevance.
run_blocking(async move {
let ctx = client
.peer_context(
&workspace_id,
PEER_USER,
&PeerRepresentationGet {
search_query: Some(query),
search_top_k: Some(10),
..Default::default()
},
)
.await
.map_err(|e| anyhow::anyhow!("honcho_search: {e}"))?;
Ok(format_conclusions(&ctx)
.unwrap_or_else(|| "No relevant context found.".to_string()))
})
}
}
/// Formats the `conclusions` array of a Honcho `peer_context` response as a
/// ranked bullet list, prefixing each fact with its `id` when present so the
/// model can target it via `honcho_conclude`. Returns `None` when empty.
fn format_conclusions(ctx: &Value) -> Option<String> {
let conclusions = ctx.get("conclusions")?.as_array()?;
let lines: Vec<String> = conclusions
.iter()
.filter_map(|c| {
let content = c.get("content").and_then(|v| v.as_str())?;
match c.get("id").and_then(|v| v.as_str()) {
Some(id) => Some(format!("- [{id}] {content}")),
None => Some(format!("- {content}")),
}
})
.collect();
(!lines.is_empty()).then(|| lines.join("\n"))
}
// ── HonchoContextTool ─────────────────────────────────────────────────────────
/// Retrieves a full context snapshot for the user (conclusions, card, summary)
/// from Honcho's `peer_context` endpoint. No LLM synthesis.
struct HonchoContextTool {
client: Arc<HonchoClient>,
workspace_id: String,
}
impl Tool for HonchoContextTool {
fn name(&self) -> &str { "honcho_context" }
fn description(&self) -> &str {
"Retrieve a full context snapshot for the user from Honcho — conclusions, \
peer card, and conversation summary. No LLM synthesis (cheaper than \
memory_query). Pass an optional `query` to focus the semantic search."
}
fn parameters_schema(&self) -> Value {
json!({
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "Optional focus query to filter context. \
Omit for a full context snapshot."
}
}
})
}
fn category(&self) -> ToolCategory { ToolCategory::Introspection }
fn execute(&self, args: Value) -> anyhow::Result<String> {
let client = Arc::clone(&self.client);
let workspace_id = self.workspace_id.clone();
let search_query = args.get("query").and_then(|v| v.as_str()).map(str::to_string);
run_blocking(async move {
let ctx = client
.peer_context(
&workspace_id,
PEER_USER,
&PeerRepresentationGet { search_query, ..Default::default() },
)
.await
.map_err(|e| anyhow::anyhow!("honcho_context: {e}"))?;
Ok(format_context(ctx).unwrap_or_else(|| "No context available yet.".to_string()))
})
}
}
// ── HonchoConcludeTool ────────────────────────────────────────────────────────
/// Writes or deletes a persistent fact (conclusion) about the user in Honcho's
/// memory. Exactly one of `conclusion` or `delete_id` must be supplied.
///
/// Written as `observer = user`, `observed = user` — matching this plugin's peer
/// model, where the `user` peer has `observe_me = true` and therefore holds the
/// self-knowledge that the read-path (`peer_context("user")`) reads back. Using
/// any other observer slot would store facts the read-path never sees.
struct HonchoConcludeTool {
client: Arc<HonchoClient>,
workspace_id: String,
}
impl Tool for HonchoConcludeTool {
fn name(&self) -> &str { "honcho_conclude" }
fn description(&self) -> &str {
"Write or delete a persistent fact about the user in Honcho's memory. \
Pass `conclusion` to create a new fact; pass `delete_id` (from \
honcho_search) to remove one. Exactly one field is required."
}
fn parameters_schema(&self) -> Value {
json!({
"type": "object",
"properties": {
"conclusion": {
"type": "string",
"description": "A factual statement about the user to persist."
},
"delete_id": {
"type": "string",
"description": "Conclusion id to delete (e.g. for PII removal)."
}
}
})
}
fn category(&self) -> ToolCategory { ToolCategory::Introspection }
fn execute(&self, args: Value) -> anyhow::Result<String> {
let conclusion = args.get("conclusion").and_then(|v| v.as_str())
.map(str::trim).filter(|s| !s.is_empty()).map(str::to_string);
let delete_id = args.get("delete_id").and_then(|v| v.as_str())
.map(str::trim).filter(|s| !s.is_empty()).map(str::to_string);
// Exactly one must be present (XOR).
if conclusion.is_some() == delete_id.is_some() {
anyhow::bail!("honcho_conclude: provide exactly one of 'conclusion' or 'delete_id'");
}
let client = Arc::clone(&self.client);
let workspace_id = self.workspace_id.clone();
run_blocking(async move {
if let Some(id) = delete_id {
client
.delete_conclusion(&workspace_id, &id)
.await
.map_err(|e| anyhow::anyhow!("honcho_conclude: {e}"))?;
Ok(format!("Conclusion {id} deleted."))
} else {
let content = conclusion.unwrap();
client
.add_conclusion(
&workspace_id,
ConclusionCreate {
content: content.clone(),
observer_id: PEER_USER.to_string(),
observed_id: PEER_USER.to_string(),
session_id: None,
},
)
.await
.map_err(|e| anyhow::anyhow!("honcho_conclude: {e}"))?;
Ok(format!("Conclusion saved: {content}"))
}
})
}
}
/// Extracts a human-readable string from the raw Honcho `session_context` /
/// `peer_context` JSON response.
///
/// Returns `None` if there is nothing *new* to inject — i.e. when the response
/// contains only raw messages (which are already present in the LLM's own
/// conversation history) or is otherwise empty.
///
/// Only synthesised knowledge is injected:
/// - `conclusions` — facts about the user derived by Honcho's background processing
/// - `summary` — a narrative summary produced by Honcho
///
/// Raw `messages` are intentionally ignored: they are redundant with the local
/// `chat_history` already sent to the LLM and would waste context tokens.
fn format_context(ctx: Value) -> Option<String> {
let mut parts: Vec<String> = Vec::new();
if let Some(conclusions) = ctx.get("conclusions").and_then(|v| v.as_array()) {
let facts: Vec<&str> = conclusions
.iter()
.filter_map(|c| c.get("content").and_then(|v| v.as_str()))
.collect();
if !facts.is_empty() {
parts.push(format!("Known facts about the user:\n- {}", facts.join("\n- ")));
}
}
if let Some(summary) = ctx.get("summary").and_then(|v| v.as_str()) {
if !summary.trim().is_empty() {
parts.push(format!("Conversation summary:\n{summary}"));
}
}
if parts.is_empty() {
return None;
}
Some(format!(
"--- Honcho memory context ---\n{}\n--- end of memory context ---",
parts.join("\n\n")
))
}
// ── HonchoPlugin ──────────────────────────────────────────────────────────────
pub struct HonchoPlugin {
config: Mutex<Option<HonchoConfig>>,
running: Arc<AtomicBool>,
cancel: Mutex<Option<CancellationToken>>,
handle: Mutex<Option<JoinHandle<()>>>,
/// Shared Memory implementation — created once, updated on start/stop.
honcho_memory: Arc<HonchoMemory>,
}
impl HonchoPlugin {
pub fn new() -> Self {
let running = Arc::new(AtomicBool::new(false));
let honcho_memory = Arc::new(HonchoMemory::new(Arc::clone(&running)));
Self {
config: Mutex::new(None),
running,
cancel: Mutex::new(None),
handle: Mutex::new(None),
honcho_memory,
}
}
}
// ── Plugin trait ──────────────────────────────────────────────────────────────
#[async_trait]
impl core_api::plugin::Plugin for HonchoPlugin {
fn id(&self) -> &str { PLUGIN_ID }
fn name(&self) -> &str { "Honcho Memory" }
fn description(&self) -> &str {
"Streams completed interactive chat turns to Honcho for long-term memory \
and injects retrieved context into every LLM turn."
}
fn is_running(&self) -> bool { self.running.load(Ordering::Relaxed) }
fn memory(&self) -> Option<Arc<dyn Memory>> {
Some(Arc::clone(&self.honcho_memory) as Arc<dyn Memory>)
}
fn config_schema(&self) -> Value {
json!({
"type": "object",
"properties": {
"base_url": {
"type": "string",
"title": "Base URL",
"description": "Honcho server URL (e.g. http://localhost:8000)",
"default": "http://localhost:8000"
},
"api_key": {
"type": "string",
"title": "API Key",
"description": "Honcho API key (leave empty for local/unauthenticated instances)",
"sensitive": true
},
"workspace_id": {
"type": "string",
"title": "Workspace ID",
"description": "Honcho workspace identifier for this agent instance",
"default": "personal-agent"
}
},
"required": ["base_url", "workspace_id"]
})
}
fn as_any(&self) -> &dyn std::any::Any { self }
fn as_arc_any(self: Arc<Self>) -> Arc<dyn std::any::Any + Send + Sync> { self }
async fn reload(&self, enabled: bool, config: Value, ctx: PluginContext) -> Result<()> {
let new_cfg = HonchoConfig {
base_url: config["base_url"].as_str().unwrap_or("http://localhost:8000").to_string(),
api_key: config["api_key"].as_str().unwrap_or("").to_string(),
workspace_id: config["workspace_id"].as_str().unwrap_or("personal-agent").to_string(),
};
let old_cfg = self.config.lock().await.clone();
let is_running = self.is_running();
let cfg_changed = old_cfg.as_ref().map_or(true, |old| old != &new_cfg);
match (enabled, is_running) {
(true, false) => {
anyhow::ensure!(
!new_cfg.base_url.is_empty(),
"honcho: cannot start — `base_url` is missing from config"
);
*self.config.lock().await = Some(new_cfg);
self.start(ctx).await?;
}
(false, true) => {
self.stop().await?;
*self.config.lock().await = None;
}
(true, true) if cfg_changed => {
info!("honcho: config changed — restarting");
self.stop().await?;
*self.config.lock().await = Some(new_cfg);
self.start(ctx).await?;
}
_ => {}
}
Ok(())
}
async fn start(&self, ctx: PluginContext) -> Result<()> {
if self.running.load(Ordering::Relaxed) {
return Ok(());
}
let cfg = self.config.lock().await.clone()
.ok_or_else(|| anyhow::anyhow!("honcho: config not set"))?;
let client = Arc::new(HonchoClient::with_base_url(&cfg.base_url, &cfg.api_key));
let workspace_id = cfg.workspace_id.clone();
self.honcho_memory.activate(Arc::clone(&client), workspace_id.clone());
let session_map = Arc::clone(&self.honcho_memory.session_map);
let mut rx = ctx.event_bus.subscribe();
let cancel = CancellationToken::new();
let cancel_clone = cancel.clone();
let running = Arc::clone(&self.running);
self.running.store(true, Ordering::Relaxed);
let task = tokio::spawn(async move {
ensure_workspace_ready(&client, &workspace_id).await;
info!("honcho plugin: listener started (workspace={workspace_id})");
loop {
tokio::select! {
_ = cancel_clone.cancelled() => {
info!("honcho plugin: cancelled");
break;
}
result = rx.recv() => {
match result {
Ok(BusEvent::UserMessage(event)) |
Ok(BusEvent::AssistantResponse(event)) => {
handle_event(
&client, &workspace_id, event, &session_map,
).await;
}
Ok(BusEvent::CompactionDone(_)) => {}
Err(RecvError::Lagged(n)) => {
warn!(
"honcho plugin: event bus lagged by {n} events \
— some turns missed"
);
}
Err(RecvError::Closed) => {
info!("honcho plugin: event bus closed");
break;
}
}
}
}
}
running.store(false, Ordering::Relaxed);
});
*self.cancel.lock().await = Some(cancel);
*self.handle.lock().await = Some(task);
Ok(())
}
async fn stop(&self) -> Result<()> {
if let Some(token) = self.cancel.lock().await.take() {
token.cancel();
}
if let Some(h) = self.handle.lock().await.take() {
let _ = h.await;
}
self.running.store(false, Ordering::Relaxed);
self.honcho_memory.deactivate();
Ok(())
}
}
// ── Internal helpers ──────────────────────────────────────────────────────────
async fn ensure_workspace_ready(client: &HonchoClient, workspace_id: &str) {
match client.create_workspace(&WorkspaceCreate {
id: workspace_id.to_string(),
metadata: None,
configuration: None,
}).await {
Ok(_) => info!("honcho: workspace '{workspace_id}' ready"),
Err(e) => warn!("honcho: workspace '{workspace_id}' create/check failed: {e}"),
}
for peer_id in [PEER_USER, PEER_ASSISTANT] {
match client.create_peer(workspace_id, &PeerCreate {
id: peer_id.to_string(),
metadata: None,
configuration: None,
}).await {
Ok(_) => debug!("honcho: peer '{peer_id}' ready"),
Err(e) => debug!("honcho: peer '{peer_id}' create/check: {e} (likely already exists)"),
}
}
}
async fn handle_event(
client: &HonchoClient,
workspace_id: &str,
event: ChatEvent,
session_map: &Arc<RwLock<HashMap<i64, String>>>,
) {
if !event.is_interactive || event.is_ephemeral || event.is_synthetic {
return;
}
let peer_id = match event.role {
ChatEventRole::User => PEER_USER,
ChatEventRole::Assistant => PEER_ASSISTANT,
ChatEventRole::Agent => return,
};
if event.content.is_empty() {
return;
}
let honcho_session_id = match get_or_create_session(
client, workspace_id, event.session_id, session_map,
).await {
Ok(id) => id,
Err(e) => {
warn!(
"honcho: failed to get/create session for local session {}: {e}",
event.session_id
);
return;
}
};
let msg = MessageCreate {
content: event.content,
peer_id: peer_id.to_string(),
metadata: Some(json!({
"local_message_id": event.message_id,
"local_stack_id": event.stack_id,
})),
configuration: None,
created_at: Some(event.created_at.to_rfc3339()),
};
match client.add_message(workspace_id, &honcho_session_id, msg).await {
Ok(_) => debug!(
"honcho: message sent (session={honcho_session_id}, peer={peer_id})"
),
Err(e) => warn!(
"honcho: add_message failed (session={honcho_session_id}): {e}"
),
}
}
async fn get_or_create_session(
client: &HonchoClient,
workspace_id: &str,
local_session_id: i64,
session_map: &Arc<RwLock<HashMap<i64, String>>>,
) -> Result<String> {
{
let map = session_map.read().await;
if let Some(id) = map.get(&local_session_id) {
return Ok(id.clone());
}
}
let mut peers = HashMap::new();
peers.insert(PEER_USER.to_string(), SessionPeerConfig {
observe_others: None,
observe_me: Some(true),
});
peers.insert(PEER_ASSISTANT.to_string(), SessionPeerConfig {
observe_me: Some(true),
observe_others: None,
});
// Use a deterministic id so the mapping survives plugin restarts without
// needing a DB column — same local_session_id always maps to the same
// Honcho session. Honcho v3 requires `id` in the creation body.
let honcho_id = format!("{workspace_id}-{local_session_id}");
let session = client.create_session(workspace_id, &SessionCreate {
id: Some(honcho_id),
metadata: Some(json!({ "local_session_id": local_session_id })),
peers: Some(peers),
configuration: None,
}).await?;
info!(
"honcho: created session {} for local session {local_session_id}",
session.id
);
let mut map = session_map.write().await;
map.entry(local_session_id).or_insert(session.id.clone());
Ok(map[&local_session_id].clone())
}
+27
View File
@@ -0,0 +1,27 @@
[package]
name = "plugin-mobile-connector"
version = "0.1.0"
edition = "2024"
description = "Mobile connector plugin — bridges Skald's Inbox to mobile apps via the relay (see data/ios-app/plugin.md)"
[dependencies]
# Application layer over `skald-relay-client` (the networking crate). All wire
# transport / E2E crypto / SQLite persistence lives there; this crate keeps only
# the Skald-specific glue (Inbox payload schemas, QR router, control tools).
core-api = { path = "../core-api" }
skald-relay-client = { path = "../skald-relay-client" }
# Still used directly: `crypto::decode_hex` in tools.rs.
skald-relay-common = { path = "../skald-relay-common" }
anyhow = "1"
async-trait = "0.1"
serde = { version = "1", features = ["derive"] }
serde_json = "1"
tokio = { version = "1", features = ["rt", "sync", "macros", "time", "net", "io-util"] }
tokio-util = { version = "0.7", features = ["rt"] }
tracing = "0.1"
axum = { version = "0.8" }
rand = "0.9"
chrono = { version = "0.4", default-features = false, features = ["clock", "std"] }
hex = "0.4"
qrcode = "0.14"
image = { version = "0.25", default-features = false, features = ["png"] }
@@ -0,0 +1,67 @@
//! The `RelayAgent` control surface (plugin.md §4). This is the domain API the
//! UI / control tools use for pairing, listing devices, and revoking. It is NOT
//! an LLM tool itself: the three LLM tools (tools.rs) call into it.
use async_trait::async_trait;
/// Returned by `start_pairing`. The `code` is a random handle distinct from the
/// `pairing_token`; it identifies the in-memory session at the QR endpoint.
pub struct PairingHandle {
/// e.g. `/api/plugin/mobile-connector/pairingqrcode?code=<random>`
pub url: String,
pub code: String,
/// Unix ms.
pub expires_at: i64,
}
/// Device authorization state, surfaced to the listing tool.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ClientState {
Pending,
Authorized,
}
/// One device, surfaced for `mobile_list_devices`.
pub struct ClientInfo {
pub ed25519_pub: [u8; 32],
pub x25519_pub: [u8; 32],
pub state: ClientState,
/// Raw device_info JSON (from `hello`), if received.
pub device_info: Option<String>,
pub platform: Option<String>,
/// Unix ms of last activity, if any.
pub last_seen: Option<i64>,
}
/// The control API exposed by the plugin. Reachable via
/// `PluginManager::get_plugin_typed::<MobileConnectorPlugin>()`.
#[async_trait]
pub trait RelayAgent: Send + Sync {
/// Open the pairing window (single-window, latest-wins) and return the
/// auto-expiring QR URL.
async fn start_pairing(&self, ttl_secs: u32) -> anyhow::Result<PairingHandle>;
/// Close the pairing window.
async fn stop_pairing(&self) -> anyhow::Result<()>;
/// ed25519 public key (namespace identity).
fn agent_ed25519_pub(&self) -> [u8; 32];
/// Derived namespace id (hex).
fn namespace_id(&self) -> String;
/// Send the current Inbox snapshot to all authorized clients.
async fn broadcast_inbox(&self) -> anyhow::Result<()>;
/// Generic push notification to all authorized clients.
async fn broadcast_notification(&self, title: &str, body: &str) -> anyhow::Result<()>;
/// List all known devices.
async fn list_clients(&self) -> Vec<ClientInfo>;
/// Authorize a Pending device by its ed25519 pubkey.
async fn authorize_client(&self, ed25519_pub: [u8; 32]) -> anyhow::Result<()>;
/// Revoke a device (lost/stolen) by its ed25519 pubkey.
async fn revoke_client(&self, ed25519_pub: [u8; 32]) -> anyhow::Result<()>;
}
+179
View File
@@ -0,0 +1,179 @@
//! The application layer on top of the payload-agnostic [`RelayClient`].
//!
//! `RelayApp` owns the Skald-specific semantics that the transport crate
//! deliberately knows nothing about: the E2E JSON payload schemas (`payloads`),
//! the `InboxApi` dispatch, and the authorization policy. It consumes
//! `client.events()` and calls `client.send(...)`; the client handles the wire,
//! crypto, counters, and device registry.
use std::sync::Arc;
use anyhow::Result;
use tokio::sync::broadcast;
use tokio_util::sync::CancellationToken;
use tracing::{debug, warn};
use core_api::inbox::InboxApi;
use skald_relay_client::{ClientState, RelayClient, RelayEvent};
use crate::PLUGIN_ID;
use crate::payloads::{self, ClientPayload};
/// Glue between the relay transport ([`RelayClient`]) and Skald's Inbox.
pub struct RelayApp {
client: Arc<RelayClient>,
inbox: Arc<dyn InboxApi>,
/// When true, a freshly paired device stays Pending until a human confirms;
/// when false, the app auto-authorizes on `ClientPaired`.
require_device_confirmation: bool,
}
impl RelayApp {
pub fn new(
client: Arc<RelayClient>,
inbox: Arc<dyn InboxApi>,
require_device_confirmation: bool,
) -> Self {
Self { client, inbox, require_device_confirmation }
}
/// The underlying transport client (used by the `RelayAgent` impl + router).
pub fn client(&self) -> &Arc<RelayClient> {
&self.client
}
// ── Inbox → clients ───────────────────────────────────────────────────────
/// Build the Inbox snapshot and send it (encrypted) to every Authorized
/// client. `live=false` so the relay stores-and-forwards + pushes to offline
/// phones.
pub async fn broadcast_inbox(&self) -> Result<()> {
let snapshot = self.inbox.list_pending().await;
let plaintext = serde_json::to_vec(&payloads::build_inbox_update(&snapshot))?;
self.broadcast_plaintext(&plaintext).await;
Ok(())
}
/// Build and send a generic notification to all Authorized clients.
pub async fn broadcast_notification(&self, title: &str, body: &str) -> Result<()> {
let plaintext = serde_json::to_vec(&payloads::build_notification(title, body))?;
self.broadcast_plaintext(&plaintext).await;
Ok(())
}
/// Send an opaque plaintext to every Authorized device (`live=false`).
async fn broadcast_plaintext(&self, plaintext: &[u8]) {
for c in self
.client
.list_clients()
.await
.into_iter()
.filter(|c| c.state == ClientState::Authorized)
{
if let Err(e) = self.client.send(&c.ed25519_pub, plaintext, false).await {
warn!(plugin = PLUGIN_ID, error = %e, "failed to send to client");
}
}
}
/// Send the current Inbox snapshot to a single client (the targeted reply to
/// `inbox_request`). `live=true`: the requester is online by construction.
async fn send_inbox_to(&self, client_ed25519_pub: &[u8; 32]) -> Result<()> {
let snapshot = self.inbox.list_pending().await;
let plaintext = serde_json::to_vec(&payloads::build_inbox_update(&snapshot))?;
self.client.send(client_ed25519_pub, &plaintext, true).await
}
// ── Clients → Inbox ───────────────────────────────────────────────────────
/// Apply a decoded client payload to the Inbox. `payload` is the clean inner
/// JSON the client already decrypted + de-framed.
async fn apply_client_payload(&self, from: &[u8; 32], payload: &[u8]) {
match payloads::parse_client_payload(payload) {
ClientPayload::ApprovalResponse { request_id, approved, reason } => {
if approved {
self.inbox.approve(request_id).await;
} else {
self.inbox.reject(request_id, reason.unwrap_or_default()).await;
}
let _ = self.broadcast_inbox().await;
}
ClientPayload::ClarificationResponse { request_id, answer } => {
self.inbox.answer(request_id, answer).await;
let _ = self.broadcast_inbox().await;
}
ClientPayload::ElicitationResponse { request_id, action, content } => {
// `content` may hold a secret (SSH/sudo password): hand it straight
// to the Inbox; never log/persist it in clear (payloads.md §3.1).
self.inbox.resolve_elicitation(request_id, action, content).await;
let _ = self.broadcast_inbox().await;
}
ClientPayload::Hello { device_info } => {
if let Err(e) = self.client.set_device_info(from, &device_info.to_string()).await {
warn!(plugin = PLUGIN_ID, error = %e, "failed to persist device_info");
}
}
ClientPayload::InboxRequest => {
if let Err(e) = self.send_inbox_to(from).await {
warn!(plugin = PLUGIN_ID, error = %e, "failed to send targeted inbox snapshot");
}
}
ClientPayload::Logout => {
if let Err(e) = self.client.revoke(from).await {
warn!(plugin = PLUGIN_ID, error = %e, "logout revoke failed");
}
}
ClientPayload::Unknown => {
debug!(plugin = PLUGIN_ID, "unknown/ignored client payload");
}
}
}
// ── Event loop ────────────────────────────────────────────────────────────
/// Consume the client's [`RelayEvent`] stream until `cancel` fires. This is
/// where the authorization policy and Inbox application live.
pub async fn run_event_loop(
self: Arc<Self>,
mut rx: broadcast::Receiver<RelayEvent>,
cancel: CancellationToken,
) {
loop {
tokio::select! {
_ = cancel.cancelled() => break,
ev = rx.recv() => match ev {
Ok(RelayEvent::Message { from, payload, .. }) => {
self.apply_client_payload(&from, &payload).await;
}
Ok(RelayEvent::ClientPaired { ed25519_pub, .. }) => {
if self.require_device_confirmation {
debug!(
plugin = PLUGIN_ID,
device = %hex::encode(ed25519_pub),
"new device paired (pending manual confirmation)"
);
let _ = self
.broadcast_notification(
"Nuovo device",
"Un nuovo device è in attesa di conferma",
)
.await;
} else if let Err(e) = self.client.authorize(&ed25519_pub).await {
warn!(plugin = PLUGIN_ID, error = %e, "auto-authorize failed");
} else {
// Send the newly-authorized device the current snapshot.
let _ = self.broadcast_inbox().await;
}
}
Ok(RelayEvent::ClientRevoked { .. })
| Ok(RelayEvent::Connected)
| Ok(RelayEvent::Disconnected) => {}
Err(broadcast::error::RecvError::Lagged(n)) => {
warn!(plugin = PLUGIN_ID, skipped = n, "relay event stream lagged");
}
Err(broadcast::error::RecvError::Closed) => break,
}
}
}
}
}
+402
View File
@@ -0,0 +1,402 @@
//! Mobile connector plugin (plugin id `mobile-connector`).
//!
//! Bridges Skald's Inbox (approvals + clarifications) to mobile apps over the
//! relay. The **networking** (v2 WS transport, E2E crypto, anti-replay counters,
//! pairing, device authorization, SQLite persistence) lives in the standalone
//! `skald-relay-client` crate; this plugin is the thin **application** layer on
//! top of it. See `data/iOS-app/v2/relay-protocol.md` for the wire contract and
//! `docs/relay/` for the client/server split.
//!
//! Module map:
//! - `payloads` — E2E JSON payload schemas (inbox_update, responses, …)
//! - `app` — `RelayApp`: Inbox dispatch, auth policy, the events() loop
//! - `router` — the QR-code HTTP endpoint
//! - `agent` — the `RelayAgent` control trait
//! - `tools` — `Tool` impls callable by the host (registered in the main crate)
mod agent;
mod app;
mod notifier;
mod payloads;
mod proxy;
mod router;
mod tools;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use anyhow::Result;
use async_trait::async_trait;
use serde_json::{json, Value};
use tokio::sync::Mutex;
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;
use tracing::{info, warn};
use core_api::plugin::{Plugin, PluginContext};
use skald_relay_client::{ClientState as RelayClientState, RelayClient, RelayClientConfig, SeedSource};
pub use agent::{ClientInfo, ClientState, PairingHandle, RelayAgent};
pub use tools::mobile_tools;
use app::RelayApp;
use notifier::{DelayedNotifier, Kind};
pub(crate) const PLUGIN_ID: &str = "mobile-connector";
const DEFAULT_TTL: u32 = 300;
const MAX_TTL: u32 = 600;
/// Default debounce before an unresolved Inbox item is pushed to the phone.
const DEFAULT_NOTIFY_DELAY_SECS: u64 = 20;
/// Seed file path (relative to the process working dir). Kept byte-identical to
/// the historical location so existing identities/devices survive the upgrade.
const SEED_PATH: &str = "data/relay/seed";
/// The mobile-connector plugin.
pub struct MobileConnectorPlugin {
running: AtomicBool,
/// Live application state — present only while running. Wrapped in `Arc` so
/// the HTTP router (built once at startup) can dynamically point to whichever
/// `RelayApp` is current after a reconfigure (plugin#reload → new state).
inner: Arc<Mutex<Option<Arc<RelayApp>>>>,
cancel: Mutex<Option<CancellationToken>>,
handles: Mutex<Vec<JoinHandle<()>>>,
/// Debounces Inbox pushes to the phone; present only while running.
notifier: Mutex<Option<Arc<DelayedNotifier>>>,
}
impl MobileConnectorPlugin {
pub fn new() -> Self {
Self {
running: AtomicBool::new(false),
inner: Arc::new(Mutex::new(None)),
cancel: Mutex::new(None),
handles: Mutex::new(Vec::new()),
notifier: Mutex::new(None),
}
}
/// Snapshot the live app, if running. Used by the `RelayAgent` impl and the
/// router accessor.
async fn app(&self) -> Option<Arc<RelayApp>> {
self.inner.lock().await.clone()
}
/// Start the runloop and the bus subscriber with the given config.
async fn start_with(&self, config: Value, ctx: &PluginContext) -> Result<()> {
let relay_url = config["relay_url"].as_str().unwrap_or("").to_string();
if relay_url.is_empty() {
warn!(plugin = PLUGIN_ID, "relay_url not configured; plugin idle");
}
let pairing_ttl = config["pairing_ttl"]
.as_u64()
.map(|v| (v as u32).min(MAX_TTL))
.unwrap_or(DEFAULT_TTL);
let require_device_confirmation = config["require_device_confirmation"]
.as_bool()
.unwrap_or(true);
let notify_delay = std::time::Duration::from_secs(
config["notify_delay_secs"]
.as_u64()
.unwrap_or(DEFAULT_NOTIFY_DELAY_SECS),
);
// Build the transport client (derives identity, inits the DB table).
let client = Arc::new(
RelayClient::new(
Arc::clone(&ctx.db),
RelayClientConfig {
relay_url,
pairing_ttl,
seed: SeedSource::Path(SEED_PATH.into()),
},
)
.await?,
);
info!(
plugin = PLUGIN_ID,
namespace = client.namespace_id_hex(),
"mobile-connector identity loaded"
);
client.start().await?;
let app = Arc::new(RelayApp::new(
Arc::clone(&client),
Arc::clone(&ctx.inbox),
require_device_confirmation,
));
let notifier = DelayedNotifier::new(Arc::clone(&app), notify_delay);
let cancel = CancellationToken::new();
let mut handles = Vec::new();
// Event loop: apply inbound payloads + authorization policy.
{
let app2 = Arc::clone(&app);
let rx = client.events();
let c = cancel.clone();
handles.push(tokio::spawn(async move {
app2.run_event_loop(rx, c).await;
}));
}
// Bus subscriber: route the four Inbox events through the debouncer.
// `*Requested` arms a delayed push; `*Resolved` cancels it (or refreshes
// the phone if the push already went out).
{
let notifier = Arc::clone(&notifier);
let c = cancel.clone();
let mut rx = ctx.chat_hub.events(PLUGIN_ID);
handles.push(tokio::spawn(async move {
use core_api::events::ServerEvent::*;
loop {
tokio::select! {
_ = c.cancelled() => break,
ev = rx.recv() => match ev {
Ok(ge) => match ge.event {
ApprovalRequested { request_id, .. } => {
notifier.on_requested((Kind::Approval, request_id)).await;
}
ApprovalResolved { request_id, .. } => {
notifier.on_resolved((Kind::Approval, request_id)).await;
}
ClarificationRequested { request_id, .. } => {
notifier.on_requested((Kind::Clarification, request_id)).await;
}
ClarificationResolved { request_id } => {
notifier.on_resolved((Kind::Clarification, request_id)).await;
}
ElicitationRequested { request_id, .. } => {
notifier.on_requested((Kind::Elicitation, request_id)).await;
}
ElicitationResolved { request_id } => {
notifier.on_resolved((Kind::Elicitation, request_id)).await;
}
_ => {}
},
Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
warn!(plugin = PLUGIN_ID, skipped = n, "event bus lagged");
}
Err(_) => break,
}
}
}
}));
}
// HTTP reverse proxy: bridge `http-local-proxy` pipes to the local web
// server so the native app can render the web UI over the relay (no NAT
// hole / Tailscale). Pinned to 127.0.0.1:<web_port>; access gated by the
// relay's pipe auth (only authorized namespace members).
{
let client2 = Arc::clone(&client);
let c = cancel.clone();
let port = ctx.web_port;
handles.push(tokio::spawn(async move {
crate::proxy::run_proxy_loop(client2, port, c).await;
}));
}
*self.notifier.lock().await = Some(notifier);
*self.inner.lock().await = Some(app);
*self.cancel.lock().await = Some(cancel);
*self.handles.lock().await = handles;
self.running.store(true, Ordering::Relaxed);
info!(plugin = PLUGIN_ID, "mobile-connector started");
Ok(())
}
async fn stop_inner(&self) {
if let Some(c) = self.cancel.lock().await.take() {
c.cancel();
}
// Cancel any armed (not-yet-fired) push timers.
if let Some(notifier) = self.notifier.lock().await.take() {
notifier.cancel_all().await;
}
// Shut down the transport (cancels + joins the WS loop) before dropping
// the app.
if let Some(app) = self.inner.lock().await.take() {
app.client().shutdown().await;
}
for h in self.handles.lock().await.drain(..) {
let _ = h.await;
}
self.running.store(false, Ordering::Relaxed);
info!(plugin = PLUGIN_ID, "mobile-connector stopped");
}
}
impl Default for MobileConnectorPlugin {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl Plugin for MobileConnectorPlugin {
fn id(&self) -> &str { PLUGIN_ID }
fn name(&self) -> &str { "Mobile Connector" }
fn description(&self) -> &str {
"Connects mobile apps to this Skald instance via the relay: bridges the \
Inbox (approvals + clarifications) to phones with end-to-end encryption."
}
fn is_running(&self) -> bool { self.running.load(Ordering::Relaxed) }
fn config_schema(&self) -> Value {
json!({
"type": "object",
"properties": {
"relay_url": {
"type": "string",
"title": "Relay URL",
"description": "wss:// URL of the relay (e.g. wss://relay.skaldagent.net/v1/ws).",
},
"pairing_ttl": {
"type": "integer",
"default": DEFAULT_TTL,
"maximum": MAX_TTL,
"title": "Pairing TTL (seconds)",
"description": "How long a pairing window stays open. Max 600.",
},
"require_device_confirmation": {
"type": "boolean",
"default": true,
"title": "Require device confirmation",
"description": "Require manual confirmation before a newly paired device is authorized (recommended).",
},
"notify_delay_secs": {
"type": "integer",
"default": DEFAULT_NOTIFY_DELAY_SECS,
"minimum": 0,
"title": "Notification delay (seconds)",
"description": "Wait this long before pushing an approval/clarification to the phone. If you answer on the computer within the window, no phone notification is sent. Set 0 to push immediately. (MCP elicitations are Inbox-only and always pushed immediately, regardless of this setting.)",
}
}
})
}
fn runtime_status(&self) -> Option<Value> {
if !self.running.load(Ordering::Relaxed) {
return None;
}
// Synchronous status: report connection flag from the live client.
let connected = self
.inner
.try_lock()
.ok()
.and_then(|g| g.as_ref().map(|app| app.client().is_connected()))
.unwrap_or(false);
Some(json!({ "connected": connected }))
}
async fn reload(&self, enabled: bool, config: Value, ctx: PluginContext) -> Result<()> {
match (enabled, self.is_running()) {
(true, false) => self.start_with(config, &ctx).await,
(false, true) => { self.stop_inner().await; Ok(()) }
(true, true) => { self.stop_inner().await; self.start_with(config, &ctx).await }
(false, false) => Ok(()),
}
}
async fn start(&self, _ctx: PluginContext) -> Result<()> {
// Lifecycle is driven by reload(enabled, ...); nothing to do here.
Ok(())
}
async fn stop(&self) -> Result<()> {
self.stop_inner().await;
Ok(())
}
fn http_router(&self) -> Option<axum::Router> {
// The router is collected once at WebFrontend startup, but the inner
// `RelayApp` may be replaced on reconfigure. We hand over the shared
// `Arc<Mutex<…>>` so the QR route always resolves the *current* app.
Some(router::build(Arc::clone(&self.inner)))
}
fn as_any(&self) -> &dyn std::any::Any { self }
fn as_arc_any(self: Arc<Self>) -> Arc<dyn std::any::Any + Send + Sync> { self }
}
// ── RelayAgent control surface ─────────────────────────────────────────────────
#[async_trait]
impl RelayAgent for MobileConnectorPlugin {
async fn start_pairing(&self, ttl_secs: u32) -> Result<PairingHandle> {
let app = self.app().await.ok_or_else(|| anyhow::anyhow!("plugin not running"))?;
let ttl = if ttl_secs == 0 { 0 } else { ttl_secs.min(MAX_TTL) };
let started = app.client().start_pairing(ttl).await?;
Ok(PairingHandle {
url: format!("/api/plugin/{PLUGIN_ID}/pairingqrcode?code={}", started.code),
code: started.code,
expires_at: started.expires_at,
})
}
async fn stop_pairing(&self) -> Result<()> {
let app = self.app().await.ok_or_else(|| anyhow::anyhow!("plugin not running"))?;
app.client().stop_pairing().await
}
fn agent_ed25519_pub(&self) -> [u8; 32] {
self.inner
.try_lock()
.ok()
.and_then(|g| g.as_ref().map(|app| app.client().agent_ed25519_pub()))
.unwrap_or([0u8; 32])
}
fn namespace_id(&self) -> String {
self.inner
.try_lock()
.ok()
.and_then(|g| g.as_ref().map(|app| app.client().namespace_id_hex()))
.unwrap_or_default()
}
async fn broadcast_inbox(&self) -> Result<()> {
let app = self.app().await.ok_or_else(|| anyhow::anyhow!("plugin not running"))?;
app.broadcast_inbox().await
}
async fn broadcast_notification(&self, title: &str, body: &str) -> Result<()> {
let app = self.app().await.ok_or_else(|| anyhow::anyhow!("plugin not running"))?;
app.broadcast_notification(title, body).await
}
async fn list_clients(&self) -> Vec<ClientInfo> {
let Some(app) = self.app().await else { return Vec::new() };
app.client()
.list_clients()
.await
.into_iter()
.map(|r| ClientInfo {
ed25519_pub: r.ed25519_pub,
x25519_pub: r.x25519_pub,
state: match r.state {
RelayClientState::Authorized => ClientState::Authorized,
RelayClientState::Pending => ClientState::Pending,
},
device_info: r.device_info,
platform: r.platform,
last_seen: r.last_seen,
})
.collect()
}
async fn authorize_client(&self, ed25519_pub: [u8; 32]) -> Result<()> {
let app = self.app().await.ok_or_else(|| anyhow::anyhow!("plugin not running"))?;
app.client().authorize(&ed25519_pub).await?;
// Send the current Inbox snapshot to the newly-authorized device
// (payload-agnostic client doesn't do this itself).
let _ = app.broadcast_inbox().await;
Ok(())
}
async fn revoke_client(&self, ed25519_pub: [u8; 32]) -> Result<()> {
let app = self.app().await.ok_or_else(|| anyhow::anyhow!("plugin not running"))?;
app.client().revoke(&ed25519_pub).await
}
}
@@ -0,0 +1,152 @@
//! Delayed push notifier.
//!
//! The mobile push for an Inbox item is only valuable when the user is *away*
//! from the computer. When they're sitting at the chat, every approval /
//! clarification would otherwise fire an instant — and pointless — phone
//! notification, since they'll answer on the computer within seconds.
//!
//! `DelayedNotifier` debounces that: when a request enters the Inbox it starts a
//! timer; only if the request is still unresolved after `delay` does it push
//! (`broadcast_inbox`). If the user resolves it on the computer first, the timer
//! is cancelled and **no** push is sent. Once a push *has* gone out, the eventual
//! resolution is broadcast so the phone clears the item.
//!
//! Elicitations are the exception: they live only in the Inbox (never inline in
//! the chat), so there is no computer-side answer to debounce against and they
//! are pushed immediately regardless of `delay`.
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Mutex;
use tokio::time::sleep;
use tokio_util::sync::CancellationToken;
use tracing::warn;
use crate::PLUGIN_ID;
use crate::app::RelayApp;
/// Which Inbox manager a `request_id` belongs to. Approvals and clarifications
/// use independent atomic counters, so the id alone is not unique.
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
pub enum Kind {
Approval,
Clarification,
Elicitation,
}
/// `(kind, request_id)` — unique across both Inbox managers.
type Key = (Kind, i64);
#[derive(Default)]
struct State {
/// Requested, timer armed, push not yet sent.
pending: HashMap<Key, CancellationToken>,
/// Push already sent, awaiting the resolution that clears the phone item.
notified: HashSet<Key>,
}
/// Debounces Inbox pushes to the phone. Cheap to clone (`Arc` inside).
pub struct DelayedNotifier {
app: Arc<RelayApp>,
delay: Duration,
state: Mutex<State>,
}
impl DelayedNotifier {
pub fn new(app: Arc<RelayApp>, delay: Duration) -> Arc<Self> {
Arc::new(Self { app, delay, state: Mutex::new(State::default()) })
}
/// A request entered the Inbox: arm a timer. If `delay` elapses before a
/// matching `on_resolved`, push the Inbox to the phone.
///
/// Exception: elicitations (e.g. MCP password prompts) live *only* in the
/// Inbox — never inline in the chat — so there is no computer-side answer to
/// debounce against. Waiting `delay` would just delay a push that is always
/// warranted, so they are pushed immediately (marked *notified* so the
/// eventual `on_resolved` refreshes the phone to clear the item).
pub async fn on_requested(self: &Arc<Self>, key: Key) {
if key.0 == Kind::Elicitation {
let newly_notified = {
let mut st = self.state.lock().await;
// Drop any stale armed timer, then mark notified.
if let Some(old) = st.pending.remove(&key) {
old.cancel();
}
st.notified.insert(key)
};
if newly_notified {
if let Err(e) = self.app.broadcast_inbox().await {
warn!(plugin = PLUGIN_ID, error = %e, "immediate elicitation push failed");
}
}
return;
}
let token = CancellationToken::new();
{
let mut st = self.state.lock().await;
// Re-arming an already-pending key: cancel the stale timer first.
if let Some(old) = st.pending.insert(key, token.clone()) {
old.cancel();
}
}
let this = Arc::clone(self);
let delay = self.delay;
tokio::spawn(async move {
tokio::select! {
_ = token.cancelled() => {} // resolved within the window — send nothing
_ = sleep(delay) => {
// Promote pending → notified under the lock, then push.
let still_armed = {
let mut st = this.state.lock().await;
if st.pending.remove(&key).is_some() {
st.notified.insert(key);
true
} else {
false
}
};
if still_armed {
if let Err(e) = this.app.broadcast_inbox().await {
warn!(plugin = PLUGIN_ID, error = %e, "delayed inbox push failed");
}
}
}
}
});
}
/// A request was resolved (approved / rejected / answered). Suppress the push
/// if it was still pending; otherwise refresh the phone so it clears the item.
pub async fn on_resolved(&self, key: Key) {
let broadcast = {
let mut st = self.state.lock().await;
if let Some(token) = st.pending.remove(&key) {
token.cancel(); // push hadn't fired yet — suppress entirely
false
} else {
// Push already went out, or key untracked: refresh the snapshot.
st.notified.remove(&key);
true
}
};
if broadcast {
if let Err(e) = self.app.broadcast_inbox().await {
warn!(plugin = PLUGIN_ID, error = %e, "inbox broadcast failed");
}
}
}
/// Cancel every armed timer (called on plugin stop).
pub async fn cancel_all(&self) {
let mut st = self.state.lock().await;
for (_, token) in st.pending.drain() {
token.cancel();
}
st.notified.clear();
}
}
@@ -0,0 +1,273 @@
//! E2E payload schemas (payloads.md). These are the JSON plaintexts sealed into
//! the `ciphertext` field of a `message` envelope; the relay never sees them.
//!
//! `request_id` travels as a decimal STRING of the i64 rowid (plugin.md §2): we
//! serialize i64 → string outbound, and parse string → i64 inbound, dropping
//! anything that does not parse.
use chrono::Utc;
use serde_json::Value;
use core_api::inbox::InboxSnapshot;
/// Generate a v4-ish UUID string for the payload `id` field. We avoid pulling in
/// the `uuid` crate: a 16-byte CSPRNG value formatted as a UUID is sufficient
/// for dedup/ack purposes (payloads.md §1).
fn new_id() -> String {
use rand::RngCore;
let mut b = [0u8; 16];
rand::rng().fill_bytes(&mut b);
// Set version (4) and variant bits for a well-formed UUID.
b[6] = (b[6] & 0x0f) | 0x40;
b[8] = (b[8] & 0x3f) | 0x80;
format!(
"{:02x}{:02x}{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}",
b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7],
b[8], b[9], b[10], b[11], b[12], b[13], b[14], b[15],
)
}
/// Convert an ISO-8601 UTC timestamp string into unix milliseconds. Falls back
/// to "now" if the string fails to parse (payloads.md wants an int ms field).
fn iso_to_ms(iso: &str) -> i64 {
chrono::DateTime::parse_from_rfc3339(iso)
.map(|dt| dt.timestamp_millis())
.unwrap_or_else(|_| Utc::now().timestamp_millis())
}
// ── Agent → Client ──────────────────────────────────────────────────────────
/// Build the `inbox_update` snapshot payload (payloads.md §3.1) from an
/// `InboxSnapshot`. Mirrors the Inbox 1:1; `badge` = total pending.
pub fn build_inbox_update(snapshot: &InboxSnapshot) -> Value {
let approvals: Vec<Value> = snapshot
.approvals
.iter()
.map(|a| {
serde_json::json!({
"request_id": a.request_id.to_string(),
"tool_name": a.tool_name,
"agent_label": "Skald",
// Short human label for card/notification; raw args for the detail
// dialog (untruncated — e.g. the full `execute_cmd` command).
"summary": a.summary,
"arguments": a.arguments,
"created_at": iso_to_ms(&a.created_at),
})
})
.collect();
let clarifications: Vec<Value> = snapshot
.clarifications
.iter()
.map(|c| {
serde_json::json!({
"request_id": c.request_id.to_string(),
"question": c.question,
"context": c.context_label,
"suggested_answers": c.suggested_answers,
"agent_label": "Skald",
"created_at": iso_to_ms(&c.created_at),
})
})
.collect();
// MCP server-initiated input requests (e.g. an SSH/sudo password). We ship
// only the prompt metadata — never the value; the value is supplied by the
// device in `elicitation_response.content` and travels E2E (payloads.md §3.1).
let elicitations: Vec<Value> = snapshot
.elicitations
.iter()
.map(|e| {
serde_json::json!({
"request_id": e.request_id.to_string(),
"server_name": e.server_name,
"message": e.message,
"field_name": e.field_name, // Option<String> → null if absent
"sensitive": e.sensitive,
"is_confirmation": e.is_confirmation,
"created_at": iso_to_ms(&e.created_at),
})
})
.collect();
serde_json::json!({
"v": 1,
"kind": "inbox_update",
"id": new_id(),
"ts": Utc::now().timestamp_millis(),
"badge": snapshot.total,
"approvals": approvals,
"clarifications": clarifications,
"elicitations": elicitations,
})
}
/// Build a generic `notification` payload (payloads.md §3.2).
pub fn build_notification(title: &str, body: &str) -> Value {
serde_json::json!({
"v": 1,
"kind": "notification",
"id": new_id(),
"ts": Utc::now().timestamp_millis(),
"title": title,
"body": body,
})
}
// ── Client → Agent ──────────────────────────────────────────────────────────
/// A decoded client→agent payload (payloads.md §4). Only the fields the agent
/// acts on are modeled; unknown kinds become [`ClientPayload::Unknown`].
#[derive(Debug)]
pub enum ClientPayload {
/// `hello`: device_info carried E2E.
Hello { device_info: Value },
/// `approval_response`.
ApprovalResponse { request_id: i64, approved: bool, reason: Option<String> },
/// `clarification_response`.
ClarificationResponse { request_id: i64, answer: String },
/// `elicitation_response`: the device's reply to an MCP elicitation. `action`
/// is `"accept"`/`"decline"`/`"cancel"`; `content` (present only for `accept`)
/// is an object keyed by `field_name` whose value may be a secret — never log it.
ElicitationResponse { request_id: i64, action: String, content: Option<Value> },
/// `inbox_request`: client asks for the current Inbox snapshot (payloads.md
/// §4.6). Sent after every `auth_ok`; the agent replies with a targeted
/// `inbox_update`. No fields beyond the common envelope.
InboxRequest,
/// `logout`: device removes itself.
Logout,
/// Anything else (ack, unknown kind, malformed request_id) — ignored.
Unknown,
}
/// Parse a decrypted client payload. Enforces `v == 1` and required-field
/// presence; on malformed input returns [`ClientPayload::Unknown`] (never panics,
/// payloads.md §6).
pub fn parse_client_payload(plaintext: &[u8]) -> ClientPayload {
let Ok(v) = serde_json::from_slice::<Value>(plaintext) else {
return ClientPayload::Unknown;
};
if v.get("v").and_then(Value::as_u64) != Some(1) {
return ClientPayload::Unknown;
}
let kind = v.get("kind").and_then(Value::as_str).unwrap_or("");
match kind {
"hello" => match v.get("device_info") {
Some(di) if di.is_object() => ClientPayload::Hello { device_info: di.clone() },
_ => ClientPayload::Unknown,
},
"approval_response" => {
let Some(rid) = parse_request_id(&v) else { return ClientPayload::Unknown };
match v.get("decision").and_then(Value::as_str) {
Some("approved") => ClientPayload::ApprovalResponse { request_id: rid, approved: true, reason: None },
Some("rejected") => ClientPayload::ApprovalResponse {
request_id: rid,
approved: false,
reason: v.get("reason").and_then(Value::as_str).map(str::to_string),
},
_ => ClientPayload::Unknown,
}
}
"clarification_response" => {
let Some(rid) = parse_request_id(&v) else { return ClientPayload::Unknown };
match v.get("answer").and_then(Value::as_str) {
Some(answer) => ClientPayload::ClarificationResponse { request_id: rid, answer: answer.to_string() },
None => ClientPayload::Unknown,
}
}
"elicitation_response" => {
let Some(rid) = parse_request_id(&v) else { return ClientPayload::Unknown };
let action = match v.get("action").and_then(Value::as_str) {
Some(a @ ("accept" | "decline" | "cancel")) => a.to_string(),
_ => return ClientPayload::Unknown,
};
// `content` is meaningful only for `accept` and must be an object
// (keyed by `field_name`); anything else is dropped.
let content = v.get("content").filter(|c| c.is_object()).cloned();
ClientPayload::ElicitationResponse { request_id: rid, action, content }
}
"inbox_request" => ClientPayload::InboxRequest,
"logout" => ClientPayload::Logout,
_ => ClientPayload::Unknown,
}
}
/// Parse the `request_id` decimal string into an i64 (plugin.md §2).
fn parse_request_id(v: &Value) -> Option<i64> {
v.get("request_id").and_then(Value::as_str)?.parse::<i64>().ok()
}
#[cfg(test)]
mod tests {
use super::*;
/// `accept` carries the secret under `content`, keyed by `field_name`.
#[test]
fn elicitation_response_accept_with_content() {
let raw = br#"{
"v": 1, "kind": "elicitation_response", "id": "abc", "ts": 1750000000000,
"request_id": "123", "action": "accept",
"content": { "password": "hunter2" }
}"#;
match parse_client_payload(raw) {
ClientPayload::ElicitationResponse { request_id, action, content } => {
assert_eq!(request_id, 123);
assert_eq!(action, "accept");
let content = content.expect("accept must carry content");
assert_eq!(content["password"], "hunter2");
}
other => panic!("expected ElicitationResponse, got {other:?}"),
}
}
/// `decline`/`cancel` have no `content`; a missing object yields `None`.
#[test]
fn elicitation_response_decline_without_content() {
let raw = br#"{
"v": 1, "kind": "elicitation_response", "id": "abc", "ts": 1750000000000,
"request_id": "7", "action": "decline"
}"#;
match parse_client_payload(raw) {
ClientPayload::ElicitationResponse { request_id, action, content } => {
assert_eq!(request_id, 7);
assert_eq!(action, "decline");
assert!(content.is_none());
}
other => panic!("expected ElicitationResponse, got {other:?}"),
}
}
/// A non-object `content` is dropped rather than forwarded.
#[test]
fn elicitation_response_non_object_content_dropped() {
let raw = br#"{
"v": 1, "kind": "elicitation_response", "id": "abc", "ts": 1750000000000,
"request_id": "9", "action": "accept", "content": "not-an-object"
}"#;
match parse_client_payload(raw) {
ClientPayload::ElicitationResponse { content, .. } => assert!(content.is_none()),
other => panic!("expected ElicitationResponse, got {other:?}"),
}
}
/// An unknown `action` is rejected as `Unknown` (no resolution attempted).
#[test]
fn elicitation_response_bad_action_is_unknown() {
let raw = br#"{
"v": 1, "kind": "elicitation_response", "id": "abc", "ts": 1750000000000,
"request_id": "1", "action": "approve"
}"#;
assert!(matches!(parse_client_payload(raw), ClientPayload::Unknown));
}
/// A missing/non-string `request_id` is rejected as `Unknown`.
#[test]
fn elicitation_response_missing_request_id_is_unknown() {
let raw = br#"{
"v": 1, "kind": "elicitation_response", "id": "abc", "ts": 1750000000000,
"action": "accept", "content": { "x": "y" }
}"#;
assert!(matches!(parse_client_payload(raw), ClientPayload::Unknown));
}
}
+204
View File
@@ -0,0 +1,204 @@
//! HTTP reverse proxy over the relay pipe (`docs/relay/pipe.md`).
//!
//! A remote client (the native app) opens a relayed byte-stream pipe with
//! `stream_type = "http-local-proxy"` and treats it as a raw TCP connection to
//! Skald's local web server. This loop accepts those pipes and splices each one,
//! byte-for-byte, to a fresh `127.0.0.1:<web_port>` connection — a transparent
//! tunnel (we never parse HTTP, so HTTP/1.1 keep-alive, parallel connections, and
//! the chat WebSocket upgrade all work). The destination is pinned to the local
//! web port: the client cannot choose host/port, so this never becomes an open
//! proxy to other local services.
//!
//! Access is already gated by the relay: only the namespace agent or an authorized
//! client can establish a pipe (`pipe.md §3.1`).
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use tokio::sync::broadcast::error::RecvError;
use tokio_util::sync::CancellationToken;
use tracing::{debug, trace, warn};
use skald_relay_client::{IncomingPipe, PipeConnection, RelayClient};
use crate::PLUGIN_ID;
/// Pipe `stream_type` this loop handles. The native app sets the same value when
/// opening the pipe.
pub(crate) const HTTP_LOCAL_PROXY_STREAM_TYPE: &str = "http-local-proxy";
/// Read buffer for the local→remote direction. Bounded so a pipe can't buffer
/// unboundedly (the relay also caps the data-plane frame size).
const READ_BUF: usize = 64 * 1024;
/// Monotonic per-connection id for log correlation across concurrent tunnels.
static CONN_SEQ: AtomicU64 = AtomicU64::new(1);
/// Subscribe to inbound pipe invites and reverse-proxy `http-local-proxy` pipes
/// to the local web server. One spawned task per accepted pipe.
pub(crate) async fn run_proxy_loop(
client: Arc<RelayClient>,
web_port: u16,
cancel: CancellationToken,
) {
let mut rx = client.incoming_pipes();
debug!(plugin = PLUGIN_ID, web_port, "http-local-proxy: listening for pipe invites");
loop {
tokio::select! {
_ = cancel.cancelled() => break,
ev = rx.recv() => match ev {
Ok(incoming) => {
trace!(
plugin = PLUGIN_ID,
stream_type = %incoming.stream_type,
from = %hex::encode(incoming.from),
"http-local-proxy: pipe invite received",
);
// Ignore (don't reject) other stream_types: `incoming_pipes` is a
// broadcast, so a future consumer may legitimately want them.
if incoming.stream_type != HTTP_LOCAL_PROXY_STREAM_TYPE {
trace!(plugin = PLUGIN_ID, stream_type = %incoming.stream_type,
"http-local-proxy: stream_type not ours, ignoring");
continue;
}
let client = Arc::clone(&client);
let child = cancel.child_token();
tokio::spawn(async move {
accept_and_proxy(client, incoming, web_port, child).await;
});
}
Err(RecvError::Lagged(n)) => {
warn!(plugin = PLUGIN_ID, skipped = n, "incoming pipes lagged");
}
Err(RecvError::Closed) => break,
}
}
}
debug!(plugin = PLUGIN_ID, "http-local-proxy: invite loop stopped");
}
/// Accept one invite, then proxy it. Runs in its own task.
async fn accept_and_proxy(
client: Arc<RelayClient>,
incoming: IncomingPipe,
web_port: u16,
cancel: CancellationToken,
) {
let conn = CONN_SEQ.fetch_add(1, Ordering::Relaxed);
debug!(plugin = PLUGIN_ID, conn, from = %hex::encode(incoming.from),
"http-local-proxy: accepting pipe");
let pipe = match client.accept_pipe(&incoming).await {
Ok(p) => p,
Err(e) => {
warn!(plugin = PLUGIN_ID, conn, error = %e, "http-local-proxy: accept_pipe failed");
return;
}
};
debug!(plugin = PLUGIN_ID, conn, "http-local-proxy: pipe accepted, opening local connection");
proxy_one(conn, pipe, web_port, cancel).await;
}
/// Splice a pipe to a fresh local TCP connection until either side closes.
///
/// Full-duplex: the pipe is `split` into independent send/receive halves, each
/// driven by its own task, so the two directions never block each other (a
/// stalled write on one side can't hold up the other). `PipeSender::send`
/// applies backpressure internally (it blocks only when the pipe's send buffer
/// is full), and `recv`/`read` are cancel-safe. When either direction ends it
/// cancels the shared token so the other unwinds; dropping both pipe halves
/// closes the socket.
async fn proxy_one(conn: u64, pipe: PipeConnection, port: u16, cancel: CancellationToken) {
let tcp = match TcpStream::connect(("127.0.0.1", port)).await {
Ok(s) => s,
Err(e) => {
warn!(plugin = PLUGIN_ID, conn, port, error = %e, "http-local-proxy: local connect failed");
pipe.close().await;
return;
}
};
debug!(plugin = PLUGIN_ID, conn, port, "http-local-proxy: local connection established");
let (mut rd, mut wr) = tcp.into_split();
let (mut tx, mut rx) = pipe.split();
let to_local = Arc::new(AtomicU64::new(0)); // remote → local bytes
let to_remote = Arc::new(AtomicU64::new(0)); // local → remote bytes
// remote → local: decrypted client bytes forwarded to the web server.
let mut rl = {
let cancel = cancel.clone();
let to_local = Arc::clone(&to_local);
tokio::spawn(async move {
loop {
tokio::select! {
_ = cancel.cancelled() => return "cancelled",
r = rx.recv() => match r {
Ok(Some(bytes)) => {
trace!(plugin = PLUGIN_ID, conn, n = bytes.len(), "http-local-proxy: remote→local");
if let Err(e) = wr.write_all(&bytes).await {
debug!(plugin = PLUGIN_ID, conn, error = %e, "http-local-proxy: local write failed");
return "local write error";
}
to_local.fetch_add(bytes.len() as u64, Ordering::Relaxed);
}
Ok(None) => return "remote closed",
Err(e) => {
debug!(plugin = PLUGIN_ID, conn, error = %e, "http-local-proxy: pipe recv error");
return "pipe recv error";
}
},
}
}
})
};
// local → remote: web-server bytes sealed back over the pipe.
let mut lr = {
let cancel = cancel.clone();
let to_remote = Arc::clone(&to_remote);
tokio::spawn(async move {
let mut buf = vec![0u8; READ_BUF];
loop {
tokio::select! {
_ = cancel.cancelled() => return "cancelled",
r = rd.read(&mut buf) => match r {
Ok(0) => return "local EOF",
Ok(n) => {
trace!(plugin = PLUGIN_ID, conn, n, "http-local-proxy: local→remote");
if let Err(e) = tx.send(&buf[..n]).await {
debug!(plugin = PLUGIN_ID, conn, error = %e, "http-local-proxy: pipe send failed");
return "pipe send error";
}
to_remote.fetch_add(n as u64, Ordering::Relaxed);
}
Err(e) => {
debug!(plugin = PLUGIN_ID, conn, error = %e, "http-local-proxy: local read error");
return "local read error";
}
},
}
}
})
};
// First direction to finish decides the reason; cancel the survivor and join
// **only** it. The handle resolved inside `select!` is already complete and
// must not be polled again (`JoinHandle polled after completion`). Joining
// the survivor lets both pipe halves drop so the socket closes cleanly.
let reason = tokio::select! {
r = &mut rl => {
cancel.cancel();
let _ = lr.await;
r.unwrap_or("remote→local task panic")
}
r = &mut lr => {
cancel.cancel();
let _ = rl.await;
r.unwrap_or("local→remote task panic")
}
};
debug!(plugin = PLUGIN_ID, conn,
to_local = to_local.load(Ordering::Relaxed),
to_remote = to_remote.load(Ordering::Relaxed),
reason, "http-local-proxy: pipe closed");
}
@@ -0,0 +1,113 @@
//! The single HTTP route the plugin contributes: the runtime QR-code endpoint
//! (plugin.md §5). Mounted by the main `WebFrontend` under
//! `/api/plugin/mobile-connector/` behind Skald's normal auth. No QR is ever
//! written to disk — the PNG is rendered on demand from the in-memory session.
//!
//! The router receives the plugin's shared state cell
//! (`Arc<Mutex<Option<Arc<RelayState>>>>`) so that every request resolves the
//! **current** `RelayState` — the same one the LLM tools use. This avoids the
//! classic stale-Arc bug when the plugin is reconfigured (reload stops the old
//! runloop + creates a fresh `RelayState`, but the router is only built once).
use std::sync::Arc;
use axum::extract::{Query, State};
use axum::http::{header, StatusCode};
use axum::response::IntoResponse;
use axum::routing::get;
use axum::Router;
use serde::Deserialize;
use tokio::sync::Mutex;
use skald_relay_client::SessionState;
use crate::app::RelayApp;
/// Shared cell type: an `Arc` to a `Mutex` holding the (optional) live app.
/// Cloned cheaply and safely shared between the plugin and the router.
type StateCell = Arc<Mutex<Option<Arc<RelayApp>>>>;
#[derive(Deserialize)]
struct QrQuery {
code: Option<String>,
}
/// Build the plugin's router. Takes the shared state cell so each request
/// resolves the *current* `RelayState` — not a snapshot from startup.
pub fn build(state_cell: StateCell) -> Router {
Router::new()
.route("/pairingqrcode", get(pairing_qr))
.with_state(state_cell)
}
/// `GET /pairingqrcode?code=<random>` → PNG of the QR while active, else a
/// placeholder PNG (plugin.md §5 table).
async fn pairing_qr(
State(cell): State<StateCell>,
Query(q): Query<QrQuery>,
) -> impl IntoResponse {
let Some(code) = q.code else {
return png_response(render_placeholder("QR non valido"));
};
// Dynamically resolve the *current* RelayApp (same one tools use).
let app = match cell.lock().await.as_ref() {
Some(s) => Arc::clone(s),
None => return png_response(render_placeholder("Plugin non attivo")),
};
match app.client().lookup_pairing(&code) {
Some((qr, SessionState::Active)) => {
// Encode the normative QrCodeData JSON into the QR.
match serde_json::to_string(&qr) {
Ok(json) => match render_qr(&json) {
Ok(png) => png_response(png),
Err(_) => png_response(render_placeholder("Errore QR")),
},
Err(_) => png_response(render_placeholder("Errore QR")),
}
}
Some((_, SessionState::Consumed)) => png_response(render_placeholder("QR già usato")),
Some((_, SessionState::Superseded)) => png_response(render_placeholder("QR scaduto")),
None => png_response(render_placeholder("QR scaduto")),
}
}
/// Wrap PNG bytes in a no-cache image response.
fn png_response(png: Vec<u8>) -> axum::response::Response {
(
StatusCode::OK,
[
(header::CONTENT_TYPE, "image/png"),
(header::CACHE_CONTROL, "no-store"),
],
png,
)
.into_response()
}
/// Render `payload` as a QR PNG (qrcode + image, all in memory).
fn render_qr(payload: &str) -> anyhow::Result<Vec<u8>> {
use image::{ImageFormat, Luma};
let code = qrcode::QrCode::new(payload.as_bytes())?;
let img = code.render::<Luma<u8>>().min_dimensions(512, 512).build();
let mut buf = std::io::Cursor::new(Vec::new());
img.write_to(&mut buf, ImageFormat::Png)?;
Ok(buf.into_inner())
}
/// Render a simple placeholder PNG carrying `msg` as a small QR (renders text so
/// a browser shows *something*; no disk I/O). Falls back to a blank image if the
/// text encode fails.
fn render_placeholder(msg: &str) -> Vec<u8> {
render_qr(msg).unwrap_or_else(|_| blank_png())
}
/// 1×1 white PNG, used only if QR rendering itself fails.
fn blank_png() -> Vec<u8> {
use image::{ImageBuffer, ImageFormat, Luma};
let img: ImageBuffer<Luma<u8>, Vec<u8>> = ImageBuffer::from_pixel(1, 1, Luma([255u8]));
let mut buf = std::io::Cursor::new(Vec::new());
let _ = img.write_to(&mut buf, ImageFormat::Png);
buf.into_inner()
}
+163
View File
@@ -0,0 +1,163 @@
//! The three LLM-callable control tools (plugin.md §11).
//!
//! These implement `core_api::tool::Tool` and close over the plugin's
//! `RelayAgent`. The host (main crate) registers them in its `ToolRegistry` via
//! [`mobile_tools`]. `mobile_start_pairing` is gated behind the approval engine
//! the same way `execute_cmd`/`restart` are: the host seeds a `require` rule for
//! the tool name (see docs), so opening a pairing window is always a deliberate
//! human action and never triggerable by prompt injection.
use std::sync::Arc;
use anyhow::Result;
use serde_json::{json, Value};
use core_api::tool::{Tool, ToolCategory};
use crate::agent::{ClientState, RelayAgent};
/// Tool name constants (also the patterns the host uses for approval rules).
pub const TOOL_START_PAIRING: &str = "mobile_start_pairing";
pub const TOOL_LIST_DEVICES: &str = "mobile_list_devices";
pub const TOOL_REVOKE_DEVICE: &str = "mobile_revoke_device";
/// Build the plugin's LLM tools, bound to a `RelayAgent`. The host calls this
/// and registers the result in its `ToolRegistry`.
pub fn mobile_tools(agent: Arc<dyn RelayAgent>) -> Vec<Arc<dyn Tool>> {
vec![
Arc::new(StartPairingTool { agent: Arc::clone(&agent) }),
Arc::new(ListDevicesTool { agent: Arc::clone(&agent) }),
Arc::new(RevokeDeviceTool { agent }),
]
}
// ── mobile_start_pairing ──────────────────────────────────────────────────────
struct StartPairingTool {
agent: Arc<dyn RelayAgent>,
}
impl Tool for StartPairingTool {
fn name(&self) -> &str { TOOL_START_PAIRING }
fn description(&self) -> &str {
"Open a pairing window so a new mobile device can be added, returning a URL \
that renders a QR code to scan. The window auto-expires. Requires user approval."
}
fn parameters_schema(&self) -> Value {
json!({
"type": "object",
"properties": {
"ttl": {
"type": "integer",
"description": "Optional window lifetime in seconds (max 600). Defaults to the configured value."
}
}
})
}
fn category(&self) -> ToolCategory { ToolCategory::Config }
fn interactive_only(&self) -> bool { true }
fn execute_async<'a>(
&'a self,
args: Value,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<String>> + Send + 'a>> {
Box::pin(async move {
let ttl = args.get("ttl").and_then(Value::as_u64).unwrap_or(0) as u32;
let handle = self.agent.start_pairing(ttl).await?;
Ok(format!(
"Pairing window open. Show this QR to the device:\n\n![pairing QR]({})\n\n\
The window expires automatically.",
handle.url
))
})
}
}
// ── mobile_list_devices ───────────────────────────────────────────────────────
struct ListDevicesTool {
agent: Arc<dyn RelayAgent>,
}
impl Tool for ListDevicesTool {
fn name(&self) -> &str { TOOL_LIST_DEVICES }
fn description(&self) -> &str {
"List paired mobile devices: state (pending/authorized), platform, device info, last seen."
}
fn parameters_schema(&self) -> Value {
json!({ "type": "object", "properties": {} })
}
fn category(&self) -> ToolCategory { ToolCategory::Introspection }
fn execute_async<'a>(
&'a self,
_args: Value,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<String>> + Send + 'a>> {
Box::pin(async move {
let clients = self.agent.list_clients().await;
if clients.is_empty() {
return Ok("No paired devices.".to_string());
}
let items: Vec<Value> = clients
.into_iter()
.map(|c| {
let device_info: Option<Value> =
c.device_info.as_deref().and_then(|s| serde_json::from_str(s).ok());
json!({
"ed25519_pub": hex::encode(c.ed25519_pub),
"state": match c.state {
ClientState::Authorized => "authorized",
ClientState::Pending => "pending",
},
"platform": c.platform,
"device_info": device_info,
"last_seen": c.last_seen,
})
})
.collect();
Ok(serde_json::to_string_pretty(&json!({ "devices": items }))?)
})
}
}
// ── mobile_revoke_device ──────────────────────────────────────────────────────
struct RevokeDeviceTool {
agent: Arc<dyn RelayAgent>,
}
impl Tool for RevokeDeviceTool {
fn name(&self) -> &str { TOOL_REVOKE_DEVICE }
fn description(&self) -> &str {
"Revoke a paired mobile device by its ed25519 public key (hex). The device loses access immediately."
}
fn parameters_schema(&self) -> Value {
json!({
"type": "object",
"properties": {
"pubkey": {
"type": "string",
"description": "The device's ed25519 public key, hex-encoded (64 chars)."
}
},
"required": ["pubkey"]
})
}
fn category(&self) -> ToolCategory { ToolCategory::Config }
fn execute_async<'a>(
&'a self,
args: Value,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<String>> + Send + 'a>> {
Box::pin(async move {
let pubkey = args
.get("pubkey")
.and_then(Value::as_str)
.ok_or_else(|| anyhow::anyhow!("mobile_revoke_device: missing `pubkey`"))?;
let ed = skald_relay_common::crypto::decode_hex::<32>(pubkey)
.ok_or_else(|| anyhow::anyhow!("mobile_revoke_device: `pubkey` is not 32-byte hex"))?;
self.agent.revoke_client(ed).await?;
Ok(format!("Device {pubkey} revoked."))
})
}
}
+24
View File
@@ -0,0 +1,24 @@
[package]
name = "plugin-tailscale-remote"
version = "0.1.0"
edition = "2024"
[dependencies]
core-api = { path = "../core-api" }
anyhow = "1"
async-trait = "0.1"
serde_json = "1"
tokio = { version = "1", features = ["full"] }
tokio-util = { version = "0.7", features = ["rt"] }
tracing = "0.1"
axum = { version = "0.8" }
tailscale = { version = "0.3.3", optional = true, features = ["axum"] }
# `remote-tailscale` pulls the pure-Rust `tailscale` crate, which internally
# forces the `aws-lc-rs` crypto backend (a C/cmake build) — the last thing that
# would keep the ring-only, C-toolchain-free static (musl) binary from building.
# It is therefore OFF by default; the recommended `tailscale_sys` provider (which
# drives the system tailscaled and needs none of this) is always compiled.
[features]
default = []
remote-tailscale = ["dep:tailscale"]
+274
View File
@@ -0,0 +1,274 @@
use std::net::Ipv4Addr;
use std::sync::{Arc, OnceLock};
use std::sync::atomic::{AtomicBool, Ordering};
use anyhow::{Result, bail};
use async_trait::async_trait;
use serde_json::{Value, json};
use tokio::sync::{Mutex, RwLock};
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;
use tracing::{error, info, warn};
use core_api::plugin::{PluginContext, RouterFactory};
use core_api::remote::RemoteAccess;
#[cfg(feature = "remote-tailscale")]
mod tailscale;
mod tailscale_sys;
#[cfg(feature = "remote-tailscale")]
use tailscale::{TailscaleConfig, TailscaleEmbeddedProvider};
pub struct RemotePlugin {
running: Arc<AtomicBool>,
cancel: Mutex<Option<CancellationToken>>,
handle: Mutex<Option<JoinHandle<()>>>,
/// Cached mesh IP — set synchronously once connected, cleared on stop.
mesh_ip: std::sync::Mutex<Option<Ipv4Addr>>,
/// Active provider name cached for synchronous runtime_status().
cached_provider: std::sync::Mutex<Option<String>>,
/// Concrete embedded provider kept alive for graceful shutdown.
#[cfg(feature = "remote-tailscale")]
provider: Mutex<Option<Arc<TailscaleEmbeddedProvider>>>,
// ── Deps extracted from AppState on first start() ─────────────────────────
// Using OnceLock so extraction is idempotent across reload() calls.
// After the first start(), the internal methods use these and never
// touch Arc<AppState> again.
port: OnceLock<u16>,
remote_slot: OnceLock<Arc<RwLock<Option<Arc<dyn RemoteAccess>>>>>,
router_factory: OnceLock<RouterFactory>,
}
impl RemotePlugin {
pub fn new() -> Self {
Self {
running: Arc::new(AtomicBool::new(false)),
cancel: Mutex::new(None),
handle: Mutex::new(None),
mesh_ip: std::sync::Mutex::new(None),
cached_provider: std::sync::Mutex::new(None),
#[cfg(feature = "remote-tailscale")]
provider: Mutex::new(None),
port: OnceLock::new(),
remote_slot: OnceLock::new(),
router_factory: OnceLock::new(),
}
}
/// Cache the three deps we need from PluginContext. Idempotent (OnceLock).
fn extract_deps(&self, ctx: &PluginContext) {
let _ = self.port.set(ctx.web_port);
let _ = self.remote_slot.set(Arc::clone(&ctx.remote_slot));
let _ = self.router_factory.set(Arc::clone(&ctx.router_factory));
}
}
#[async_trait]
impl core_api::plugin::Plugin for RemotePlugin {
fn id(&self) -> &str { "remote_connectivity" }
fn name(&self) -> &str { "Remote Connectivity" }
fn description(&self) -> &str {
"Exposes the web app on a mesh network (Tailscale) so remote clients can connect \
without port forwarding or internet exposure."
}
fn is_running(&self) -> bool { self.running.load(Ordering::Relaxed) }
fn config_schema(&self) -> Value {
json!({
"type": "object",
"properties": {
"provider": {
"type": "string",
"enum": ["tailscale_sys", "tailscale"],
"default": "tailscale_sys",
"title": "Provider",
"description": "tailscale_sys: uses the system Tailscale daemon (recommended). tailscale: experimental embedded Tailscale, no daemon required."
},
"auth_key": {
"type": "string",
"title": "Auth Key",
"description": "Tailscale auth key (tskey-auth-...). Only required for the embedded 'tailscale' provider on first join.",
"sensitive": true
},
"hostname": {
"type": "string",
"default": "personal-agent",
"title": "Hostname",
"description": "Hostname this node requests on the tailnet. Only used by the embedded 'tailscale' provider."
},
"key_file": {
"type": "string",
"default": "data/tailscale_keys.json",
"title": "Key File",
"description": "Path for persisting node identity between restarts. Only used by the embedded 'tailscale' provider."
}
}
})
}
fn runtime_status(&self) -> Option<Value> {
if !self.running.load(Ordering::Relaxed) {
return None;
}
let ip = self.mesh_ip.lock().ok()
.and_then(|g| g.map(|ip| ip.to_string()))
.unwrap_or_default();
let provider = self.cached_provider.lock().ok()
.and_then(|g| g.clone())
.unwrap_or_else(|| "unknown".to_string());
Some(json!({ "provider": provider, "ip": ip, "connected": true }))
}
fn as_any(&self) -> &dyn std::any::Any { self }
fn as_arc_any(self: Arc<Self>) -> Arc<dyn std::any::Any + Send + Sync> { self }
async fn reload(&self, enabled: bool, config: Value, ctx: PluginContext) -> Result<()> {
self.extract_deps(&ctx);
match (enabled, self.is_running()) {
(true, false) => self.start_with_config(config).await,
(false, true) => self.stop().await,
(true, true) => { self.stop().await?; self.start_with_config(config).await }
(false, false) => Ok(()),
}
}
async fn start(&self, ctx: PluginContext) -> Result<()> {
self.extract_deps(&ctx);
Ok(())
}
async fn stop(&self) -> Result<()> {
if let Some(token) = self.cancel.lock().await.take() {
token.cancel();
}
if let Some(h) = self.handle.lock().await.take() {
let _ = h.await;
}
#[cfg(feature = "remote-tailscale")]
if let Some(provider) = self.provider.lock().await.take() {
provider.shutdown().await;
}
if let Ok(mut g) = self.mesh_ip.lock() { *g = None; }
if let Ok(mut g) = self.cached_provider.lock() { *g = None; }
self.running.store(false, Ordering::Relaxed);
info!(plugin = "remote_connectivity", "remote plugin stopped");
Ok(())
}
}
// ── Internal start helpers ────────────────────────────────────────────────────
impl RemotePlugin {
async fn start_with_config(&self, config: Value) -> Result<()> {
let provider_name = config["provider"].as_str().unwrap_or("tailscale_sys");
match provider_name {
"tailscale_sys" => self.start_tailscale_sys().await,
#[cfg(feature = "remote-tailscale")]
"tailscale" => self.start_tailscale(config).await,
other => bail!("remote_connectivity: unknown provider '{other}'"),
}
}
#[cfg(feature = "remote-tailscale")]
async fn start_tailscale(&self, config: Value) -> Result<()> {
let auth_key = config["auth_key"].as_str().unwrap_or("").to_string();
let hostname = config["hostname"].as_str().unwrap_or("personal-agent").to_string();
let key_file = config["key_file"]
.as_str()
.unwrap_or("data/tailscale_keys.json")
.to_string();
let port = *self.port.get().expect("extract_deps not called");
let remote_slot = self.remote_slot.get().expect("extract_deps not called");
let router_factory = self.router_factory.get().expect("extract_deps not called");
let provider = Arc::new(
TailscaleEmbeddedProvider::new(TailscaleConfig { auth_key, hostname, key_file }).await?
);
let ip = provider.device_ip().await?;
if let Ok(mut g) = self.mesh_ip.lock() { *g = Some(ip); }
if let Ok(mut g) = self.cached_provider.lock() { *g = Some("tailscale".into()); }
*remote_slot.write().await = Some(Arc::clone(&provider) as Arc<dyn RemoteAccess>);
let listener = provider.axum_listener(port).await?;
let router = router_factory();
let cancel = CancellationToken::new();
let cancel_c = cancel.clone();
let running = Arc::clone(&self.running);
self.running.store(true, Ordering::Relaxed);
let handle = tokio::spawn(async move {
info!(provider = "tailscale", %ip, port, "mesh server listening");
tokio::select! {
_ = cancel_c.cancelled() => {
info!(provider = "tailscale", "mesh server cancelled");
}
result = axum::serve(listener, router) => {
match result {
Ok(()) => warn!(provider = "tailscale", "mesh server exited unexpectedly"),
Err(e) => error!(provider = "tailscale", error = %e, "mesh server error"),
}
}
}
running.store(false, Ordering::Relaxed);
});
*self.cancel.lock().await = Some(cancel);
*self.handle.lock().await = Some(handle);
*self.provider.lock().await = Some(provider);
Ok(())
}
async fn start_tailscale_sys(&self) -> Result<()> {
use tailscale_sys::TailscaleSystemProvider;
let port = *self.port.get().expect("extract_deps not called");
let remote_slot = self.remote_slot.get().expect("extract_deps not called");
let router_factory = self.router_factory.get().expect("extract_deps not called");
let provider = Arc::new(TailscaleSystemProvider::new().await?);
let ip = provider.device_ip().await?;
if let Ok(mut g) = self.mesh_ip.lock() { *g = Some(ip); }
if let Ok(mut g) = self.cached_provider.lock() { *g = Some("tailscale_sys".into()); }
*remote_slot.write().await = Some(Arc::clone(&provider) as Arc<dyn RemoteAccess>);
let listener = tokio::net::TcpListener::bind((ip, port)).await?;
let router = router_factory();
let cancel = CancellationToken::new();
let cancel_c = cancel.clone();
let running = Arc::clone(&self.running);
let provider_c = Arc::clone(&provider);
self.running.store(true, Ordering::Relaxed);
let handle = tokio::spawn(async move {
info!(provider = "tailscale_sys", %ip, port, "mesh server listening");
tokio::select! {
_ = cancel_c.cancelled() => {
info!(provider = "tailscale_sys", "mesh server cancelled");
}
result = axum::serve(listener, router) => {
match result {
Ok(()) => warn!(provider = "tailscale_sys", "mesh server exited unexpectedly"),
Err(e) => error!(provider = "tailscale_sys", error = %e, "mesh server error"),
}
}
}
provider_c.shutdown().await;
running.store(false, Ordering::Relaxed);
});
*self.cancel.lock().await = Some(cancel);
*self.handle.lock().await = Some(handle);
Ok(())
}
}
@@ -0,0 +1,89 @@
use std::net::{Ipv4Addr, SocketAddr};
use std::sync::OnceLock;
use anyhow::{Context, Result, anyhow};
use async_trait::async_trait;
use tokio::sync::Mutex;
use tracing::{info, warn};
use core_api::remote::RemoteAccess;
pub struct TailscaleConfig {
/// Tailscale auth key (tskey-auth-...). Required for first join; reused from key file afterwards.
pub auth_key: String,
/// Hostname this node will request on the tailnet.
pub hostname: String,
/// Path where tailscale-rs persists node keys between restarts.
pub key_file: String,
}
pub struct TailscaleEmbeddedProvider {
device: Mutex<Option<tailscale::Device>>,
ip: OnceLock<Ipv4Addr>,
hostname: String,
}
impl TailscaleEmbeddedProvider {
pub async fn new(cfg: TailscaleConfig) -> Result<Self> {
let mut ts_cfg = tailscale::Config::default_with_key_file(&cfg.key_file)
.await
.context("loading tailscale key file")?;
ts_cfg.requested_hostname = Some(cfg.hostname.clone());
let auth_key = if cfg.auth_key.is_empty() { None } else { Some(cfg.auth_key) };
let device = tailscale::Device::new(&ts_cfg, auth_key)
.await
.context("creating tailscale device")?;
let ip = device.ipv4_addr().await.context("getting tailscale IPv4")?;
info!(provider = "tailscale", %ip, hostname = %cfg.hostname, "mesh device ready");
Ok(Self {
device: Mutex::new(Some(device)),
ip: OnceLock::from(ip),
hostname: cfg.hostname,
})
}
/// Create an Axum-compatible listener bound to the mesh IP on `port`.
/// The caller is responsible for spawning an `axum::serve` task with the result.
pub async fn axum_listener(&self, port: u16) -> Result<tailscale::axum::Listener> {
let ip = self.ip.get().copied().ok_or_else(|| anyhow!("device not ready"))?;
let addr = SocketAddr::from((ip, port));
let guard = self.device.lock().await;
let device = guard.as_ref().ok_or_else(|| anyhow!("device shut down"))?;
let net_listener = device
.tcp_listen(addr)
.await
.with_context(|| format!("tcp_listen on {addr}"))?;
Ok(tailscale::axum::Listener::from(net_listener))
}
}
#[async_trait]
impl RemoteAccess for TailscaleEmbeddedProvider {
fn provider_name(&self) -> &str { "tailscale" }
async fn device_ip(&self) -> Result<Ipv4Addr> {
self.ip.get().copied().ok_or_else(|| anyhow!("device not ready"))
}
fn is_connected(&self) -> bool { self.ip.get().is_some() }
async fn shutdown(&self) {
let device = self.device.lock().await.take();
if let Some(dev) = device {
let clean = dev.shutdown(Some(std::time::Duration::from_secs(5))).await;
if clean {
info!(provider = "tailscale", hostname = %self.hostname, "mesh device shut down cleanly");
} else {
warn!(provider = "tailscale", hostname = %self.hostname, "mesh device shutdown timed out");
}
}
}
}
@@ -0,0 +1,71 @@
use std::net::Ipv4Addr;
use std::sync::OnceLock;
use std::sync::atomic::{AtomicBool, Ordering};
use anyhow::{Context, Result};
use async_trait::async_trait;
use tokio::process::Command;
use core_api::remote::RemoteAccess;
/// Remote-access provider that reads the Tailscale IP from the system daemon.
///
/// Requires `tailscaled` to be installed and running; the `tailscale` CLI must
/// be on PATH. No embedded netstack — zero experimental dependencies.
pub struct TailscaleSystemProvider {
ip: OnceLock<Ipv4Addr>,
connected: AtomicBool,
}
impl TailscaleSystemProvider {
/// Runs `tailscale ip -4`, caches the result, and returns `Self` on success.
pub async fn new() -> Result<Self> {
let provider = Self {
ip: OnceLock::new(),
connected: AtomicBool::new(false),
};
let ip = provider.fetch_ip().await?;
let _ = provider.ip.set(ip);
provider.connected.store(true, Ordering::Relaxed);
Ok(provider)
}
async fn fetch_ip(&self) -> Result<Ipv4Addr> {
let output = Command::new("tailscale")
.args(["ip", "-4"])
.output()
.await
.context("tailscale command not found — is the Tailscale daemon installed and running?")?;
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr);
anyhow::bail!("tailscale ip -4 failed: {stderr}");
}
let stdout = String::from_utf8_lossy(&output.stdout);
let ip_str = stdout.trim();
ip_str
.parse::<Ipv4Addr>()
.with_context(|| format!("failed to parse tailscale IP: '{ip_str}'"))
}
}
#[async_trait]
impl RemoteAccess for TailscaleSystemProvider {
fn provider_name(&self) -> &str { "tailscale_sys" }
async fn device_ip(&self) -> Result<Ipv4Addr> {
self.ip
.get()
.copied()
.ok_or_else(|| anyhow::anyhow!("tailscale_sys: not connected"))
}
fn is_connected(&self) -> bool {
self.connected.load(Ordering::Relaxed)
}
async fn shutdown(&self) {
self.connected.store(false, Ordering::Relaxed);
}
}
+18
View File
@@ -0,0 +1,18 @@
[package]
name = "plugin-telegram-bot"
version = "0.1.0"
edition = "2024"
[dependencies]
core-api = { path = "../core-api" }
anyhow = "1"
async-trait = "0.1"
serde = { version = "1", features = ["derive"] }
serde_json = "1"
tokio = { version = "1", features = ["full"] }
tokio-util = { version = "0.7", features = ["rt"] }
tracing = "0.1"
teloxide = { version = "0.17.0", default-features = false, features = ["macros", "rustls", "ctrlc_handler"] }
chrono = { version = "0.4", default-features = false, features = ["clock", "std"] }
rand = "0.10.1"
regex = "1"
@@ -0,0 +1,132 @@
use std::path::Path;
use anyhow::Result;
use core_api::message_meta::Attachment;
use teloxide::net::Download;
use teloxide::prelude::*;
/// A media item sent by the user via Telegram.
///
/// # Extending
/// Add a new variant here, then handle it in:
/// 1. `handlers::classify_message` — detect the message type and build the variant
/// 2. `TelegramAttachment::download_and_save` — fetch bytes and persist to disk,
/// returning an [`Attachment`]
/// (return `Ok(None)` if no file is involved)
/// 3. `TelegramAttachment::system_info_message` — describe a file-less variant
/// (Location) for the LLM
pub(crate) enum TelegramAttachment {
Document {
file_id: String,
file_name: String,
mime_type: Option<String>,
caption: Option<String>,
},
Photo {
file_id: String,
caption: Option<String>,
},
Location {
latitude: f64,
longitude: f64,
accuracy: Option<f64>,
/// True when the user shared a live location (continuously updated).
is_live: bool,
},
}
impl TelegramAttachment {
/// Downloads the attachment from Telegram, writes it to `base_dir/<chat_id>/<name>`,
/// and returns the saved [`Attachment`] (shared with the web/mobile path so the
/// copilot UI renders it identically). Returns `None` for attachment types that
/// carry no binary content (e.g. Location).
///
/// The returned `path` is made relative to the process working directory (the
/// project root) when possible, so it is both servable under `/data/…` and
/// resolvable by the filesystem tools — matching web uploads.
pub(crate) async fn download_and_save(
&self,
bot: &Bot,
base_dir: &Path,
chat_id: i64,
) -> Result<Option<Attachment>> {
let (file_id, file_name, mimetype): (&str, String, Option<String>) = match self {
Self::Document { file_id, file_name, mime_type, .. } =>
(file_id, file_name.clone(), mime_type.clone()),
Self::Photo { file_id, .. } =>
(file_id, format!("{file_id}.jpg"), Some("image/jpeg".to_string())),
Self::Location { .. } => return Ok(None),
};
let dir = base_dir.join(chat_id.to_string());
tokio::fs::create_dir_all(&dir).await?;
let tg_file = bot.get_file(teloxide::types::FileId(file_id.to_string())).await?;
let mut bytes = Vec::new();
bot.download_file(&tg_file.path, &mut bytes).await?;
let path = dir.join(&file_name);
tokio::fs::write(&path, &bytes).await?;
// Prefer a project-root-relative path so `/data/…` serving works.
let rel = std::env::current_dir()
.ok()
.and_then(|cwd| path.strip_prefix(&cwd).ok().map(Path::to_path_buf))
.unwrap_or_else(|| path.clone());
Ok(Some(Attachment {
path: rel.to_string_lossy().to_string(),
name: file_name,
mimetype,
filesize: Some(bytes.len() as u64),
}))
}
/// Builds the `[TELEGRAM SYSTEM INFO]` message injected into the conversation history.
/// `saved_path` is `None` for attachment types that produce no file on disk.
pub(crate) fn system_info_message(&self, saved_path: Option<&Path>) -> String {
match self {
Self::Document { file_name, mime_type, caption, .. } => {
let mime = mime_type.as_deref().unwrap_or("application/octet-stream");
let path = saved_path.map(|p| p.display().to_string()).unwrap_or_default();
format!(
"[TELEGRAM SYSTEM INFO]\n\
The user has sent a file attachment.\n\
File name: {file_name}\n\
MIME type: {mime}\n\
Saved at: {path}{}",
caption_line(caption.as_deref()),
)
}
Self::Photo { caption, .. } => {
let path = saved_path.map(|p| p.display().to_string()).unwrap_or_default();
format!(
"[TELEGRAM SYSTEM INFO]\n\
The user has sent a photo.\n\
Saved at: {path}{}",
caption_line(caption.as_deref()),
)
}
Self::Location { latitude, longitude, accuracy, is_live } => {
let maps_url = format!("https://maps.google.com/?q={latitude},{longitude}");
let accuracy_line = accuracy
.map(|a| format!("\nAccuracy: ±{a:.0} m"))
.unwrap_or_default();
let kind = if *is_live { "live location (snapshot at time of receipt)" } else { "location" };
format!(
"[TELEGRAM SYSTEM INFO]\n\
The user has shared a {kind}.\n\
Latitude: {latitude}\n\
Longitude: {longitude}{accuracy_line}\n\
Maps URL: {maps_url}"
)
}
}
}
}
fn caption_line(caption: Option<&str>) -> String {
caption
.map(|c| format!("\nCaption: {c}"))
.unwrap_or_default()
}
+170
View File
@@ -0,0 +1,170 @@
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::SystemTime;
use anyhow::Result;
use chrono::{DateTime, Local, Utc};
use rand::RngExt;
use serde::{Deserialize, Serialize};
use teloxide::prelude::*;
use teloxide::types::ParseMode;
use tokio::time::Duration;
use tokio_util::sync::CancellationToken;
use tracing::{error, info};
use super::TgShared;
// ── Whitelist file schema ─────────────────────────────────────────────────────
//
// Written to secrets/telegram_whitelist.json.
// The main agent edits this file directly to authorise users.
#[derive(Debug, Serialize, Deserialize, Default)]
pub struct WhitelistFile {
#[serde(default)]
pub whitelist: Vec<i64>,
#[serde(default)]
pub pending_pairings: Vec<PairingEntry>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct PairingEntry {
pub code: String,
pub chat_id: i64,
pub issued_at: String,
}
pub(crate) async fn load_wl(secrets_dir: &Path) -> WhitelistFile {
let path = secrets_dir.join("telegram_whitelist.json");
match tokio::fs::read_to_string(&path).await {
Ok(s) => serde_json::from_str(&s).unwrap_or_default(),
Err(_) => WhitelistFile::default(),
}
}
pub(crate) async fn save_wl(secrets_dir: &Path, wl: &WhitelistFile) -> Result<()> {
tokio::fs::create_dir_all(secrets_dir).await?;
let path = secrets_dir.join("telegram_whitelist.json");
tokio::fs::write(&path, serde_json::to_string_pretty(wl)?).await?;
Ok(())
}
// ── Pairing ───────────────────────────────────────────────────────────────────
/// Pairing codes older than this are considered abandoned and pruned, so the
/// whitelist file does not accumulate stale `pending_pairings` entries.
const PAIRING_TTL_HOURS: i64 = 24;
pub(crate) async fn handle_pairing(bot: &Bot, chat_id: ChatId, shared: &Arc<TgShared>) {
let mut wl = load_wl(&shared.secrets_dir).await;
// Drop pairing codes past their TTL. Entries with an unparseable timestamp
// are kept (don't silently lose data on a format change).
let cutoff = Utc::now() - chrono::Duration::hours(PAIRING_TTL_HOURS);
let before = wl.pending_pairings.len();
wl.pending_pairings.retain(|e| match DateTime::parse_from_rfc3339(&e.issued_at) {
Ok(ts) => ts.with_timezone(&Utc) > cutoff,
Err(_) => true,
});
let pruned = wl.pending_pairings.len() != before;
// Re-use an existing (non-expired) code if one is already pending for this chat.
let (code, added) = if let Some(entry) = wl.pending_pairings.iter().find(|e| e.chat_id == chat_id.0) {
(entry.code.clone(), false)
} else {
let code = generate_code();
wl.pending_pairings.push(PairingEntry {
code: code.clone(),
chat_id: chat_id.0,
issued_at: Local::now().format("%Y-%m-%dT%H:%M:%S%:z").to_string(),
});
(code, true)
};
// Persist if we added a new code or pruned expired ones.
if added || pruned {
if let Err(e) = save_wl(&shared.secrets_dir, &wl).await {
error!(error = %e, "telegram: failed to write whitelist file");
}
}
if added {
info!(chat_id = chat_id.0, code = %code, "TELEGRAM PAIRING: code written to telegram_whitelist.json");
}
bot.send_message(
chat_id,
format!(
"🔐 <b>Pairing required.</b>\n\n\
Code: <code>{code}</code>\n\n\
Provide this code to the web agent to authorize access.",
),
)
.parse_mode(ParseMode::Html)
.await
.ok();
}
pub(crate) fn generate_code() -> String {
const CHARS: &[u8] = b"ABCDEFGHJKLMNPQRSTUVWXYZ23456789";
let mut rng = rand::rng();
(0..6).map(|_| CHARS[rng.random_range(0..CHARS.len())] as char).collect()
}
// ── Whitelist watchdog ────────────────────────────────────────────────────────
//
// Polls telegram_whitelist.json every 10 s for mtime changes.
// When a new chat_id appears in `whitelist` (agent moved it from pending),
// sends a welcome message so the user knows they are authorized.
pub(crate) async fn whitelist_watchdog(bot: Bot, secrets_dir: PathBuf, cancel: CancellationToken) {
let path = secrets_dir.join("telegram_whitelist.json");
let mut last_mtime: Option<SystemTime> = tokio::fs::metadata(&path).await.ok()
.and_then(|m| m.modified().ok());
let mut known_wl = load_wl(&secrets_dir).await.whitelist;
let mut interval = tokio::time::interval(Duration::from_secs(10));
interval.tick().await; // skip the immediate first tick
loop {
tokio::select! {
_ = cancel.cancelled() => break,
_ = interval.tick() => {
let new_mtime = tokio::fs::metadata(&path).await.ok()
.and_then(|m| m.modified().ok());
if new_mtime.is_none() || new_mtime == last_mtime {
continue;
}
last_mtime = new_mtime;
let wl = load_wl(&secrets_dir).await;
let newly_authorized: Vec<i64> = wl.whitelist.iter()
.filter(|id| !known_wl.contains(id))
.cloned()
.collect();
if !newly_authorized.is_empty() {
info!(users = ?newly_authorized, "telegram: new users authorized — sending welcome");
for &chat_id in &newly_authorized {
bot.send_message(
ChatId(chat_id),
"✅ <b>Access granted!</b>\n\
You can now talk to your agent.\n\n\
/help for available commands.",
)
.parse_mode(ParseMode::Html)
.await
.ok();
}
}
known_wl = wl.whitelist;
info!(
whitelist = known_wl.len(),
pending = wl.pending_pairings.len(),
"telegram: whitelist file reloaded"
);
}
}
}
}
+405
View File
@@ -0,0 +1,405 @@
use std::sync::Arc;
use teloxide::prelude::*;
use teloxide::types::{InlineKeyboardButton, InlineKeyboardMarkup, ParseMode};
use tokio::sync::broadcast;
use tokio_util::sync::CancellationToken;
use tracing::{error, info, warn};
use core_api::events::{GlobalEvent, ServerEvent};
use super::TgShared;
use super::auth::load_wl;
use super::helpers::{escape_html, label_to_html, send_long};
/// Sends an inline keyboard for an approval request and records the request_id.
async fn send_approval_keyboard(
bot: &Bot,
chat_id: ChatId,
text: String,
request_id: i64,
shared: &Arc<TgShared>,
) {
let keyboard = InlineKeyboardMarkup::new(vec![
vec![
InlineKeyboardButton::callback("✅ Approve", format!("approve:{request_id}")),
InlineKeyboardButton::callback("❌ Reject", format!("reject:{request_id}")),
],
vec![
InlineKeyboardButton::callback("⏱ 15 min", format!("bypass_time:900:{request_id}")),
InlineKeyboardButton::callback("🔄 Session", format!("bypass_session:{request_id}")),
],
]);
match bot
.send_message(chat_id, text)
.parse_mode(ParseMode::Html)
.reply_markup(keyboard)
.await
{
Ok(m) => { shared.pending_approvals.lock().await.insert(m.id, request_id); }
Err(e) => error!(error = %e, "telegram: failed to send approval message"),
}
}
// ── Persistent background forwarder ──────────────────────────────────────────
/// Spawned once when the plugin starts.
/// Stays subscribed to the "telegram" broadcast channel forever, forwarding
/// events to the home chat_id. This is the only subscriber — per-message
/// subscriptions are not used — so it also catches background notifications
/// that arrive without a user message triggering them.
///
/// Re-subscribes immediately after each `Done`/`Error` so no events from the
/// next turn are missed. Safe because Tokio's cooperative scheduler guarantees
/// no other task runs between the re-subscription point and the next `await`,
/// and the processing mutex in `ChatSessionHandler` serialises turns.
pub(crate) async fn persistent_forwarder(
bot: Bot,
shared: Arc<TgShared>,
cancel: CancellationToken,
) {
info!("telegram: persistent forwarder started");
let mut rx = shared.chat_hub.events("telegram");
// Single loop: rx is updated in-place on Done/Error so we never miss events
// from the next turn (re-subscription happens before the async send).
loop {
let ge: GlobalEvent = tokio::select! {
_ = cancel.cancelled() => {
info!("telegram: persistent forwarder stopped");
return;
}
result = rx.recv() => match result {
Ok(e) => e,
Err(broadcast::error::RecvError::Lagged(n)) => {
warn!(skipped = n, "telegram: persistent forwarder lagged");
continue;
}
Err(broadcast::error::RecvError::Closed) => return,
},
};
// ApprovalResolved is handled regardless of source so Telegram removes its
// keyboard even when the approval was resolved via web or REST.
if let ServerEvent::ApprovalResolved { request_id, approved, .. } = ge.event {
let label = if approved { "✅ Approved" } else { "❌ Rejected" };
let mut pending = shared.pending_approvals.lock().await;
if let Some((&msg_id, _)) = pending.iter().find(|(_, rid)| **rid == request_id) {
let msg_id = msg_id;
pending.remove(&msg_id);
drop(pending);
if let Some(cid) = resolve_chat_id(&shared).await {
bot.delete_message(cid, msg_id).await.ok();
}
}
continue;
}
// All other events: only process if they belong to the "telegram" source.
if ge.source.as_deref() != Some("telegram") {
tracing::debug!(event_type = ge.event.type_name(), source = ?ge.source, "persistent_forwarder: skipping non-telegram event");
continue;
}
let event = ge.event;
tracing::debug!(event_type = event.type_name(), "persistent_forwarder: processing telegram event");
// Resolve the destination chat_id (last known user, or first in whitelist).
// For terminal events (Done/Error) with no known chat, still re-subscribe.
let chat_id = match resolve_chat_id(&shared).await {
Some(id) => id,
None => {
warn!(event_type = %event.type_name(), "telegram: persistent_forwarder — no chat_id resolved, dropping event");
if matches!(event, ServerEvent::Done { .. } | ServerEvent::Error { .. }) {
rx = shared.chat_hub.events("telegram");
}
continue;
}
};
match event {
ServerEvent::Done { content, .. } => {
// Re-subscribe BEFORE any await so we don't miss the next turn.
rx = shared.chat_hub.events("telegram");
if !content.trim().is_empty() {
send_long(&bot, chat_id, &content, Some(ParseMode::Html)).await;
}
}
ServerEvent::Error { message } => {
rx = shared.chat_hub.events("telegram");
bot.send_message(
chat_id,
format!("⚠️ <b>Error:</b> {}", escape_html(&message)),
)
.parse_mode(ParseMode::Html)
.await
.ok();
}
ServerEvent::ToolStart { label_short, .. } => {
bot.send_message(chat_id, format!("🔧 <i>{}</i>…", label_to_html(&label_short)))
.parse_mode(ParseMode::Html)
.await
.ok();
}
ServerEvent::Thinking { content, .. } => {
if !content.trim().is_empty() {
send_long(&bot, chat_id, &content, Some(ParseMode::Html)).await;
}
}
ServerEvent::AgentStart { agent_id, parent_agent_id, prompt_preview, .. } => {
let preview = prompt_preview.chars().take(300).collect::<String>();
let ellipsis = if prompt_preview.len() > 300 { "" } else { "" };
bot.send_message(
chat_id,
format!(
"🤖 <b>{}</b> → <b>{}</b>\n<blockquote>{}{ellipsis}</blockquote>",
escape_html(&parent_agent_id),
escape_html(&agent_id),
escape_html(&preview),
),
)
.parse_mode(ParseMode::Html)
.await
.ok();
}
ServerEvent::AgentDone { agent_id, parent_agent_id, result_preview, .. } => {
let preview = result_preview.chars().take(300).collect::<String>();
let ellipsis = if result_preview.len() > 300 { "" } else { "" };
bot.send_message(
chat_id,
format!(
"✅ <b>{}</b> finished → <b>{}</b>\n<blockquote>{}{ellipsis}</blockquote>",
escape_html(&agent_id),
escape_html(&parent_agent_id),
escape_html(&preview),
),
)
.parse_mode(ParseMode::Html)
.await
.ok();
}
ServerEvent::PendingWrite { request_id, path, new_content, .. } => {
let preview = truncate_chars(&new_content, 800);
let text = format!(
"🔐 <b>Approval required</b>\n\
<b>Operation:</b> <code>{}</code>\n\n\
<b>Content:</b>\n<pre>{}</pre>",
escape_html(&path),
escape_html(&preview),
);
send_approval_keyboard(&bot, chat_id, text, request_id, &shared).await;
}
ServerEvent::ApprovalRequired { request_id, tool_name, arguments, .. } => {
let args_str = serde_json::to_string_pretty(&arguments)
.unwrap_or_else(|_| arguments.to_string());
let args_preview = truncate_chars(&args_str, 600);
let text = format!(
"🔐 <b>Approval required</b>\n\
<b>Tool:</b> <code>{}</code>\n\n\
<b>Arguments:</b>\n<pre>{}</pre>",
escape_html(&tool_name),
escape_html(&args_preview),
);
send_approval_keyboard(&bot, chat_id, text, request_id, &shared).await;
}
ServerEvent::AgentQuestion { request_id, tool_call_id, title, question, suggested_answers, .. } => {
info!(request_id, tool_call_id, %question, "telegram: persistent_forwarder received AgentQuestion");
// If a previous question is still pending, disable its (now-dead)
// buttons so tapping them doesn't silently no-op.
if let Some(prev) = shared.pending_question.lock().await.take() {
bot.edit_message_reply_markup(chat_id, prev.message_id)
.reply_markup(InlineKeyboardMarkup::new(vec![vec![
InlineKeyboardButton::callback("⏭ Superseded by a newer question", "noop"),
]]))
.await
.ok();
}
let header = format!("❓ <b>{}</b>\n{}", escape_html(&title), escape_html(&question));
let keyboard = if suggested_answers.is_empty() {
None
} else {
let buttons: Vec<Vec<InlineKeyboardButton>> = suggested_answers
.iter()
.enumerate()
.map(|(i, s)| vec![InlineKeyboardButton::callback(
s.clone(),
format!("ansidx:{request_id}:{i}"),
)])
.collect();
Some(InlineKeyboardMarkup::new(buttons))
};
let mut req = bot.send_message(chat_id, header).parse_mode(ParseMode::Html);
if let Some(kb) = keyboard {
req = req.reply_markup(kb);
}
match req.await {
Ok(m) => {
info!(request_id, msg_id = m.id.0, "telegram: AgentQuestion sent to user, pending_question set");
*shared.pending_question.lock().await = Some(super::PendingQuestion {
request_id,
message_id: m.id,
suggested_answers,
});
}
Err(e) => error!(error = %e, request_id, "telegram: failed to send AgentQuestion to user"),
}
}
ServerEvent::LlmFailed { tried, last_error } => {
let models = tried.join(", ");
bot.send_message(
chat_id,
format!(
"⚠️ <b>LLM unavailable</b>\nTried: <code>{}</code>\n{}",
escape_html(&models),
escape_html(&last_error),
),
)
.parse_mode(ParseMode::Html)
.await
.ok();
}
// ToolDone, ToolError, FileChanged, Truncated, ModelFallback,
// NewSession, ApprovalResolved (handled above) — silenced.
_ => {}
}
}
}
/// Resolves the Telegram chat_id to use for outbound messages.
/// Prefers the last chat_id that sent a message; falls back to the first
/// whitelisted user.
async fn resolve_chat_id(shared: &TgShared) -> Option<ChatId> {
if let Some(id) = *shared.home_chat_id.lock().await {
return Some(id);
}
let wl = load_wl(&shared.secrets_dir).await;
wl.whitelist.first().map(|&id| ChatId(id))
}
// ── Helpers ───────────────────────────────────────────────────────────────────
/// Truncates `s` to at most `max_chars` Unicode scalar values.
/// Appends `…` if truncated. Never panics on multibyte UTF-8 content.
fn truncate_chars(s: &str, max_chars: usize) -> String {
let mut chars = s.chars();
let truncated: String = chars.by_ref().take(max_chars).collect();
if chars.next().is_some() {
format!("{truncated}")
} else {
truncated
}
}
// ── Callback query handler (button presses) ───────────────────────────────────
pub(crate) async fn callback_handler(
bot: Bot,
q: CallbackQuery,
shared: Arc<TgShared>,
) -> ResponseResult<()> {
let approval_msg = q
.message
.as_ref()
.and_then(|m| m.regular_message())
.map(|m| (m.chat.id, m.id));
let Some((msg_chat_id, msg_id)) = approval_msg else {
bot.answer_callback_query(q.id.clone()).await.ok();
return Ok(());
};
let Some(data) = q.data.as_deref() else {
bot.answer_callback_query(q.id.clone()).await.ok();
return Ok(());
};
// ── Suggested-answer button (ask_user_clarification) ─────────────────────
if let Some(rest) = data.strip_prefix("ansidx:") {
let mut parts = rest.splitn(2, ':');
let req_id = parts.next().and_then(|s| s.parse::<i64>().ok());
let idx_str = parts.next().and_then(|s| s.parse::<usize>().ok());
if let (Some(request_id), Some(idx)) = (req_id, idx_str) {
let mut pq = shared.pending_question.lock().await;
if let Some(pq_inner) = pq.as_ref() {
if pq_inner.request_id == request_id {
let answer = pq_inner.suggested_answers.get(idx).cloned().unwrap_or_default();
drop(pq);
*shared.pending_question.lock().await = None;
shared.chat_hub.resolve_question("telegram", request_id, answer.clone()).await;
info!(request_id, %answer, "telegram: clarification answered via button");
bot.edit_message_reply_markup(msg_chat_id, msg_id)
.reply_markup(InlineKeyboardMarkup::new(vec![vec![
InlineKeyboardButton::callback(format!("{answer}"), "noop"),
]]))
.await
.ok();
}
}
}
bot.answer_callback_query(q.id.clone()).await.ok();
return Ok(());
}
// ── Approval buttons ──────────────────────────────────────────────────────
enum ApprovalAction {
Approve,
Reject,
BypassTime(u64),
BypassSession,
}
let parsed: Option<(i64, ApprovalAction, &str)> =
if let Some(id_str) = data.strip_prefix("approve:") {
id_str.parse::<i64>().ok().map(|id| (id, ApprovalAction::Approve, "✅ Approved"))
} else if let Some(id_str) = data.strip_prefix("reject:") {
id_str.parse::<i64>().ok().map(|id| (id, ApprovalAction::Reject, "❌ Rejected"))
} else if let Some(rest) = data.strip_prefix("bypass_time:") {
let mut parts = rest.splitn(2, ':');
let secs = parts.next().and_then(|s| s.parse::<u64>().ok());
let id = parts.next().and_then(|s| s.parse::<i64>().ok());
secs.zip(id).map(|(s, id)| (id, ApprovalAction::BypassTime(s), "⏱ Bypass (timed)"))
} else if let Some(id_str) = data.strip_prefix("bypass_session:") {
id_str.parse::<i64>().ok().map(|id| (id, ApprovalAction::BypassSession, "🔄 Bypass (session)"))
} else {
None
};
if let Some((request_id, action, label)) = parsed {
let stored = shared.pending_approvals.lock().await.remove(&msg_id);
if let Some(stored_id) = stored {
if stored_id == request_id {
match action {
ApprovalAction::Approve =>
shared.approval.approve(request_id).await,
ApprovalAction::Reject =>
shared.approval.reject(request_id, String::new()).await,
ApprovalAction::BypassTime(secs) =>
shared.approval.approve_with_bypass(request_id, Some(secs)).await,
ApprovalAction::BypassSession =>
shared.approval.approve_with_bypass(request_id, None).await,
}
info!(request_id, label, "telegram: approval resolved");
bot.delete_message(msg_chat_id, msg_id).await.ok();
}
} else {
warn!(request_id, "telegram: approval not found (already resolved?)");
}
}
bot.answer_callback_query(q.id.clone()).await.ok();
Ok(())
}
+541
View File
@@ -0,0 +1,541 @@
use std::sync::Arc;
use teloxide::prelude::*;
use teloxide::types::{ChatAction, ParseMode};
use tracing::{error, info};
use core_api::chat_hub::{ModelCommandOutcome, SendMessageOptions};
use core_api::command::expand_template;
use core_api::location::GpsCoord;
use core_api::message_meta::{CommandRef, MessageMetadata};
use super::TELEGRAM_FORMAT_CONTEXT;
use super::TgShared;
use super::attachments::TelegramAttachment;
use super::auth::{handle_pairing, load_wl};
// ── Available commands help text (shared by /help and unknown-command replies) ──
const HELP_TEXT: &str = "<b>Available commands</b>\n\n\
/clear — start a new conversation\n\
/new — alias for /clear\n\
/stop — interrupt the agent mid-turn\n\
/models — list available LLM models, ordered by priority\n\
/model &lt;N|name|auto&gt; — select the model for this chat\n\
/context — show last turn's token usage\n\
/cost — show total spend for this session (USD)\n\
/compact — force context compaction\n\
/resettools — remove all activated tool groups (MCP + config) from the session\n\
/sethome — receive agent notifications here\n\
/help — this message";
/// Builds the `/help` text: the static system-command list plus a dynamically
/// discovered "Custom commands" section (`commands/<name>/`). Descriptions are
/// HTML-escaped since the message is sent with `ParseMode::Html`.
fn help_text(command: &dyn core_api::command::CommandApi) -> String {
let mut out = String::from(HELP_TEXT);
let cmds = command.list_enabled();
if !cmds.is_empty() {
out.push_str("\n\n<b>Custom commands</b>");
for c in cmds {
out.push_str(&format!(
"\n/{}{}",
c.name,
super::helpers::escape_html(&c.description)
));
}
}
out
}
// ── Incoming message classification ───────────────────────────────────────────
//
// To add a new media type: add a variant to IncomingEvent, handle it in
// classify_message, then dispatch it in message_handler.
pub(crate) enum IncomingEvent {
Text(String),
Command { name: String, args: Vec<String> },
Voice { file_id: String },
Attachment(TelegramAttachment),
}
pub(crate) fn classify_message(msg: &Message) -> Option<IncomingEvent> {
if let Some(voice) = msg.voice() {
return Some(IncomingEvent::Voice { file_id: voice.file.id.to_string() });
}
if let Some(doc) = msg.document() {
return Some(IncomingEvent::Attachment(TelegramAttachment::Document {
file_id: doc.file.id.to_string(),
file_name: doc.file_name.clone().unwrap_or_else(|| "attachment".to_string()),
mime_type: doc.mime_type.as_ref().map(|m| m.to_string()),
caption: msg.caption().map(str::to_string),
}));
}
if let Some(photos) = msg.photo() {
if let Some(largest) = photos.last() {
return Some(IncomingEvent::Attachment(TelegramAttachment::Photo {
file_id: largest.file.id.to_string(),
caption: msg.caption().map(str::to_string),
}));
}
}
if let Some(loc) = msg.location() {
return Some(IncomingEvent::Attachment(TelegramAttachment::Location {
latitude: loc.latitude,
longitude: loc.longitude,
accuracy: loc.horizontal_accuracy,
is_live: loc.live_period.is_some(),
}));
}
let text = msg.text()?;
// A command is any message that *starts* with '/'. We deliberately do NOT
// rely on teloxide's BotCommand entities: those are emitted for every
// "/token" anywhere in the text, so a normal sentence containing a "/path"
// (e.g. "stop /usr/bin/foo") would be misclassified as a command. A leading
// slash is the only signal. Arguments are parsed from the message `text`
// (not `entity.text()`, which spans only "/model" and would drop the arg).
// An unknown command is handled by the dispatcher, which replies with help.
if text.starts_with('/') {
let full = text.trim_start_matches('/');
let mut parts = full.splitn(2, ' ');
let name = parts.next().unwrap_or("").to_ascii_lowercase();
let name = name.split('@').next().unwrap_or(&name).to_string();
let rest = parts.next().unwrap_or("").trim().to_string();
let args: Vec<String> = if rest.is_empty() {
vec![]
} else {
rest.split_whitespace().map(str::to_string).collect()
};
return Some(IncomingEvent::Command { name, args });
}
Some(IncomingEvent::Text(text.to_string()))
}
// ── Message handler ───────────────────────────────────────────────────────────
pub(crate) async fn message_handler(
bot: Bot,
msg: Message,
shared: Arc<TgShared>,
) -> ResponseResult<()> {
let chat_id = msg.chat.id;
// Whitelist check — re-read the file on every message so agent edits are
// picked up without a plugin restart.
let wl = load_wl(&shared.secrets_dir).await;
if !wl.whitelist.contains(&chat_id.0) {
handle_pairing(&bot, chat_id, &shared).await;
return Ok(());
}
// Track the last active chat_id so the persistent forwarder knows
// where to send background notifications.
*shared.home_chat_id.lock().await = Some(chat_id);
let Some(incoming) = classify_message(&msg) else {
bot.send_message(chat_id, "Unsupported message format.").await.ok();
return Ok(());
};
match incoming {
IncomingEvent::Command { ref name, .. } if name == "clear" || name == "new" => {
handle_clear(&bot, chat_id, &shared).await;
}
IncomingEvent::Command { ref name, .. } if name == "sethome" => {
match shared.chat_hub.set_home("telegram").await {
Ok(_) => {
info!("telegram: set as home source");
bot.send_message(chat_id, "🏠 Telegram set as <b>home</b>. Agent notifications will be delivered here.")
.parse_mode(ParseMode::Html)
.await
.ok();
}
Err(e) => {
bot.send_message(chat_id, format!("⚠️ Error: {e}")).await.ok();
}
}
}
IncomingEvent::Command { ref name, .. } if name == "help" => {
bot.send_message(chat_id, help_text(&*shared.command))
.parse_mode(ParseMode::Html)
.await
.ok();
}
IncomingEvent::Command { ref name, .. } if name == "stop" => {
handle_stop(&bot, chat_id, &shared).await;
}
IncomingEvent::Command { ref name, .. } if name == "context" => {
handle_context(&bot, chat_id, &shared).await;
}
IncomingEvent::Command { ref name, .. } if name == "cost" => {
handle_cost(&bot, chat_id, &shared).await;
}
IncomingEvent::Command { ref name, .. } if name == "compact" => {
handle_compact(&bot, chat_id, &shared).await;
}
IncomingEvent::Command { ref name, .. } if name == "resettools" => {
handle_reset_mcp(&bot, chat_id, &shared).await;
}
IncomingEvent::Command { ref name, .. } if name == "models" => {
handle_list_models(&bot, chat_id, &shared).await;
}
IncomingEvent::Command { ref name, ref args, .. } if name == "model" => {
handle_set_model(&bot, chat_id, args, &shared).await;
}
// A recognised custom `/command` expands its `COMMAND.md` template into a
// normal user message (fully interactive: the model can then ask questions,
// iterate, dispatch sub-agents). Any other `/...` is an unknown command and
// is never forwarded to the LLM — reply with a not-found notice + help.
IncomingEvent::Command { ref name, ref args, .. } => {
if let Some(resolved) = shared.command.resolve(name) {
let args_str = args.join(" ");
let display = if args_str.is_empty() {
format!("/{name}")
} else {
format!("/{name} {args_str}")
};
let content = expand_template(&resolved.template, &args_str);
let metadata = MessageMetadata {
command: Some(CommandRef {
name: resolved.name,
display,
}),
..Default::default()
};
handle_llm_message(bot, chat_id, content, Some(metadata), shared).await;
} else {
bot.send_message(
chat_id,
format!("Unknown command: /{name}\n\n{}", help_text(&*shared.command)),
)
.parse_mode(ParseMode::Html)
.await
.ok();
}
}
IncomingEvent::Voice { file_id } => {
handle_voice(&bot, chat_id, file_id, &shared).await;
}
IncomingEvent::Attachment(attachment) => {
handle_attachment(bot, chat_id, attachment, shared).await;
}
_ => {
let text = match &incoming {
IncomingEvent::Text(t) => t.clone(),
IncomingEvent::Command { .. }
| IncomingEvent::Voice { .. }
| IncomingEvent::Attachment(_) => unreachable!(),
};
// If a clarification question is pending, treat any text as the answer.
{
let mut pq = shared.pending_question.lock().await;
if let Some(pq_inner) = pq.take() {
let request_id = pq_inner.request_id;
let question_msg_id = pq_inner.message_id;
drop(pq);
shared.chat_hub.resolve_question("telegram", request_id, text.clone()).await;
tracing::info!(request_id, %text, "telegram: clarification answered via text");
bot.edit_message_reply_markup(chat_id, question_msg_id)
.reply_markup(teloxide::types::InlineKeyboardMarkup::new(vec![vec![
teloxide::types::InlineKeyboardButton::callback(
format!("{}", super::helpers::escape_html(&text)),
"noop",
),
]]))
.await
.ok();
return Ok(());
}
}
handle_llm_message(bot, chat_id, text, None, shared).await;
}
}
Ok(())
}
// ── /clear command ────────────────────────────────────────────────────────────
async fn handle_clear(bot: &Bot, chat_id: ChatId, shared: &Arc<TgShared>) {
match shared.chat_hub.clear("telegram").await {
Ok(_) => {
info!("telegram: session cleared via /clear");
bot.send_message(chat_id, "🆕 New conversation started.").await.ok();
}
Err(e) => {
error!(error = %e, "telegram: failed to clear session");
bot.send_message(chat_id, format!("⚠️ Error: {e}")).await.ok();
}
}
}
// ── /context command ──────────────────────────────────────────────────────────
async fn handle_context(bot: &Bot, chat_id: ChatId, shared: &Arc<TgShared>) {
match shared.chat_hub.context_info("telegram").await {
Ok((input, output)) => {
let input_str = input.map_or("?".to_string(), |t| t.to_string());
let output_str = output.map_or("?".to_string(), |t| t.to_string());
bot.send_message(
chat_id,
format!("<i>↑{input_str} tok · ↓{output_str} tok</i>"),
)
.parse_mode(ParseMode::Html)
.await
.ok();
}
Err(e) => {
bot.send_message(chat_id, format!("⚠️ Error: {e}")).await.ok();
}
}
}
// ── /cost command ─────────────────────────────────────────────────────────────
async fn handle_cost(bot: &Bot, chat_id: ChatId, shared: &Arc<TgShared>) {
match shared.chat_hub.cost_info("telegram").await {
Ok(Some(c)) => {
bot.send_message(chat_id, format!("💰 Session cost: ${c:.4}")).await.ok();
}
Ok(None) => {
bot.send_message(chat_id, "💰 No cost recorded for this session.").await.ok();
}
Err(e) => {
bot.send_message(chat_id, format!("⚠️ Error: {e}")).await.ok();
}
}
}
// ── /compact command ──────────────────────────────────────────────────────────
async fn handle_compact(bot: &Bot, chat_id: ChatId, shared: &Arc<TgShared>) {
match shared.chat_hub.force_compact("telegram").await {
Ok(true) => {
info!("telegram: manual compaction succeeded");
bot.send_message(chat_id, "✅ Context compacted.").await.ok();
}
Ok(false) => {
bot.send_message(chat_id, "⏩ Compaction skipped (no messages to summarise or compaction disabled).").await.ok();
}
Err(e) => {
error!(error = %e, "telegram: manual compaction failed");
bot.send_message(chat_id, format!("⚠️ Compaction failed: {e}")).await.ok();
}
}
}
// ── /resettools command ───────────────────────────────────────────────────────
async fn handle_reset_mcp(bot: &Bot, chat_id: ChatId, shared: &Arc<TgShared>) {
match shared.chat_hub.reset_mcp("telegram").await {
Ok(()) => {
info!("telegram: tool-group grants reset via /resettools");
bot.send_message(chat_id, "✅ Activated tool groups removed from the session.").await.ok();
}
Err(e) => {
error!(error = %e, "telegram: /resettools failed");
bot.send_message(chat_id, format!("⚠️ Error: {e}")).await.ok();
}
}
}
// ── /stop command ────────────────────────────────────────────────────────────
async fn handle_stop(bot: &Bot, chat_id: ChatId, shared: &Arc<TgShared>) {
shared.chat_hub.cancel("telegram").await;
info!("telegram: agent cancelled via /stop");
bot.send_message(chat_id, "⏹ Agent stopped.").await.ok();
}
// ── /models and /model commands ──────────────────────────────────────────────
//
// Business logic (resolve arg, mutate pin, broadcast) lives in
// `ChatHub::apply_model_command` / `ChatHub::list_clients_marked`. Here we only
// format for Telegram (HTML) and send via the bot — same pattern the web WS
// handler uses with Markdown.
async fn handle_list_models(bot: &Bot, chat_id: ChatId, shared: &Arc<TgShared>) {
let items = shared.chat_hub.list_clients_marked("telegram").await;
let mut text = String::from("<b>Available models</b>\n\n");
for (i, name, is_current) in &items {
let marker = if *is_current { "" } else { "" };
text.push_str(&format!(
"{} <code>{:2}</code> {}\n",
marker,
i,
super::helpers::escape_html(name)
));
}
text.push_str("\nUse <code>/model N</code>, <code>/model name</code>, or <code>/model auto</code>.");
bot.send_message(chat_id, text)
.parse_mode(ParseMode::Html)
.await
.ok();
}
async fn handle_set_model(bot: &Bot, chat_id: ChatId, args: &[String], shared: &Arc<TgShared>) {
let arg = args.first().cloned().unwrap_or_default();
let outcome = shared.chat_hub.apply_model_command("telegram", &arg).await;
let text = match outcome {
ModelCommandOutcome::Set(name) => format!("✅ Model set: <b>{}</b>", super::helpers::escape_html(&name)),
ModelCommandOutcome::Cleared => "✅ Model reset to <b>auto</b>.".to_string(),
ModelCommandOutcome::Error(msg) => format!("⚠️ {}", super::helpers::escape_html(&msg)),
};
bot.send_message(chat_id, text)
.parse_mode(ParseMode::Html)
.await
.ok();
}
// ── LLM dispatch ─────────────────────────────────────────────────────────────
async fn handle_llm_message(
bot: Bot,
chat_id: ChatId,
text: String,
metadata: Option<MessageMetadata>,
shared: Arc<TgShared>,
) {
bot.send_chat_action(chat_id, ChatAction::Typing).await.ok();
// The persistent_forwarder (spawned once in start()) is always subscribed
// to the "telegram" broadcast channel and will pick up all events for this
// turn — including Done → send to Telegram. No per-message subscription needed.
let client_name = shared.chat_hub.get_selected_client("telegram").await;
let opts = SendMessageOptions {
client_name,
extra_system_context: Some(TELEGRAM_FORMAT_CONTEXT.to_string()),
tail_reminder: Some(super::TELEGRAM_FORMAT_REMINDER.to_string()),
interface_tools: super::tools::interface_tools(bot, chat_id, &*shared.tts).await,
metadata,
..Default::default()
};
// send_message only enqueues — the turn runs on ChatHub's per-source consumer —
// so awaiting inline keeps this message handler responsive.
if let Err(e) = shared.chat_hub.send_message("telegram", &text, opts).await {
error!(error = %e, "telegram: enqueue error");
}
}
// ── Voice message → transcribe → LLM ─────────────────────────────────────────
async fn handle_voice(
bot: &Bot,
chat_id: ChatId,
file_id: String,
shared: &Arc<TgShared>,
) {
use teloxide::net::Download;
let transcriber = match shared.transcriber().await {
Some(t) => t,
None => {
bot.send_message(chat_id, "⚠️ Transcription not available (no transcription provider configured).").await.ok();
return;
}
};
let file = match bot.get_file(teloxide::types::FileId(file_id)).await {
Ok(f) => f,
Err(e) => {
error!(error = %e, "telegram: get_file failed");
bot.send_message(chat_id, "⚠️ Could not download audio file.").await.ok();
return;
}
};
let mut audio_bytes = Vec::new();
if let Err(e) = bot.download_file(&file.path, &mut audio_bytes).await {
error!(error = %e, "telegram: download_file failed");
bot.send_message(chat_id, "⚠️ Audio download failed.").await.ok();
return;
}
bot.send_chat_action(chat_id, ChatAction::Typing).await.ok();
let text = match transcriber.transcribe(audio_bytes, "ogg").await {
Ok(t) => t,
Err(e) => {
error!(error = %e, "telegram: transcription failed");
bot.send_message(chat_id, format!("⚠️ Transcription failed: {e}")).await.ok();
return;
}
};
info!(chat_id = chat_id.0, "telegram: voice transcribed, forwarding to LLM");
let message = format!(
"[TELEGRAM SYSTEM INFO]\n\
The user sent a voice message. The following is the audio transcript:\n\n\
{text}"
);
handle_llm_message(bot.clone(), chat_id, message, None, Arc::clone(shared)).await;
}
// ── Edited message (live location updates) ────────────────────────────────────
pub(crate) async fn edited_message_handler(
msg: Message,
shared: Arc<TgShared>,
) -> ResponseResult<()> {
if let Some(loc) = msg.location() {
let coord = GpsCoord { latitude: loc.latitude, longitude: loc.longitude };
shared.location.update("telegram", coord, loc.horizontal_accuracy, true);
}
Ok(())
}
// ── File / media attachment ───────────────────────────────────────────────────
async fn handle_attachment(
bot: Bot,
chat_id: ChatId,
attachment: TelegramAttachment,
shared: Arc<TgShared>,
) {
// Update LocationManager immediately, before any LLM dispatch.
if let TelegramAttachment::Location { latitude, longitude, accuracy, is_live } = &attachment {
let coord = GpsCoord { latitude: *latitude, longitude: *longitude };
shared.location.update("telegram", coord, *accuracy, *is_live);
}
bot.send_chat_action(chat_id, ChatAction::UploadDocument).await.ok();
let saved = match attachment.download_and_save(&bot, &shared.uploads_dir, chat_id.0).await {
Ok(s) => s,
Err(e) => {
error!(error = %e, "telegram: failed to save attachment");
bot.send_message(chat_id, "⚠️ Could not save the attachment.").await.ok();
return;
}
};
match saved {
// Document / Photo: carry the file as structured metadata (rendered as a
// chip in the copilot UI; the LLM gets the shared [SYSTEM INFO] block).
// The caption, if any, becomes the user's text for this turn.
Some(att) => {
info!(chat_id = chat_id.0, path = %att.path, "telegram: attachment saved, forwarding to LLM");
let caption = match &attachment {
TelegramAttachment::Document { caption, .. } => caption.clone(),
TelegramAttachment::Photo { caption, .. } => caption.clone(),
TelegramAttachment::Location { .. } => None,
}.unwrap_or_default();
let metadata = MessageMetadata { attachments: vec![att], ..Default::default() };
handle_llm_message(bot, chat_id, caption, Some(metadata), shared).await;
}
// Location (no file): keep the textual system-info block.
None => {
let message = attachment.system_info_message(None);
handle_llm_message(bot, chat_id, message, None, shared).await;
}
}
}
+184
View File
@@ -0,0 +1,184 @@
use regex::Regex;
use std::sync::OnceLock;
use teloxide::prelude::*;
use teloxide::types::ParseMode;
pub(crate) fn escape_html(s: &str) -> String {
s.replace('&', "&amp;").replace('<', "&lt;").replace('>', "&gt;")
}
/// Strip HTML tags for plain-text fallback.
fn strip_html(s: &str) -> String {
static RE: OnceLock<Regex> = OnceLock::new();
let re = RE.get_or_init(|| Regex::new(r"<[^>]+>").unwrap());
re.replace_all(s, "").into_owned()
}
// ── Markdown → Telegram HTML sanitizer ───────────────────────────────────────
/// Convert a Markdown table block (slice of raw lines) into bullet list lines.
fn table_to_bullets(rows: &[&str]) -> String {
let mut out = String::new();
let mut header_seen = false;
for &line in rows {
let trimmed = line.trim();
// Separator row (e.g. |---|---|) → skip
if trimmed.starts_with('|')
&& trimmed.chars().all(|c| matches!(c, '|' | '-' | ':' | ' '))
{
header_seen = true; // next rows are data
continue;
}
// Table row: split on '|', drop empty outer segments
if trimmed.starts_with('|') && trimmed.ends_with('|') {
let cells: Vec<&str> = trimmed
.trim_matches('|')
.split('|')
.map(str::trim)
.filter(|c| !c.is_empty())
.collect();
if cells.is_empty() { continue; }
// First row before separator = header → emit as bold label, not bullet
if !header_seen {
out.push_str(&format!("<b>{}</b>\n", cells.join("")));
} else {
out.push_str(&format!("{}\n", cells.join("")));
}
}
}
out
}
/// Safety-net: convert residual Markdown bold (**text**) to <b>text</b>.
fn md_bold_to_html(text: &str) -> String {
static RE: OnceLock<Regex> = OnceLock::new();
let re = RE.get_or_init(|| Regex::new(r"\*\*(.+?)\*\*").unwrap());
re.replace_all(text, "<b>$1</b>").into_owned()
}
/// Safety-net: convert residual Markdown headers (## text) to <b>text</b>.
fn md_headers_to_html(text: &str) -> String {
static RE: OnceLock<Regex> = OnceLock::new();
let re = RE.get_or_init(|| Regex::new(r"(?m)^#{1,6} +(.+)$").unwrap());
re.replace_all(text, "<b>$1</b>").into_owned()
}
/// Safety-net: convert residual inline `` `code` `` to <code>code</code>,
/// HTML-escaping the inner text so it renders verbatim.
fn md_inline_code_to_html(text: &str) -> String {
static RE: OnceLock<Regex> = OnceLock::new();
let re = RE.get_or_init(|| Regex::new(r"`([^`\n]+?)`").unwrap());
re.replace_all(text, |caps: &regex::Captures| {
format!("<code>{}</code>", escape_html(&caps[1]))
})
.into_owned()
}
/// Sanitize LLM output for Telegram HTML rendering:
/// 1. Convert fenced code blocks (```) → `<pre>…</pre>` (inner text escaped).
/// 2. Convert Markdown tables → bullet lists (Telegram has no `<table>` support).
/// 3. Convert residual `**bold**` → `<b>bold</b>`.
/// 4. Convert residual `## headers` → `<b>text</b>`.
/// 5. Convert residual inline `` `code` `` → `<code>code</code>`.
fn sanitize_for_telegram(text: &str) -> String {
// Pass 1: block conversion (line-by-line state machine for fences & tables)
let lines: Vec<&str> = text.lines().collect();
let mut out = String::with_capacity(text.len());
let mut i = 0;
while i < lines.len() {
let line = lines[i];
let trimmed = line.trim();
// Fenced code block: ``` … ``` → <pre>escaped</pre>
if trimmed.starts_with("```") {
i += 1; // skip the opening fence (and any language tag)
let start = i;
while i < lines.len() && !lines[i].trim().starts_with("```") {
i += 1;
}
let inner = lines[start..i].join("\n");
if i < lines.len() { i += 1; } // skip the closing fence
out.push_str("<pre>");
out.push_str(&escape_html(&inner));
out.push_str("</pre>\n");
continue;
}
// Markdown table block
if trimmed.starts_with('|') && trimmed.ends_with('|') && trimmed.len() > 1 {
let start = i;
while i < lines.len() && {
let t = lines[i].trim();
(t.starts_with('|') && t.ends_with('|') && t.len() > 1)
|| (t.starts_with('|')
&& t.chars().all(|c| matches!(c, '|' | '-' | ':' | ' ')))
} {
i += 1;
}
out.push_str(&table_to_bullets(&lines[start..i]));
} else {
out.push_str(line);
out.push('\n');
i += 1;
}
}
// Passes 2-4: residual Markdown
let out = md_bold_to_html(&out);
let out = md_headers_to_html(&out);
md_inline_code_to_html(&out)
}
/// Convert a tool label (our internal format with backtick-wrapped args) to
/// Telegram HTML: plain text is HTML-escaped, `` `code` `` becomes `<code>code</code>`.
pub(crate) fn label_to_html(s: &str) -> String {
let mut out = String::with_capacity(s.len() + 16);
let mut rest = s;
while let Some(open) = rest.find('`') {
out.push_str(&escape_html(&rest[..open]));
rest = &rest[open + 1..];
if let Some(close) = rest.find('`') {
out.push_str("<code>");
out.push_str(&escape_html(&rest[..close]));
out.push_str("</code>");
rest = &rest[close + 1..];
} else {
// Unmatched backtick — emit as-is and stop.
out.push('`');
break;
}
}
out.push_str(&escape_html(rest));
out
}
// ── send_long ─────────────────────────────────────────────────────────────────
pub(crate) async fn send_long(bot: &Bot, chat_id: ChatId, text: &str, parse_mode: Option<ParseMode>) {
const MAX: usize = 4000;
if text.is_empty() { return; }
// Sanitize before chunking so table blocks are never split mid-row.
let sanitized;
let text = if parse_mode == Some(ParseMode::Html) {
sanitized = sanitize_for_telegram(text);
&sanitized
} else {
text
};
let chars: Vec<char> = text.chars().collect();
let mut start = 0;
while start < chars.len() {
let end = (start + MAX).min(chars.len());
let chunk: String = chars[start..end].iter().collect();
let mut req = bot.send_message(chat_id, &chunk);
if let Some(pm) = parse_mode { req = req.parse_mode(pm); }
if req.await.is_err() {
// Retry without parse_mode so the text reaches the user even if
// the markup was malformed. Strip HTML tags first so we don't
// display raw `<b>…</b>` to the user.
let plain = strip_html(&chunk);
bot.send_message(chat_id, plain).await.ok();
}
start = end;
}
}
+276
View File
@@ -0,0 +1,276 @@
/// Telegram plugin — connects the Skald LLM to a private Telegram bot.
///
/// # Pairing
/// Unknown users receive a pairing code in chat. The code is also written to
/// `secrets/telegram_whitelist.json` under `pending_pairings`. The main agent
/// (via `read_file` / `write_file`) can inspect that file and move the
/// `chat_id` into the `whitelist` array to complete the authorisation — no
/// code changes required, just a file edit.
///
/// # Human-in-the-loop approvals
/// Tool calls requiring approval emit a `PendingWrite` event; the plugin
/// forwards it to Telegram as an inline-keyboard message with
/// [✅ Approve] [❌ Reject] / [⏱ 15 min] [🔄 Session] buttons.
///
/// # Adding new message types
/// 1. Add a variant to `IncomingEvent` in `handlers.rs`.
/// 2. Handle it in `classify_message` (same file).
/// 3. Dispatch it in `message_handler` (same file).
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use anyhow::Result;
use async_trait::async_trait;
use serde_json::{Value, json};
use teloxide::prelude::*;
use teloxide::types::MessageId;
use tokio::sync::Mutex;
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;
use tracing::{info, warn};
use core_api::approval::ApprovalApi;
use core_api::chat_hub::ChatHubApi;
use core_api::command::CommandApi;
use core_api::location::LocationUpdater;
use core_api::plugin::{Plugin, PluginContext};
use core_api::transcribe::{Transcribe, TranscribeProvider};
use core_api::tts::TtsProvider;
mod attachments;
mod auth;
mod events;
mod handlers;
mod helpers;
mod tools;
/// Injected as extra system context for every Telegram turn.
/// Kept compact to minimise token overhead.
pub(crate) const TELEGRAM_FORMAT_CONTEXT: &str = "\
OUTPUT FORMAT — TELEGRAM HTML ONLY.\n\
Allowed tags: <b> <i> <u> <s> <code> <pre> <a> <blockquote>. \
Telegram supports NO other HTML and NO Markdown.\n\
FORBIDDEN (will appear as raw symbols): ** * _ ` # | and Markdown tables.\n\
• Headers → <b>text</b>\n\
• Structured data → bullet lists with •, never | tables\n\
• Escape & < > as &amp; &lt; &gt;";
/// Short reminder injected near the end of the message list to counter
/// instruction drift in long conversations.
pub(crate) const TELEGRAM_FORMAT_REMINDER: &str = "\
[FORMAT] Telegram HTML only: <b> <i> <code> <pre>. \
No Markdown: no ** * _ ` # |. No tables — use bullet lists.";
// ── Shared state injected into every teloxide handler ─────────────────────────
/// A pending `ask_user_clarification` question waiting for the user's reply.
pub(crate) struct PendingQuestion {
pub(crate) request_id: i64,
pub(crate) message_id: MessageId,
/// Suggested answers (used to resolve the selection when the user taps a button).
pub(crate) suggested_answers: Vec<String>,
}
pub(crate) struct TgShared {
pub(crate) chat_hub: Arc<dyn ChatHubApi>,
/// Custom slash-command resolver (`commands/<name>/`). Read-only: lets the
/// bot expand a recognised `/command` into a template before forwarding it to
/// the LLM, mirroring the WS handler.
pub(crate) command: Arc<dyn CommandApi>,
pub(crate) approval: Arc<dyn ApprovalApi>,
pub(crate) transcribe: Arc<dyn TranscribeProvider>,
pub(crate) tts: Arc<dyn TtsProvider>,
pub(crate) location: Arc<dyn LocationUpdater>,
/// MessageId of the approval message → request_id.
pub(crate) pending_approvals: Mutex<HashMap<MessageId, i64>>,
/// Currently active clarification question (at most one at a time per session).
pub(crate) pending_question: Mutex<Option<PendingQuestion>>,
pub(crate) secrets_dir: PathBuf,
/// Base directory for file attachments: `<data_root>/uploads/telegram/`.
pub(crate) uploads_dir: PathBuf,
/// Last chat_id that sent a message — used as the target for background notifications.
/// Set on every incoming message; read by the persistent event forwarder.
pub(crate) home_chat_id: Mutex<Option<ChatId>>,
}
impl TgShared {
pub(crate) async fn transcriber(&self) -> Option<Arc<dyn Transcribe>> {
self.transcribe.get().await
}
}
// ── Plugin struct ─────────────────────────────────────────────────────────────
pub struct TelegramPlugin {
secrets_dir: PathBuf,
/// Bot token — set by reload() before start() is called.
token: Mutex<String>,
running: Arc<AtomicBool>,
cancel: Mutex<Option<CancellationToken>>,
handle: Mutex<Option<JoinHandle<()>>>,
}
impl TelegramPlugin {
pub fn new(secrets_dir: impl Into<PathBuf>) -> Self {
Self {
secrets_dir: secrets_dir.into(),
token: Mutex::new(String::new()),
running: Arc::new(AtomicBool::new(false)),
cancel: Mutex::new(None),
handle: Mutex::new(None),
}
}
}
#[async_trait]
impl Plugin for TelegramPlugin {
fn id(&self) -> &str { "telegram" }
fn name(&self) -> &str { "Telegram Bot" }
fn description(&self) -> &str {
"Private Telegram bot. Forwards messages to the LLM; supports HITL approval via inline keyboards."
}
fn is_running(&self) -> bool { self.running.load(Ordering::Relaxed) }
fn config_schema(&self) -> Value {
json!({
"type": "object",
"properties": {
"token": {
"type": "string",
"title": "Bot Token",
"description": "Telegram bot token from @BotFather",
"sensitive": true
}
},
"required": ["token"]
})
}
fn as_any(&self) -> &dyn std::any::Any { self }
fn as_arc_any(self: Arc<Self>) -> Arc<dyn std::any::Any + Send + Sync> { self }
async fn reload(&self, enabled: bool, config: Value, ctx: PluginContext) -> Result<()> {
let new_token = config["token"].as_str().unwrap_or("").to_string();
let old_token = self.token.lock().await.clone();
let is_running = self.is_running();
match (enabled, is_running) {
(true, false) => {
anyhow::ensure!(!new_token.is_empty(),
"telegram: cannot start — `token` is missing from config");
*self.token.lock().await = new_token;
self.start(ctx).await?;
}
(false, true) => {
self.stop().await?;
}
(true, true) => {
if new_token != old_token {
info!("telegram: token changed — restarting");
self.stop().await?;
*self.token.lock().await = new_token;
self.start(ctx).await?;
}
}
(false, false) => {}
}
Ok(())
}
async fn start(&self, ctx: PluginContext) -> Result<()> {
if self.running.load(Ordering::Relaxed) {
return Ok(());
}
let token = self.token.lock().await.clone();
if token.is_empty() {
anyhow::bail!("telegram: token is empty — set it via the plugins API");
}
let uploads_dir = self.secrets_dir
.parent()
.unwrap_or(std::path::Path::new("."))
.join("uploads")
.join("telegram");
// Register "telegram" source with ChatHub (idempotent).
// ChatHub restores the active session from the sources table automatically.
ctx.chat_hub.register("telegram").await;
info!("telegram: registered with ChatHub");
let shared = Arc::new(TgShared {
chat_hub: Arc::clone(&ctx.chat_hub),
command: Arc::clone(&ctx.command),
approval: Arc::clone(&ctx.approval),
transcribe: Arc::clone(&ctx.transcribe),
tts: Arc::clone(&ctx.tts_provider),
location: Arc::clone(&ctx.location),
pending_approvals: Mutex::new(HashMap::new()),
pending_question: Mutex::new(None),
secrets_dir: self.secrets_dir.clone(),
uploads_dir,
home_chat_id: Mutex::new(None),
});
let bot = Bot::new(&token);
let cancel = CancellationToken::new();
tokio::spawn(events::persistent_forwarder(
bot.clone(),
Arc::clone(&shared),
cancel.clone(),
));
let hub_clone = Arc::clone(&ctx.chat_hub);
tokio::spawn(async move {
if let Err(e) = hub_clone.resume("telegram").await {
tracing::warn!(error = %e, "telegram: startup resume failed");
}
});
let cancel_clone = cancel.clone();
let cancel_wdg = cancel.clone();
let running_clone = Arc::clone(&self.running);
self.running.store(true, Ordering::Relaxed);
let handler = dptree::entry()
.branch(Update::filter_message().endpoint(handlers::message_handler))
.branch(Update::filter_edited_message().endpoint(handlers::edited_message_handler))
.branch(Update::filter_callback_query().endpoint(events::callback_handler));
let secrets_dir_wdg = self.secrets_dir.clone();
let bot_wdg = bot.clone();
let task = tokio::spawn(async move {
let mut dispatcher = Dispatcher::builder(bot, handler)
.dependencies(dptree::deps![shared])
.build();
info!("telegram plugin: dispatcher starting");
tokio::select! {
_ = cancel_clone.cancelled() => info!("telegram plugin: cancellation received"),
_ = dispatcher.dispatch() => warn!("telegram plugin: dispatcher exited unexpectedly"),
_ = auth::whitelist_watchdog(bot_wdg, secrets_dir_wdg, cancel_wdg) => {}
}
running_clone.store(false, Ordering::Relaxed);
info!("telegram plugin: stopped");
});
*self.cancel.lock().await = Some(cancel);
*self.handle.lock().await = Some(task);
Ok(())
}
async fn stop(&self) -> Result<()> {
if let Some(token) = self.cancel.lock().await.take() {
token.cancel();
}
if let Some(h) = self.handle.lock().await.take() {
let _ = h.await;
}
self.running.store(false, Ordering::Relaxed);
Ok(())
}
}
+227
View File
@@ -0,0 +1,227 @@
use std::sync::Arc;
use serde_json::json;
use teloxide::prelude::*;
use teloxide::types::InputFile;
use core_api::interface_tool::InterfaceTool;
use core_api::tts::{TextToSpeech, TtsProvider};
/// Returns all LLM-callable tools available in a Telegram session.
///
/// Each tool captures `bot` and `chat_id` so its handler can send content
/// back to the user without any additional context.
///
/// `send_voice_message` is included only when at least one TTS provider is active.
///
/// # Adding a new tool
/// Implement a private `fn <name>_tool(bot: Bot, chat_id: ChatId, ...) -> InterfaceTool`
/// and push it into the vec returned by this function.
pub(crate) async fn interface_tools(
bot: Bot,
chat_id: ChatId,
tts: &dyn TtsProvider,
) -> Vec<InterfaceTool> {
let mut tools = vec![send_attachment_tool(bot.clone(), chat_id)];
if let Some(synth) = tts.get().await {
tools.push(send_voice_tool(bot, chat_id, synth));
}
tools
}
// ── send_attachment ───────────────────────────────────────────────────────────
fn send_attachment_tool(bot: Bot, chat_id: ChatId) -> InterfaceTool {
InterfaceTool {
definition: json!({
"type": "function",
"function": {
"name": "send_attachment",
"description": "Send a file from the local filesystem to the user on Telegram. Images (jpg/png/webp) and videos (mp4/mov/webm) are sent inline by default; any other type is sent as a document. Set as_document=true to force sending as a downloadable file.",
"parameters": {
"type": "object",
"properties": {
"file_path": {
"type": "string",
"description": "Absolute or relative path to the file to send."
},
"caption": {
"type": "string",
"description": "Optional caption shown below the file."
},
"as_document": {
"type": "boolean",
"description": "Force sending as a downloadable file instead of an inline photo/video (default false)."
}
},
"required": ["file_path"]
}
}
}),
handler: Arc::new(move |args| {
let bot = bot.clone();
Box::pin(async move {
let file_path = args["file_path"]
.as_str()
.ok_or_else(|| anyhow::anyhow!("send_attachment: missing `file_path`"))?;
let caption = args["caption"].as_str().map(str::to_string);
let as_document = args["as_document"].as_bool().unwrap_or(false);
let path = std::path::Path::new(file_path);
if !path.exists() {
anyhow::bail!("send_attachment: file not found: {file_path}");
}
// Present images/videos inline by default; everything else (and
// anything when as_document=true) as a downloadable document.
let ext = path.extension()
.and_then(|e| e.to_str())
.unwrap_or("")
.to_ascii_lowercase();
let kind = if as_document {
"document"
} else {
match ext.as_str() {
"jpg" | "jpeg" | "png" | "webp" => "photo",
"mp4" | "mov" | "webm" => "video",
_ => "document",
}
};
let file = InputFile::file(path);
let result = match kind {
"photo" => {
let mut req = bot.send_photo(chat_id, file);
if let Some(cap) = caption { req = req.caption(cap); }
req.await.map(|_| ())
}
"video" => {
let mut req = bot.send_video(chat_id, file);
if let Some(cap) = caption { req = req.caption(cap); }
req.await.map(|_| ())
}
_ => {
let mut req = bot.send_document(chat_id, file);
if let Some(cap) = caption { req = req.caption(cap); }
req.await.map(|_| ())
}
};
result.map_err(|e| anyhow::anyhow!("send_attachment: Telegram error: {e}"))?;
Ok(format!("File sent ({kind}): {file_path}"))
})
}),
}
}
// ── send_voice_message ────────────────────────────────────────────────────────
fn send_voice_tool(bot: Bot, chat_id: ChatId, synth: Arc<dyn TextToSpeech>) -> InterfaceTool {
let instructions_hint = synth
.instructions()
.map(|i| format!("\n\nVoice instructions: {i}"))
.unwrap_or_default();
InterfaceTool {
definition: json!({
"type": "function",
"function": {
"name": "send_voice_message",
"description": format!(
"Synthesise text to speech and send it to the user as a Telegram voice message. \
Use when audio is a better medium than text — e.g. short answers, \
confirmations, or when the user asks you to speak.{instructions_hint}"
),
"parameters": {
"type": "object",
"properties": {
"text": {
"type": "string",
"description": "The text to synthesise and send as audio."
}
},
"required": ["text"]
}
}
}),
handler: Arc::new(move |args| {
let bot = bot.clone();
let synth = Arc::clone(&synth);
Box::pin(async move {
let text = args["text"]
.as_str()
.ok_or_else(|| anyhow::anyhow!("send_voice_message: missing `text`"))?;
let audio = synth
.synthesize(text, None)
.await
.map_err(|e| anyhow::anyhow!("send_voice_message: TTS error: {e}"))?;
// Telegram only renders Ogg/Opus as a playable voice message, so
// transcode whatever the synthesiser produced (mp3, wav, raw pcm…).
let audio = to_ogg_opus(audio, synth.output_format())
.await
.map_err(|e| anyhow::anyhow!("send_voice_message: audio conversion failed: {e}"))?;
bot.send_voice(chat_id, InputFile::memory(audio).file_name("voice.ogg"))
.await
.map_err(|e| anyhow::anyhow!("send_voice_message: Telegram error: {e}"))?;
Ok("Voice message sent.".to_string())
})
}),
}
}
/// Transcode synthesised audio to Ogg/Opus — the only format Telegram renders as
/// a playable voice message — using ffmpeg over stdin/stdout pipes (no temp files).
///
/// `format` is the synthesiser's [`TextToSpeech::output_format`]. Ogg/Opus input
/// is passed through untouched. Raw `pcm` is headerless, so it is described to
/// ffmpeg as the 24 kHz / mono / s16le stream OpenAI and Gemini TTS emit; every
/// other (self-describing) container is auto-detected by ffmpeg.
async fn to_ogg_opus(audio: Vec<u8>, format: &str) -> anyhow::Result<Vec<u8>> {
use std::process::Stdio;
use tokio::io::AsyncWriteExt;
// Already a Telegram-native container — nothing to do.
if matches!(format, "opus" | "ogg") {
return Ok(audio);
}
let mut cmd = tokio::process::Command::new("ffmpeg");
cmd.args(["-hide_banner", "-loglevel", "error"]);
if format == "pcm" {
cmd.args(["-f", "s16le", "-ar", "24000", "-ac", "1"]);
}
cmd.args(["-i", "pipe:0", "-c:a", "libopus", "-b:a", "32k", "-f", "ogg", "pipe:1"])
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped());
let mut child = cmd.spawn().map_err(|e| anyhow::anyhow!(
"ffmpeg not available (required to convert {format} audio to Telegram Ogg/Opus): {e}"
))?;
// Feed stdin from a separate task so a full stdout pipe can't deadlock the write.
let mut stdin = child.stdin.take().expect("stdin piped");
let feeder = tokio::spawn(async move {
let _ = stdin.write_all(&audio).await;
let _ = stdin.shutdown().await;
});
let out = child.wait_with_output().await
.map_err(|e| anyhow::anyhow!("ffmpeg execution failed: {e}"))?;
let _ = feeder.await;
if !out.status.success() {
anyhow::bail!(
"ffmpeg ({format} → Ogg/Opus) exited with {}: {}",
out.status,
String::from_utf8_lossy(&out.stderr).trim(),
);
}
Ok(out.stdout)
}
@@ -0,0 +1,17 @@
[package]
name = "plugin-transcribe-whisper-local"
version = "0.1.0"
edition = "2024"
[dependencies]
core-api = { path = "../core-api" }
anyhow = "1"
async-trait = "0.1"
serde_json = "1"
tokio = { version = "1", features = ["full"] }
tracing = "0.1"
whisper-rs = { version = "0.16.0" }
[target.'cfg(target_os = "macos")'.dependencies]
whisper-rs = { version = "0.16.0", features = ["metal"] }
hound = "3"
@@ -0,0 +1,457 @@
/// WhisperLocalPlugin — local Speech-to-Text via whisper.cpp (Metal-accelerated on Apple Silicon).
///
/// Audio (any format) is first converted to 16 kHz mono WAV by ffmpeg, then fed to
/// whisper.cpp through the `whisper-rs` crate. No Python involved.
///
/// The ~2 GB model is **lazily loaded**: by default it is loaded into memory only on
/// the first transcription and unloaded again after a configurable idle period
/// (`idle_timeout_secs`, default 20 min). Set `load_at_startup: true` to load eagerly.
///
/// The model must be a GGML `.bin` file. Download with:
/// curl -L -o models/ggml-large-v3.bin \
/// https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-large-v3.bin
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::time::Duration;
use anyhow::Result;
use async_trait::async_trait;
use serde_json::{Value, json};
use tokio::task::JoinHandle;
use tracing::{debug, info, warn};
use whisper_rs::{FullParams, SamplingStrategy, WhisperContext, WhisperContextParameters};
use core_api::plugin::{Plugin, PluginContext};
use core_api::transcribe::{Transcribe, TranscribeRegistry};
/// Default idle timeout before the model is unloaded from memory (20 minutes).
const DEFAULT_IDLE_TIMEOUT_SECS: u64 = 1200;
/// How often the eviction task checks for idleness.
const EVICTION_TICK: Duration = Duration::from_secs(60);
// ── LazyModel ─────────────────────────────────────────────────────────────────
//
// Shared, droppable home of the GGML weights. Both the plugin and the registered
// transcriber hold an `Arc<LazyModel>`; the eviction task holds a `Weak`. The
// weights live behind `ctx` and can be dropped at any time to reclaim ~2 GB — the
// transcriber keeps working because it reloads on demand via `ensure_loaded()`.
struct LazyModel {
/// Path to the GGML `.bin` file. Behind a Mutex so a config change can swap it.
path: tokio::sync::Mutex<PathBuf>,
/// Loaded model context — `None` until first use (or after eviction).
ctx: tokio::sync::Mutex<Option<Arc<WhisperContext>>>,
/// Timestamp of the last transcription, used to decide idle eviction.
last_used: std::sync::Mutex<std::time::Instant>,
}
impl LazyModel {
fn new() -> Self {
Self {
path: tokio::sync::Mutex::new(PathBuf::new()),
ctx: tokio::sync::Mutex::new(None),
last_used: std::sync::Mutex::new(std::time::Instant::now()),
}
}
/// Reset the idle timer to now.
fn touch(&self) {
*self.last_used.lock().unwrap() = std::time::Instant::now();
}
async fn is_loaded(&self) -> bool {
self.ctx.lock().await.is_some()
}
/// Ensure the model is resident in memory and return a handle to it. Loads
/// from `path` on first use (or after eviction); a no-op clone otherwise.
/// The `ctx` lock is held across the load so concurrent first-callers wait
/// for a single load instead of loading the weights twice — but the lock is
/// released before the returned handle is used for inference.
async fn ensure_loaded(&self) -> Result<Arc<WhisperContext>> {
let mut guard = self.ctx.lock().await;
if let Some(ctx) = guard.as_ref() {
self.touch();
return Ok(Arc::clone(ctx));
}
let model_path = self.path.lock().await.clone();
anyhow::ensure!(
model_path.exists(),
"whisper_local: model file not found at '{}'. \
Download a GGML model and set the path via the plugins API.",
model_path.display()
);
let path_str = model_path
.to_str()
.ok_or_else(|| anyhow::anyhow!("model path is not valid UTF-8"))?
.to_string();
info!(model = %model_path.display(), "whisper_local: loading model into memory");
let whisper_ctx = tokio::task::spawn_blocking(move || {
let mut params = WhisperContextParameters::default();
params.use_gpu(true);
WhisperContext::new_with_params(&path_str, params)
.map_err(|e| anyhow::anyhow!("failed to load whisper model: {e:?}"))
})
.await??;
let ctx = Arc::new(whisper_ctx);
*guard = Some(Arc::clone(&ctx));
self.touch();
Ok(ctx)
}
/// Drop the in-memory model. The actual free runs on a blocking thread because
/// whisper.cpp cleanup may touch the GPU. In-flight transcriptions hold their
/// own `Arc`, so memory is only reclaimed once they finish.
async fn unload(&self) {
if let Some(ctx) = self.ctx.lock().await.take() {
let in_flight = Arc::strong_count(&ctx) - 1;
tokio::task::spawn_blocking(move || drop(ctx));
debug!(in_flight, "whisper_local: model unloaded from memory");
}
}
/// Unload the model if it has been idle for at least `timeout`.
async fn evict_if_idle(&self, timeout: Duration) {
if !self.is_loaded().await {
return;
}
let idle = self.last_used.lock().unwrap().elapsed();
if idle >= timeout {
info!(idle_secs = idle.as_secs(), "whisper_local: idle timeout reached — unloading model");
self.unload().await;
}
}
}
// ── WhisperLocalPlugin ────────────────────────────────────────────────────────
pub struct WhisperLocalPlugin {
/// Shared, lazily-loaded model weights.
model: Arc<LazyModel>,
/// BCP-47 language code. Shared with the registered transcriber so runtime
/// changes take effect without re-registering.
language: Arc<tokio::sync::Mutex<String>>,
/// Idle seconds before unload. `0` = never unload.
idle_timeout_secs: AtomicU64,
/// Load the model eagerly in `start()` instead of on first use.
load_at_startup: AtomicBool,
running: AtomicBool,
/// Kept so stop() can deregister without needing the full context.
transcribe_registry: tokio::sync::Mutex<Option<Arc<dyn TranscribeRegistry>>>,
/// Background idle-eviction task; aborted on stop()/reconfigure.
evictor: tokio::sync::Mutex<Option<JoinHandle<()>>>,
}
impl WhisperLocalPlugin {
pub fn new() -> Self {
Self {
model: Arc::new(LazyModel::new()),
language: Arc::new(tokio::sync::Mutex::new("auto".to_string())),
idle_timeout_secs: AtomicU64::new(DEFAULT_IDLE_TIMEOUT_SECS),
load_at_startup: AtomicBool::new(false),
running: AtomicBool::new(false),
transcribe_registry: tokio::sync::Mutex::new(None),
evictor: tokio::sync::Mutex::new(None),
}
}
/// (Re)start the idle-eviction task to match the current `idle_timeout_secs`.
/// A timeout of `0` means "never unload" — any existing task is dropped and
/// none is spawned.
async fn respawn_evictor(&self) {
if let Some(handle) = self.evictor.lock().await.take() {
handle.abort();
}
let secs = self.idle_timeout_secs.load(Ordering::Relaxed);
if secs == 0 {
return;
}
let timeout = Duration::from_secs(secs);
let model = Arc::downgrade(&self.model);
let handle = tokio::spawn(async move {
let mut ticker = tokio::time::interval(EVICTION_TICK);
ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
loop {
ticker.tick().await;
match model.upgrade() {
Some(m) => m.evict_if_idle(timeout).await,
None => break, // plugin dropped
}
}
});
*self.evictor.lock().await = Some(handle);
}
}
impl Default for WhisperLocalPlugin {
fn default() -> Self { Self::new() }
}
#[async_trait]
impl Plugin for WhisperLocalPlugin {
fn id(&self) -> &str { "whisper_local" }
fn name(&self) -> &str { "Whisper Local" }
fn description(&self) -> &str { "Local STT via whisper.cpp (Metal-accelerated)" }
fn is_running(&self) -> bool { self.running.load(Ordering::Relaxed) }
fn config_schema(&self) -> Value {
json!({
"type": "object",
"properties": {
"model": {
"type": "string",
"title": "Model path",
"description": "Path to GGML .bin file (e.g. models/ggml-large-v3.bin)"
},
"language": {
"type": "string",
"title": "Language",
"description": "BCP-47 code (e.g. 'it', 'en') or 'auto' for auto-detect",
"default": "auto"
},
"load_at_startup": {
"type": "boolean",
"title": "Load at startup",
"description": "Load the model into memory when the plugin starts. \
If false (default), the model is loaded lazily on the first transcription.",
"default": false
},
"idle_timeout_secs": {
"type": "integer",
"title": "Idle unload timeout (seconds)",
"description": "Unload the model from memory after this many seconds of inactivity. 0 = never unload.",
"default": 1200
}
},
"required": ["model"]
})
}
fn as_any(&self) -> &dyn std::any::Any { self }
fn as_arc_any(self: Arc<Self>) -> Arc<dyn std::any::Any + Send + Sync> { self }
async fn reload(&self, enabled: bool, config: Value, ctx: PluginContext) -> Result<()> {
let new_model = config["model"].as_str().unwrap_or("").to_string();
let new_lang = config["language"].as_str().unwrap_or("auto").to_string();
let new_eager = config["load_at_startup"].as_bool().unwrap_or(false);
let new_timeout = config["idle_timeout_secs"].as_u64().unwrap_or(DEFAULT_IDLE_TIMEOUT_SECS);
let old_model = self.model.path.lock().await.to_string_lossy().to_string();
let is_running = self.is_running();
// Scalar settings that don't depend on lifecycle. Language is shared with
// the live transcriber, so this takes effect immediately.
self.load_at_startup.store(new_eager, Ordering::Relaxed);
self.idle_timeout_secs.store(new_timeout, Ordering::Relaxed);
*self.language.lock().await = new_lang;
match (enabled, is_running) {
(true, false) => {
anyhow::ensure!(!new_model.is_empty(),
"whisper_local: cannot start — `model` is missing from config");
*self.model.path.lock().await = PathBuf::from(&new_model);
self.start(ctx).await?;
}
(false, true) => {
self.stop().await?;
}
(true, true) => {
if new_model != old_model {
info!("whisper_local: model path changed — clearing cached model");
*self.model.path.lock().await = PathBuf::from(&new_model);
self.model.unload().await;
}
// Pick up a possibly-changed idle timeout.
self.respawn_evictor().await;
// Honour load_at_startup turning on (or a fresh model path) eagerly.
if self.load_at_startup.load(Ordering::Relaxed) && !self.model.is_loaded().await {
self.model.ensure_loaded().await?;
}
}
(false, false) => {}
}
Ok(())
}
async fn start(&self, ctx: PluginContext) -> Result<()> {
if self.running.load(Ordering::Relaxed) { return Ok(()); }
// Redirect all whisper.cpp / ggml C-level log output through Rust callbacks.
// Since we don't enable the `log_backend` or `tracing_backend` features,
// the trampolines are no-ops — this silences all init/model-load chatter.
// The call is idempotent (backed by std::sync::Once).
whisper_rs::install_logging_hooks();
// Validate the model path up front so config errors surface immediately,
// even though the weights are not loaded until first use (lazy mode).
let model_path = self.model.path.lock().await.clone();
anyhow::ensure!(
model_path.exists(),
"whisper_local: model file not found at '{}'. \
Download a GGML model and set the path via the plugins API.",
model_path.display()
);
self.running.store(true, Ordering::Relaxed);
// Register a lightweight transcriber that shares the lazy model + language;
// it holds no strong reference to the weights, so eviction can free them.
*self.transcribe_registry.lock().await = Some(Arc::clone(&ctx.transcribe_registry));
ctx.transcribe_registry.register(Arc::new(WhisperLocalTranscriber {
model: Arc::clone(&self.model),
language: Arc::clone(&self.language),
})).await;
if self.load_at_startup.load(Ordering::Relaxed) {
self.model.ensure_loaded().await?;
info!(model = %model_path.display(), "whisper_local plugin started (model preloaded)");
} else {
info!(model = %model_path.display(), "whisper_local plugin started (lazy — loads on first use)");
}
self.respawn_evictor().await;
Ok(())
}
async fn stop(&self) -> Result<()> {
self.running.store(false, Ordering::Relaxed);
if let Some(handle) = self.evictor.lock().await.take() {
handle.abort();
}
self.model.unload().await;
info!("whisper_local plugin stopped");
if let Some(reg) = self.transcribe_registry.lock().await.take() {
reg.unregister("whisper_local").await;
}
Ok(())
}
}
// ── WhisperLocalTranscriber ───────────────────────────────────────────────────
//
// Lightweight handle registered in TranscribeManager. Shares the plugin's lazy
// model and language via Arc, so it never pins the weights in memory and always
// reflects the current language. Loads the model on demand at transcription time.
struct WhisperLocalTranscriber {
model: Arc<LazyModel>,
language: Arc<tokio::sync::Mutex<String>>,
}
#[async_trait]
impl Transcribe for WhisperLocalTranscriber {
fn id(&self) -> &str { "whisper_local" }
async fn transcribe(&self, audio: Vec<u8>, format: &str) -> Result<String> {
debug!(bytes = audio.len(), format, "whisper_local: transcribing audio");
// Load on demand (no-op if already resident) and refresh the idle timer.
let ctx = self.model.ensure_loaded().await?;
let language = self.language.lock().await.clone();
let pcm = audio_to_pcm_f32(audio, format).await?;
let text = tokio::task::spawn_blocking(move || {
let mut state = ctx.create_state()
.map_err(|e| anyhow::anyhow!("whisper state error: {e:?}"))?;
let mut params = FullParams::new(SamplingStrategy::Greedy { best_of: 1 });
params.set_language(Some(&language));
params.set_print_special(false);
params.set_print_progress(false);
params.set_print_realtime(false);
params.set_print_timestamps(false);
state.full(params, &pcm)
.map_err(|e| anyhow::anyhow!("whisper inference error: {e:?}"))?;
let n = state.full_n_segments();
let text = (0..n)
.map(|i| {
state.get_segment(i)
.ok_or_else(|| anyhow::anyhow!("whisper: segment {i} out of range"))
.and_then(|s| s.to_str()
.map(str::to_string)
.map_err(|e| anyhow::anyhow!("whisper segment text error: {e:?}")))
})
.collect::<Result<Vec<_>>>()?
.join(" ")
.trim()
.to_string();
info!(chars = text.len(), segments = n, "whisper_local: transcription complete");
Ok::<String, anyhow::Error>(text)
})
.await??;
// Reset the idle timer again so a long inference isn't evicted right after.
self.model.touch();
Ok(text)
}
}
// ── Audio conversion ──────────────────────────────────────────────────────────
//
// Uses ffmpeg (assumed installed) to decode any audio format to 16 kHz mono WAV,
// then reads it with hound. whisper.cpp requires f32 PCM at exactly 16 kHz mono.
async fn audio_to_pcm_f32(audio: Vec<u8>, format: &str) -> Result<Vec<f32>> {
let pid = std::process::id();
let tmp_in = std::env::temp_dir().join(format!("whisper_in_{pid}.{format}"));
let tmp_out = std::env::temp_dir().join(format!("whisper_out_{pid}.wav"));
tokio::fs::write(&tmp_in, &audio).await?;
let result = tokio::process::Command::new("ffmpeg")
.args([
"-y", "-i", tmp_in.to_str().unwrap(),
"-ar", "16000",
"-ac", "1",
tmp_out.to_str().unwrap(),
])
.output()
.await;
tokio::fs::remove_file(&tmp_in).await.ok();
let output = result?;
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr);
anyhow::bail!("ffmpeg audio conversion failed: {stderr}");
}
let wav_bytes = tokio::fs::read(&tmp_out).await?;
tokio::fs::remove_file(&tmp_out).await.ok();
tokio::task::spawn_blocking(move || {
let mut reader = hound::WavReader::new(std::io::Cursor::new(wav_bytes))?;
let spec = reader.spec();
let pcm: Vec<f32> = match spec.sample_format {
hound::SampleFormat::Int => {
let bits = spec.bits_per_sample;
reader.samples::<i32>()
.map(|s| s.map(|v| v as f32 / (1i32 << (bits - 1)) as f32))
.collect::<Result<_, _>>()?
}
hound::SampleFormat::Float => {
reader.samples::<f32>()
.collect::<Result<_, _>>()?
}
};
if pcm.is_empty() {
warn!("whisper_local: decoded PCM is empty");
}
Ok::<Vec<f32>, anyhow::Error>(pcm)
})
.await?
}
+13
View File
@@ -0,0 +1,13 @@
[package]
name = "plugin-tts-kokoro"
version = "0.1.0"
edition = "2024"
[dependencies]
core-api = { path = "../core-api" }
anyhow = "1"
async-trait = "0.1"
serde_json = "1"
tokio = { version = "1", features = ["full"] }
tracing = "0.1"
reqwest = { version = "0.13", default-features = false, features = ["rustls-no-provider", "charset", "http2", "system-proxy", "json"] }
@@ -0,0 +1,145 @@
#!/usr/bin/env python3
"""Kokoro-ONNX TTS inference server.
Started by the plugin-tts-kokoro Rust plugin. Prints "PORT:<n>" to stdout
once the HTTP server is bound so the plugin knows which port to connect to.
Model files (kokoro-v1.0.onnx + voices-v1.0.bin) are downloaded from GitHub
releases on first run and cached in --model-dir.
Endpoints
---------
POST /synthesize
Body: {"text": "...", "voice": "if_sara", "lang": "it", "speed": 1.0}
Returns: audio/wav bytes
GET /health
Returns: {"status": "ok"}
"""
import argparse
import io
import os
import socket
import sys
import threading
import numpy as np
import soundfile as sf
import urllib.request
from fastapi import FastAPI, HTTPException
from fastapi.responses import Response
from pydantic import BaseModel
from kokoro_onnx import Kokoro
import uvicorn
# ── Model URLs ────────────────────────────────────────────────────────────────
MODEL_BASE_URL = "https://github.com/thewh1teagle/kokoro-onnx/releases/download/model-files-v1.0"
MODEL_ONNX_FILE = "kokoro-v1.0.onnx"
MODEL_VOICES_FILE = "voices-v1.0.bin"
VALID_VOICES = {
"af_heart", "af_bella", "af_nicole", "af_sarah", "af_sky",
"am_adam", "am_michael",
"bf_emma", "bf_isabella", "bm_george", "bm_lewis",
"if_sara", "im_nicola",
"jf_alpha", "jf_gongitsune", "jm_kumo",
"zf_xiaobei", "zm_yunxi",
}
# ── Globals set at startup ────────────────────────────────────────────────────
kokoro_model: Kokoro | None = None
default_voice = "if_sara"
default_lang = "it"
default_speed = 1.0
# ── Model loading ─────────────────────────────────────────────────────────────
def download_if_missing(model_dir: str, filename: str) -> str:
path = os.path.join(model_dir, filename)
if not os.path.exists(path):
url = f"{MODEL_BASE_URL}/{filename}"
print(f"[kokoro] downloading {filename} from {url}", flush=True)
urllib.request.urlretrieve(url, path)
print(f"[kokoro] downloaded {filename}", flush=True)
return path
def load_model(model_dir: str) -> None:
global kokoro_model
os.makedirs(model_dir, exist_ok=True)
onnx_path = download_if_missing(model_dir, MODEL_ONNX_FILE)
voices_path = download_if_missing(model_dir, MODEL_VOICES_FILE)
print(f"[kokoro] loading model from {model_dir}", flush=True)
kokoro_model = Kokoro(onnx_path, voices_path)
print("[kokoro] model ready", flush=True)
# ── FastAPI server ────────────────────────────────────────────────────────────
app = FastAPI()
class SynthesizeRequest(BaseModel):
text: str
voice: str | None = None
lang: str | None = None
speed: float | None = None
@app.get("/health")
def health():
return {"status": "ok"}
@app.post("/synthesize")
def synthesize(req: SynthesizeRequest):
if kokoro_model is None:
raise HTTPException(status_code=503, detail="model not loaded")
voice = req.voice if req.voice in VALID_VOICES else default_voice
lang = req.lang or default_lang
speed = req.speed or default_speed
try:
samples, sample_rate = kokoro_model.create(req.text, voice=voice, speed=speed, lang=lang)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
buf = io.BytesIO()
sf.write(buf, samples, sample_rate, format="WAV")
return Response(content=buf.getvalue(), media_type="audio/wav")
# ── Entry point ───────────────────────────────────────────────────────────────
def find_free_port() -> int:
with socket.socket() as s:
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--model-dir", default="models/kokoro")
parser.add_argument("--default-voice", default="if_sara")
parser.add_argument("--default-lang", default="it")
parser.add_argument("--default-speed", type=float, default=1.0)
args = parser.parse_args()
global default_voice, default_lang, default_speed
default_voice = args.default_voice
default_lang = args.default_lang
default_speed = args.default_speed
load_model(args.model_dir)
port = find_free_port()
# Signal to the Rust parent that we are ready.
print(f"PORT:{port}", flush=True)
uvicorn.run(app, host="127.0.0.1", port=port, log_level="warning")
if __name__ == "__main__":
main()
+318
View File
@@ -0,0 +1,318 @@
//! Kokoro ONNX TTS plugin.
//!
//! On start, writes the embedded `kokoro_server.py` to `models/kokoro/`,
//! spawns it as a subprocess, reads the bound port from its stdout, then
//! registers itself as a [`TextToSpeech`] provider with the TTS manager.
//!
//! The subprocess downloads the Kokoro ONNX model files from GitHub releases
//! on first run (~310 MB + 27 MB) and caches them in `models/kokoro/`.
//! No API token required.
//!
//! # Config (stored in `plugins` SQLite table)
//!
//! ```json
//! {
//! "voice": "if_sara",
//! "lang": "it",
//! "speed": 1.0
//! }
//! ```
//!
//! | Field | Values | Default |
//! |---------|-----------------------------------------|------------|
//! | `voice` | `"if_sara"` \| `"im_nicola"` \| … | `"if_sara"` |
//! | `lang` | `"it"` \| `"en-us"` \| `"en-gb"` \| … | `"it"` |
//! | `speed` | `0.5` `2.0` | `1.0` |
use std::sync::{Arc, atomic::{AtomicBool, Ordering}};
use anyhow::{Context, Result, anyhow};
use async_trait::async_trait;
use serde_json::{Value, json};
use tokio::io::{AsyncBufReadExt, BufReader};
use tokio::process::{Child, Command};
use tokio::sync::Mutex;
use tracing::{info, warn};
const KOKORO_SERVER_PY: &str = include_str!("kokoro_server.py");
use core_api::plugin::{Plugin, PluginContext};
use core_api::tts::TextToSpeech;
const PLUGIN_ID: &str = "kokoro_tts";
const MODEL_DIR: &str = "models/kokoro";
const PROVIDER_ID: &str = "kokoro_tts";
const SERVER_PY_NAME: &str = "kokoro_server.py";
// ── Config ────────────────────────────────────────────────────────────────────
#[derive(Clone, PartialEq, Debug)]
struct KokoroConfig {
voice: String,
lang: String,
speed: f64,
}
impl KokoroConfig {
fn from_value(v: &Value) -> Self {
Self {
voice: v["voice"].as_str().unwrap_or("if_sara").to_string(),
lang: v["lang"].as_str().unwrap_or("it").to_string(),
speed: v["speed"].as_f64().unwrap_or(1.0),
}
}
}
// ── KokoroSynthesiser ─────────────────────────────────────────────────────────
struct KokoroSynthesiser {
port: u16,
config: KokoroConfig,
http: reqwest::Client,
}
impl KokoroSynthesiser {
fn new(port: u16, config: KokoroConfig) -> Self {
Self { port, config, http: reqwest::Client::new() }
}
}
#[async_trait]
impl TextToSpeech for KokoroSynthesiser {
fn id(&self) -> &str { PROVIDER_ID }
fn name(&self) -> &str { "Kokoro TTS" }
fn description(&self) -> Option<&str> {
Some("Local Kokoro ONNX TTS — lightweight, fast, multilingual. \
Runs on CPU, no GPU required. ~310 MB model, auto-downloaded on first use.")
}
fn instructions(&self) -> Option<&str> {
Some("\
Kokoro TTS produces natural spoken audio. \
Write as you would speak: short sentences, no markdown, no bullet points, no symbols. \
For Italian, use standard Italian orthography — accented letters are fine. \
Do not include URLs, file paths, or code snippets; describe them in words instead.")
}
async fn synthesize(&self, text: &str, instructions: Option<&str>) -> Result<Vec<u8>> {
// instructions may carry a voice override as the first word (same convention as Orpheus).
let voice = instructions
.and_then(|s| s.split_whitespace().next())
.map(str::to_owned);
let url = format!("http://127.0.0.1:{}/synthesize", self.port);
let resp = self.http
.post(&url)
.json(&json!({
"text": text,
"voice": voice.as_deref().unwrap_or(&self.config.voice),
"lang": self.config.lang,
"speed": self.config.speed,
}))
.send()
.await
.map_err(|e| anyhow!("kokoro_tts: request failed: {e}"))?;
let status = resp.status();
if !status.is_success() {
let msg = resp.text().await.unwrap_or_default();
anyhow::bail!("kokoro_tts: server error {status}: {msg}");
}
Ok(resp.bytes().await.map(|b| b.to_vec())
.map_err(|e| anyhow!("kokoro_tts: failed to read bytes: {e}"))?)
}
}
// ── Plugin inner state ────────────────────────────────────────────────────────
struct Inner {
child: Child,
port: u16,
config: KokoroConfig,
script_path: std::path::PathBuf,
}
// ── KokoroTtsPlugin ───────────────────────────────────────────────────────────
pub struct KokoroTtsPlugin {
running: AtomicBool,
inner: Mutex<Option<Inner>>,
}
impl KokoroTtsPlugin {
pub fn new() -> Self {
Self {
running: AtomicBool::new(false),
inner: Mutex::new(None),
}
}
async fn do_start(&self, config: &KokoroConfig, ctx: &PluginContext) -> Result<()> {
std::fs::create_dir_all(MODEL_DIR)
.context("kokoro_tts: failed to create model dir")?;
let script_path = std::path::Path::new(MODEL_DIR).join(SERVER_PY_NAME);
std::fs::write(&script_path, KOKORO_SERVER_PY)
.context("kokoro_tts: failed to write embedded server script")?;
let mut child = Command::new("python3")
.args([
script_path.to_str().unwrap(),
"--model-dir", MODEL_DIR,
"--default-voice", &config.voice,
"--default-lang", &config.lang,
"--default-speed", &config.speed.to_string(),
])
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.spawn()
.context("kokoro_tts: failed to spawn python3")?;
// Log stderr in background so Python errors appear in Skald's log file.
if let Some(stderr) = child.stderr.take() {
tokio::spawn(async move {
let mut lines = BufReader::new(stderr).lines();
while let Ok(Some(line)) = lines.next_line().await {
tracing::warn!("kokoro_tts(py/err): {line}");
}
});
}
let stdout = child.stdout.take()
.ok_or_else(|| anyhow!("kokoro_tts: no stdout from subprocess"))?;
let mut lines = BufReader::new(stdout).lines();
let port = loop {
match lines.next_line().await? {
None => anyhow::bail!("kokoro_tts: subprocess exited before printing port"),
Some(line) => {
if let Some(p) = line.strip_prefix("PORT:") {
break p.trim().parse::<u16>()
.context("kokoro_tts: invalid port from subprocess")?;
}
info!("kokoro_tts(py): {line}");
}
}
};
info!(port, "kokoro_tts: python server ready");
let synthesiser = Arc::new(KokoroSynthesiser::new(port, config.clone()));
ctx.tts_registry.register(Arc::clone(&synthesiser) as _).await;
self.running.store(true, Ordering::Relaxed);
*self.inner.lock().await = Some(Inner {
child,
port,
config: config.clone(),
script_path,
});
Ok(())
}
async fn do_stop(&self, ctx: &PluginContext) {
ctx.tts_registry.unregister(PROVIDER_ID).await;
if let Some(mut inner) = self.inner.lock().await.take() {
let _ = inner.child.kill().await;
let _ = std::fs::remove_file(&inner.script_path);
}
self.running.store(false, Ordering::Relaxed);
info!("kokoro_tts: stopped");
}
}
#[async_trait]
impl Plugin for KokoroTtsPlugin {
fn id(&self) -> &str { PLUGIN_ID }
fn name(&self) -> &str { "Kokoro TTS" }
fn description(&self) -> &str {
"Local text-to-speech using Kokoro ONNX. Lightweight (~310 MB), fast on CPU, \
multilingual (Italian, English, Japanese, Chinese, Spanish, French…). \
No API token required — model is downloaded automatically on first use."
}
fn is_running(&self) -> bool { self.running.load(Ordering::Relaxed) }
fn config_schema(&self) -> Value {
json!({
"type": "object",
"properties": {
"voice": {
"type": "string",
"enum": [
"if_sara", "im_nicola",
"af_heart", "af_bella", "af_nicole", "af_sarah", "af_sky",
"am_adam", "am_michael",
"bf_emma", "bf_isabella", "bm_george", "bm_lewis"
],
"default": "if_sara",
"description": "Voice ID. Prefix: a=American, b=British, i=Italian, j=Japanese, z=Chinese; f=female, m=male."
},
"lang": {
"type": "string",
"enum": ["it", "en-us", "en-gb", "ja", "zh", "es", "fr", "hi", "pt-br", "ko"],
"default": "it",
"description": "Language code for phonemisation."
},
"speed": {
"type": "number",
"minimum": 0.5,
"maximum": 2.0,
"default": 1.0,
"description": "Speech speed multiplier."
}
}
})
}
async fn reload(&self, enabled: bool, config: Value, ctx: PluginContext) -> Result<()> {
let new_cfg = KokoroConfig::from_value(&config);
let is_running = self.is_running();
let config_changed = self.inner.lock().await
.as_ref()
.map(|i| i.config != new_cfg)
.unwrap_or(false);
match (enabled, is_running) {
(true, false) => self.do_start(&new_cfg, &ctx).await?,
(false, true) => self.do_stop(&ctx).await,
(true, true) if config_changed => {
info!("kokoro_tts: config changed — restarting");
self.do_stop(&ctx).await;
self.do_start(&new_cfg, &ctx).await?;
}
_ => {}
}
Ok(())
}
async fn start(&self, ctx: PluginContext) -> Result<()> {
let _ = ctx;
Ok(())
}
async fn stop(&self) -> Result<()> {
warn!("kokoro_tts: stop() called without ctx — cannot unregister from TtsManager");
if let Some(mut inner) = self.inner.lock().await.take() {
let _ = inner.child.kill().await;
let _ = inner.child.wait().await;
}
self.running.store(false, Ordering::Relaxed);
Ok(())
}
fn runtime_status(&self) -> Option<Value> {
let inner = self.inner.try_lock().ok()?;
let inner = inner.as_ref()?;
Some(json!({
"port": inner.port,
"voice": inner.config.voice,
"lang": inner.config.lang,
"speed": inner.config.speed,
}))
}
fn as_any(&self) -> &dyn std::any::Any { self }
fn as_arc_any(self: Arc<Self>) -> Arc<dyn std::any::Any + Send + Sync> { self }
}
+13
View File
@@ -0,0 +1,13 @@
[package]
name = "plugin-tts-orpheus-3b"
version = "0.1.0"
edition = "2024"
[dependencies]
core-api = { path = "../core-api" }
anyhow = "1"
async-trait = "0.1"
serde_json = "1"
tokio = { version = "1", features = ["full"] }
tracing = "0.1"
reqwest = { version = "0.13", default-features = false, features = ["rustls-no-provider", "charset", "http2", "system-proxy", "json"] }
+314
View File
@@ -0,0 +1,314 @@
//! Orpheus TTS 3B plugin.
//!
//! On start, writes the embedded `orpheus_server.py` (bundled via
//! `include_str!`) to `models/orpheus-3b/`, spawns it as a subprocess, reads
//! the bound port from its stdout, then registers itself as a [`TextToSpeech`]
//! provider with the TTS manager.
//!
//! The subprocess loads the Orpheus 3B model from HuggingFace (auto-download
//! on first run, cached in `models/orpheus-3b/`) and exposes a minimal HTTP
//! server on a random OS-assigned port.
//!
//! # Required secret
//!
//! Set before enabling the plugin:
//! ```
//! set_secret("HUGGINGFACE_TOKEN", "hf_...")
//! ```
//! Get a token at <https://huggingface.co/settings/tokens>.
//!
//! # Config (stored in `plugins` SQLite table)
//!
//! ```json
//! {
//! "quantization": "int8",
//! "voice": "tara"
//! }
//! ```
//!
//! | Field | Values | Default |
//! |-------|--------|---------|
//! | `quantization` | `"none"` \| `"int8"` \| `"int4"` | `"int8"` |
//! | `voice` | `"tara"` \| `"dan"` \| `"leah"` \| `"zac"` \| `"zoe"` \| `"mia"` \| `"julia"` \| `"leo"` | `"tara"` |
use std::sync::{Arc, atomic::{AtomicBool, Ordering}};
use anyhow::{Context, Result, anyhow};
use async_trait::async_trait;
use serde_json::{Value, json};
use tokio::io::{AsyncBufReadExt, BufReader};
use tokio::process::{Child, Command};
use tokio::sync::Mutex;
use tracing::{info, warn};
const ORPHEUS_SERVER_PY: &str = include_str!("orpheus_server.py");
use core_api::plugin::{Plugin, PluginContext};
use core_api::secrets;
use core_api::tts::TextToSpeech;
const PLUGIN_ID: &str = "orpheus_tts_3b";
const MODEL_DIR: &str = "models/orpheus-3b";
const PROVIDER_ID: &str = "orpheus_tts_3b";
const SERVER_PY_NAME: &str = "orpheus_server.py";
// ── Config ────────────────────────────────────────────────────────────────────
#[derive(Clone, PartialEq, Debug)]
struct OrpheusTtsConfig {
quantization: String,
voice: String,
}
impl OrpheusTtsConfig {
fn from_value(v: &Value) -> Self {
Self {
quantization: v["quantization"].as_str().unwrap_or("int8").to_string(),
voice: v["voice"].as_str().unwrap_or("tara").to_string(),
}
}
}
// ── OrpheusSynthesiser ────────────────────────────────────────────────────────
/// Calls the local Orpheus Python server to synthesise audio.
struct OrpheusSynthesiser {
port: u16,
default_voice: String,
http: reqwest::Client,
}
impl OrpheusSynthesiser {
fn new(port: u16, default_voice: impl Into<String>) -> Self {
Self {
port,
default_voice: default_voice.into(),
http: reqwest::Client::new(),
}
}
}
#[async_trait]
impl TextToSpeech for OrpheusSynthesiser {
fn id(&self) -> &str { PROVIDER_ID }
fn name(&self) -> &str { "Orpheus TTS 3B" }
fn description(&self) -> Option<&str> {
Some("Local Orpheus TTS 3B — high-quality expressive speech, runs on-device.")
}
fn instructions(&self) -> Option<&str> {
Some("\
Orpheus TTS supports inline emotion tags. Insert them directly in the text where the effect should occur.\n\
Supported tags: <laugh>, <chuckle>, <sigh>, <cough>, <sniffle>, <groan>, <yawn>, <gasp>\n\
Example: \"I told him the meeting was at nine, not eleven. <sigh> He showed up at noon. <chuckle> Classic.\"\
")
}
async fn synthesize(&self, text: &str, instructions: Option<&str>) -> Result<Vec<u8>> {
let voice = instructions
.and_then(|s| s.split_whitespace().next()) // first word as voice override
.unwrap_or(&self.default_voice);
let url = format!("http://127.0.0.1:{}/synthesize", self.port);
let resp = self.http
.post(&url)
.json(&json!({
"text": text,
"voice": voice,
"instructions": instructions,
}))
.send()
.await
.map_err(|e| anyhow!("orpheus_tts: request failed: {e}"))?;
let status = resp.status();
if !status.is_success() {
let msg = resp.text().await.unwrap_or_default();
anyhow::bail!("orpheus_tts: server error {status}: {msg}");
}
Ok(resp.bytes().await.map(|b| b.to_vec())
.map_err(|e| anyhow!("orpheus_tts: failed to read bytes: {e}"))?)
}
}
// ── Plugin inner state ────────────────────────────────────────────────────────
struct Inner {
child: Child,
port: u16,
config: OrpheusTtsConfig,
script_path: std::path::PathBuf,
}
// ── OrpheusTtsPlugin ──────────────────────────────────────────────────────────
pub struct OrpheusTtsPlugin {
running: AtomicBool,
inner: Mutex<Option<Inner>>,
}
impl OrpheusTtsPlugin {
pub fn new() -> Self {
Self {
running: AtomicBool::new(false),
inner: Mutex::new(None),
}
}
async fn do_start(&self, config: &OrpheusTtsConfig, ctx: &PluginContext) -> Result<()> {
std::fs::create_dir_all(MODEL_DIR)
.context("orpheus_tts: failed to create model dir")?;
// Write the embedded script to the model dir so it can be executed.
let script_path = std::path::Path::new(MODEL_DIR).join(SERVER_PY_NAME);
std::fs::write(&script_path, ORPHEUS_SERVER_PY)
.context("orpheus_tts: failed to write embedded server script")?;
// HuggingFace token — required for gated repos. Passed as env var so
// transformers/huggingface_hub pick it up automatically.
let hf_token = secrets::require(&ctx.secrets, "HUGGINGFACE_TOKEN").await?;
let mut child = Command::new("python3")
.args([
script_path.to_str().unwrap(),
"--model-dir", MODEL_DIR,
"--quantization", &config.quantization,
"--default-voice", &config.voice,
])
.env("HF_TOKEN", &hf_token)
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::inherit())
.spawn()
.context("orpheus_tts: failed to spawn python3")?;
// Read stdout until we see "PORT:<n>" — the server prints it once bound.
let stdout = child.stdout.take()
.ok_or_else(|| anyhow!("orpheus_tts: no stdout from subprocess"))?;
let mut lines = BufReader::new(stdout).lines();
let port = loop {
match lines.next_line().await? {
None => anyhow::bail!("orpheus_tts: subprocess exited before printing port"),
Some(line) => {
if let Some(p) = line.strip_prefix("PORT:") {
break p.trim().parse::<u16>()
.context("orpheus_tts: invalid port from subprocess")?;
}
// Forward other startup lines as info.
info!("orpheus_tts(py): {line}");
}
}
};
info!(port, "orpheus_tts: python server ready");
let synthesiser = Arc::new(OrpheusSynthesiser::new(port, &config.voice));
ctx.tts_registry.register(Arc::clone(&synthesiser) as _).await;
self.running.store(true, Ordering::Relaxed);
*self.inner.lock().await = Some(Inner {
child,
port,
config: config.clone(),
script_path,
});
Ok(())
}
async fn do_stop(&self, ctx: &PluginContext) {
ctx.tts_registry.unregister(PROVIDER_ID).await;
if let Some(mut inner) = self.inner.lock().await.take() {
let _ = inner.child.kill().await;
let _ = std::fs::remove_file(&inner.script_path);
}
self.running.store(false, Ordering::Relaxed);
info!("orpheus_tts: stopped");
}
}
#[async_trait]
impl Plugin for OrpheusTtsPlugin {
fn id(&self) -> &str { PLUGIN_ID }
fn name(&self) -> &str { "Orpheus TTS 3B" }
fn description(&self) -> &str {
"Local text-to-speech using Orpheus 3B. Expressive, high-quality, runs fully on-device. \
Requires ~7 GB VRAM (fp16), ~4 GB (int8), or ~2.5 GB (int4). \
Requires secret: HUGGINGFACE_TOKEN (HuggingFace access token — \
get one at https://huggingface.co/settings/tokens, then call \
set_secret(\"HUGGINGFACE_TOKEN\", \"hf_...\")). \
See docs/tts-providers.md for full setup instructions."
}
fn is_running(&self) -> bool { self.running.load(Ordering::Relaxed) }
fn config_schema(&self) -> Value {
json!({
"type": "object",
"properties": {
"quantization": {
"type": "string",
"enum": ["none", "int8", "int4"],
"default": "int8",
"description": "bitsandbytes precision: none=fp16 (~7 GB VRAM), int8 (~4 GB), int4 (~2.5 GB)"
},
"voice": {
"type": "string",
"enum": ["tara", "dan", "leah", "zac", "zoe", "mia", "julia", "leo"],
"default": "tara",
"description": "Default voice. Can be overridden per synthesis call via instructions."
}
}
})
}
async fn reload(&self, enabled: bool, config: Value, ctx: PluginContext) -> Result<()> {
let new_cfg = OrpheusTtsConfig::from_value(&config);
let is_running = self.is_running();
let config_changed = self.inner.lock().await
.as_ref()
.map(|i| i.config != new_cfg)
.unwrap_or(false);
match (enabled, is_running) {
(true, false) => self.do_start(&new_cfg, &ctx).await?,
(false, true) => self.do_stop(&ctx).await,
(true, true) if config_changed => {
info!("orpheus_tts: config changed — restarting");
self.do_stop(&ctx).await;
self.do_start(&new_cfg, &ctx).await?;
}
_ => {}
}
Ok(())
}
async fn start(&self, ctx: PluginContext) -> Result<()> {
// start() is called by the plugin manager; reload() handles the normal path.
// This is a no-op here — reload() does the real work.
let _ = ctx;
Ok(())
}
async fn stop(&self) -> Result<()> {
warn!("orpheus_tts: stop() called without ctx — cannot unregister from TtsManager");
if let Some(mut inner) = self.inner.lock().await.take() {
let _ = inner.child.kill().await;
}
self.running.store(false, Ordering::Relaxed);
Ok(())
}
fn runtime_status(&self) -> Option<Value> {
let inner = self.inner.try_lock().ok()?;
let inner = inner.as_ref()?;
Some(json!({
"port": inner.port,
"quantization": inner.config.quantization,
"voice": inner.config.voice,
}))
}
fn as_any(&self) -> &dyn std::any::Any { self }
fn as_arc_any(self: Arc<Self>) -> Arc<dyn std::any::Any + Send + Sync> { self }
}
@@ -0,0 +1,228 @@
#!/usr/bin/env python3
"""Orpheus TTS 3B inference server.
Started by the plugin-tts-orpheus-3b Rust plugin. Prints "PORT:<n>" to stdout
once the HTTP server is bound so the plugin knows which port to connect to.
The model is downloaded from HuggingFace on first run and cached in --model-dir.
Endpoints
---------
POST /synthesize
Body: {"text": "...", "voice": "tara", "instructions": "..."}
Returns: audio/wav bytes
GET /health
Returns: {"status": "ok"}
"""
import argparse
import io
import json
import os
import socket
import sys
import threading
import numpy as np
import scipy.io.wavfile as wavfile
import torch
from fastapi import FastAPI, HTTPException
from fastapi.responses import Response
from huggingface_hub import snapshot_download
from pydantic import BaseModel
from snac import SNAC
from transformers import AutoModelForCausalLM, AutoTokenizer
import uvicorn
# ── Model IDs ────────────────────────────────────────────────────────────────
ORPHEUS_MODEL_ID = "canopylabs/orpheus-3b-0.1-ft"
SNAC_MODEL_ID = "hubertsiuzdak/snac_24khz"
SAMPLE_RATE = 24000
VALID_VOICES = {"tara", "dan", "leah", "zac", "zoe", "mia", "julia", "leo"}
# ── Globals set at startup ────────────────────────────────────────────────────
model = None
tokenizer = None
snac_model = None
default_voice = "tara"
if torch.cuda.is_available():
device = "cuda"
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
device = "mps"
# Some transformer ops are not yet implemented on MPS; fall back to CPU.
os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")
else:
device = "cpu"
# ── Model loading ─────────────────────────────────────────────────────────────
def load_model(model_dir: str, quantization: str) -> None:
global model, tokenizer, snac_model
print(f"[orpheus] loading model (quantization={quantization}, device={device})", flush=True)
hf_cache = os.path.join(model_dir, "hf_cache")
os.makedirs(hf_cache, exist_ok=True)
tokenizer = AutoTokenizer.from_pretrained(
ORPHEUS_MODEL_ID,
cache_dir=hf_cache,
)
load_kwargs: dict = {
"cache_dir": hf_cache,
"torch_dtype": torch.float16 if device in ("cuda", "mps") else torch.float32,
"device_map": "auto" if device == "cuda" else None,
"low_cpu_mem_usage": True,
}
# bitsandbytes quantization is CUDA-only — skip on MPS and CPU.
if device == "cuda":
if quantization == "int8":
load_kwargs["load_in_8bit"] = True
elif quantization == "int4":
load_kwargs["load_in_4bit"] = True
elif quantization != "none":
print(f"[orpheus] quantization '{quantization}' not supported on {device}, running fp16", flush=True)
model = AutoModelForCausalLM.from_pretrained(ORPHEUS_MODEL_ID, **load_kwargs)
if device != "cuda": # for cuda, device_map="auto" already handles placement
model = model.to(device)
model.eval()
snac_model = SNAC.from_pretrained(SNAC_MODEL_ID, cache_dir=hf_cache).to(device)
snac_model.eval()
print("[orpheus] model loaded", flush=True)
# ── Inference ─────────────────────────────────────────────────────────────────
def _tokens_to_audio(token_ids: list[int]) -> np.ndarray:
"""Decode Orpheus audio token stream via SNAC to a float32 waveform."""
# Orpheus uses a 7-level SNAC codec; tokens are interleaved in groups of 7.
# Filter to valid audio token range (typically 128266129290 for 24 kHz SNAC).
audio_token_start = 128266
audio_tokens = [t - audio_token_start for t in token_ids if t >= audio_token_start]
if len(audio_tokens) < 7:
return np.zeros(0, dtype=np.float32)
# Trim to multiple of 7.
n = (len(audio_tokens) // 7) * 7
audio_tokens = audio_tokens[:n]
layers = [[] for _ in range(7)]
for i, tok in enumerate(audio_tokens):
layers[i % 7].append(tok)
with torch.no_grad():
codes = [
torch.tensor(layer, dtype=torch.long, device=device).unsqueeze(0)
for layer in layers
]
audio = snac_model.decode(codes)
return audio.squeeze().cpu().float().numpy()
def synthesize_text(text: str, voice: str, instructions: str | None) -> bytes:
voice = voice if voice in VALID_VOICES else default_voice
# Build prompt in Orpheus format.
prompt = f"<|audio|>{voice}: {text}<|eot_id|>"
if instructions:
prompt = f"<|audio|>{voice}: {text} [style: {instructions}]<|eot_id|>"
inputs = tokenizer(prompt, return_tensors="pt").to(device)
with torch.no_grad():
output_ids = model.generate(
**inputs,
max_new_tokens=4096,
do_sample=True,
temperature=0.7,
repetition_penalty=1.1,
eos_token_id=tokenizer.eos_token_id,
)
# Strip the prompt tokens; keep only newly generated tokens.
new_tokens = output_ids[0][inputs["input_ids"].shape[1]:].tolist()
waveform = _tokens_to_audio(new_tokens)
if waveform.size == 0:
raise RuntimeError("orpheus: decoding produced no audio samples")
# Encode to 16-bit WAV in memory.
pcm = (waveform * 32767).astype(np.int16)
buf = io.BytesIO()
wavfile.write(buf, SAMPLE_RATE, pcm)
return buf.getvalue()
# ── FastAPI app ───────────────────────────────────────────────────────────────
app = FastAPI()
class SynthesizeRequest(BaseModel):
text: str
voice: str | None = None
instructions: str | None = None
@app.post("/synthesize")
def synthesize(req: SynthesizeRequest):
if not req.text.strip():
raise HTTPException(status_code=400, detail="text is empty")
try:
audio = synthesize_text(
req.text,
req.voice or default_voice,
req.instructions,
)
return Response(content=audio, media_type="audio/wav")
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.get("/health")
def health():
return {"status": "ok"}
# ── Entry point ───────────────────────────────────────────────────────────────
def main() -> None:
global default_voice
parser = argparse.ArgumentParser()
parser.add_argument("--model-dir", default="models/orpheus-3b")
parser.add_argument("--quantization", default="int8", choices=["none", "int8", "int4"])
parser.add_argument("--default-voice", default="tara")
args = parser.parse_args()
default_voice = args.default_voice
load_model(args.model_dir, args.quantization)
# Bind on port 0 — OS assigns a free port.
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.bind(("127.0.0.1", 0))
port = sock.getsockname()[1]
sock.close()
# Print port for the Rust plugin to read.
print(f"PORT:{port}", flush=True)
config = uvicorn.Config(app, host="127.0.0.1", port=port, log_level="warning")
server = uvicorn.Server(config)
server.run()
if __name__ == "__main__":
main()
+54
View File
@@ -0,0 +1,54 @@
[package]
name = "skald-relay-client"
version = "0.1.0"
edition = "2024"
description = "Standalone relay client (agent role): WS v2 transport, E2E crypto, pairing, device authorization, SQLite persistence. Payload-agnostic — emits decoded bytes via RelayEvent. See docs/relay/"
[dependencies]
# --- shared frame types + crypto (frames + derive/namespace/ecdh/seal/open) ---
skald-relay-common = { path = "../skald-relay-common" }
# --- async runtime ---
tokio = { version = "1", features = ["full"] }
tokio-util = { version = "0.7", features = ["rt"] }
# --- WS transport (v2 binary) ---
# `tokio-tungstenite` pinned to 0.29 to match `skald-relay-server`'s
# `[dev-dependencies]`: the integration test links both crates in the same
# binary and the `WsMessage`/`connect_async` types must be the same compile
# units on both sides.
tokio-tungstenite = { version = "0.29", features = ["rustls-tls-webpki-roots"] }
futures-util = "0.3"
# --- persistence (counters MUST survive restarts; nonce never reused) ---
# `sqlx` 0.9.0 pinned to match the server (the integration test shares the
# `SqlitePool` / `row` types).
sqlx = { version = "0.9.0", features = ["runtime-tokio", "sqlite"] }
# --- v2 wire format (prost 0.13 + bytes 1 match skald-relay-common/server) ---
prost = "0.13"
bytes = "1"
# --- identity / crypto helpers used directly here ---
ed25519-dalek = "2"
# `rand` 0.9 matches the plugin's usage of `rand::rng()` (RngCore) in
# `identity.rs`/`pairing.rs`. `skald-relay-common` keeps its own `rand = "0.8"`
# in its build — the two crates coexist fine in the workspace.
rand = "0.9"
# --- misc ---
anyhow = "1"
chrono = { version = "0.4", default-features = false, features = ["clock", "std"] }
hex = "0.4"
serde = { version = "1", features = ["derive"] } # QrCodeData
tracing = "0.1"
[dev-dependencies]
# The integration test boots the relay in-process (`AppState::build` + `router`)
# and speaks v2 protobuf against it over a real WS on `127.0.0.1:0`.
skald-relay-server = { path = "../skald-relay-server" }
serde_json = "1"
# `axum` to serve the in-process relay; `tokio` `full` for `tokio::test` +
# the `net`/`time` features the harness needs.
axum = { version = "0.8", features = ["ws"] }
tokio = { version = "1", features = ["full"] }
+238
View File
@@ -0,0 +1,238 @@
//! [`RelayClient`] — the public façade over the networking layer.
//!
//! Concrete struct with inherent async methods (no trait): there is exactly one
//! implementation and the consumer wants a thin, direct handle. The client owns
//! the WS loop lifecycle and the broadcast event channel; all transport/crypto
//! logic lives in [`crate::state::RelayState`], shared behind an `Arc`.
use std::sync::Arc;
use anyhow::Result;
use sqlx::SqlitePool;
use tokio::sync::{broadcast, mpsc, Mutex};
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;
use tracing::info;
use crate::config::RelayClientConfig;
use crate::db::{self, ClientRow};
use crate::events::RelayEvent;
use crate::identity::Identity;
use crate::pairing::{QrCodeData, SessionState, StartedPairing};
use crate::state::{RelayState, StateConfig};
use crate::ws;
/// How many events the broadcast channel buffers before lagging slow consumers.
const EVENT_CHANNEL_CAP: usize = 256;
/// A standalone, payload-agnostic relay client (agent role).
///
/// Lifecycle: [`new`](Self::new) derives the identity and initializes the DB but
/// does **not** connect; [`start`](Self::start) spawns the reconnecting WS loop;
/// [`shutdown`](Self::shutdown) cancels it and joins. Inbound traffic and
/// lifecycle transitions are delivered via [`events`](Self::events).
pub struct RelayClient {
state: Arc<RelayState>,
/// Token cancelling the WS loop; `Some` only while started.
cancel: Mutex<Option<CancellationToken>>,
handle: Mutex<Option<JoinHandle<()>>>,
}
impl RelayClient {
/// Derive the identity from the seed source, ensure the `relay_clients`
/// table exists, and build the client. Does NOT connect — call
/// [`start`](Self::start).
pub async fn new(db: Arc<SqlitePool>, config: RelayClientConfig) -> Result<Self> {
db::init(&db).await?;
let identity = Identity::from_source(&config.seed)?;
info!(
crate_name = "skald-relay-client",
namespace = identity.namespace_id_hex(),
"relay client identity loaded"
);
let (events_tx, _) = broadcast::channel(EVENT_CHANNEL_CAP);
let state = Arc::new(RelayState::new(
identity,
db,
StateConfig { relay_url: config.relay_url, pairing_ttl: config.pairing_ttl },
events_tx,
));
Ok(Self {
state,
cancel: Mutex::new(None),
handle: Mutex::new(None),
})
}
/// Spawn the reconnecting WS loop. No-op (stays idle) if `relay_url` is
/// empty. Wires a fresh outbound channel into the state. Calling `start`
/// while already started replaces the loop (the caller should `shutdown`
/// first; this guards by cancelling any prior token).
pub async fn start(&self) -> Result<()> {
// Cancel any previous loop defensively.
if let Some(c) = self.cancel.lock().await.take() {
c.cancel();
}
if let Some(h) = self.handle.lock().await.take() {
let _ = h.await;
}
let cancel = CancellationToken::new();
let (out_tx, out_rx) = mpsc::unbounded_channel::<Vec<u8>>();
self.state.set_outbound(out_tx);
if self.state.relay_url().is_empty() {
// Idle: no WS loop, but the outbound sender is set so pairing/send
// calls fail loudly ("WS not started") rather than panic.
*self.cancel.lock().await = Some(cancel);
return Ok(());
}
let st = Arc::clone(&self.state);
let c = cancel.clone();
let handle = tokio::spawn(async move {
ws::run_loop(st, out_rx, c).await;
});
*self.cancel.lock().await = Some(cancel);
*self.handle.lock().await = Some(handle);
Ok(())
}
/// Cancel the WS loop, clear the outbound sender, and join the task.
pub async fn shutdown(&self) {
if let Some(c) = self.cancel.lock().await.take() {
c.cancel();
}
self.state.clear_outbound();
self.state.set_connected(false);
if let Some(h) = self.handle.lock().await.take() {
let _ = h.await;
}
}
/// Subscribe to the client's [`RelayEvent`] stream. Each call returns a new
/// receiver; a slow consumer lags (`RecvError::Lagged`) rather than blocking
/// the WS loop.
pub fn events(&self) -> broadcast::Receiver<RelayEvent> {
self.state.subscribe()
}
/// Seal `payload` to one authorized client and queue the `message` frame.
/// `live=true` routes-or-fails (peer online by construction); `live=false`
/// stores-and-forwards + pushes for offline phones.
pub async fn send(&self, dest: &[u8; 32], payload: &[u8], live: bool) -> Result<()> {
self.state.send_to_client(dest, payload, live).await
}
// ── Pipe (relayed byte-stream, docs/relay/pipe.md) ─────────────────────────
/// Open an end-to-end-encrypted byte pipe to `peer` (a namespace member).
/// Brokers the rendezvous over the E2E channel (`pipe_invite`/`pipe_accept`,
/// ephemeral DH → PFS) and returns the live data-plane channel.
pub async fn open_pipe(
&self,
peer: &[u8; 32],
stream_type: &str,
headers: std::collections::BTreeMap<String, String>,
) -> Result<crate::pipe::PipeConnection> {
self.state.open_pipe(peer, stream_type, headers).await
}
/// Subscribe to inbound pipe invites (responder side). Each invite is an
/// [`IncomingPipe`](crate::pipe::IncomingPipe); call [`accept_pipe`](Self::accept_pipe)
/// or [`reject_pipe`](Self::reject_pipe) on it. Single-consumer expected.
pub fn incoming_pipes(&self) -> broadcast::Receiver<crate::pipe::IncomingPipe> {
self.state.incoming_pipes()
}
/// Accept an inbound invite → returns the live data-plane channel.
pub async fn accept_pipe(
&self,
incoming: &crate::pipe::IncomingPipe,
) -> Result<crate::pipe::PipeConnection> {
self.state.accept_pipe(incoming).await
}
/// Decline an inbound invite.
pub async fn reject_pipe(
&self,
incoming: &crate::pipe::IncomingPipe,
reason: &str,
) -> Result<()> {
self.state.reject_pipe(incoming, reason).await
}
// ── Pairing ───────────────────────────────────────────────────────────────
/// Open the pairing window (single-window, latest-wins). `ttl_secs == 0`
/// uses the configured default.
pub async fn start_pairing(&self, ttl_secs: u32) -> Result<StartedPairing> {
let ttl = if ttl_secs == 0 { self.state.default_pairing_ttl() } else { ttl_secs };
self.state.start_pairing(ttl).await
}
/// Close the pairing window locally and tell the relay.
pub async fn stop_pairing(&self) -> Result<()> {
self.state.stop_pairing().await
}
/// Resolve a pairing `code` to its QR payload + lifecycle state (QR router).
pub fn lookup_pairing(&self, code: &str) -> Option<(QrCodeData, SessionState)> {
self.state.lookup_pairing(code)
}
/// The configured default pairing TTL (seconds).
pub fn default_pairing_ttl(&self) -> u32 {
self.state.default_pairing_ttl()
}
// ── Device registry / authorization ───────────────────────────────────────
/// Mark a Pending device Authorized and push the updated authorize set.
/// Payload-agnostic: it does not broadcast any application snapshot — the
/// consumer does that after authorizing if needed.
pub async fn authorize(&self, ed25519_pub: &[u8; 32]) -> Result<()> {
self.state.authorize(ed25519_pub).await
}
/// Revoke a device (delete keys/counters, re-push the authorize set without
/// it). Emits [`RelayEvent::ClientRevoked`].
pub async fn revoke(&self, ed25519_pub: &[u8; 32]) -> Result<()> {
self.state.revoke(ed25519_pub).await
}
/// Remove every device and push an empty authorize set. Emits one
/// `ClientRevoked` per removed device.
pub async fn clear_all(&self) -> Result<()> {
self.state.clear_all().await
}
/// All known devices (pending + authorized), ordered by `authorized_at`.
pub async fn list_clients(&self) -> Vec<ClientRow> {
self.state.list_clients().await
}
/// Persist the `device_info` JSON for a device (the consumer decodes the
/// `hello` payload and hands the raw JSON here).
pub async fn set_device_info(&self, ed25519_pub: &[u8; 32], json: &str) -> Result<()> {
self.state.set_device_info(ed25519_pub, json).await
}
// ── Identity accessors ────────────────────────────────────────────────────
pub fn agent_ed25519_pub(&self) -> [u8; 32] {
self.state.identity().ed25519_pub()
}
pub fn agent_x25519_pub(&self) -> [u8; 32] {
self.state.identity().x25519_pub()
}
pub fn namespace_id_hex(&self) -> String {
self.state.identity().namespace_id_hex().to_string()
}
pub fn is_connected(&self) -> bool {
self.state.is_connected()
}
}
+31
View File
@@ -0,0 +1,31 @@
//! Configuration for [`crate::RelayClient`].
use std::path::PathBuf;
/// Configuration snapshot passed to `RelayClient::new`.
///
/// `relay_url` may be empty: the client then stays idle (no WS loop is
/// spawned), which keeps the plugin toggleable without a relay configured.
pub struct RelayClientConfig {
/// `wss://` URL of the relay (e.g. `wss://relay.skaldagent.net/v1/ws`).
/// Empty => the client is idle.
pub relay_url: String,
/// Default pairing TTL in seconds (used when `start_pairing(0)` is called).
pub pairing_ttl: u32,
/// Where the agent's 32-byte identity seed comes from (crypto.md §9).
pub seed: SeedSource,
}
/// Source of the persistent identity seed (crypto.md §9).
///
/// `Path` preserves an existing on-disk identity: the plugin passes
/// `Path("data/relay/seed")` — the same relative path as today — so no device
/// is orphaned on upgrade and the namespace id is unchanged. `Bytes` is for
/// tests / in-memory identities.
pub enum SeedSource {
/// A raw 32-byte seed (tests, in-memory).
Bytes([u8; 32]),
/// Load (or generate + persist `0600`) the seed at the given path. The
/// parent directory is created on first use.
Path(PathBuf),
}
+305
View File
@@ -0,0 +1,305 @@
//! Persistence for authorized devices and their anti-replay counters
//! (crypto.md §9). The client does NOT open its own SQLite file: it reuses
//! Skald's shared `SqlitePool` (passed into `RelayClient::new`) and namespaces
//! its one table with the `relay_` prefix.
//!
//! Counters MUST survive restarts (crypto.md §9 "⚠️"): a `send_counter` reset
//! to 0 would reuse an AES-GCM nonce under the same key, and a `recv_counter`
//! reset would re-open the replay window. So both are columns here, not
//! in-memory.
use anyhow::Result;
use chrono::Utc;
use sqlx::{Row, SqlitePool};
/// Authorization state of a paired device.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ClientState {
/// Paired but not yet confirmed by the human (relay-protocol.md §6).
Pending,
/// Confirmed — receives Inbox snapshots and may answer.
Authorized,
}
impl ClientState {
#[allow(dead_code)] // mirrors from_str; kept for completeness/debugging
pub fn as_str(self) -> &'static str {
match self {
ClientState::Pending => "pending",
ClientState::Authorized => "authorized",
}
}
#[allow(clippy::should_implement_trait)] // small internal mapper, not the std trait
pub fn from_str(s: &str) -> ClientState {
match s {
"authorized" => ClientState::Authorized,
_ => ClientState::Pending,
}
}
}
/// One row of `relay_clients`.
///
/// `send_counter` / `authorized_at` are part of the persisted schema (read back
/// for diagnostics / future use) even though the hot paths use the dedicated
/// counter helpers.
#[allow(dead_code)]
#[derive(Debug, Clone)]
pub struct ClientRow {
pub ed25519_pub: [u8; 32],
pub x25519_pub: [u8; 32],
pub state: ClientState,
pub platform: Option<String>,
/// Raw JSON of the `device_info` object received in `hello`.
pub device_info: Option<String>,
pub send_counter: u64,
pub recv_counter: u64,
pub authorized_at: Option<i64>,
pub last_seen: Option<i64>,
}
/// Create the `relay_clients` table if missing (idempotent — called on start).
pub async fn init(pool: &SqlitePool) -> Result<()> {
sqlx::query(
"CREATE TABLE IF NOT EXISTS relay_clients (
ed25519_pub BLOB PRIMARY KEY,
x25519_pub BLOB NOT NULL,
state TEXT NOT NULL,
platform TEXT,
device_info TEXT,
send_counter INTEGER NOT NULL DEFAULT 0,
recv_counter INTEGER NOT NULL DEFAULT 0,
authorized_at INTEGER,
last_seen INTEGER
)",
)
.execute(pool)
.await?;
Ok(())
}
/// Insert (or replace) a freshly paired client with counters reset to 0 and
/// state = Pending (relay-protocol.md §6 step 7c).
pub async fn upsert_paired(
pool: &SqlitePool,
ed25519_pub: &[u8; 32],
x25519_pub: &[u8; 32],
platform: Option<&str>,
) -> Result<()> {
sqlx::query(
"INSERT INTO relay_clients
(ed25519_pub, x25519_pub, state, platform, send_counter, recv_counter)
VALUES (?, ?, 'pending', ?, 0, 0)
ON CONFLICT(ed25519_pub) DO UPDATE SET
x25519_pub = excluded.x25519_pub,
state = 'pending',
platform = excluded.platform,
send_counter = 0,
recv_counter = 0",
)
.bind(ed25519_pub.as_slice())
.bind(x25519_pub.as_slice())
.bind(platform)
.execute(pool)
.await?;
Ok(())
}
/// Mark a client Authorized, stamping `authorized_at` with the current time.
pub async fn set_authorized(pool: &SqlitePool, ed25519_pub: &[u8; 32]) -> Result<()> {
sqlx::query("UPDATE relay_clients SET state = 'authorized', authorized_at = ? WHERE ed25519_pub = ?")
.bind(Utc::now().timestamp_millis())
.bind(ed25519_pub.as_slice())
.execute(pool)
.await?;
Ok(())
}
/// Persist the device_info JSON received in a `hello` payload.
pub async fn set_device_info(pool: &SqlitePool, ed25519_pub: &[u8; 32], device_info_json: &str) -> Result<()> {
sqlx::query("UPDATE relay_clients SET device_info = ?, last_seen = ? WHERE ed25519_pub = ?")
.bind(device_info_json)
.bind(Utc::now().timestamp_millis())
.bind(ed25519_pub.as_slice())
.execute(pool)
.await?;
Ok(())
}
/// Atomically reserve the next send counter for a client and return it.
///
/// The new value is persisted BEFORE the caller seals/sends a message
/// (crypto.md §8): even if the process dies right after, the counter never
/// regresses, so no AES-GCM nonce is ever reused. Returns the counter value to
/// embed in the nonce.
///
/// A single `UPDATE … RETURNING` (not SELECT-then-UPDATE in a deferred
/// transaction): the latter starts as a reader, takes a WAL snapshot, then tries
/// to upgrade to a writer — and if another connection committed to the same row
/// meanwhile it fails with `SQLITE_BUSY_SNAPSHOT` (517), which `busy_timeout`
/// does **not** retry. Concurrent `accept_pipe`/`send` for one peer hit the same
/// row at once (e.g. a WebView opening many connections), so the snapshot upgrade
/// loses constantly. A lone `UPDATE` starts directly as a write, so callers
/// serialize on the write lock (which `busy_timeout` *does* cover).
pub async fn next_send_counter(pool: &SqlitePool, ed25519_pub: &[u8; 32]) -> Result<u64> {
let next: i64 = sqlx::query_scalar(
"UPDATE relay_clients SET send_counter = send_counter + 1 \
WHERE ed25519_pub = ? RETURNING send_counter",
)
.bind(ed25519_pub.as_slice())
.fetch_optional(pool)
.await?
.ok_or_else(|| anyhow::anyhow!("next_send_counter: client not found"))?;
Ok(next as u64)
}
/// Persist a newly-seen receive counter after a valid `open` (crypto.md §8).
pub async fn set_recv_counter(pool: &SqlitePool, ed25519_pub: &[u8; 32], counter: u64) -> Result<()> {
sqlx::query("UPDATE relay_clients SET recv_counter = ?, last_seen = ? WHERE ed25519_pub = ?")
.bind(counter as i64)
.bind(Utc::now().timestamp_millis())
.bind(ed25519_pub.as_slice())
.execute(pool)
.await?;
Ok(())
}
/// Delete a client and all its derived state (keys/counters/device_info) on
/// revoke (relay-protocol.md §7).
pub async fn delete(pool: &SqlitePool, ed25519_pub: &[u8; 32]) -> Result<()> {
sqlx::query("DELETE FROM relay_clients WHERE ed25519_pub = ?")
.bind(ed25519_pub.as_slice())
.execute(pool)
.await?;
Ok(())
}
/// Delete every client row (used by `clear_all`). Does NOT drop the table.
pub async fn delete_all(pool: &SqlitePool) -> Result<()> {
sqlx::query("DELETE FROM relay_clients")
.execute(pool)
.await?;
Ok(())
}
/// Fetch one client by pubkey.
pub async fn get(pool: &SqlitePool, ed25519_pub: &[u8; 32]) -> Result<Option<ClientRow>> {
let row = sqlx::query(
"SELECT ed25519_pub, x25519_pub, state, platform, device_info,
send_counter, recv_counter, authorized_at, last_seen
FROM relay_clients WHERE ed25519_pub = ?",
)
.bind(ed25519_pub.as_slice())
.fetch_optional(pool)
.await?;
Ok(row.map(row_to_client))
}
/// List all clients.
pub async fn list_all(pool: &SqlitePool) -> Result<Vec<ClientRow>> {
let rows = sqlx::query(
"SELECT ed25519_pub, x25519_pub, state, platform, device_info,
send_counter, recv_counter, authorized_at, last_seen
FROM relay_clients ORDER BY authorized_at",
)
.fetch_all(pool)
.await?;
Ok(rows.into_iter().map(row_to_client).collect())
}
/// Hex pubkeys of all Authorized clients — the `authorize` set sent to the relay.
pub async fn authorized_pubkeys_hex(pool: &SqlitePool) -> Result<Vec<String>> {
let rows = sqlx::query("SELECT ed25519_pub FROM relay_clients WHERE state = 'authorized'")
.fetch_all(pool)
.await?;
Ok(rows
.into_iter()
.map(|r| {
let pk: Vec<u8> = r.get("ed25519_pub");
hex::encode(pk)
})
.collect())
}
fn row_to_client(row: sqlx::sqlite::SqliteRow) -> ClientRow {
let ed: Vec<u8> = row.get("ed25519_pub");
let x: Vec<u8> = row.get("x25519_pub");
let state: String = row.get("state");
ClientRow {
ed25519_pub: to_array(&ed),
x25519_pub: to_array(&x),
state: ClientState::from_str(&state),
platform: row.get("platform"),
device_info: row.get("device_info"),
send_counter: row.get::<i64, _>("send_counter") as u64,
recv_counter: row.get::<i64, _>("recv_counter") as u64,
authorized_at: row.get("authorized_at"),
last_seen: row.get("last_seen"),
}
}
/// Convert a byte slice into a 32-byte array (zero-padded / truncated defensively).
fn to_array(bytes: &[u8]) -> [u8; 32] {
let mut out = [0u8; 32];
let n = bytes.len().min(32);
out[..n].copy_from_slice(&bytes[..n]);
out
}
#[cfg(test)]
mod tests {
use super::*;
async fn mem_pool() -> SqlitePool {
let pool = SqlitePool::connect("sqlite::memory:").await.expect("pool");
init(&pool).await.expect("init");
pool
}
#[tokio::test]
async fn next_send_counter_is_monotonic() {
let pool = mem_pool().await;
let ed = [1u8; 32];
let x = [2u8; 32];
upsert_paired(&pool, &ed, &x, None).await.expect("upsert");
let c1 = next_send_counter(&pool, &ed).await.expect("next1");
let c2 = next_send_counter(&pool, &ed).await.expect("next2");
let c3 = next_send_counter(&pool, &ed).await.expect("next3");
assert_eq!(c1, 1);
assert_eq!(c2, 2);
assert_eq!(c3, 3, "send counter must be strictly monotonic");
// The persisted value survives a fresh connection to the same DB file
// is not testable with :memory:; instead assert the in-DB value.
let row = get(&pool, &ed).await.expect("get").expect("row");
assert_eq!(row.send_counter, 3);
}
#[tokio::test]
async fn upsert_resets_counters_on_repair() {
let pool = mem_pool().await;
let ed = [3u8; 32];
upsert_paired(&pool, &ed, &[4u8; 32], None).await.expect("upsert");
next_send_counter(&pool, &ed).await.expect("bump");
next_send_counter(&pool, &ed).await.expect("bump");
// Re-pairing the same device resets counters to 0.
upsert_paired(&pool, &ed, &[5u8; 32], Some("ios")).await.expect("re-upsert");
let c = next_send_counter(&pool, &ed).await.expect("next after re-pair");
assert_eq!(c, 1, "re-pairing must reset the send counter");
}
#[tokio::test]
async fn delete_all_clears_rows() {
let pool = mem_pool().await;
upsert_paired(&pool, &[1u8; 32], &[2u8; 32], None).await.expect("upsert");
upsert_paired(&pool, &[3u8; 32], &[4u8; 32], None).await.expect("upsert");
assert_eq!(list_all(&pool).await.unwrap().len(), 2);
delete_all(&pool).await.expect("delete_all");
assert_eq!(list_all(&pool).await.unwrap().len(), 0, "delete_all must clear every row");
// Table still usable afterwards (init not required again).
upsert_paired(&pool, &[1u8; 32], &[2u8; 32], None).await.expect("upsert post-clear");
assert_eq!(list_all(&pool).await.unwrap().len(), 1);
}
}
+33
View File
@@ -0,0 +1,33 @@
//! Events broadcast by the relay client. Consumers (the plugin's
//! `RelayApp::run_event_loop`) subscribe via `RelayClient::events()`.
/// One event emitted by the relay client.
///
/// `Message.payload` is already **decrypted AND decompressed**: the client
/// peels off the v2 `version‖comp‖json` framing before emission, so the
/// consumer sees only clean bytes and applies its own JSON semantics. This is
/// the payload-agnostic boundary (see the "Crate split" section of
/// `docs/plugins/mobile-connector.md`).
#[derive(Debug, Clone)]
pub enum RelayEvent {
/// The WS handshake completed and the relay verified our namespace_id.
Connected,
/// The WS connection dropped. The client will reconnect with backoff.
Disconnected,
/// An inbound `message` from a client, decoded end-to-end. `from` is the
/// sender's ed25519 pubkey; `live` mirrors the wire flag.
Message {
from: [u8; 32],
payload: Vec<u8>,
live: bool,
},
/// A device paired (pending authorization). The consumer decides whether
/// to call `client.authorize(ed)` (auto) or to notify the human (manual).
ClientPaired {
ed25519_pub: [u8; 32],
x25519_pub: [u8; 32],
platform: String,
},
/// A device was revoked via `client.revoke` or removed by `client.clear_all`.
ClientRevoked { ed25519_pub: [u8; 32] },
}
+152
View File
@@ -0,0 +1,152 @@
//! Agent identity: the 32-byte seed and the keys derived from it.
//!
//! The seed is the only persistent secret (crypto.md §9). When sourced from a
//! path it lives on the filesystem with `0600` permissions and is generated on
//! first use. All Ed25519 / X25519 key material and the `namespace_id` are
//! derived from it at runtime via `skald-relay-common` (crypto.md §3-7), never
//! persisted — so the byte-for-byte interop with the reference vectors is
//! inherited from the shared crate.
//!
//! The seed location is supplied by the caller via [`SeedSource`] (no more
//! CWD-coupled `const`). The plugin passes `SeedSource::Path("data/relay/seed")`
//! to preserve the existing identity on upgrade.
use std::path::{Path, PathBuf};
use anyhow::{Context, Result};
use rand::RngCore;
use skald_relay_common::crypto::{self, DerivedKeys};
use crate::config::SeedSource;
/// The agent's cryptographic identity, derived from the persistent seed.
pub struct Identity {
keys: DerivedKeys,
/// `namespace_id` raw 32 bytes (used to build the AEAD AAD).
ns_raw: [u8; 32],
/// `namespace_id` lowercase hex (64 chars), used on the wire.
ns_hex: String,
}
impl Identity {
/// Build an identity from a [`SeedSource`]: for `Bytes` this is pure
/// in-memory derivation; for `Path` it loads (or generates + persists
/// `0600`) the seed at the given path, preserving any existing identity.
pub fn from_source(source: &SeedSource) -> Result<Self> {
match source {
SeedSource::Bytes(seed) => Ok(Self::from_seed(seed)),
SeedSource::Path(p) => {
let seed = load_or_create_seed(p)?;
Ok(Self::from_seed(&seed))
}
}
}
/// Build an identity from a raw seed (used by `from_source` and tests).
pub fn from_seed(seed: &[u8; 32]) -> Self {
let keys = crypto::derive_keys(seed);
let (ns_raw, ns_hex) = crypto::namespace_id(&keys.ed25519_pub);
Self { keys, ns_raw, ns_hex }
}
pub fn ed25519_pub(&self) -> [u8; 32] {
self.keys.ed25519_pub
}
pub fn x25519_pub(&self) -> [u8; 32] {
self.keys.x25519_pub
}
pub fn signing_key(&self) -> ed25519_dalek::SigningKey {
self.keys.signing_key()
}
pub fn namespace_id_raw(&self) -> [u8; 32] {
self.ns_raw
}
pub fn namespace_id_hex(&self) -> &str {
&self.ns_hex
}
/// Derive the per-client AES-256-GCM key from this agent's X25519 private key
/// and the peer's X25519 public key (crypto.md §4-5).
pub fn derive_aes_key(&self, client_x25519_pub: &[u8; 32]) -> [u8; 32] {
let shared = crypto::ecdh(&self.keys.x25519_priv, client_x25519_pub);
crypto::derive_aes_key(&shared)
}
}
/// Read the seed, or generate a fresh 32-byte CSPRNG seed and persist it `0600`.
fn load_or_create_seed(path: &Path) -> Result<[u8; 32]> {
if let Ok(bytes) = std::fs::read(path) {
if bytes.len() == 32 {
let mut seed = [0u8; 32];
seed.copy_from_slice(&bytes);
return Ok(seed);
}
anyhow::bail!(
"relay seed at {} has wrong length ({}, expected 32) — refusing to overwrite",
path.display(),
bytes.len()
);
}
// First run: create the parent dir (defaulting to "." if the path has no
// parent, e.g. a bare filename) and persist a fresh seed.
let dir = path.parent().map(Path::to_path_buf).unwrap_or_else(|| PathBuf::from("."));
std::fs::create_dir_all(&dir)
.with_context(|| format!("creating seed dir {}", dir.display()))?;
let mut seed = [0u8; 32];
rand::rng().fill_bytes(&mut seed);
write_secret_file(path, &seed)
.with_context(|| format!("writing seed file {}", path.display()))?;
tracing::info!(
crate_name = "skald-relay-client",
"generated new relay seed at {}",
path.display()
);
Ok(seed)
}
/// Write `bytes` to `path` with `0600` permissions on Unix.
fn write_secret_file(path: &Path, bytes: &[u8]) -> Result<()> {
std::fs::write(path, bytes)?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let perms = std::fs::Permissions::from_mode(0o600);
std::fs::set_permissions(path, perms)?;
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn identity_matches_reference_vectors() {
// Same seed as skald-relay-common's pinned vector (bytes 0..32).
let seed: [u8; 32] = (0u8..32).collect::<Vec<_>>().try_into().unwrap();
let id = Identity::from_seed(&seed);
assert_eq!(
hex::encode(id.ed25519_pub()),
"b3e202f4ac99fd9929da47df20adedd5b2598411a466a229f086eda3467ffa7b"
);
assert_eq!(
id.namespace_id_hex(),
"f7d340d3c3f0b0052fa904ba60ebd38a0f7e7d10672ac80648991a2c632c9e58"
);
}
#[test]
fn from_source_bytes_matches_from_seed() {
let seed: [u8; 32] = (0u8..32).collect::<Vec<_>>().try_into().unwrap();
let a = Identity::from_source(&SeedSource::Bytes(seed)).expect("bytes");
let b = Identity::from_seed(&seed);
assert_eq!(a.ed25519_pub(), b.ed25519_pub());
assert_eq!(a.namespace_id_hex(), b.namespace_id_hex());
}
}
+39
View File
@@ -0,0 +1,39 @@
//! skald-relay-client — standalone agent-role relay client.
//!
//! Payload-agnostic: it speaks the v2 WebSocket protocol (challenge/auth,
//! `Authorize`, store-and-forward `Message`), performs E2E seal/open with
//! per-client AES-256-GCM keys, manages anti-replay counters, pairing windows,
//! and device authorization. It never interprets the decrypted bytes: it emits
//! them via [`events::RelayEvent`] and lets the application layer (the
//! `plugin-mobile-connector` crate) apply JSON semantics. See the
//! "Crate split" section of `docs/plugins/mobile-connector.md` and
//! `docs/relay/`.
//!
//! Module map:
//! - `config` — `RelayClientConfig` + `SeedSource`
//! - `events` — `RelayEvent` (broadcast)
//! - `identity` — seed + derived keys + namespace_id
//! - `db` — `relay_clients` table (devices + anti-replay counters)
//! - `pairing` — in-memory pairing sessions + QR payload
//! - `state` — `RelayState` (networking-only: seal/open, counters, send/recv)
//! - `ws` — the permanent reconnecting agent WebSocket (v2 binary)
//! - `pipe` — pipe data plane: `PipeConnection` + `IncomingPipe` (pipe.md)
//! - `client` — `RelayClient`, the public façade
pub mod client;
pub mod config;
pub mod db;
pub mod events;
pub mod identity;
pub mod pairing;
pub mod pipe;
mod state;
mod ws;
pub use client::RelayClient;
pub use config::{RelayClientConfig, SeedSource};
pub use db::{ClientRow, ClientState};
pub use events::RelayEvent;
pub use identity::Identity;
pub use pairing::{QrCodeData, SessionState, StartedPairing};
pub use pipe::{IncomingPipe, PipeConnection, PipeReceiver, PipeRole, PipeSender};
+228
View File
@@ -0,0 +1,228 @@
//! In-memory pairing window (relay-protocol.md §5, §6).
//!
//! A pairing session is transient (≤ TTL) and is NEVER persisted: it holds a
//! live `pairing_token` whose bytes must not touch disk (crypto.md §9). The map
//! `code → session` is single-window, latest-wins: a new `start_pairing`
//! supersedes the previous session.
//!
//! The `code` is a random, non-enumerable handle distinct from the
//! `pairing_token`; it is what travels in the QR endpoint URL, so a URL that
//! leaks into `chat_history` is only a capability that self-revokes once the
//! window closes.
use std::collections::HashMap;
use std::sync::Mutex;
use chrono::Utc;
use rand::RngCore;
use serde::Serialize;
/// Lifecycle state of a pairing session, used to pick the QR-endpoint response.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SessionState {
/// Window open, token not yet consumed.
Active,
/// A device paired against this session (token consumed).
Consumed,
/// A newer `start_pairing` replaced this session.
Superseded,
}
/// The JSON object encoded INSIDE the QR (relay-protocol.md §5). Field
/// order/encoding is normative: all 32-byte values are lowercase hex (64 chars).
#[derive(Debug, Clone, Serialize)]
pub struct QrCodeData {
pub v: u8,
pub relay_url: String,
pub namespace_id: String,
pub agent_ed25519_pub: String,
pub agent_x25519_pub: String,
pub pairing_token: String,
}
/// One in-memory pairing session keyed by its `code`.
#[derive(Debug, Clone)]
pub struct PairingSession {
/// The QR payload to render while the window is active.
pub qr: QrCodeData,
/// Raw pairing token (kept to correlate, never serialized except in `qr`).
pub token: [u8; 32],
/// Unix-ms expiry.
pub expires_at: i64,
pub state: SessionState,
}
impl PairingSession {
/// Effective state at `now`, folding in TTL expiry on top of the stored state.
pub fn effective_state(&self, now_ms: i64) -> SessionState {
match self.state {
SessionState::Active if now_ms >= self.expires_at => SessionState::Superseded, // expired ⇒ placeholder
other => other,
}
}
}
/// Result of opening a pairing window.
pub struct StartedPairing {
pub code: String,
pub token: [u8; 32],
pub expires_at: i64,
}
/// Single-window registry of pairing sessions.
#[derive(Default)]
pub struct PairingStore {
inner: Mutex<HashMap<String, PairingSession>>,
}
impl PairingStore {
pub fn new() -> Self {
Self::default()
}
/// Open a new window (latest-wins): every existing session is marked
/// Superseded, then the new one is inserted Active. Returns the handle.
pub fn start(
&self,
relay_url: &str,
namespace_id: &str,
agent_ed25519_pub: &[u8; 32],
agent_x25519_pub: &[u8; 32],
ttl_secs: u32,
) -> StartedPairing {
let mut token = [0u8; 32];
rand::rng().fill_bytes(&mut token);
let mut code_bytes = [0u8; 16];
rand::rng().fill_bytes(&mut code_bytes);
let code = hex::encode(code_bytes);
let expires_at = Utc::now().timestamp_millis() + (ttl_secs as i64) * 1000;
let qr = QrCodeData {
v: 1,
relay_url: relay_url.to_string(),
namespace_id: namespace_id.to_string(),
agent_ed25519_pub: hex::encode(agent_ed25519_pub),
agent_x25519_pub: hex::encode(agent_x25519_pub),
pairing_token: hex::encode(token),
};
let mut map = self.inner.lock().unwrap();
for s in map.values_mut() {
if s.state == SessionState::Active {
s.state = SessionState::Superseded;
}
}
map.insert(
code.clone(),
PairingSession { qr, token, expires_at, state: SessionState::Active },
);
StartedPairing { code, token, expires_at }
}
/// Mark every active session Superseded (used by `stop_pairing`).
pub fn supersede_all(&self) {
let mut map = self.inner.lock().unwrap();
for s in map.values_mut() {
if s.state == SessionState::Active {
s.state = SessionState::Superseded;
}
}
}
/// Mark the active session whose token matches as Consumed (after a device
/// pairs). Returns true if one was found.
pub fn consume_by_token(&self, token: &[u8; 32]) -> bool {
let mut map = self.inner.lock().unwrap();
for s in map.values_mut() {
if s.state == SessionState::Active
&& skald_relay_common::crypto::ct_eq(&s.token, token)
{
s.state = SessionState::Consumed;
return true;
}
}
false
}
/// The token of the single currently-active session, if any. The relay echoes
/// no token in `client_paired`, so the active session's token is what we
/// consume on pairing.
pub fn active_token(&self) -> Option<[u8; 32]> {
let now = Utc::now().timestamp_millis();
let map = self.inner.lock().unwrap();
map.values()
.find(|s| s.effective_state(now) == SessionState::Active)
.map(|s| s.token)
}
/// Look up a session by code, returning its QR payload and effective state.
pub fn lookup(&self, code: &str) -> Option<(QrCodeData, SessionState)> {
let now = Utc::now().timestamp_millis();
let map = self.inner.lock().unwrap();
map.get(code).map(|s| (s.qr.clone(), s.effective_state(now)))
}
}
#[cfg(test)]
mod tests {
use super::*;
fn store_with_active() -> (PairingStore, StartedPairing) {
let s = PairingStore::new();
let started = s.start("wss://relay", "ns", &[1u8; 32], &[2u8; 32], 300);
(s, started)
}
#[test]
fn single_window_latest_wins() {
let (s, first) = store_with_active();
let _second = s.start("wss://relay", "ns", &[1u8; 32], &[2u8; 32], 300);
let (_, state) = s.lookup(&first.code).expect("first session present");
assert_eq!(state, SessionState::Superseded, "opening a new window supersedes the old one");
}
#[test]
fn consume_by_token_marks_consumed() {
let (s, started) = store_with_active();
let token = started.token;
let (_, state) = s.lookup(&started.code).expect("present");
assert_eq!(state, SessionState::Active);
assert!(s.consume_by_token(&token));
let (_, state) = s.lookup(&started.code).expect("present");
assert_eq!(state, SessionState::Consumed);
}
#[test]
fn consume_with_wrong_token_is_noop() {
let (s, _started) = store_with_active();
assert!(!s.consume_by_token(&[0u8; 32]));
}
#[test]
fn stop_pairing_supersedes_active() {
let (s, started) = store_with_active();
s.supersede_all();
let (_, state) = s.lookup(&started.code).expect("present");
assert_eq!(state, SessionState::Superseded);
}
#[test]
fn lookup_unknown_code_is_none() {
let s = PairingStore::new();
assert!(s.lookup("nope").is_none());
}
#[test]
fn expired_active_session_reports_superseded() {
let (s, started) = store_with_active();
// Construct a session already in the past by inserting manually.
let mut session = s.inner.lock().unwrap();
let entry = session.get_mut(&started.code).expect("present");
entry.expires_at = Utc::now().timestamp_millis() - 1;
drop(session);
let (_, state) = s.lookup(&started.code).expect("present");
assert_eq!(state, SessionState::Superseded, "expiry folds into Superseded");
}
}
+537
View File
@@ -0,0 +1,537 @@
//! Pipe data-plane client (docs/relay/pipe.md §2-3): the [`PipeConnection`]
//! secure byte channel over a `/v1/pipe` WebSocket.
//!
//! The control plane (invite/accept signaling, ephemeral DH) lives in
//! [`crate::state`]; by the time `PipeConnection::connect` runs, both peers have
//! derived the same per-pipe `pipe_key`. This module only does the data plane:
//! dial `/v1/pipe`, prove identity to the relay (`pipe_auth`), then seal/open
//! every frame with AES-256-GCM keyed by `pipe_key`, using a per-direction
//! counter nonce (the relay forwards opaque ciphertext, pipe.md §2.2).
use std::collections::BTreeMap;
use std::sync::Arc;
use anyhow::{anyhow, Result};
use futures_util::stream::{SplitSink, SplitStream};
use futures_util::{SinkExt, StreamExt};
use skald_relay_common::crypto;
use skald_relay_common::pipe::{self, PipeAuth, PipeChallenge, PipeSuite};
use tokio::sync::{mpsc, OwnedSemaphorePermit, Semaphore};
use tokio::task::JoinHandle;
use tokio_tungstenite::tungstenite::Message as WsMessage;
use tokio_util::sync::CancellationToken;
/// An inbound pipe invite surfaced to the application (responder side). The app
/// inspects `from` / `stream_type` / `headers` and then calls
/// `RelayClient::accept_pipe` or `reject_pipe`. The remaining fields carry the
/// handshake state the accept path needs.
#[derive(Debug, Clone)]
pub struct IncomingPipe {
/// The initiator's ed25519 pubkey.
pub from: [u8; 32],
/// App-defined purpose discriminator.
pub stream_type: String,
/// Arbitrary app-defined headers from the invite.
pub headers: BTreeMap<String, String>,
/// Rendezvous key (echoed in the accept + data-plane auth).
pub(crate) connection_id: [u8; 32],
/// Negotiated suite (v1: only `X25519Sealed`).
pub(crate) suite: PipeSuite,
/// The initiator's opaque handshake material (its ephemeral X25519 pubkey).
pub(crate) peer_handshake: Vec<u8>,
}
type ClientWs =
tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
type ClientSink = SplitSink<ClientWs, WsMessage>;
type ClientStream = SplitStream<ClientWs>;
/// Soft cap on **unflushed** outbound bytes buffered inside the pipe before
/// [`PipeSender::send`] blocks (backpressure). The background writer task
/// releases each reservation once the corresponding frame is flushed to the
/// socket, so this bounds in-flight memory while letting `send` and `recv`
/// proceed independently. ~10 MiB.
const SEND_BUFFER_BYTES: usize = 10 * 1024 * 1024;
/// Which end of the pipe this peer is — selects the send/receive nonce
/// directions so the two AES-GCM streams never collide.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PipeRole {
/// Sent `pipe_invite`. Sends on the INITIATOR direction.
Initiator,
/// Replied with `pipe_accept`. Sends on the RESPONDER direction.
Responder,
}
/// Write half of a [`PipeConnection`]: seals plaintext and queues it for the
/// background writer task that owns the socket sink. Single-writer (`&mut self`)
/// so the per-direction counter nonce stays strictly ordered on the wire.
pub struct PipeSender {
/// Sealed frames + their byte reservation, drained by the writer task in
/// FIFO order (preserving counter order). Count-unbounded but byte-bounded
/// by `buffer`.
data_tx: mpsc::UnboundedSender<(Vec<u8>, OwnedSemaphorePermit)>,
/// ~[`SEND_BUFFER_BYTES`] of permits; `send` blocks when exhausted.
buffer: Arc<Semaphore>,
key: [u8; 32],
send_dir: [u8; 4],
send_ctr: u64,
/// AAD binding every frame to the rendezvous (the 32-byte connection_id).
aad: [u8; 32],
}
/// Read half of a [`PipeConnection`]: reads sealed frames off the socket stream
/// and opens them. WS-level pings are answered via the shared writer task.
pub struct PipeReceiver {
stream: ClientStream,
/// Forwards `Pong` replies to the writer task (which owns the sink).
ctrl_tx: mpsc::UnboundedSender<WsMessage>,
key: [u8; 32],
recv_dir: [u8; 4],
recv_ctr: u64,
aad: [u8; 32],
}
/// An end-to-end-encrypted byte channel to a namespace peer, relayed through
/// `/v1/pipe`. The relay never sees plaintext. Full-duplex: a background writer
/// task owns the socket sink and drains a byte-bounded buffer, so `send` and
/// `recv` never block each other at the socket. Use the unified `send`/`recv`
/// on `&mut self`, or [`split`](Self::split) into independent halves to drive
/// the two directions from separate tasks.
pub struct PipeConnection {
sender: PipeSender,
receiver: PipeReceiver,
/// Cancels the writer task on explicit [`close`](Self::close).
cancel: CancellationToken,
writer: JoinHandle<()>,
}
impl PipeConnection {
/// Dial `/v1/pipe`, complete the relay auth handshake, and return the ready
/// channel. `pipe_key` must already be derived from the signaling ephemeral
/// DH (same value on both peers).
#[allow(clippy::too_many_arguments)]
pub(crate) async fn connect(
relay_url: &str,
signing_key: &ed25519_dalek::SigningKey,
my_ed_pub: &[u8; 32],
peer_ed_pub: &[u8; 32],
namespace_id_raw: &[u8; 32],
connection_id: &[u8; 32],
pipe_key: &[u8; 32],
role: PipeRole,
) -> Result<PipeConnection> {
let url = pipe_url(relay_url);
let (mut ws, _) = tokio_tungstenite::connect_async(&url).await?;
// Relay speaks first: PipeChallenge.
let nonce = read_challenge(&mut ws).await?;
// Reply with a signature over PIPE_AUTH_DOMAIN ‖ 0x00 ‖ nonce ‖ cid.
let sig = crypto::sign_pipe_auth(signing_key, &nonce, connection_id);
let auth = PipeAuth {
connection_id: connection_id.to_vec(),
pubkey: my_ed_pub.to_vec(),
dest: crypto::sha256(peer_ed_pub).to_vec(),
namespace_id: namespace_id_raw.to_vec(),
signature: sig.to_vec(),
};
ws.send(WsMessage::Binary(pipe::encode(&auth).into())).await?;
let (send_dir, recv_dir) = match role {
PipeRole::Initiator => (crypto::DIR_PIPE_INITIATOR, crypto::DIR_PIPE_RESPONDER),
PipeRole::Responder => (crypto::DIR_PIPE_RESPONDER, crypto::DIR_PIPE_INITIATOR),
};
// Split the socket so the two directions are independent: the writer task
// owns the sink and drains the byte-bounded buffer; the read half owns
// the stream. Counters start at 1 (pipe.md §4).
let (sink, stream) = ws.split();
let buffer = Arc::new(Semaphore::new(SEND_BUFFER_BYTES));
let (data_tx, data_rx) = mpsc::unbounded_channel::<(Vec<u8>, OwnedSemaphorePermit)>();
let (ctrl_tx, ctrl_rx) = mpsc::unbounded_channel::<WsMessage>();
let cancel = CancellationToken::new();
let writer = tokio::spawn(writer_loop(sink, data_rx, ctrl_rx, cancel.clone()));
Ok(PipeConnection {
sender: PipeSender {
data_tx,
buffer,
key: *pipe_key,
send_dir,
send_ctr: 1,
aad: *connection_id,
},
receiver: PipeReceiver {
stream,
ctrl_tx,
key: *pipe_key,
recv_dir,
recv_ctr: 1,
aad: *connection_id,
},
cancel,
writer,
})
}
/// Seal and queue one application chunk (delegates to the write half).
pub async fn send(&mut self, plaintext: &[u8]) -> Result<()> {
self.sender.send(plaintext).await
}
/// Receive and open the next application chunk (delegates to the read half).
/// `Ok(None)` on a clean close.
pub async fn recv(&mut self) -> Result<Option<Vec<u8>>> {
self.receiver.recv().await
}
/// Split into independent write/read halves for full-duplex use: move each
/// into its own task and the two directions run concurrently. Dropping
/// **both** halves tears the pipe down (the writer task closes the socket
/// once both of its channels are gone).
pub fn split(self) -> (PipeSender, PipeReceiver) {
// Detach the writer (drop its JoinHandle); teardown is drop-driven for
// this path. `cancel` is dropped without firing — a dropped token does
// not cancel — so the writer lives until both halves drop.
(self.sender, self.receiver)
}
/// Cancel the writer task and close the underlying socket.
pub async fn close(self) {
self.cancel.cancel();
let _ = self.writer.await;
}
}
impl PipeSender {
/// Seal and queue one application chunk. Returns once the frame is buffered
/// (or, when the ~[`SEND_BUFFER_BYTES`] buffer is full, once space frees up)
/// — **not** once it is flushed. The 12-byte nonce is implicit
/// (per-direction counter), so it is not transmitted.
pub async fn send(&mut self, plaintext: &[u8]) -> Result<()> {
let nonce = crypto::build_nonce(self.send_dir, self.send_ctr);
let sealed = crypto::seal(&self.key, &nonce, &self.aad, plaintext)
.map_err(|e| anyhow!("pipe seal failed: {e}"))?;
self.send_ctr += 1;
// Reserve the frame's bytes; block here when the buffer is full. Clamp so
// a single frame larger than the whole buffer can never deadlock.
let want = sealed.len().min(SEND_BUFFER_BYTES) as u32;
let permit = Arc::clone(&self.buffer)
.acquire_many_owned(want)
.await
.map_err(|_| anyhow!("pipe send buffer closed"))?;
self.data_tx
.send((sealed, permit))
.map_err(|_| anyhow!("pipe writer stopped"))?;
Ok(())
}
}
impl PipeReceiver {
/// Receive and open the next application chunk. `Ok(None)` on a clean close.
/// WS-level pings are answered transparently (forwarded to the writer task).
pub async fn recv(&mut self) -> Result<Option<Vec<u8>>> {
loop {
let Some(msg) = self.stream.next().await else { return Ok(None) };
match msg? {
WsMessage::Binary(data) => {
let nonce = crypto::build_nonce(self.recv_dir, self.recv_ctr);
let pt = crypto::open(&self.key, &nonce, &self.aad, &data)
.map_err(|_| anyhow!("pipe open failed (tag mismatch / desync)"))?;
self.recv_ctr += 1;
return Ok(Some(pt));
}
// The writer owns the sink; hand it the Pong. A send error means
// the pipe is tearing down anyway — ignore it.
WsMessage::Ping(p) => {
let _ = self.ctrl_tx.send(WsMessage::Pong(p));
}
WsMessage::Pong(_) => {}
WsMessage::Close(_) => return Ok(None),
WsMessage::Text(_) | WsMessage::Frame(_) => {} // pipe is binary-only
}
}
}
}
/// Background writer: owns the socket sink and flushes queued frames. Control
/// frames (Pong) are prioritized (`biased`) over data so keep-alives are never
/// starved by a full data buffer. Each data permit is released **after** the
/// frame is flushed, so the buffer measures unflushed bytes. Exits on `cancel`
/// or once both channels close (both halves dropped), then closes the socket.
async fn writer_loop(
mut sink: ClientSink,
mut data_rx: mpsc::UnboundedReceiver<(Vec<u8>, OwnedSemaphorePermit)>,
mut ctrl_rx: mpsc::UnboundedReceiver<WsMessage>,
cancel: CancellationToken,
) {
let mut data_open = true;
let mut ctrl_open = true;
loop {
if !data_open && !ctrl_open {
break;
}
tokio::select! {
biased;
_ = cancel.cancelled() => break,
msg = ctrl_rx.recv(), if ctrl_open => match msg {
Some(m) => {
if sink.send(m).await.is_err() {
break;
}
}
None => ctrl_open = false,
},
msg = data_rx.recv(), if data_open => match msg {
Some((bytes, permit)) => {
let r = sink.send(WsMessage::Binary(bytes.into())).await;
drop(permit); // release the reservation after the flush
if r.is_err() {
break;
}
}
None => data_open = false,
},
}
}
let _ = sink.send(WsMessage::Close(None)).await;
let _ = sink.close().await;
}
/// Derive the data-plane URL from the control URL by swapping the path
/// `/v1/ws` → `/v1/pipe` (config stores the full control-plane URL).
fn pipe_url(relay_url: &str) -> String {
if relay_url.contains("/v1/ws") {
relay_url.replace("/v1/ws", "/v1/pipe")
} else {
format!("{}/v1/pipe", relay_url.trim_end_matches('/'))
}
}
/// Read frames until the relay's [`PipeChallenge`]; return the 32-byte nonce.
async fn read_challenge(ws: &mut ClientWs) -> Result<[u8; 32]> {
while let Some(msg) = ws.next().await {
match msg? {
WsMessage::Binary(data) => {
let c: PipeChallenge = pipe::decode(&data)
.map_err(|e| anyhow!("malformed pipe challenge: {e}"))?;
return pipe::to_array::<32>(&c.nonce)
.ok_or_else(|| anyhow!("pipe challenge nonce is not 32 bytes"));
}
WsMessage::Ping(p) => ws.send(WsMessage::Pong(p)).await?,
WsMessage::Close(_) => return Err(anyhow!("relay closed before pipe challenge")),
_ => {}
}
}
Err(anyhow!("connection closed before pipe challenge"))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn pipe_url_swaps_path() {
assert_eq!(pipe_url("wss://r.example/v1/ws"), "wss://r.example/v1/pipe");
assert_eq!(pipe_url("ws://127.0.0.1:8080/v1/ws"), "ws://127.0.0.1:8080/v1/pipe");
assert_eq!(pipe_url("wss://r.example"), "wss://r.example/v1/pipe");
}
}
/// Data-plane E2E against the **real** relay server (booted in-process): two
/// `PipeConnection`s (initiator + responder, sharing a pre-derived key) dial
/// `/v1/pipe`, get matched, and stream sealed bytes both ways through a relay
/// that only ever sees ciphertext.
#[cfg(test)]
mod net_tests {
use super::*;
use std::net::SocketAddr;
use std::time::Duration;
use skald_relay_server::config::{Config, PipeConfig};
use skald_relay_server::{router, AppState};
async fn spawn_relay() -> (SocketAddr, AppState) {
use std::sync::atomic::{AtomicU64, Ordering};
static C: AtomicU64 = AtomicU64::new(0);
let n = C.fetch_add(1, Ordering::Relaxed);
let db = std::env::temp_dir().join(format!("relay-pipe-cli-{}-{n}.db", std::process::id()));
let cfg = Config {
bind: "127.0.0.1:0".parse().unwrap(),
db_path: db.to_string_lossy().into(),
pipe: PipeConfig::default(),
};
let state = AppState::build(cfg).await.unwrap();
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let serve = state.clone();
tokio::spawn(async move {
axum::serve(listener, router(serve).into_make_service_with_connect_info::<SocketAddr>())
.await
.unwrap();
});
(addr, state)
}
fn id(seed: u8) -> (ed25519_dalek::SigningKey, [u8; 32]) {
let dk = crypto::derive_keys(&[seed; 32]);
(dk.signing_key(), dk.ed25519_pub)
}
#[tokio::test]
async fn pipe_connection_streams_through_real_relay() {
let (addr, state) = spawn_relay().await;
let (agent_sk, agent_ed) = id(0xA1);
let (client_sk, client_ed) = id(0xB2);
// Seed: agent owns the namespace, client is authorized in it.
let (ns_raw, ns_hex) = crypto::namespace_id(&agent_ed);
state.store.upsert_namespace(&ns_hex, &agent_ed).await.unwrap();
let cx = crypto::derive_keys(&[0xC3; 32]).x25519_pub;
state.store.upsert_pending_client(&ns_hex, &client_ed, &cx, "", "ios").await.unwrap();
state.store.apply_authorize(&ns_hex, &[client_ed]).await.unwrap();
// A pre-shared per-pipe key (in production: ephemeral DH from signaling).
let key = crypto::derive_pipe_key(&[0x07; 32]);
let cid = [0x9C; 32];
let url = format!("ws://{addr}/v1/ws");
let mut a = PipeConnection::connect(
&url, &agent_sk, &agent_ed, &client_ed, &ns_raw, &cid, &key, PipeRole::Initiator,
)
.await
.expect("initiator connect");
tokio::time::sleep(Duration::from_millis(50)).await; // A pending before B
let mut b = PipeConnection::connect(
&url, &client_sk, &client_ed, &agent_ed, &ns_raw, &cid, &key, PipeRole::Responder,
)
.await
.expect("responder connect");
// Bytes both ways.
a.send(b"ping").await.unwrap();
assert_eq!(b.recv().await.unwrap().as_deref(), Some(&b"ping"[..]));
b.send(b"pong").await.unwrap();
assert_eq!(a.recv().await.unwrap().as_deref(), Some(&b"pong"[..]));
// A larger blob round-trips intact (seal/open + relay splice).
let blob = vec![0x5A_u8; 200_000];
a.send(&blob).await.unwrap();
assert_eq!(b.recv().await.unwrap(), Some(blob));
// Closing one tears down the other.
a.close().await;
assert_eq!(b.recv().await.unwrap(), None);
}
#[tokio::test]
async fn pipe_wrong_key_fails_to_open() {
let (addr, state) = spawn_relay().await;
let (agent_sk, agent_ed) = id(0xD4);
let (client_sk, client_ed) = id(0xE5);
let (ns_raw, ns_hex) = crypto::namespace_id(&agent_ed);
state.store.upsert_namespace(&ns_hex, &agent_ed).await.unwrap();
let cx = crypto::derive_keys(&[0xF6; 32]).x25519_pub;
state.store.upsert_pending_client(&ns_hex, &client_ed, &cx, "", "ios").await.unwrap();
state.store.apply_authorize(&ns_hex, &[client_ed]).await.unwrap();
let cid = [0x1A; 32];
let url = format!("ws://{addr}/v1/ws");
// Mismatched keys: the relay still splices, but `open` must fail (the
// relay never had the plaintext — confidentiality holds end to end).
let ka = crypto::derive_pipe_key(&[0x01; 32]);
let kb = crypto::derive_pipe_key(&[0x02; 32]);
let mut a = PipeConnection::connect(
&url, &agent_sk, &agent_ed, &client_ed, &ns_raw, &cid, &ka, PipeRole::Initiator,
)
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
let mut b = PipeConnection::connect(
&url, &client_sk, &client_ed, &agent_ed, &ns_raw, &cid, &kb, PipeRole::Responder,
)
.await
.unwrap();
a.send(b"secret").await.unwrap();
assert!(b.recv().await.is_err(), "wrong key must fail AEAD open");
}
/// Full-duplex: `split` both ends, then stream a large blob **both ways at
/// once** — each side sends while it receives. Proves send and recv run
/// concurrently (neither direction blocks the other) through the real relay.
#[tokio::test]
async fn pipe_split_streams_both_directions_concurrently() {
let (addr, state) = spawn_relay().await;
let (agent_sk, agent_ed) = id(0x11);
let (client_sk, client_ed) = id(0x22);
let (ns_raw, ns_hex) = crypto::namespace_id(&agent_ed);
state.store.upsert_namespace(&ns_hex, &agent_ed).await.unwrap();
let cx = crypto::derive_keys(&[0x33; 32]).x25519_pub;
state.store.upsert_pending_client(&ns_hex, &client_ed, &cx, "", "ios").await.unwrap();
state.store.apply_authorize(&ns_hex, &[client_ed]).await.unwrap();
let key = crypto::derive_pipe_key(&[0x44; 32]);
let cid = [0x55; 32];
let url = format!("ws://{addr}/v1/ws");
let a = PipeConnection::connect(
&url, &agent_sk, &agent_ed, &client_ed, &ns_raw, &cid, &key, PipeRole::Initiator,
)
.await
.expect("initiator connect");
tokio::time::sleep(Duration::from_millis(50)).await; // A pending before B
let b = PipeConnection::connect(
&url, &client_sk, &client_ed, &agent_ed, &ns_raw, &cid, &key, PipeRole::Responder,
)
.await
.expect("responder connect");
let (mut a_tx, mut a_rx) = a.split();
let (mut b_tx, mut b_rx) = b.split();
const CHUNKS: usize = 64;
const CHUNK: usize = 16 * 1024; // 64 × 16 KiB = 1 MiB each way
const TOTAL: usize = CHUNKS * CHUNK;
// Both directions send at the same time…
let a_send = tokio::spawn(async move {
for i in 0..CHUNKS {
a_tx.send(&vec![i as u8; CHUNK]).await.unwrap();
}
});
let b_send = tokio::spawn(async move {
for i in 0..CHUNKS {
b_tx.send(&vec![0xFF - i as u8; CHUNK]).await.unwrap();
}
});
// …while both directions receive concurrently.
let a_recv = tokio::spawn(async move {
let mut got = 0usize;
while got < TOTAL {
match a_rx.recv().await.unwrap() {
Some(p) => got += p.len(),
None => break,
}
}
got
});
let b_recv = tokio::spawn(async move {
let mut got = 0usize;
while got < TOTAL {
match b_rx.recv().await.unwrap() {
Some(p) => got += p.len(),
None => break,
}
}
got
});
a_send.await.unwrap();
b_send.await.unwrap();
assert_eq!(a_recv.await.unwrap(), TOTAL, "A must receive all of B's bytes");
assert_eq!(b_recv.await.unwrap(), TOTAL, "B must receive all of A's bytes");
}
}
+673
View File
@@ -0,0 +1,673 @@
//! Networking-only shared state, owned behind an `Arc` and shared by the WS
//! loop, the pairing/authorization surface, and the QR lookup. Everything here
//! is transport + crypto + the device registry: there is **no** knowledge of
//! what the decrypted bytes mean (the payload-agnostic boundary). Decoded
//! inbound bytes and lifecycle transitions are surfaced via [`RelayEvent`].
//!
//! The wire transport is **v2 protobuf** (docs/relay/relay-protocol.md): every
//! frame queued onto the WS outbound channel is the
//! `prost::Message::encode_to_vec()` of a `RelayFrame`. E2E plaintexts are
//! wrapped in the v2 framing (`compress_payload`) before sealing, and peeled
//! (`decompress_payload`) before being emitted, so consumers only ever see the
//! clean inner payload.
use std::collections::{BTreeMap, HashMap};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use anyhow::{anyhow, Result};
use prost::Message as _;
use rand::RngCore;
use skald_relay_common::crypto::{self, DIR_AGENT_TO_CLIENT, DIR_CLIENT_TO_AGENT};
use skald_relay_common::pipe::{PipeAccept, PipeInvite, PipeReject, PipeSignal, PipeSuite, to_array};
use skald_relay_common::proto::v2::*;
use skald_relay_common::proto::v2::relay_frame::Frame;
use sqlx::SqlitePool;
use tokio::sync::{broadcast, mpsc, oneshot};
use tracing::{debug, warn};
use crate::db::{self, ClientRow, ClientState};
use crate::events::RelayEvent;
use crate::identity::Identity;
use crate::pairing::{PairingStore, QrCodeData, SessionState, StartedPairing};
use crate::pipe::{IncomingPipe, PipeConnection, PipeRole};
/// How many inbound pipe invites the broadcast buffers before lagging.
const INCOMING_PIPE_CHANNEL_CAP: usize = 64;
/// How long `open_pipe` waits for a `pipe_accept` before giving up.
const PIPE_ACCEPT_TIMEOUT: Duration = Duration::from_secs(30);
/// Networking config snapshot the runtime needs.
pub(crate) struct StateConfig {
pub relay_url: String,
pub pairing_ttl: u32,
}
/// Everything the runloop and surfaces share. Payload-agnostic.
pub(crate) struct RelayState {
identity: Identity,
db: Arc<SqlitePool>,
pairing: PairingStore,
config: StateConfig,
/// Sender into the WS outbound queue. `None` until the loop is started.
/// Carries **encoded protobuf bytes** ready to be wrapped in
/// `Message::Binary` by the WS layer (v2 transport).
outbound: Mutex<Option<mpsc::UnboundedSender<Vec<u8>>>>,
/// Cache of per-client aes_key, keyed by ed25519 pubkey (crypto.md §8).
/// Derived from the seed + the client's x25519 pubkey; never persisted.
aes_cache: Mutex<HashMap<[u8; 32], [u8; 32]>>,
connected: AtomicBool,
/// Broadcast sink for [`RelayEvent`]s consumed by the application layer.
events_tx: broadcast::Sender<RelayEvent>,
/// Pending `open_pipe` waiters: connection_id → accept/reject delivery
/// (docs/relay/pipe.md §1). The initiator parks here until the peer replies.
pipe_waiters: Mutex<HashMap<[u8; 32], oneshot::Sender<Result<PipeAccept, String>>>>,
/// Broadcast of inbound `pipe_invite`s (responder side). The consumer calls
/// `accept_pipe`/`reject_pipe`. Single-consumer expected.
incoming_pipes_tx: broadcast::Sender<IncomingPipe>,
}
impl RelayState {
pub(crate) fn new(
identity: Identity,
db: Arc<SqlitePool>,
config: StateConfig,
events_tx: broadcast::Sender<RelayEvent>,
) -> Self {
let (incoming_pipes_tx, _) = broadcast::channel(INCOMING_PIPE_CHANNEL_CAP);
Self {
identity,
db,
pairing: PairingStore::new(),
config,
outbound: Mutex::new(None),
aes_cache: Mutex::new(HashMap::new()),
connected: AtomicBool::new(false),
events_tx,
pipe_waiters: Mutex::new(HashMap::new()),
incoming_pipes_tx,
}
}
// ── Accessors ─────────────────────────────────────────────────────────────
pub(crate) fn identity(&self) -> &Identity {
&self.identity
}
pub(crate) fn relay_url(&self) -> String {
self.config.relay_url.clone()
}
pub(crate) fn default_pairing_ttl(&self) -> u32 {
self.config.pairing_ttl
}
/// Emit a [`RelayEvent`]; ignores the "no subscribers" case.
pub(crate) fn emit(&self, ev: RelayEvent) {
let _ = self.events_tx.send(ev);
}
pub(crate) fn subscribe(&self) -> broadcast::Receiver<RelayEvent> {
self.events_tx.subscribe()
}
pub(crate) fn set_connected(&self, v: bool) {
let was = self.connected.swap(v, Ordering::Relaxed);
if was != v {
self.emit(if v { RelayEvent::Connected } else { RelayEvent::Disconnected });
}
}
pub(crate) fn is_connected(&self) -> bool {
self.connected.load(Ordering::Relaxed)
}
pub(crate) fn set_outbound(&self, tx: mpsc::UnboundedSender<Vec<u8>>) {
*self.outbound.lock().unwrap() = Some(tx);
}
pub(crate) fn clear_outbound(&self) {
*self.outbound.lock().unwrap() = None;
}
/// Queue an already-encoded `RelayFrame` onto the WS outbound channel.
fn send_frame(&self, bytes: Vec<u8>) -> Result<()> {
let guard = self.outbound.lock().unwrap();
match guard.as_ref() {
Some(tx) => tx
.send(bytes)
.map_err(|_| anyhow::anyhow!("WS outbound channel closed")),
None => Err(anyhow::anyhow!("WS not started")),
}
}
pub(crate) async fn authorized_pubkeys_hex(&self) -> Result<Vec<String>> {
db::authorized_pubkeys_hex(&self.db).await
}
/// Re-send the full authorize set (replacement semantics,
/// relay-protocol.md §7). v2: each client pubkey travels as a raw 32-byte
/// `bytes` field.
async fn send_authorize(&self) -> Result<()> {
let clients_hex = db::authorized_pubkeys_hex(&self.db).await?;
let clients: Vec<prost::bytes::Bytes> = clients_hex
.iter()
.filter_map(|h| hex::decode(h).ok())
.map(prost::bytes::Bytes::from)
.collect();
let frame = RelayFrame {
frame: Some(Frame::Authorize(Authorize { clients })),
};
self.send_frame(frame.encode_to_vec())
}
// ── Pairing ───────────────────────────────────────────────────────────────
/// Open a pairing window: generate a token, send `pairing_start`, register
/// the in-memory session (latest-wins). Returns the handle for the QR URL.
pub(crate) async fn start_pairing(&self, ttl_secs: u32) -> Result<StartedPairing> {
let started = self.pairing.start(
&self.config.relay_url,
self.identity.namespace_id_hex(),
&self.identity.ed25519_pub(),
&self.identity.x25519_pub(),
ttl_secs,
);
let frame = RelayFrame {
frame: Some(Frame::PairingStart(PairingStart {
pairing_token: prost::bytes::Bytes::copy_from_slice(&started.token),
ttl: ttl_secs,
})),
};
self.send_frame(frame.encode_to_vec())?;
debug!(crate_name = "skald-relay-client", ttl_secs, "pairing window opened");
Ok(started)
}
/// Close the pairing window locally and tell the relay.
pub(crate) async fn stop_pairing(&self) -> Result<()> {
self.pairing.supersede_all();
let frame = RelayFrame {
frame: Some(Frame::PairingStop(PairingStop {})),
};
self.send_frame(frame.encode_to_vec())
}
/// Look up a pairing session for the QR endpoint.
pub(crate) fn lookup_pairing(&self, code: &str) -> Option<(QrCodeData, SessionState)> {
self.pairing.lookup(code)
}
/// Handle `client_paired` (relay-protocol.md §6 step 7): derive aes_key,
/// persist the client as Pending, consume the pairing session, then emit
/// [`RelayEvent::ClientPaired`]. The **authorization policy is the
/// consumer's** — this layer never auto-authorizes.
pub(crate) async fn handle_client_paired(
&self,
client_ed25519_pub: &[u8; 32],
client_x25519_pub: &[u8; 32],
platform: &str,
) {
let ed = *client_ed25519_pub;
let x = *client_x25519_pub;
// Derive + cache the per-client aes_key.
let aes_key = self.identity.derive_aes_key(&x);
self.aes_cache.lock().unwrap().insert(ed, aes_key);
// Persist as Pending with counters at 0.
if let Err(e) = db::upsert_paired(&self.db, &ed, &x, Some(platform)).await {
warn!(crate_name = "skald-relay-client", error = %e, "failed to persist paired client");
return;
}
// Mark the active pairing session as consumed.
if let Some(tok) = self.pairing.active_token() {
self.pairing.consume_by_token(&tok);
}
self.emit(RelayEvent::ClientPaired {
ed25519_pub: ed,
x25519_pub: x,
platform: platform.to_string(),
});
}
/// Mark a client Authorized and push the updated authorize set. Does NOT
/// broadcast any application payload — that is the consumer's job after
/// authorizing (the client is payload-agnostic).
pub(crate) async fn authorize(&self, ed25519_pub: &[u8; 32]) -> Result<()> {
db::set_authorized(&self.db, ed25519_pub).await?;
self.send_authorize().await?;
debug!(crate_name = "skald-relay-client", device = %hex::encode(ed25519_pub), "device authorized");
Ok(())
}
/// Revoke a client (relay-protocol.md §7): drop from the set, re-authorize
/// without it, delete its keys/counters/device_info, emit `ClientRevoked`.
pub(crate) async fn revoke(&self, ed25519_pub: &[u8; 32]) -> Result<()> {
db::delete(&self.db, ed25519_pub).await?;
self.aes_cache.lock().unwrap().remove(ed25519_pub);
self.send_authorize().await?;
debug!(crate_name = "skald-relay-client", device = %hex::encode(ed25519_pub), "device revoked");
self.emit(RelayEvent::ClientRevoked { ed25519_pub: *ed25519_pub });
Ok(())
}
/// Remove every device, clear the aes cache, and push an empty authorize
/// set. Emits one `ClientRevoked` per removed device.
pub(crate) async fn clear_all(&self) -> Result<()> {
let removed = db::list_all(&self.db).await.unwrap_or_default();
db::delete_all(&self.db).await?;
self.aes_cache.lock().unwrap().clear();
self.send_authorize().await?;
for c in removed {
self.emit(RelayEvent::ClientRevoked { ed25519_pub: c.ed25519_pub });
}
Ok(())
}
/// Persist the device_info JSON for a client (from a `hello` payload, decoded
/// by the consumer).
pub(crate) async fn set_device_info(&self, ed25519_pub: &[u8; 32], json: &str) -> Result<()> {
db::set_device_info(&self.db, ed25519_pub, json).await
}
pub(crate) async fn list_clients(&self) -> Vec<ClientRow> {
db::list_all(&self.db).await.unwrap_or_default()
}
// ── E2E: aes_key cache ────────────────────────────────────────────────────
/// Resolve (and cache) the aes_key for a client, deriving from the stored
/// x25519 pubkey on a cache miss.
async fn aes_key_for(&self, ed25519_pub: &[u8; 32]) -> Option<[u8; 32]> {
if let Some(k) = self.aes_cache.lock().unwrap().get(ed25519_pub) {
return Some(*k);
}
let row = db::get(&self.db, ed25519_pub).await.ok().flatten()?;
let key = self.identity.derive_aes_key(&row.x25519_pub);
self.aes_cache.lock().unwrap().insert(*ed25519_pub, key);
Some(key)
}
// ── Send ──────────────────────────────────────────────────────────────────
/// Seal an opaque `payload` to one client and queue the `message` frame.
///
/// v2 transport: the payload is wrapped in the `version ‖ comp ‖ payload`
/// framing (`compress_payload`) before sealing, then wrapped in
/// `RelayFrame{Message{ciphertext, nonce, peer, live}}`. `live=true` routes
/// or fails (the peer is online by construction); `live=false` stores-and-
/// forwards + pushes for offline phones.
pub(crate) async fn send_to_client(
&self,
client_ed25519_pub: &[u8; 32],
payload: &[u8],
live: bool,
) -> Result<()> {
// v2 framing: version(1B) ‖ comp(1B) ‖ payload (compresses over threshold).
let framed = crypto::compress_payload(payload);
self.seal_and_queue(client_ed25519_pub, &framed, live).await
}
/// Seal an already-framed plaintext to `dest` and queue the `message` frame.
/// Shared by [`send_to_client`](Self::send_to_client) (v2 app framing) and
/// [`send_pipe_signal`](Self::send_pipe_signal) (pipe framing).
async fn seal_and_queue(&self, dest: &[u8; 32], framed: &[u8], live: bool) -> Result<()> {
let aes_key = self
.aes_key_for(dest)
.await
.ok_or_else(|| anyhow!("no aes_key for client"))?;
// Persist the send counter BEFORE sealing/sending (crypto.md §8/§9):
// a crash after this point never reuses a nonce.
let counter = db::next_send_counter(&self.db, dest).await?;
let nonce = crypto::build_nonce(DIR_AGENT_TO_CLIENT, counter);
let aad = crypto::build_aad(
&self.identity.namespace_id_raw(),
&self.identity.ed25519_pub(),
dest,
);
let sealed = crypto::seal(&aes_key, &nonce, &aad, framed)
.map_err(|e| anyhow!("seal failed: {e}"))?;
let frame = RelayFrame {
frame: Some(Frame::Message(Message {
ciphertext: prost::bytes::Bytes::from(sealed),
nonce: prost::bytes::Bytes::copy_from_slice(&nonce),
peer: prost::bytes::Bytes::copy_from_slice(dest),
live,
})),
};
self.send_frame(frame.encode_to_vec())
}
/// Seal + queue a pipe-signaling message (docs/relay/pipe.md §1) over the E2E
/// channel, wrapped in the reserved pipe framing so the peer routes it to its
/// pipe layer. Always `live` (a stale invite is useless, pipe.md §1).
async fn send_pipe_signal(&self, dest: &[u8; 32], signal: &PipeSignal) -> Result<()> {
let framed = crypto::frame_pipe_signal(&skald_relay_common::pipe::encode(signal));
self.seal_and_queue(dest, &framed, true).await
}
// ── Receive ───────────────────────────────────────────────────────────────
/// Handle an inbound `message` (relay-protocol.md §3.1): authorize the
/// sender, check nonce direction + counter monotonicity, open, advance the
/// recv counter, peel the v2 framing, then emit [`RelayEvent::Message`] with
/// the clean inner payload. The client never inspects the payload contents.
pub(crate) async fn handle_inbound_message(
&self,
from: &[u8; 32],
nonce: &[u8; 12],
ciphertext: &[u8],
live: bool,
) {
// `from` must be an Authorized client.
let row = match db::get(&self.db, from).await {
Ok(Some(r)) if r.state == ClientState::Authorized => r,
_ => {
warn!(crate_name = "skald-relay-client", "message from non-authorized sender dropped");
return;
}
};
// Extract the counter from the nonce and check direction + monotonicity.
if nonce[..4] != DIR_CLIENT_TO_AGENT {
warn!(crate_name = "skald-relay-client", "message with wrong nonce direction dropped");
return;
}
let counter = u64::from_be_bytes(nonce[4..].try_into().unwrap());
if counter <= row.recv_counter {
warn!(crate_name = "skald-relay-client", "replayed/old counter dropped");
return;
}
let Some(aes_key) = self.aes_key_for(from).await else { return };
let aad = crypto::build_aad(
&self.identity.namespace_id_raw(),
from,
&self.identity.ed25519_pub(),
);
let framed = match crypto::open(&aes_key, nonce, &aad, ciphertext) {
Ok(pt) => pt,
Err(_) => {
// No content logging on decrypt failure (crypto.md §8).
warn!(crate_name = "skald-relay-client", "decrypt failed, message dropped");
return;
}
};
// Valid open → advance recv_counter.
if let Err(e) = db::set_recv_counter(&self.db, from, counter).await {
warn!(crate_name = "skald-relay-client", error = %e, "failed to persist recv_counter");
}
// Pipe signaling rides this same E2E channel under a reserved framing
// version (crypto::FRAMING_VERSION_PIPE). Route it to the pipe layer
// instead of emitting a Message; all other payloads stay pass-through.
if crypto::is_pipe_signal(&framed) {
match crypto::unframe_pipe_signal(&framed) {
Some(body) => self.handle_pipe_signal(from, body),
None => warn!(crate_name = "skald-relay-client", "malformed pipe signal framing dropped"),
}
return;
}
// Peel the v2 framing so the consumer sees the clean inner payload.
let payload = match crypto::decompress_payload(&framed) {
Ok(p) => p,
Err(e) => {
warn!(crate_name = "skald-relay-client", error = %e, "framing decompress failed");
return;
}
};
self.emit(RelayEvent::Message { from: *from, payload, live });
}
// ── Pipe control plane (docs/relay/pipe.md §1, §3) ────────────────────────
/// Subscribe to inbound `pipe_invite`s (responder side). Single-consumer
/// expected: the consumer accepts/rejects each pipe exactly once.
pub(crate) fn incoming_pipes(&self) -> broadcast::Receiver<IncomingPipe> {
self.incoming_pipes_tx.subscribe()
}
/// Route a decoded pipe-signaling message: invites surface to the app via the
/// incoming-pipes broadcast; accept/reject wake the matching `open_pipe`
/// waiter. This is the only payload kind the otherwise payload-agnostic client
/// interprets (it owns the pipe control plane end-to-end).
fn handle_pipe_signal(&self, from: &[u8; 32], body: &[u8]) {
let signal: PipeSignal = match skald_relay_common::pipe::decode(body) {
Ok(s) => s,
Err(e) => {
warn!(crate_name = "skald-relay-client", error = %e, "malformed pipe signal dropped");
return;
}
};
match signal {
PipeSignal::Invite(inv) => {
let Some(connection_id) = to_array::<32>(&inv.connection_id) else {
warn!(crate_name = "skald-relay-client", "pipe invite with bad connection_id");
return;
};
let _ = self.incoming_pipes_tx.send(IncomingPipe {
from: *from,
stream_type: inv.stream_type,
headers: inv.headers,
connection_id,
suite: inv.suite,
peer_handshake: inv.handshake,
});
}
PipeSignal::Accept(acc) => {
if let Some(cid) = to_array::<32>(&acc.connection_id)
&& let Some(tx) = self.pipe_waiters.lock().unwrap().remove(&cid)
{
let _ = tx.send(Ok(acc));
}
}
PipeSignal::Reject(rej) => {
if let Some(cid) = to_array::<32>(&rej.connection_id)
&& let Some(tx) = self.pipe_waiters.lock().unwrap().remove(&cid)
{
let _ = tx.send(Err(rej.reason));
}
}
}
}
/// Initiator: open a pipe to `peer`. Generates an ephemeral X25519, sends
/// `pipe_invite`, waits for `pipe_accept`, derives the per-pipe key (PFS),
/// then dials the data plane.
pub(crate) async fn open_pipe(
&self,
peer: &[u8; 32],
stream_type: &str,
headers: BTreeMap<String, String>,
) -> Result<PipeConnection> {
let mut eph_priv = [0u8; 32];
rand::rng().fill_bytes(&mut eph_priv);
let eph_pub = crypto::x25519_pubkey(&eph_priv);
let mut connection_id = [0u8; 32];
rand::rng().fill_bytes(&mut connection_id);
let rx = self.register_pipe_waiter(connection_id);
let invite = PipeSignal::Invite(PipeInvite {
connection_id: connection_id.to_vec(),
suite: PipeSuite::X25519Sealed,
handshake: eph_pub.to_vec(),
stream_type: stream_type.to_string(),
compress: vec![skald_relay_common::pipe::PipeCompress::None],
headers,
});
if let Err(e) = self.send_pipe_signal(peer, &invite).await {
self.pipe_waiters.lock().unwrap().remove(&connection_id);
return Err(e);
}
let accept = match tokio::time::timeout(PIPE_ACCEPT_TIMEOUT, rx).await {
Ok(Ok(Ok(acc))) => acc,
Ok(Ok(Err(reason))) => return Err(anyhow!("pipe rejected by peer: {reason}")),
_ => {
self.pipe_waiters.lock().unwrap().remove(&connection_id);
return Err(anyhow!("pipe accept timed out"));
}
};
let peer_eph = to_array::<32>(&accept.handshake)
.ok_or_else(|| anyhow!("pipe accept has a bad ephemeral key"))?;
let pipe_key = crypto::derive_pipe_key(&crypto::ecdh(&eph_priv, &peer_eph));
PipeConnection::connect(
&self.relay_url(),
&self.identity.signing_key(),
&self.identity.ed25519_pub(),
peer,
&self.identity.namespace_id_raw(),
&connection_id,
&pipe_key,
PipeRole::Initiator,
)
.await
}
/// Responder: accept an inbound invite. Replies with `pipe_accept`, derives
/// the per-pipe key, then dials the data plane.
pub(crate) async fn accept_pipe(&self, incoming: &IncomingPipe) -> Result<PipeConnection> {
// v1 supports only the X25519Sealed suite; a future Noise suite is a new
// arm here (the wire shape is unchanged — pipe.md forward-compat).
if incoming.suite != PipeSuite::X25519Sealed {
return Err(anyhow!("unsupported pipe suite"));
}
let peer_eph = to_array::<32>(&incoming.peer_handshake)
.ok_or_else(|| anyhow!("pipe invite has a bad ephemeral key"))?;
let mut eph_priv = [0u8; 32];
rand::rng().fill_bytes(&mut eph_priv);
let eph_pub = crypto::x25519_pubkey(&eph_priv);
let pipe_key = crypto::derive_pipe_key(&crypto::ecdh(&eph_priv, &peer_eph));
let accept = PipeSignal::Accept(PipeAccept {
connection_id: incoming.connection_id.to_vec(),
suite: PipeSuite::X25519Sealed,
handshake: eph_pub.to_vec(),
compress: skald_relay_common::pipe::PipeCompress::None,
});
self.send_pipe_signal(&incoming.from, &accept).await?;
PipeConnection::connect(
&self.relay_url(),
&self.identity.signing_key(),
&self.identity.ed25519_pub(),
&incoming.from,
&self.identity.namespace_id_raw(),
&incoming.connection_id,
&pipe_key,
PipeRole::Responder,
)
.await
}
/// Decline an inbound invite (sends `pipe_reject`).
pub(crate) async fn reject_pipe(&self, incoming: &IncomingPipe, reason: &str) -> Result<()> {
let reject = PipeSignal::Reject(PipeReject {
connection_id: incoming.connection_id.to_vec(),
reason: reason.to_string(),
});
self.send_pipe_signal(&incoming.from, &reject).await
}
/// Register an `open_pipe` waiter keyed by `connection_id`; the inbound
/// `pipe_accept`/`pipe_reject` resolves it.
fn register_pipe_waiter(
&self,
connection_id: [u8; 32],
) -> oneshot::Receiver<Result<PipeAccept, String>> {
let (tx, rx) = oneshot::channel();
self.pipe_waiters.lock().unwrap().insert(connection_id, tx);
rx
}
}
#[cfg(test)]
mod pipe_signal_tests {
use super::*;
use skald_relay_common::pipe::PipeCompress;
async fn make_state() -> RelayState {
let db = std::env::temp_dir().join(format!("relay-cli-state-{}.db", std::process::id()));
let pool = SqlitePool::connect(&format!("sqlite://{}?mode=rwc", db.display()))
.await
.unwrap();
db::init(&pool).await.unwrap();
let (events_tx, _) = broadcast::channel(16);
RelayState::new(
Identity::from_seed(&[1u8; 32]),
Arc::new(pool),
StateConfig { relay_url: String::new(), pairing_ttl: 300 },
events_tx,
)
}
#[tokio::test]
async fn invite_surfaces_on_incoming_pipes() {
let st = make_state().await;
let mut rx = st.incoming_pipes();
let invite = PipeSignal::Invite(PipeInvite {
connection_id: vec![7; 32],
suite: PipeSuite::X25519Sealed,
handshake: vec![8; 32],
stream_type: "log".into(),
compress: vec![PipeCompress::None],
headers: BTreeMap::from([("k".into(), "v".into())]),
});
st.handle_pipe_signal(&[2u8; 32], &skald_relay_common::pipe::encode(&invite));
let got = rx.try_recv().expect("invite surfaced");
assert_eq!(got.from, [2u8; 32]);
assert_eq!(got.stream_type, "log");
assert_eq!(got.connection_id, [7u8; 32]);
assert_eq!(got.headers.get("k").map(String::as_str), Some("v"));
}
#[tokio::test]
async fn accept_resolves_the_waiter() {
let st = make_state().await;
let cid = [3u8; 32];
let rx = st.register_pipe_waiter(cid);
let accept = PipeSignal::Accept(PipeAccept {
connection_id: cid.to_vec(),
suite: PipeSuite::X25519Sealed,
handshake: vec![9; 32],
compress: PipeCompress::None,
});
st.handle_pipe_signal(&[2u8; 32], &skald_relay_common::pipe::encode(&accept));
let resolved = rx.await.expect("waiter not dropped");
assert_eq!(resolved.expect("accept ok").handshake, vec![9; 32]);
}
#[tokio::test]
async fn reject_resolves_waiter_with_reason() {
let st = make_state().await;
let cid = [4u8; 32];
let rx = st.register_pipe_waiter(cid);
let reject = PipeSignal::Reject(PipeReject { connection_id: cid.to_vec(), reason: "busy".into() });
st.handle_pipe_signal(&[2u8; 32], &skald_relay_common::pipe::encode(&reject));
assert_eq!(rx.await.expect("waiter").unwrap_err(), "busy");
}
#[tokio::test]
async fn unknown_connection_id_is_ignored() {
let st = make_state().await;
// An accept for a connection_id with no waiter must not panic.
let accept = PipeSignal::Accept(PipeAccept {
connection_id: vec![0xEE; 32],
suite: PipeSuite::X25519Sealed,
handshake: vec![0; 32],
compress: PipeCompress::None,
});
st.handle_pipe_signal(&[2u8; 32], &skald_relay_common::pipe::encode(&accept));
}
}
+361
View File
@@ -0,0 +1,361 @@
//! The permanent agent WebSocket toward the relay, speaking **v2 protobuf**
//! (docs/relay/relay-protocol.md).
//!
//! A single WS carries everything: challenge-response auth, the `Authorize` set,
//! outbound E2E `Message`s, and inbound `Message` / `ClientPaired` frames. v2
//! transport is **binary-only**: every wire frame is a `RelayFrame` protobuf
//! message wrapped in `Message::Binary`; WS-level `Ping`/`Pong`/`Close` are
//! their own `WsMessage` variants and never appear as protobuf.
//!
//! Reconnection uses exponential backoff (1,2,4,…,60 s) with jitter, and the
//! whole loop is cancellable on stop.
use std::sync::Arc;
use std::time::Duration;
use anyhow::{anyhow, Result};
use futures_util::{SinkExt, StreamExt};
use prost::Message as _;
use rand::Rng;
use skald_relay_common::crypto;
use skald_relay_common::proto::v2::*;
use skald_relay_common::proto::v2::relay_frame::Frame;
use tokio::sync::mpsc;
use tokio_tungstenite::tungstenite::Message as WsMessage;
use tokio_util::sync::CancellationToken;
use tracing::{debug, info, warn};
use crate::state::RelayState;
/// Run the reconnecting WS loop until `cancel` fires (relay-protocol.md §8).
pub(crate) async fn run_loop(
state: Arc<RelayState>,
mut outbound_rx: mpsc::UnboundedReceiver<Vec<u8>>,
cancel: CancellationToken,
) {
let mut backoff_step: u32 = 0;
loop {
if cancel.is_cancelled() {
return;
}
match connect_once(&state, &mut outbound_rx, &cancel).await {
Ok(()) => {
// Clean disconnect (cancelled or graceful): reset backoff.
backoff_step = 0;
}
Err(e) => {
warn!(crate_name = "skald-relay-client", error = %e, "relay connection ended");
}
}
if cancel.is_cancelled() {
return;
}
let delay = backoff_delay(backoff_step);
backoff_step = backoff_step.saturating_add(1);
state.set_connected(false);
debug!(crate_name = "skald-relay-client", secs = delay.as_secs_f64(), "reconnect backoff");
tokio::select! {
_ = cancel.cancelled() => return,
_ = tokio::time::sleep(delay) => {}
}
}
}
/// Backoff schedule 1,2,4,…,60 s plus up to 50% jitter (relay-protocol.md §8).
fn backoff_delay(step: u32) -> Duration {
let base = 1u64.checked_shl(step).unwrap_or(60).min(60);
let jitter_ms = rand::rng().random_range(0..=(base * 500));
Duration::from_millis(base * 1000 + jitter_ms)
}
/// One full connection lifecycle: connect → challenge → auth → authorize → loop.
async fn connect_once(
state: &Arc<RelayState>,
outbound_rx: &mut mpsc::UnboundedReceiver<Vec<u8>>,
cancel: &CancellationToken,
) -> Result<()> {
let url = state.relay_url();
info!(crate_name = "skald-relay-client", %url, "connecting to relay");
let (ws_stream, _resp) = tokio::select! {
_ = cancel.cancelled() => return Ok(()),
r = tokio_tungstenite::connect_async(&url) => r?,
};
let (mut sink, mut stream) = ws_stream.split();
// 1. Wait for the relay's challenge (it speaks first, relay-protocol.md §4).
let challenge_nonce = wait_for_challenge(&mut stream).await?;
// 2. Sign AUTH_DOMAIN ‖ 0x00 ‖ nonce and send the agent Auth frame.
let sig = crypto::sign_challenge(&state.identity().signing_key(), &challenge_nonce);
let auth = RelayFrame {
frame: Some(Frame::Auth(Auth {
role: Some(auth::Role::Agent(AuthAgent {
agent_ed25519_pub: prost::bytes::Bytes::copy_from_slice(
&state.identity().ed25519_pub(),
),
})),
signature: prost::bytes::Bytes::copy_from_slice(&sig),
})),
};
sink.send(WsMessage::Binary(auth.encode_to_vec().into())).await?;
// 3. Expect AuthOk and verify the namespace_id locally.
let ns_raw = wait_for_auth_ok(&mut stream).await?;
if ns_raw != state.identity().namespace_id_raw() {
return Err(anyhow!(
"relay returned mismatched namespace_id (got {}, expected {})",
hex::encode(ns_raw),
hex::encode(state.identity().namespace_id_raw())
));
}
info!(crate_name = "skald-relay-client", "relay auth ok, namespace verified");
state.set_connected(true);
// 4. Send the current authorize set from the DB (empty on first run).
// We push it directly via the sink rather than through `outbound_rx` so it
// lands immediately — the queue is only drained inside the main loop below.
let authorized = state.authorized_pubkeys_hex().await.unwrap_or_default();
let clients: Vec<prost::bytes::Bytes> = authorized
.iter()
.filter_map(|h| hex::decode(h).ok())
.map(prost::bytes::Bytes::from)
.collect();
let authorize = RelayFrame {
frame: Some(Frame::Authorize(Authorize { clients })),
};
sink.send(WsMessage::Binary(authorize.encode_to_vec().into())).await?;
// 5. Main dispatch loop: outbound queue, inbound frames, WS-level Ping/Pong.
loop {
tokio::select! {
_ = cancel.cancelled() => {
let _ = sink.send(WsMessage::Close(None)).await;
return Ok(());
}
// Outbound: already-encoded protobuf frames queued by pairing / send
// / revoke. The channel carries `Vec<u8>` ready to be shipped as a
// binary WS frame.
maybe = outbound_rx.recv() => {
match maybe {
Some(bytes) => sink.send(WsMessage::Binary(bytes.into())).await?,
None => return Ok(()), // channel closed → client stopping
}
}
// Inbound: relay → agent frames.
maybe = stream.next() => {
let Some(msg) = maybe else { return Ok(()) }; // stream ended
match msg? {
WsMessage::Binary(data) => {
handle_incoming(state, &data).await;
}
WsMessage::Ping(p) => sink.send(WsMessage::Pong(p)).await?,
WsMessage::Pong(_) => {}
WsMessage::Close(_) => return Ok(()),
WsMessage::Text(_) | WsMessage::Frame(_) => {
// v2 transport is binary-only; ignore text/frame
// variants (forward-compat, no protocol-defined reaction).
}
}
}
}
}
}
/// Read binary frames until `Challenge` arrives; returns the raw 32-byte nonce.
async fn wait_for_challenge<S>(stream: &mut S) -> Result<[u8; 32]>
where
S: StreamExt<Item = Result<WsMessage, tokio_tungstenite::tungstenite::Error>> + Unpin,
{
while let Some(msg) = stream.next().await {
match msg? {
WsMessage::Binary(data) => {
let frame = RelayFrame::decode(&data[..])?;
if let Some(Frame::Challenge(c)) = frame.frame {
if c.nonce.len() != 32 {
return Err(anyhow!("challenge nonce is not 32 bytes"));
}
let mut out = [0u8; 32];
out.copy_from_slice(&c.nonce);
return Ok(out);
}
}
WsMessage::Close(_) => return Err(anyhow!("closed before challenge")),
_ => {}
}
}
Err(anyhow!("connection closed before challenge"))
}
/// Read binary frames until `AuthOk`; returns the raw 32-byte namespace_id.
async fn wait_for_auth_ok<S>(stream: &mut S) -> Result<[u8; 32]>
where
S: StreamExt<Item = Result<WsMessage, tokio_tungstenite::tungstenite::Error>> + Unpin,
{
while let Some(msg) = stream.next().await {
match msg? {
WsMessage::Binary(data) => {
let frame = RelayFrame::decode(&data[..])?;
match frame.frame {
Some(Frame::AuthOk(AuthOk { namespace_id })) => {
if namespace_id.len() != 32 {
return Err(anyhow!("namespace_id is not 32 bytes"));
}
let mut out = [0u8; 32];
out.copy_from_slice(&namespace_id);
return Ok(out);
}
Some(Frame::AuthError(AuthError { code, message })) => {
return Err(anyhow!("auth_error from relay: {code} ({message})"));
}
_ => {}
}
}
WsMessage::Close(_) => return Err(anyhow!("closed before auth_ok")),
_ => {}
}
}
Err(anyhow!("connection closed before auth_ok"))
}
/// Dispatch one decoded relay→agent `RelayFrame`. WS-level Ping/Pong are
/// handled at the transport layer above; everything that arrives as a binary
/// frame is decoded to `RelayFrame` and matched on the `Frame` oneof here.
async fn handle_incoming(state: &Arc<RelayState>, data: &[u8]) {
let frame = match RelayFrame::decode(data) {
Ok(f) => f,
Err(e) => {
warn!(crate_name = "skald-relay-client", error = %e, "malformed protobuf frame dropped");
return;
}
};
let Some(f) = frame.frame else {
debug!(crate_name = "skald-relay-client", "empty relay frame dropped");
return;
};
match f {
Frame::Message(m) => {
// Validate lengths before handing off to the E2E layer.
if m.peer.len() != 32 || m.nonce.len() != 12 {
warn!(crate_name = "skald-relay-client", "message with wrong peer/nonce length dropped");
return;
}
let mut from = [0u8; 32];
from.copy_from_slice(&m.peer);
let mut nonce = [0u8; 12];
nonce.copy_from_slice(&m.nonce);
state.handle_inbound_message(&from, &nonce, &m.ciphertext, m.live).await;
}
Frame::ClientPaired(cp) => {
if cp.client_ed25519_pub.len() != 32 || cp.client_x25519_pub.len() != 32 {
warn!(crate_name = "skald-relay-client", "client_paired with wrong pubkey length dropped");
return;
}
let mut ed = [0u8; 32];
ed.copy_from_slice(&cp.client_ed25519_pub);
let mut x = [0u8; 32];
x.copy_from_slice(&cp.client_x25519_pub);
// Decode the protobuf `Platform` enum to the lowercase string the DB
// expects. The wire value defaults to `0` (`UNSPECIFIED`) — the helper
// maps that to `"unknown"`.
let platform = platform_i32_to_str(cp.platform);
state.handle_client_paired(&ed, &x, platform).await;
}
Frame::AuthorizeOk(aok) => {
debug!(crate_name = "skald-relay-client", authorized = aok.authorized, "authorize_ok");
}
Frame::PairingReady(_) | Frame::PairingStopOk(_) => {}
Frame::PresenceEvent(pe) => {
debug!(
crate_name = "skald-relay-client",
pubkey = %hex::encode(&pe.pubkey),
status = pe.status,
"presence event"
);
}
Frame::PresenceList(pl) => {
debug!(crate_name = "skald-relay-client", online = pl.online.len(), "presence list");
}
Frame::PeerOffline(po) => {
// Expected backstop for route-or-fail live sends (relay-protocol.md
// §3): a `live=true` send found the peer gone. A normal protocol
// event, not an error.
debug!(
crate_name = "skald-relay-client",
peer = %hex::encode(&po.peer),
"peer offline for live send; dropping"
);
}
Frame::Error(e) => {
warn!(crate_name = "skald-relay-client", code = %e.code, message = %e.message, "relay error frame");
}
// Server-to-client or handshake frames the agent never expects inbound.
Frame::Challenge(_)
| Frame::Auth(_)
| Frame::AuthOk(_)
| Frame::AuthError(_)
| Frame::Authorize(_)
| Frame::PairingStart(_)
| Frame::PairingStop(_)
| Frame::PresenceRequest(_) => {
warn!(crate_name = "skald-relay-client", "unexpected relay→agent frame dropped");
}
}
}
/// Map a protobuf `Platform` enum wire value to the lowercase string the DB
/// stores in the `platform` column. Unknown values become `"unknown"`.
fn platform_i32_to_str(v: i32) -> &'static str {
if v == Platform::Ios as i32 {
"ios"
} else if v == Platform::Android as i32 {
"android"
} else {
"unknown"
}
}
#[cfg(test)]
mod tests {
use super::*;
/// `platform_i32_to_str` is total on the wire values the relay emits and
/// never panics on bogus inputs (relay-protocol.md §11 forward-compat).
#[test]
fn platform_conversion() {
assert_eq!(platform_i32_to_str(0), "unknown");
assert_eq!(platform_i32_to_str(1), "ios");
assert_eq!(platform_i32_to_str(2), "android");
assert_eq!(platform_i32_to_str(99), "unknown");
}
/// A minimal `Message` frame round-trips through `prost` so the wire
/// encoding we emit is the same one the relay will decode.
#[test]
fn message_frame_round_trip() {
let frame = RelayFrame {
frame: Some(Frame::Message(Message {
ciphertext: vec![0xAA; 64].into(),
nonce: vec![0x01; 12].into(),
peer: vec![0x02; 32].into(),
live: false,
})),
};
let bytes = frame.encode_to_vec();
let decoded = RelayFrame::decode(&bytes[..]).expect("decode");
match decoded.frame {
Some(Frame::Message(m)) => {
assert_eq!(m.ciphertext.len(), 64);
assert_eq!(m.nonce.len(), 12);
assert_eq!(m.peer.len(), 32);
assert!(!m.live);
}
other => panic!("expected Message, got {other:?}"),
}
}
}
@@ -0,0 +1,355 @@
//! End-to-end integration test for `skald-relay-client` against the **real**
//! relay server (`skald-relay-server`) booted in-process on an ephemeral port.
//!
//! It drives the full agent-role flow through the public `RelayClient` API while
//! a hand-rolled "mobile client" (raw `tokio-tungstenite` speaking v2 protobuf +
//! the shared E2E crypto) plays the counterpart:
//!
//! start → Connected → start_pairing → pair → ClientPaired → authorize →
//! send (mobile decrypts) → mobile reply → Message → replay dropped →
//! revoke → ClientRevoked.
//!
//! This exercises the persist-before-seal counter path, the nonce direction /
//! AAD construction, the pairing window, and the events channel end to end.
use std::net::SocketAddr;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use bytes::Bytes;
use ed25519_dalek::SigningKey;
use futures_util::{SinkExt, StreamExt};
use prost::Message as _;
use skald_relay_common::crypto::{self, DIR_CLIENT_TO_AGENT};
use skald_relay_common::proto::v2::{
self, Auth, AuthClient, AuthPairing, Message as ProtoMessage, RelayFrame,
};
use skald_relay_common::proto::v2::auth::Role as AuthRole;
use skald_relay_common::proto::v2::relay_frame::Frame;
use sqlx::SqlitePool;
use tokio::sync::broadcast;
use tokio_tungstenite::tungstenite::Message as WsMessage;
use skald_relay_client::{RelayClient, RelayClientConfig, RelayEvent, SeedSource};
use skald_relay_server::config::Config;
use skald_relay_server::{AppState, router};
type Ws =
tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
static COUNTER: AtomicU64 = AtomicU64::new(0);
fn unique_suffix() -> String {
let nanos = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_nanos();
let seq = COUNTER.fetch_add(1, Ordering::Relaxed);
format!("{nanos}-{}-{seq}", std::process::id())
}
/// Boot a relay on a random port with a throwaway SQLite file. Returns its addr.
async fn spawn_relay() -> SocketAddr {
let db = std::env::temp_dir().join(format!("relay-srv-{}.db", unique_suffix()));
let cfg = Config {
bind: "127.0.0.1:0".parse().unwrap(),
db_path: db.to_string_lossy().into(),
pipe: skald_relay_server::config::PipeConfig::default(),
};
let state = AppState::build(cfg).await.expect("build relay state");
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move {
axum::serve(
listener,
router(state).into_make_service_with_connect_info::<SocketAddr>(),
)
.await
.unwrap();
});
addr
}
/// A fresh on-disk SQLite pool for the client (a temp-file DB, not `:memory:`,
/// so the WS-loop task and the test task share the same database across the
/// pool's connections).
async fn client_pool() -> Arc<SqlitePool> {
let path = std::env::temp_dir().join(format!("relay-cli-{}.db", unique_suffix()));
let _ = std::fs::remove_file(&path);
let url = format!("sqlite://{}?mode=rwc", path.display());
Arc::new(SqlitePool::connect(&url).await.expect("client pool"))
}
// ── raw WS helpers (mobile side) ────────────────────────────────────────────
async fn connect(addr: SocketAddr) -> Ws {
let url = format!("ws://{addr}/v1/ws");
let (ws, _) = tokio_tungstenite::connect_async(url).await.expect("connect");
ws
}
async fn send(ws: &mut Ws, frame: &RelayFrame) {
ws.send(WsMessage::Binary(frame.encode_to_vec().into()))
.await
.expect("send binary");
}
async fn recv(ws: &mut Ws) -> RelayFrame {
loop {
let m = ws.next().await.expect("stream open").expect("ws frame");
match m {
WsMessage::Binary(b) => return RelayFrame::decode(b.as_ref()).expect("decode"),
WsMessage::Ping(_) | WsMessage::Pong(_) => continue,
WsMessage::Close(f) => panic!("unexpected ws close: {f:?}"),
other => panic!("unexpected ws frame: {other:?}"),
}
}
}
async fn read_challenge(ws: &mut Ws) -> [u8; 32] {
match recv(ws).await.frame {
Some(Frame::Challenge(c)) => c.nonce.as_ref().try_into().expect("32B challenge"),
other => panic!("expected Challenge, got {other:?}"),
}
}
fn auth_pairing_frame(
sk: &SigningKey,
challenge: &[u8; 32],
ns_raw: &[u8; 32],
token: &[u8; 32],
x25519_pub: &[u8; 32],
) -> RelayFrame {
let sig = crypto::sign_challenge(sk, challenge);
RelayFrame {
frame: Some(Frame::Auth(Auth {
signature: Bytes::copy_from_slice(&sig),
role: Some(AuthRole::Pairing(AuthPairing {
namespace_id: Bytes::copy_from_slice(ns_raw),
client_ed25519_pub: Bytes::copy_from_slice(&sk.verifying_key().to_bytes()),
client_x25519_pub: Bytes::copy_from_slice(x25519_pub),
pairing_token: Bytes::copy_from_slice(token),
device_token: "devtok".into(),
platform: v2::Platform::Ios as i32,
})),
})),
}
}
fn auth_client_frame(sk: &SigningKey, challenge: &[u8; 32], ns_raw: &[u8; 32]) -> RelayFrame {
let sig = crypto::sign_challenge(sk, challenge);
RelayFrame {
frame: Some(Frame::Auth(Auth {
signature: Bytes::copy_from_slice(&sig),
role: Some(AuthRole::Client(AuthClient {
namespace_id: Bytes::copy_from_slice(ns_raw),
client_ed25519_pub: Bytes::copy_from_slice(&sk.verifying_key().to_bytes()),
device_token: "devtok".into(),
platform: v2::Platform::Ios as i32,
})),
})),
}
}
/// Pair on a short-lived side connection (challenge → auth(pairing) → AuthOk).
async fn pair(addr: SocketAddr, sk: &SigningKey, ns_raw: &[u8; 32], token: &[u8; 32], x_pub: &[u8; 32]) {
let mut ws = connect(addr).await;
let c = read_challenge(&mut ws).await;
send(&mut ws, &auth_pairing_frame(sk, &c, ns_raw, token, x_pub)).await;
match recv(&mut ws).await.frame {
Some(Frame::AuthOk(_)) => {}
other => panic!("pairing expected AuthOk, got {other:?}"),
}
drop(ws);
}
/// Connect as the authorized client role; returns the live socket.
async fn auth_client(addr: SocketAddr, sk: &SigningKey, ns_raw: &[u8; 32]) -> Ws {
let mut ws = connect(addr).await;
let c = read_challenge(&mut ws).await;
send(&mut ws, &auth_client_frame(sk, &c, ns_raw)).await;
match recv(&mut ws).await.frame {
Some(Frame::AuthOk(_)) => {}
other => panic!("client expected AuthOk, got {other:?}"),
}
ws
}
// ── event helpers ───────────────────────────────────────────────────────────
async fn next_event(rx: &mut broadcast::Receiver<RelayEvent>) -> RelayEvent {
tokio::time::timeout(Duration::from_secs(3), rx.recv())
.await
.expect("timed out waiting for event")
.expect("event recv")
}
/// Next event that is not a `Connected`/`Disconnected` heartbeat.
async fn next_significant(rx: &mut broadcast::Receiver<RelayEvent>) -> RelayEvent {
loop {
match next_event(rx).await {
RelayEvent::Connected | RelayEvent::Disconnected => continue,
other => return other,
}
}
}
#[tokio::test]
async fn full_round_trip() {
let addr = spawn_relay().await;
// Build the client (agent role) pointed at the in-process relay.
let pool = client_pool().await;
let config = RelayClientConfig {
relay_url: format!("ws://{addr}/v1/ws"),
pairing_ttl: 300,
seed: SeedSource::Bytes([1u8; 32]),
};
let client = RelayClient::new(pool, config).await.expect("new client");
let agent_ed = client.agent_ed25519_pub();
let agent_x = client.agent_x25519_pub();
let ns_raw: [u8; 32] = hex::decode(client.namespace_id_hex()).unwrap().try_into().unwrap();
// Mobile identity.
let mobile = crypto::derive_keys(&[7u8; 32]);
let mobile_sk = mobile.signing_key();
let mobile_ed = mobile.ed25519_pub;
// Shared AES key both sides derive independently.
let aes = crypto::derive_aes_key(&crypto::ecdh(&mobile.x25519_priv, &agent_x));
let mut rx = client.events();
client.start().await.expect("start");
// Connected handshake completes.
match next_event(&mut rx).await {
RelayEvent::Connected => {}
other => panic!("expected Connected, got {other:?}"),
}
// 1) Open a pairing window and pair the mobile. `start_pairing` only queues
// the frame; let the relay register the token before the mobile pairs.
let started = client.start_pairing(0).await.expect("start_pairing");
tokio::time::sleep(Duration::from_millis(150)).await;
pair(addr, &mobile_sk, &ns_raw, &started.token, &mobile.x25519_pub).await;
match next_significant(&mut rx).await {
RelayEvent::ClientPaired { ed25519_pub, platform, .. } => {
assert_eq!(ed25519_pub, mobile_ed);
assert_eq!(platform, "ios");
}
other => panic!("expected ClientPaired, got {other:?}"),
}
// The device is Pending until we authorize it.
let clients = client.list_clients().await;
assert_eq!(clients.len(), 1);
assert_eq!(clients[0].state, skald_relay_client::ClientState::Pending);
// 2) Authorize, then the mobile connects as the authorized client role.
client.authorize(&mobile_ed).await.expect("authorize");
// Give the relay a moment to process the Authorize set before connecting.
tokio::time::sleep(Duration::from_millis(150)).await;
let mut mobile_ws = auth_client(addr, &mobile_sk, &ns_raw).await;
// 3) Agent → mobile: send an opaque payload; the mobile decrypts it.
let agent_payload = b"hello-from-agent";
client.send(&mobile_ed, agent_payload, false).await.expect("send");
let frame = recv(&mut mobile_ws).await;
let m = match frame.frame {
Some(Frame::Message(m)) => m,
other => panic!("mobile expected Message, got {other:?}"),
};
assert_eq!(m.peer.as_ref(), &agent_ed[..], "relay rewrites peer=from");
let nonce: [u8; 12] = m.nonce.as_ref().try_into().unwrap();
let aad = crypto::build_aad(&ns_raw, &agent_ed, &mobile_ed);
let framed = crypto::open(&aes, &nonce, &aad, &m.ciphertext).expect("mobile open");
let got = crypto::decompress_payload(&framed).expect("decompress");
assert_eq!(got, agent_payload);
// 4) Mobile → agent: seal a reply (counter 1, client→agent direction).
let reply = b"hi-from-mobile";
let reply_nonce = crypto::build_nonce(DIR_CLIENT_TO_AGENT, 1);
let reply_aad = crypto::build_aad(&ns_raw, &mobile_ed, &agent_ed);
let reply_framed = crypto::compress_payload(reply);
let reply_ct = crypto::seal(&aes, &reply_nonce, &reply_aad, &reply_framed).expect("mobile seal");
let reply_frame = RelayFrame {
frame: Some(Frame::Message(ProtoMessage {
ciphertext: Bytes::copy_from_slice(&reply_ct),
nonce: Bytes::copy_from_slice(&reply_nonce),
peer: Bytes::copy_from_slice(&agent_ed),
live: false,
})),
};
send(&mut mobile_ws, &reply_frame).await;
match next_significant(&mut rx).await {
RelayEvent::Message { from, payload, .. } => {
assert_eq!(from, mobile_ed);
assert_eq!(payload, reply);
}
other => panic!("expected Message, got {other:?}"),
}
// 5) Replay the exact same frame (counter 1 again) → dropped, no event.
send(&mut mobile_ws, &reply_frame).await;
let replayed = tokio::time::timeout(Duration::from_millis(400), rx.recv()).await;
assert!(
!matches!(replayed, Ok(Ok(RelayEvent::Message { .. }))),
"a replayed counter must not surface a Message event"
);
// 6) Revoke the device → ClientRevoked event + empty registry.
client.revoke(&mobile_ed).await.expect("revoke");
match next_significant(&mut rx).await {
RelayEvent::ClientRevoked { ed25519_pub } => assert_eq!(ed25519_pub, mobile_ed),
other => panic!("expected ClientRevoked, got {other:?}"),
}
assert!(client.list_clients().await.is_empty(), "registry empty after revoke");
client.shutdown().await;
}
/// `clear_all` removes every device and emits one `ClientRevoked` per device.
#[tokio::test]
async fn clear_all_wipes_devices() {
let addr = spawn_relay().await;
let pool = client_pool().await;
let client = RelayClient::new(
pool,
RelayClientConfig {
relay_url: format!("ws://{addr}/v1/ws"),
pairing_ttl: 300,
seed: SeedSource::Bytes([2u8; 32]),
},
)
.await
.expect("new client");
let ns_raw: [u8; 32] = hex::decode(client.namespace_id_hex()).unwrap().try_into().unwrap();
let mut rx = client.events();
client.start().await.expect("start");
match next_event(&mut rx).await {
RelayEvent::Connected => {}
other => panic!("expected Connected, got {other:?}"),
}
// Pair + authorize one device.
let mobile = crypto::derive_keys(&[9u8; 32]);
let started = client.start_pairing(0).await.expect("pairing");
tokio::time::sleep(Duration::from_millis(150)).await;
pair(addr, &mobile.signing_key(), &ns_raw, &started.token, &mobile.x25519_pub).await;
match next_significant(&mut rx).await {
RelayEvent::ClientPaired { .. } => {}
other => panic!("expected ClientPaired, got {other:?}"),
}
client.authorize(&mobile.ed25519_pub).await.expect("authorize");
assert_eq!(client.list_clients().await.len(), 1);
client.clear_all().await.expect("clear_all");
match next_significant(&mut rx).await {
RelayEvent::ClientRevoked { ed25519_pub } => assert_eq!(ed25519_pub, mobile.ed25519_pub),
other => panic!("expected ClientRevoked, got {other:?}"),
}
assert!(client.list_clients().await.is_empty());
client.shutdown().await;
}
+44
View File
@@ -0,0 +1,44 @@
[package]
name = "skald-relay-common"
version = "0.1.0"
edition = "2024"
description = "Shared frame types + crypto for the Skald Remote Control relay and the mobile-connector plugin (see data/ios-app/plugin.md §1.1)"
license = "MIT"
# Library shared byte-for-byte between the relay server and the mobile-connector
# plugin. Keep it lightweight: NO axum / tokio here. The `gen-vectors` binary
# (src/bin/gen-vectors.rs) reproduces the crypto interop vectors
# (data/ios-app/test-vectors.md §3) using only the library functions below.
[dependencies]
# --- serde / frames ---
serde = { version = "1", features = ["derive"] }
serde_json = "1"
# MsgPack wire format for the pipe control/data-plane messages (docs/relay/pipe.md).
# `rmp-serde` keeps the same serde derives as the JSON path so a struct maps to
# either format; the pipe layer is binary-first per the transport direction.
rmp-serde = "1"
# --- crypto ---
ed25519-dalek = "2"
x25519-dalek = { version = "2", features = ["static_secrets"] }
hkdf = "0.12"
aes-gcm = "0.10"
sha2 = "0.10"
subtle = "2"
hex = "0.4"
base64 = "0.22"
rand = "0.8"
flate2 = "1"
# --- protobuf (v2 transport) ---
prost = "0.13"
bytes = "1"
[build-dependencies]
# Generates the Rust types for `proto/skald/relay/v2/relay_frame.proto`.
# The output lands in OUT_DIR and is `include!`-ed by `src/proto.rs`.
prost-build = "0.13"
# Vendored `protoc` binary so the build needs no system protoc — keeps the musl
# / container distribution build self-contained (see docs/build-and-distribution.md).
protoc-bin-vendored = "3"
+41
View File
@@ -0,0 +1,41 @@
// Build script for skald-relay-common.
//
// Compiles the v2 relay protocol schema (proto/skald/relay/v2/relay_frame.proto)
// into Rust types using prost-build. The generated module is named after the
// proto package (`skald.relay.v2`) and lands in OUT_DIR; `src/proto.rs`
// `include!`s it under a `v2` submodule.
//
// We intentionally use `Config::new()` (no `extern_path`, no `include_file`)
// so the generated code is self-contained — `proto.rs` is the single
// integration point and downstream crates depend on `skald_relay_common::proto`
// without having to plumb any extra paths through their own `build.rs`.
fn main() -> std::io::Result<()> {
// Point prost-build at a vendored `protoc` so the build needs no system
// protoc installed (this is what lets the musl / container distribution
// build run in a bare toolchain image — see docs/build-and-distribution.md).
// An explicit PROTOC in the environment still wins.
if std::env::var_os("PROTOC").is_none() {
let protoc = protoc_bin_vendored::protoc_bin_path()
.expect("vendored protoc unavailable for this build host");
// SAFETY: build scripts run single-threaded before any other code, so
// there is no concurrent env access. `set_var` is `unsafe` in edition 2024.
unsafe { std::env::set_var("PROTOC", protoc); }
}
let mut config = prost_build::Config::new();
// Default for `bytes` fields is `Vec<u8>`, but state it explicitly so a
// future prost default change can't silently flip the wire encoding.
config.bytes(["."]);
config.compile_protos(
&["proto/skald/relay/v2/relay_frame.proto"],
&["proto"],
)?;
// Rerun if the schema (or this build script) changes.
println!("cargo:rerun-if-changed=proto/skald/relay/v2/relay_frame.proto");
println!("cargo:rerun-if-changed=build.rs");
Ok(())
}
@@ -0,0 +1,82 @@
// NORMATIVE schema for the Skald Relay v2 transport.
// Mirror copy of data/iOS-app/v2/relay-protocol.md §2.
// Do NOT edit field numbers / names — clients and servers depend on byte
// compatibility.
syntax = "proto3";
package skald.relay.v2;
message RelayFrame {
oneof frame {
Challenge challenge = 1;
Auth auth = 2;
AuthOk auth_ok = 3;
AuthError auth_error = 4;
Authorize authorize = 5;
AuthorizeOk authorize_ok = 6;
PairingStart pairing_start = 7;
PairingReady pairing_ready = 8;
PairingStop pairing_stop = 9;
PairingStopOk pairing_stop_ok = 10;
ClientPaired client_paired = 11;
Message message = 12;
PeerOffline peer_offline = 13;
PresenceRequest presence_request = 14;
PresenceList presence_list = 15;
PresenceEvent presence_event = 16;
Error error = 17;
}
reserved 18, 19;
}
message Message {
bytes ciphertext = 1;
bytes nonce = 2;
bytes peer = 3;
bool live = 4;
}
message PeerOffline { bytes peer = 1; }
message PresenceRequest {}
message PresenceList { repeated bytes online = 1; }
message PresenceEvent { bytes pubkey = 1; Status status = 2; }
enum Status { STATUS_UNSPECIFIED = 0; STATUS_ONLINE = 1; STATUS_OFFLINE = 2; }
message Challenge { bytes nonce = 1; }
message Auth {
oneof role {
AuthAgent agent = 1;
AuthClient client = 2;
AuthPairing pairing = 3;
}
bytes signature = 4;
}
message AuthAgent { bytes agent_ed25519_pub = 1; }
message AuthClient {
bytes namespace_id = 1;
bytes client_ed25519_pub = 2;
string device_token = 3;
Platform platform = 4;
}
message AuthPairing {
bytes namespace_id = 1;
bytes client_ed25519_pub = 2;
bytes client_x25519_pub = 3;
bytes pairing_token = 4;
string device_token = 5;
Platform platform = 6;
}
enum Platform { PLATFORM_UNSPECIFIED = 0; PLATFORM_IOS = 1; PLATFORM_ANDROID = 2; }
message AuthOk { bytes namespace_id = 1; }
message AuthError { string code = 1; string message = 2; }
message Authorize { repeated bytes clients = 1; }
message AuthorizeOk { uint32 authorized = 1; }
message PairingStart { bytes pairing_token = 1; uint32 ttl = 2; }
message PairingReady { uint32 ttl = 1; }
message PairingStop {}
message PairingStopOk {}
message ClientPaired { bytes client_ed25519_pub = 1; bytes client_x25519_pub = 2; Platform platform = 3; }
message Error { string code = 1; string message = 2; }
@@ -0,0 +1,151 @@
//! Reference generator for the crypto interop test vectors
//! (data/ios-app/test-vectors.md §3). This is the *source of truth*: run it once
//! and paste the output into test-vectors.md §4, then commit.
//!
//! cargo run -p skald-relay-common --bin gen-vectors
//!
//! The relay itself only verifies Ed25519 and derives namespace_id; the full
//! E2E suite (X25519/HKDF/AES-GCM) lives in `skald-relay-common::crypto` so
//! independent implementations (Swift app, Kotlin app) can assert byte-for-byte
//! equality. This binary is a thin driver over those library functions, so the
//! plugin, relay tests and the generator all share the exact same code path.
//!
//! ## v2 framing (data/iOS-app/v2/framing.md §1)
//!
//! In v2 the bytes that get fed to AES-GCM are not the raw JSON plaintext: the
//! sender wraps them in a tiny envelope
//!
//! ```text
//! plaintext = version (0x01) ‖ comp (1B) ‖ payload(JSON)
//! ```
//!
//! where `comp = 0x00` (no compression) when `len(payload) <= 1024`, else
//! `comp = 0x01` (zlib). See [`crypto::compress_payload`] for the implementation.
//! The reference vectors below seal that framed byte stream; interop consumers
//! must do the same.
use base64::{Engine, engine::general_purpose::STANDARD as B64};
use skald_relay_common::crypto::{
DIR_AGENT_TO_CLIENT, DIR_CLIENT_TO_AGENT, build_aad, build_nonce, compress_payload,
decompress_payload, derive_aes_key, derive_keys, ecdh, namespace_id, seal, sign_challenge,
};
// Exact UTF-8 plaintexts from test-vectors.md §1 (no spaces, no field reorder).
const PLAINTEXT_A2C: &[u8] = br#"{"v":1,"kind":"inbox_update","id":"00000000-0000-4000-8000-000000000001","ts":1750000000000,"badge":1,"approvals":[{"request_id":"appr_test_1","tool_name":"send_email","agent_label":"Skald","summary":"Test","created_at":1750000000000}],"clarifications":[]}"#;
const PLAINTEXT_C2A: &[u8] = br#"{"v":1,"kind":"approval_response","id":"00000000-0000-4000-8000-000000000002","ts":1750000000000,"request_id":"appr_test_1","decision":"approved"}"#;
fn main() {
let seed_a: [u8; 32] = (0u8..32).collect::<Vec<_>>().try_into().unwrap();
let seed_c: [u8; 32] = (32u8..64).collect::<Vec<_>>().try_into().unwrap();
// --- key derivation (crypto.md §3) ---
let a = derive_keys(&seed_a);
let c = derive_keys(&seed_c);
// --- namespace_id (crypto.md §7) ---
let (ns_raw, ns_hex) = namespace_id(&a.ed25519_pub);
// --- ECDH + AEAD key (crypto.md §4-5), with the mandatory self-check ---
let s1 = ecdh(&a.x25519_priv, &c.x25519_pub);
let s2 = ecdh(&c.x25519_priv, &a.x25519_pub);
assert_eq!(s1, s2, "ECDH mismatch");
let shared = s1;
let aes_key = derive_aes_key(&shared);
// --- A2C (agent → client) ---
let n_a2c = build_nonce(DIR_AGENT_TO_CLIENT, 1);
let aad_a2c = build_aad(&ns_raw, &a.ed25519_pub, &c.ed25519_pub);
// v2 framing: wrap the JSON in version+comp+payload before sealing.
let framed_a2c = compress_payload(PLAINTEXT_A2C);
// Self-check: decompress_payload is the exact inverse of compress_payload.
assert_eq!(
decompress_payload(&framed_a2c).unwrap(),
PLAINTEXT_A2C,
"A2C framing round-trip mismatch"
);
// In v2 the AES-GCM input is the *framed* plaintext (framing.md §1).
let sealed_a2c = seal(&aes_key, &n_a2c, &aad_a2c, &framed_a2c).unwrap();
// --- C2A (client → agent) ---
let n_c2a = build_nonce(DIR_CLIENT_TO_AGENT, 1);
let aad_c2a = build_aad(&ns_raw, &c.ed25519_pub, &a.ed25519_pub);
let framed_c2a = compress_payload(PLAINTEXT_C2A);
assert_eq!(
decompress_payload(&framed_c2a).unwrap(),
PLAINTEXT_C2A,
"C2A framing round-trip mismatch"
);
let sealed_c2a = seal(&aes_key, &n_c2a, &aad_c2a, &framed_c2a).unwrap();
// --- auth signature (client), crypto.md §8 ---
let challenge: [u8; 32] =
hex::decode("aabbccddeeff00112233445566778899aabbccddeeff00112233445566778899")
.unwrap()
.try_into()
.unwrap();
let sig = sign_challenge(&c.signing_key(), &challenge);
// Round-trip self-check: decrypt what we sealed. In v2 this recovers the
// *framed* plaintext (not the raw JSON); a real receiver would then call
// decompress_payload to peel off the version+comp header.
let dec =
skald_relay_common::crypto::open(&aes_key, &n_a2c, &aad_a2c, &sealed_a2c).unwrap();
assert_eq!(dec, framed_a2c, "A2C round-trip mismatch");
assert_eq!(
decompress_payload(&dec).unwrap(),
PLAINTEXT_A2C,
"A2C decompress after open mismatch"
);
println!("# Generated by `cargo run -p skald-relay-common --bin gen-vectors`");
println!("# v2 framing (data/ios-app/v2/framing.md §1): the bytes fed to AES-GCM");
println!("# are plaintext = version(0x01) ‖ comp(1B) ‖ payload(JSON). Below the");
println!("# threshold (1024B), comp = 0x00 and payload is the raw JSON. V14/V17");
println!("# below are computed by sealing the FRAMED plaintext, not the raw JSON.");
println!("V1 agent_x25519_priv = {}", hex::encode(a.x25519_priv));
println!("V2 agent_x25519_pub = {}", hex::encode(a.x25519_pub));
println!("V3 agent_ed25519_priv = {}", hex::encode(a.ed25519_priv));
println!("V4 agent_ed25519_pub = {}", hex::encode(a.ed25519_pub));
println!("V5 client_x25519_priv = {}", hex::encode(c.x25519_priv));
println!("V6 client_x25519_pub = {}", hex::encode(c.x25519_pub));
println!("V7 client_ed25519_priv= {}", hex::encode(c.ed25519_priv));
println!("V8 client_ed25519_pub = {}", hex::encode(c.ed25519_pub));
println!("V9 namespace_id = {}", ns_hex);
println!("V10 shared_secret = {}", hex::encode(shared));
println!("V11 aes_key = {}", hex::encode(aes_key));
println!("V12 nonce_a2c = {}", hex::encode(n_a2c));
println!("V13 aad_a2c = {}", hex::encode(&aad_a2c));
println!("V14 sealed_a2c (b64) = {}", B64.encode(&sealed_a2c));
println!("V15 nonce_c2a = {}", hex::encode(n_c2a));
println!("V16 aad_c2a = {}", hex::encode(&aad_c2a));
println!("V17 sealed_c2a (b64) = {}", B64.encode(&sealed_c2a));
println!("V18 auth_sig_client = {}", hex::encode(sig));
println!();
println!("# v2 framed plaintexts (input to AES-GCM, framing.md §1):");
println!(
"PT_FRAMED_A2C = {}",
hex::encode(&framed_a2c)
);
println!(
"PT_FRAMED_C2A = {}",
hex::encode(&framed_c2a)
);
println!(
"# framed_a2c[:2] = {:02x}{:02x} (version=01, comp=00 = none for <1024 B)",
framed_a2c[0], framed_a2c[1]
);
println!(
"# framed_c2a[:2] = {:02x}{:02x} (version=01, comp=00 = none for <1024 B)",
framed_c2a[0], framed_c2a[1]
);
println!(
"# PT_FRAMED_A2C.len = {} (PLAINTEXT_A2C.len + 2 framing header)",
framed_a2c.len()
);
println!(
"# PT_FRAMED_C2A.len = {} (PLAINTEXT_C2A.len + 2 framing header)",
framed_c2a.len()
);
}

Some files were not shown because too many files have changed in this diff Show More