use std::{collections::HashMap, time::Duration};
use rig::{
agent::Agent,
completion::{Message, Prompt},
providers::anthropic::completion::CompletionModel,
};
use serde::Deserialize;
use crate::SharedRcon;
const DIM: &str = "\x1b[2m";
const RESET: &str = "\x1b[0m";
#[derive(Debug, Deserialize)]
struct SenseiMessage {
player: String,
message: String,
}
pub async fn run(rcon: SharedRcon, sensei: Agent<CompletionModel>, poll_interval: Duration) {
let mut histories: HashMap<String, Vec<Message>> = HashMap::new();
let mut consecutive_errors: u32 = 0;
loop {
match poll_messages(&rcon).await {
Ok(messages) => {
consecutive_errors = 0;
for msg in messages {
handle_message(&rcon, &sensei, &mut histories, &msg).await;
}
}
Err(e) => {
consecutive_errors += 1;
if consecutive_errors <= 3 {
eprintln!("{DIM}[Bridge] Poll error: {e}{RESET}");
} else if consecutive_errors == 4 {
eprintln!("{DIM}[Bridge] Repeated errors, backing off to 10s intervals{RESET}");
}
}
}
let sleep_duration = if consecutive_errors > 3 {
Duration::from_secs(10)
} else {
poll_interval
};
tokio::time::sleep(sleep_duration).await;
}
}
async fn poll_messages(rcon: &SharedRcon) -> Result<Vec<SenseiMessage>, BridgeError> {
let response = rcon.lock().await.execute("/sensei_poll").await?;
let trimmed = response.trim();
if trimmed.is_empty() || trimmed == "[]" {
return Ok(Vec::new());
}
if trimmed.contains("\"error\"") {
return Err(BridgeError::Lua(trimmed.to_string()));
}
let messages: Vec<SenseiMessage> = serde_json::from_str(trimmed)?;
Ok(messages)
}
async fn handle_message(
rcon: &SharedRcon,
sensei: &Agent<CompletionModel>,
histories: &mut HashMap<String, Vec<Message>>,
msg: &SenseiMessage,
) {
eprintln!("{DIM}[Bridge] {}: {}{RESET}", msg.player, msg.message);
let history = histories.entry(msg.player.clone()).or_default();
let prompt = format!(
"[In-game message from player {}] {}",
msg.player, msg.message
);
match sensei.prompt(&prompt).with_history(history).await {
Ok(response) => {
let sanitized = sanitize_for_game(&response);
if let Err(e) = send_response(rcon, &sanitized).await {
eprintln!("{DIM}[Bridge] Failed to send response: {e}{RESET}");
}
}
Err(e) => {
eprintln!("{DIM}[Bridge] Agent error: {e}{RESET}");
let _ = send_response(
rcon,
"Sorry, I encountered an error processing your question.",
)
.await;
}
}
}
async fn send_response(rcon: &SharedRcon, message: &str) -> Result<(), BridgeError> {
let command = format!("/sensei_respond {message}");
rcon.lock().await.execute(&command).await?;
Ok(())
}
fn sanitize_for_game(response: &str) -> String {
let mut result = response.to_string();
result = result.replace("```", "");
result = result.replace("**", "");
result = result.replace("__", "");
result = result.replace('`', "");
for prefix in ["#### ", "### ", "## ", "# "] {
result = result.replace(prefix, "");
}
result = result.replace('[', "(");
result = result.replace(']', ")");
result = result
.lines()
.map(str::trim)
.filter(|l| !l.is_empty())
.collect::<Vec<_>>()
.join(" | ");
while result.contains(" ") {
result = result.replace(" ", " ");
}
if result.len() > 1000 {
let mut end = 1000;
while !result.is_char_boundary(end) {
end -= 1;
}
result.truncate(end);
result.push_str("...");
}
result.trim().to_string()
}
#[derive(Debug)]
enum BridgeError {
Rcon(factorio_rcon::RconError),
Json(serde_json::Error),
Lua(String),
}
impl std::fmt::Display for BridgeError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Rcon(e) => write!(f, "RCON: {e}"),
Self::Json(e) => write!(f, "JSON parse: {e}"),
Self::Lua(e) => write!(f, "Lua: {e}"),
}
}
}
impl From<factorio_rcon::RconError> for BridgeError {
fn from(e: factorio_rcon::RconError) -> Self {
Self::Rcon(e)
}
}
impl From<serde_json::Error> for BridgeError {
fn from(e: serde_json::Error) -> Self {
Self::Json(e)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn strips_bold_markers() {
assert_eq!(sanitize_for_game("use **bold** text"), "use bold text");
}
#[test]
fn strips_code_fences() {
let input = "before\n```lua\nprint('hi')\n```\nafter";
let result = sanitize_for_game(input);
assert!(result.contains("print('hi')"));
assert!(!result.contains("```"));
}
#[test]
fn strips_inline_backticks() {
assert_eq!(sanitize_for_game("use `iron-plate`"), "use iron-plate");
}
#[test]
fn strips_markdown_headers() {
assert_eq!(sanitize_for_game("## Analysis"), "Analysis");
}
#[test]
fn flattens_newlines() {
let input = "line one\nline two\nline three";
assert_eq!(sanitize_for_game(input), "line one | line two | line three");
}
#[test]
fn truncates_long_text() {
let long = "a".repeat(1500);
let result = sanitize_for_game(&long);
assert!(result.len() <= 1003); assert!(result.ends_with("..."));
}
#[test]
fn collapses_whitespace() {
assert_eq!(sanitize_for_game("too many spaces"), "too many spaces");
}
#[test]
fn escapes_brackets() {
assert_eq!(
sanitize_for_game("use [item=iron-plate]"),
"use (item=iron-plate)"
);
}
#[test]
fn empty_input() {
assert_eq!(sanitize_for_game(""), "");
}
#[test]
fn strips_blank_lines() {
let input = "first\n\n\nsecond";
assert_eq!(sanitize_for_game(input), "first | second");
}
#[test]
fn truncates_multibyte_safely() {
let mut input = "a".repeat(999);
input.push('€'); input.push_str("tail");
let result = sanitize_for_game(&input);
assert!(result.ends_with("..."));
assert!(result.len() <= 1003);
}
#[test]
fn combined_markdown() {
let input = "## Tips\n\n**First**, use `assembler`.\n\n### Next\n\nBuild more.";
let result = sanitize_for_game(input);
assert!(!result.contains("**"));
assert!(!result.contains("##"));
assert!(!result.contains('`'));
assert!(result.contains("First"));
assert!(result.contains("assembler"));
}
#[test]
fn deserialize_single_message() {
let json = r#"[{"player":"nick","message":"what now?","tick":12345}]"#;
let msgs: Vec<SenseiMessage> = serde_json::from_str(json).unwrap();
assert_eq!(msgs.len(), 1);
assert_eq!(msgs[0].player, "nick");
assert_eq!(msgs[0].message, "what now?");
}
#[test]
fn deserialize_multiple_messages() {
let json = r#"[
{"player":"alice","message":"help","tick":100},
{"player":"bob","message":"tips?","tick":200}
]"#;
let msgs: Vec<SenseiMessage> = serde_json::from_str(json).unwrap();
assert_eq!(msgs.len(), 2);
assert_eq!(msgs[1].player, "bob");
}
#[test]
fn deserialize_empty_array() {
let json = "[]";
let msgs: Vec<SenseiMessage> = serde_json::from_str(json).unwrap();
assert!(msgs.is_empty());
}
}