huskies: merge 1247 refactor Deduplicate chat-transport LLM command handling (discord <-> whatsapp 45-line clone)
This commit is contained in:
Huskies Agent
@@ -0,0 +1,250 @@
|
||||
//! 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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user