ai/log
2
0
Fork 0
Code Activity

fix bot --user

Browse Source
This commit is contained in:
syui
2026-03-23 03:49:07 +09:00
parent 2c867eaf5c
commit 2bf780cea0
2 changed files with 99 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

109
src/commands/bot.rs
View File

@@ -1,7 +1,7 @@
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashSet;
use std::collections::{HashMap, HashSet};
use std::fs;
use std::path::PathBuf;
use std::process::Stdio;
@@ -27,7 +27,7 @@ struct ClaudeSession {
impl ClaudeSession {
/// Spawn a persistent claude process with stream-json I/O
async fn spawn() -> Result<Self> {
async fn spawn(user_mode: bool) -> Result<Self> {
// Run claude inside a dedicated directory under config
let work_dir = dirs::config_dir()
.context("Could not find config directory")?
@@ -39,15 +39,50 @@ impl ClaudeSession {
let rules_path = work_dir.join("CLAUDE.md");
fs::write(&rules_path, BOT_RULES)?;
eprintln!("bot: claude working directory = {}", work_dir.display());
// Manage settings.json: create in user mode, remove otherwise
let settings_dir = work_dir.join(".claude");
let settings_path = settings_dir.join("settings.json");
if user_mode {
fs::create_dir_all(&settings_dir)?;
let settings = serde_json::json!({
"permissions": {
"deny": [
"Bash",
"Edit(.claude/**)",
"Write(.claude/**)",
"Edit(.mcp.json)",
"Write(.mcp.json)"
],
"allow": [
"Read",
"Glob",
"Grep"
]
}
});
fs::write(&settings_path, serde_json::to_string_pretty(&settings)?)?;
} else {
// Remove restrictive settings from previous --user run
let _ = fs::remove_file(&settings_path);
}
let mut child = tokio::process::Command::new("claude")
.arg("--input-format")
eprintln!("bot: claude working directory = {}", work_dir.display());
eprintln!("bot: user_mode = {}", user_mode);
let mut cmd = tokio::process::Command::new("claude");
cmd.arg("--input-format")
.arg("stream-json")
.arg("--output-format")
.arg("stream-json")
.arg("--verbose")
.arg("--dangerously-skip-permissions")
.arg("--verbose");
if user_mode {
cmd.arg("--permission-mode").arg("dontAsk");
} else {
cmd.arg("--dangerously-skip-permissions");
}
let mut child = cmd
.current_dir(&work_dir)
.stdin(Stdio::piped())
.stdout(Stdio::piped())
@@ -187,6 +222,9 @@ struct BotState {
/// Timestamp of the last seen notification
#[serde(default)]
last_seen: Option<String>,
/// Per-user daily message counts: DID -> (date "YYYY-MM-DD", count)
#[serde(default)]
user_counts: HashMap<String, (String, u32)>,
}
/// Parsed notification for processing
@@ -242,6 +280,33 @@ fn load_admin_did() -> Result<String> {
.context("config.json missing 'did' field")
}
/// Load daily rate limit from config.json (bot.limit, default 3)
fn load_user_limit() -> u32 {
token::load_config()
.ok()
.and_then(|c| c["bot"]["limit"].as_u64())
.unwrap_or(3) as u32
}
/// Check if a user is within daily rate limit. Returns true if allowed.
/// Increments count on success.
fn check_rate_limit(state: &mut BotState, did: &str, limit: u32) -> bool {
let today = chrono::Utc::now().format("%Y-%m-%d").to_string();
let entry = state.user_counts.entry(did.to_string()).or_insert_with(|| (today.clone(), 0));
// Reset if new day
if entry.0 != today {
*entry = (today, 0);
}
if entry.1 >= limit {
false
} else {
entry.1 += 1;
true
}
}
/// Fetch notifications using the bot session
async fn fetch_notifications(
client: &XrpcClient,
@@ -340,20 +405,22 @@ async fn post_reply(
}
/// Main bot entry point
pub async fn start(interval_secs: u64) -> Result<()> {
pub async fn start(interval_secs: u64, user_mode: bool) -> Result<()> {
let admin_did = load_admin_did()?;
let user_limit = if user_mode { load_user_limit() } else { 0 };
eprintln!("bot: admin DID = {}", admin_did);
eprintln!("bot: user_mode = {}, limit = {}/day", user_mode, user_limit);
eprintln!("bot: polling interval = {}s", interval_secs);
let mut state = load_state();
// Spawn persistent Claude session
let mut session = ClaudeSession::spawn().await?;
let mut session = ClaudeSession::spawn(user_mode).await?;
eprintln!("bot: claude session started");
eprintln!("bot: starting notification loop...");
loop {
if let Err(e) = poll_once(&admin_did, &mut state, &mut session).await {
if let Err(e) = poll_once(&admin_did, &mut state, &mut session, user_mode, user_limit).await {
eprintln!("bot: poll error: {}", e);
}
tokio::time::sleep(std::time::Duration::from_secs(interval_secs)).await;
@@ -365,6 +432,8 @@ async fn poll_once(
admin_did: &str,
state: &mut BotState,
claude: &mut ClaudeSession,
user_mode: bool,
user_limit: u32,
) -> Result<()> {
// Refresh bot session
let session = auth::refresh_bot_session().await?;
@@ -403,12 +472,22 @@ async fn poll_once(
}
}
// Admin filter: only respond to admin
let author_did = match notif["author"]["did"].as_str() {
Some(d) => d,
None => continue,
};
if author_did != admin_did {
if author_did == admin_did {
// Admin: always pass
} else if user_mode {
// User mode: check rate limit
if !check_rate_limit(state, author_did, user_limit) {
eprintln!("bot: rate limited user {} (limit {}/day)", author_did, user_limit);
state.processed.insert(uri.to_string());
continue;
}
} else {
// Admin-only mode: skip non-admin
continue;
}
@@ -437,7 +516,7 @@ async fn poll_once(
Err(e) => {
eprintln!("bot: claude error: {}, respawning session...", e);
// Try to respawn session
match ClaudeSession::spawn().await {
match ClaudeSession::spawn(user_mode).await {
Ok(new_session) => {
*claude = new_session;
eprintln!("bot: claude session respawned");
@@ -584,7 +663,7 @@ pub async fn start_chat(interval_secs: u64) -> Result<()> {
let mut state = load_chat_state();
// Spawn persistent Claude session
let mut claude = ClaudeSession::spawn().await?;
let mut claude = ClaudeSession::spawn(false).await?;
eprintln!("chat-bot: claude session started");
eprintln!("chat-bot: starting chat poll loop...");
@@ -692,7 +771,7 @@ async fn chat_poll_once(
Ok(r) => r,
Err(e) => {
eprintln!("chat-bot: claude error: {}, respawning...", e);
match ClaudeSession::spawn().await {
match ClaudeSession::spawn(false).await {
Ok(new_session) => {
*claude = new_session;
match claude.send(text).await {

7
src/main.rs
View File

@@ -287,6 +287,9 @@ enum BotCommands {
/// Poll interval in seconds
#[arg(short, long, default_value = "30")]
interval: u64,
/// Allow all users (with daily rate limit from config bot.limit)
#[arg(long)]
user: bool,
},
/// Start the DM chat bot (poll chat messages and reply)
Chat {
@@ -412,8 +415,8 @@ async fn main() -> Result<()> {
}
Commands::Bot { command } => {
match command {
BotCommands::Start { interval } => {
commands::bot::start(interval).await?;
BotCommands::Start { interval, user } => {
commands::bot::start(interval, user).await?;
}
BotCommands::Chat { interval } => {
commands::bot::start_chat(interval).await?;