251 lines
10 KiB
Rust
251 lines
10 KiB
Rust
|
//! Transport-agnostic LLM turn driver shared by every chat transport that
|
||
|
|
//! forwards a non-command message to Claude Code (currently Discord and
|
||
|
|
//! WhatsApp). Formatting, chunking, and error-message differences between
|
||
|
|
//! transports are captured by the small [`TransportFormat`] trait; the
|
||
|
|
//! streaming/session/history plumbing lives once in [`run_llm_turn`].
|
||
|
|
|
||
|
|
use std::collections::HashMap;
|
||
|
|
use std::sync::Arc;
|
||
|
|
|
||
|
|
use tokio::sync::Mutex as TokioMutex;
|
||
|
|
|
||
|
|
use crate::chat::ChatTransport;
|
||
|
|
use crate::chat::transport::matrix::{ConversationEntry, ConversationRole, RoomConversation};
|
||
|
|
use crate::http::context::PermissionDecision;
|
||
|
|
use crate::services::Services;
|
||
|
|
use crate::slog;
|
||
|
|
|
||
|
|
/// Per-transport formatting hooks needed to run a shared LLM turn.
|
||
|
|
pub trait TransportFormat: Send + Sync {
|
||
|
|
/// Convert markdown into the transport's native formatting and split it
|
||
|
|
/// into messages that respect the transport's size limit. Transports
|
||
|
|
/// without a size limit (or that don't chunk) return a single-element
|
||
|
|
/// vector.
|
||
|
|
fn format_and_chunk(&self, markdown: &str) -> Vec<String>;
|
||
|
|
|
||
|
|
/// Prefix used in `slog!` calls for this transport, e.g. `"[discord]"`.
|
||
|
|
fn log_prefix(&self) -> &'static str;
|
||
|
|
|
||
|
|
/// Convert an LLM error into a user-facing message.
|
||
|
|
fn format_error(&self, err: &str) -> String;
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Build the prompt for an LLM turn, prepending any pending CRDT
|
||
|
|
/// pipeline-transition events as a `<system-reminder>` block.
|
||
|
|
pub fn build_llm_prompt(persona: &str, bot_name: &str, user: &str, user_message: &str) -> String {
|
||
|
|
let event_ctx = crate::llm_session::assemble_prompt_context(persona);
|
||
|
|
format!(
|
||
|
|
"{event_ctx}[Your name is {bot_name}. Refer to yourself as {bot_name}, not Claude.]\n\n{user}: {user_message}"
|
||
|
|
)
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Everything a shared LLM turn needs beyond formatting: transport handle,
|
||
|
|
/// routing/history key, display name, and shared services/history state.
|
||
|
|
pub struct LlmTurnContext<'a> {
|
||
|
|
/// Chat transport used to send/post messages for this turn.
|
||
|
|
pub transport: Arc<dyn ChatTransport>,
|
||
|
|
/// Routing and history key (Discord channel id, WhatsApp phone number).
|
||
|
|
pub key: &'a str,
|
||
|
|
/// Display name stored in conversation history and used in the prompt.
|
||
|
|
pub user: &'a str,
|
||
|
|
/// Shared services bundle (project root, permissions, agents).
|
||
|
|
pub services: &'a Arc<Services>,
|
||
|
|
/// Per-key conversation history.
|
||
|
|
pub history: &'a Arc<TokioMutex<HashMap<String, RoomConversation>>>,
|
||
|
|
/// Maximum number of conversation entries to keep per key.
|
||
|
|
pub history_size: usize,
|
||
|
|
/// Persists conversation history to disk.
|
||
|
|
pub save_history: fn(&std::path::Path, &HashMap<String, RoomConversation>),
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Forward a message to Claude Code and stream the response back through the
|
||
|
|
/// transport, handling permission prompts and conversation history.
|
||
|
|
///
|
||
|
|
/// `F` is taken by value (and cloned into the posting task) rather than by
|
||
|
|
/// reference because the posting task is a `tokio::spawn`ed future, which
|
||
|
|
/// requires `'static` — the formatter types are zero-sized markers, so
|
||
|
|
/// cloning is free.
|
||
|
|
pub async fn run_llm_turn<F: TransportFormat + Clone + 'static>(
|
||
|
|
ctx: LlmTurnContext<'_>,
|
||
|
|
format: F,
|
||
|
|
user_message: &str,
|
||
|
|
) {
|
||
|
|
use crate::chat::util::drain_complete_paragraphs;
|
||
|
|
use crate::llm::providers::claude_code::{ClaudeCodeProvider, ClaudeCodeResult};
|
||
|
|
use std::sync::atomic::{AtomicBool, Ordering};
|
||
|
|
use tokio::sync::watch;
|
||
|
|
|
||
|
|
let prefix = format.log_prefix();
|
||
|
|
|
||
|
|
// Look up existing session ID for this key.
|
||
|
|
let resume_session_id: Option<String> = {
|
||
|
|
let guard = ctx.history.lock().await;
|
||
|
|
guard.get(ctx.key).and_then(|conv| conv.session_id.clone())
|
||
|
|
};
|
||
|
|
|
||
|
|
let bot_name = &ctx.services.bot_name;
|
||
|
|
let persona = bot_name.to_lowercase();
|
||
|
|
let prompt = build_llm_prompt(&persona, bot_name, ctx.user, user_message);
|
||
|
|
|
||
|
|
let provider = ClaudeCodeProvider::new();
|
||
|
|
let (_cancel_tx, mut cancel_rx) = watch::channel(false);
|
||
|
|
|
||
|
|
// Channel for sending complete chunks to the posting task.
|
||
|
|
let (msg_tx, mut msg_rx) = tokio::sync::mpsc::unbounded_channel::<String>();
|
||
|
|
let msg_tx_for_callback = msg_tx.clone();
|
||
|
|
|
||
|
|
// Spawn a task to post messages as they arrive.
|
||
|
|
let post_transport = Arc::clone(&ctx.transport);
|
||
|
|
let post_key = ctx.key.to_string();
|
||
|
|
let post_format = format.clone();
|
||
|
|
let post_task = tokio::spawn(async move {
|
||
|
|
while let Some(chunk) = msg_rx.recv().await {
|
||
|
|
for part in post_format.format_and_chunk(&chunk) {
|
||
|
|
let _ = post_transport.send_message(&post_key, &part, "").await;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
});
|
||
|
|
|
||
|
|
// Shared buffer between the sync token callback and the async scope.
|
||
|
|
let buffer = Arc::new(std::sync::Mutex::new(String::new()));
|
||
|
|
let buffer_for_callback = Arc::clone(&buffer);
|
||
|
|
let sent_any_chunk = Arc::new(AtomicBool::new(false));
|
||
|
|
let sent_any_chunk_for_callback = Arc::clone(&sent_any_chunk);
|
||
|
|
|
||
|
|
let project_root_str = ctx.services.project_root.to_string_lossy().to_string();
|
||
|
|
let chat_fut = provider.chat_stream(
|
||
|
|
&prompt,
|
||
|
|
&project_root_str,
|
||
|
|
resume_session_id.as_deref(),
|
||
|
|
None,
|
||
|
|
None,
|
||
|
|
&mut cancel_rx,
|
||
|
|
move |token| {
|
||
|
|
let mut buf = buffer_for_callback.lock().unwrap();
|
||
|
|
buf.push_str(token);
|
||
|
|
let paragraphs = drain_complete_paragraphs(&mut buf);
|
||
|
|
for chunk in paragraphs {
|
||
|
|
sent_any_chunk_for_callback.store(true, Ordering::Relaxed);
|
||
|
|
let _ = msg_tx_for_callback.send(chunk);
|
||
|
|
}
|
||
|
|
},
|
||
|
|
|_thinking| {},
|
||
|
|
|_activity| {},
|
||
|
|
);
|
||
|
|
tokio::pin!(chat_fut);
|
||
|
|
|
||
|
|
// Register as a permission responder for the duration of this chat turn.
|
||
|
|
let (_perm_guard, mut perm_rx) = ctx.services.permission_registry.register();
|
||
|
|
|
||
|
|
let result = loop {
|
||
|
|
tokio::select! {
|
||
|
|
r = &mut chat_fut => break r,
|
||
|
|
|
||
|
|
Some(perm_fwd) = perm_rx.recv() => {
|
||
|
|
let prompt_msg = format!(
|
||
|
|
"**Permission Request**\n\nTool: `{}`\n```json\n{}\n```\n\nReply **yes** to approve or **no** to deny.",
|
||
|
|
perm_fwd.tool_name,
|
||
|
|
serde_json::to_string_pretty(&perm_fwd.tool_input)
|
||
|
|
.unwrap_or_else(|_| perm_fwd.tool_input.to_string()),
|
||
|
|
);
|
||
|
|
for part in format.format_and_chunk(&prompt_msg) {
|
||
|
|
let _ = ctx.transport.send_message(ctx.key, &part, "").await;
|
||
|
|
}
|
||
|
|
|
||
|
|
// Keyed by request_id (not just key) so a second
|
||
|
|
// concurrent request doesn't drop the first's sender.
|
||
|
|
ctx.services
|
||
|
|
.pending_perm_replies
|
||
|
|
.insert(ctx.key.to_string(), perm_fwd.request_id.clone(), perm_fwd.response_tx)
|
||
|
|
.await;
|
||
|
|
|
||
|
|
// Spawn a timeout task: auto-deny if the user does not respond.
|
||
|
|
let pending = Arc::clone(&ctx.services.pending_perm_replies);
|
||
|
|
let timeout_key = ctx.key.to_string();
|
||
|
|
let timeout_request_id = perm_fwd.request_id.clone();
|
||
|
|
let timeout_transport = Arc::clone(&ctx.transport);
|
||
|
|
let timeout_secs = ctx.services.permission_timeout_secs;
|
||
|
|
tokio::spawn(async move {
|
||
|
|
tokio::time::sleep(std::time::Duration::from_secs(timeout_secs)).await;
|
||
|
|
if let Some(tx) = pending.remove_by_request_id(&timeout_key, &timeout_request_id).await {
|
||
|
|
let _ = tx.send(PermissionDecision::Deny);
|
||
|
|
let msg = "Permission request timed out — denied (fail-closed).";
|
||
|
|
let _ = timeout_transport.send_message(&timeout_key, msg, "").await;
|
||
|
|
}
|
||
|
|
});
|
||
|
|
}
|
||
|
|
}
|
||
|
|
};
|
||
|
|
|
||
|
|
// Flush remaining text.
|
||
|
|
let remaining = buffer.lock().unwrap().trim().to_string();
|
||
|
|
let did_send_any = sent_any_chunk.load(Ordering::Relaxed);
|
||
|
|
|
||
|
|
let (assistant_reply, new_session_id) = match result {
|
||
|
|
Ok(ClaudeCodeResult {
|
||
|
|
messages,
|
||
|
|
session_id,
|
||
|
|
..
|
||
|
|
}) => {
|
||
|
|
let reply = if !remaining.is_empty() {
|
||
|
|
let _ = msg_tx.send(remaining.clone());
|
||
|
|
remaining
|
||
|
|
} else if !did_send_any {
|
||
|
|
let last_text = messages
|
||
|
|
.iter()
|
||
|
|
.rev()
|
||
|
|
.find(|m| m.role == crate::llm::types::Role::Assistant && !m.content.is_empty())
|
||
|
|
.map(|m| m.content.clone())
|
||
|
|
.unwrap_or_default();
|
||
|
|
if !last_text.is_empty() {
|
||
|
|
let _ = msg_tx.send(last_text.clone());
|
||
|
|
}
|
||
|
|
last_text
|
||
|
|
} else {
|
||
|
|
remaining
|
||
|
|
};
|
||
|
|
slog!("{prefix} session_id from chat_stream: {:?}", session_id);
|
||
|
|
(reply, session_id)
|
||
|
|
}
|
||
|
|
Err(e) => {
|
||
|
|
slog!("{prefix} LLM error: {e}");
|
||
|
|
let err_msg = format.format_error(&e);
|
||
|
|
let _ = msg_tx.send(err_msg.clone());
|
||
|
|
(err_msg, None)
|
||
|
|
}
|
||
|
|
};
|
||
|
|
|
||
|
|
// Signal the posting task to finish and wait for it.
|
||
|
|
drop(msg_tx);
|
||
|
|
let _ = post_task.await;
|
||
|
|
|
||
|
|
// Record this exchange in conversation history.
|
||
|
|
if !assistant_reply.starts_with("Error processing") {
|
||
|
|
let mut guard = ctx.history.lock().await;
|
||
|
|
let conv = guard.entry(ctx.key.to_string()).or_default();
|
||
|
|
|
||
|
|
if new_session_id.is_some() {
|
||
|
|
conv.session_id = new_session_id;
|
||
|
|
}
|
||
|
|
|
||
|
|
conv.entries.push(ConversationEntry {
|
||
|
|
role: ConversationRole::User,
|
||
|
|
sender: ctx.user.to_string(),
|
||
|
|
content: user_message.to_string(),
|
||
|
|
});
|
||
|
|
conv.entries.push(ConversationEntry {
|
||
|
|
role: ConversationRole::Assistant,
|
||
|
|
sender: String::new(),
|
||
|
|
content: assistant_reply,
|
||
|
|
});
|
||
|
|
|
||
|
|
// Trim to configured maximum.
|
||
|
|
if conv.entries.len() > ctx.history_size {
|
||
|
|
let excess = conv.entries.len() - ctx.history_size;
|
||
|
|
conv.entries.drain(..excess);
|
||
|
|
}
|
||
|
|
|
||
|
|
(ctx.save_history)(&ctx.services.project_root, &guard);
|
||
|
|
}
|
||
|
|
}
|