First Version
This commit is contained in:
@@ -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"] }
|
||||
@@ -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>);
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
@@ -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(", ")
|
||||
)),
|
||||
}
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
@@ -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>;
|
||||
}
|
||||
@@ -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>,
|
||||
}
|
||||
@@ -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",
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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("")
|
||||
}
|
||||
}
|
||||
@@ -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};
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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![]
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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>;
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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,
|
||||
)
|
||||
})
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
@@ -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])
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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"
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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>;
|
||||
@@ -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};
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
/// 1–65535 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 {
|
||||
/// 1–100 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.0–1.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>,
|
||||
/// 1–100
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub search_top_k: Option<u32>,
|
||||
/// 0.0–1.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>,
|
||||
/// 1–100
|
||||
#[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>,
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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"
|
||||
@@ -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)))
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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 })
|
||||
}
|
||||
}
|
||||
@@ -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)))
|
||||
}
|
||||
}
|
||||
@@ -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"
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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", ¶ms);
|
||||
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");
|
||||
}
|
||||
}
|
||||
@@ -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) = ¬ification_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 }
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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"] }
|
||||
@@ -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 ¤t {
|
||||
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 })) }
|
||||
}
|
||||
@@ -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"
|
||||
@@ -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 }
|
||||
}
|
||||
@@ -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"
|
||||
@@ -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())
|
||||
}
|
||||
@@ -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<()>;
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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(¬ifier);
|
||||
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));
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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\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."))
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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"]
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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 <N|name|auto> — 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;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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('&', "&").replace('<', "<").replace('>', ">")
|
||||
}
|
||||
|
||||
/// 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: ®ex::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;
|
||||
}
|
||||
}
|
||||
@@ -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 & < >";
|
||||
|
||||
/// 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(())
|
||||
}
|
||||
}
|
||||
@@ -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?
|
||||
}
|
||||
@@ -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()
|
||||
@@ -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 }
|
||||
}
|
||||
@@ -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"] }
|
||||
@@ -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 128266–129290 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()
|
||||
@@ -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"] }
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
@@ -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),
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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] },
|
||||
}
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -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};
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
@@ -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));
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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"
|
||||
@@ -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
Reference in New Issue
Block a user