use crate::command::chat::storage::ChatMessage;
use crate::util::log::write_info_log;
use std::cell::RefCell;
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::{
Arc, Mutex, OnceLock,
atomic::{AtomicBool, Ordering},
};
use tokio_util::sync::CancellationToken;
thread_local! {
static CURRENT_AGENT_NAME: RefCell<String> = RefCell::new("Main".to_string());
static THREAD_CWD: RefCell<Option<PathBuf>> = const { RefCell::new(None) };
}
pub fn set_current_agent_name(name: &str) {
CURRENT_AGENT_NAME.with(|cell| {
*cell.borrow_mut() = name.to_string();
});
}
pub fn current_agent_name() -> String {
CURRENT_AGENT_NAME.with(|cell| cell.borrow().clone())
}
pub fn set_thread_cwd(path: &std::path::Path) {
THREAD_CWD.with(|cell| {
*cell.borrow_mut() = Some(path.to_path_buf());
});
}
pub fn thread_cwd() -> Option<PathBuf> {
THREAD_CWD.with(|cell| cell.borrow().clone())
}
pub fn clear_thread_cwd() {
THREAD_CWD.with(|cell| {
*cell.borrow_mut() = None;
});
}
static GLOBAL_FILE_LOCKS: OnceLock<Mutex<HashMap<PathBuf, String>>> = OnceLock::new();
fn global_file_locks() -> &'static Mutex<HashMap<PathBuf, String>> {
GLOBAL_FILE_LOCKS.get_or_init(|| Mutex::new(HashMap::new()))
}
pub fn acquire_global_file_lock(
path: &std::path::Path,
agent_name: &str,
) -> Result<FileLockGuard, String> {
let canonical = path.to_path_buf();
let mut map = global_file_locks()
.lock()
.map_err(|_| "file_locks mutex poisoned".to_string())?;
if let Some(holder) = map.get(&canonical)
&& holder != agent_name
{
return Err(holder.clone());
}
map.insert(canonical.clone(), agent_name.to_string());
Ok(FileLockGuard {
path: canonical,
agent_name: agent_name.to_string(),
})
}
#[allow(dead_code)]
pub struct TeammateHandle {
pub name: String,
pub role: String,
pub pending_user_messages: Arc<Mutex<Vec<ChatMessage>>>,
pub streaming_content: Arc<Mutex<String>>,
pub cancel_token: CancellationToken,
pub is_running: Arc<AtomicBool>,
pub thread_handle: Option<std::thread::JoinHandle<()>>,
pub system_prompt_snapshot: Arc<Mutex<String>>,
pub messages_snapshot: Arc<Mutex<Vec<ChatMessage>>>,
}
#[allow(dead_code)]
impl TeammateHandle {
pub fn running(&self) -> bool {
self.is_running.load(Ordering::Relaxed)
}
pub fn cancel(&self) {
self.cancel_token.cancel();
}
}
pub struct FileLockGuard {
path: PathBuf,
agent_name: String,
}
impl Drop for FileLockGuard {
fn drop(&mut self) {
if let Ok(mut map) = global_file_locks().lock()
&& map.get(&self.path).map(|s| s.as_str()) == Some(self.agent_name.as_str())
{
map.remove(&self.path);
}
}
}
#[allow(dead_code)]
pub struct TeammateManager {
pub teammates: HashMap<String, TeammateHandle>,
pub main_pending: Arc<Mutex<Vec<ChatMessage>>>,
pub shared_messages: Arc<Mutex<Vec<ChatMessage>>>,
}
#[allow(dead_code)]
impl TeammateManager {
pub fn new(
main_pending: Arc<Mutex<Vec<ChatMessage>>>,
shared_messages: Arc<Mutex<Vec<ChatMessage>>>,
) -> Self {
Self {
teammates: HashMap::new(),
main_pending,
shared_messages,
}
}
pub fn broadcast(&self, from: &str, text: &str, at_target: Option<&str>) {
let formatted = if let Some(target) = at_target {
format!("<{}> @{} {}", from, target, text)
} else {
format!("<{}> {}", from, text)
};
write_info_log(
"TeammateManager",
&format!(
"broadcast from={}: {}",
from,
&formatted[..{
let mut b = formatted.len().min(100);
while b > 0 && !formatted.is_char_boundary(b) {
b -= 1;
}
b
}]
),
);
if from != "Main"
&& let Ok(mut pending) = self.main_pending.lock()
{
pending.push(ChatMessage::text("user", &formatted));
}
for (name, handle) in &self.teammates {
if name == from {
continue; }
if let Ok(mut pending) = handle.pending_user_messages.lock() {
pending.push(ChatMessage::text("user", &formatted));
}
}
if from != "Main"
&& let Ok(mut shared) = self.shared_messages.lock()
{
shared.push(ChatMessage::text("assistant", &formatted));
}
}
pub fn team_summary(&self) -> String {
if self.teammates.is_empty() {
return String::new();
}
let mut summary = String::from("## Teammates\n\n当前团队成员:\n");
summary.push_str("- Main (主协调者)\n");
for (name, handle) in &self.teammates {
let status = if handle.running() {
"工作中"
} else {
"空闲"
};
summary.push_str(&format!("- {} ({}) [{}]\n", name, handle.role, status));
}
summary.push_str(
"\n使用 SendMessage 工具向其他 agent 发送消息。可以用 @AgentName 指定目标。\n",
);
summary
}
pub fn all_names(&self) -> Vec<String> {
let mut names = vec!["Main".to_string()];
names.extend(self.teammates.keys().cloned());
names
}
pub fn stop_teammate(&mut self, name: &str) {
if let Some(handle) = self.teammates.get(name) {
handle.cancel();
write_info_log("TeammateManager", &format!("stopped teammate: {}", name));
}
}
pub fn stop_all(&mut self) {
for (name, handle) in &self.teammates {
handle.cancel();
write_info_log("TeammateManager", &format!("stopping teammate: {}", name));
}
}
pub fn cleanup_finished(&mut self) {
let finished: Vec<String> = self
.teammates
.iter()
.filter(|(_, h)| {
!h.running()
&& h.thread_handle
.as_ref()
.map(|t| t.is_finished())
.unwrap_or(true)
})
.map(|(name, _)| name.clone())
.collect();
for name in finished {
if let Some(mut handle) = self.teammates.remove(&name) {
if let Some(th) = handle.thread_handle.take() {
let _ = th.join();
}
write_info_log("TeammateManager", &format!("cleaned up teammate: {}", name));
}
}
}
pub fn register_teammate(&mut self, handle: TeammateHandle) {
write_info_log(
"TeammateManager",
&format!("registered teammate: {} ({})", handle.name, handle.role),
);
self.teammates.insert(handle.name.clone(), handle);
}
}
impl Default for TeammateManager {
fn default() -> Self {
Self {
teammates: HashMap::new(),
main_pending: Arc::new(Mutex::new(Vec::new())),
shared_messages: Arc::new(Mutex::new(Vec::new())),
}
}
}