mod api_auth;
mod api_codingplan;
mod api_config;
mod api_provider;
pub(crate) mod live_api;
pub use live_api::current_live_session;
pub use live_api::ensure_live_session;
pub use live_api::ensure_live_session_seeded;
pub use live_api::live_set_provider;
mod telemetry_scope;
pub mod auth_token;
pub mod permission_bridge;
pub mod webui;
pub(crate) use telemetry_scope::daemon_scope;
use axum::{
extract::{Path, Query, State},
http::{header, request::Parts as RequestParts, HeaderValue, Method, StatusCode},
response::{sse::Sse, IntoResponse, Json},
routing::{delete, get, post},
Router,
};
use futures::stream::StreamExt;
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{mpsc, watch, RwLock};
use tokio_stream::wrappers::UnboundedReceiverStream;
use tokio_util::sync::CancellationToken;
use tower_http::cors::{AllowOrigin, CorsLayer};
use atomcode_core::config::Config;
use atomcode_core::conversation::Conversation;
use atomcode_core::lsp::manager::build_lsp_manager;
use atomcode_core::mcp::{register_mcp_tools, McpRegistry};
use atomcode_core::provider;
use atomcode_core::session::{Session, SessionId, SessionManager, SessionMeta};
use atomcode_core::tool::diagnostics::DiagnosticsTool;
use atomcode_core::tool::{ToolContext, ToolRegistry};
use atomcode_core::turn::event::{TurnEvent, TurnResult};
use atomcode_core::turn::permission::{
ApprovalRequest, AutoPermissionDecider, AutoPermissionMode, InteractivePermissionDecider,
PermissionDecider,
};
use atomcode_core::turn::runner::TurnRunner;
use atomcode_telemetry::{
config::{resolve, ProcessEnv},
CliOverride, CurrentContext, Event, RepoOrigin, SessionMode,
Telemetry, TelemetryState,
};
use atomcode_core::auth;
use atomcode_core::telemetry_bootstrap::detect_repo_origin;
#[derive(Debug, Serialize)]
pub(crate) struct ApiError {
pub success: bool,
pub error: String,
}
#[derive(Debug, Serialize)]
pub(crate) struct ConfigResponse {
pub path: PathBuf,
pub default_provider: String,
pub default_workdir: Option<String>,
pub providers: Vec<ProviderInfo>,
}
#[derive(Debug, Serialize)]
pub(crate) struct ProviderInfo {
pub name: String,
#[serde(rename = "type")]
pub provider_type: String,
pub model: String,
pub base_url: Option<String>,
pub has_api_key: bool,
pub is_default: bool,
pub context_window: usize,
pub max_tokens: Option<usize>,
pub thinking_enabled: Option<bool>,
pub thinking_budget: Option<u32>,
pub thinking_type: Option<String>,
pub thinking_keep: Option<String>,
pub reasoning_history: Option<String>,
pub skip_tls_verify: bool,
pub ephemeral: bool,
}
pub struct LoginSessionEntry {
pub session: atomcode_core::auth::LoginSession,
pub created_at: std::time::Instant,
}
pub(crate) type LoginSessionsStore = Arc<RwLock<HashMap<String, LoginSessionEntry>>>;
pub(crate) fn json_error(
status: StatusCode,
message: impl Into<String>,
) -> (StatusCode, Json<ApiError>) {
(
status,
Json(ApiError {
success: false,
error: message.into(),
}),
)
}
#[derive(Debug, Clone, Serialize)]
pub struct ProjectInfo {
pub hash: String,
pub name: String,
pub working_dir: PathBuf,
#[serde(skip_serializing_if = "Option::is_none")]
pub description: Option<String>,
pub session_count: usize,
pub created_at: u64,
pub last_updated: u64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProjectState {
pub working_dir: PathBuf,
#[serde(skip_serializing_if = "Option::is_none")]
pub previous_dir: Option<PathBuf>,
pub recent_dirs: Vec<PathBuf>,
pub name: String,
}
#[derive(Debug, Deserialize)]
pub struct ChangeDirRequest {
pub path: String,
#[serde(default)]
pub set_default: bool,
}
#[derive(Debug, Serialize)]
pub struct ChangeDirResponse {
pub success: bool,
pub message: String,
pub current_dir: PathBuf,
pub project_hash: String,
}
#[derive(Debug, Deserialize)]
pub struct SearchQuery {
pub q: String,
}
#[derive(Debug, Deserialize)]
pub struct CreateSessionRequest {
#[serde(default)]
pub working_dir: Option<PathBuf>,
#[serde(default)]
pub title: Option<String>,
}
#[derive(Debug, Serialize)]
pub struct CreateSessionResponse {
pub id: String,
pub name: String,
pub working_dir: PathBuf,
pub project_hash: String,
pub created_at: u64,
}
#[derive(Debug, Serialize)]
pub struct SessionDetail {
pub id: String,
pub name: String,
pub working_dir: PathBuf,
pub created_at: u64,
pub updated_at: u64,
pub message_count: usize,
pub messages: Vec<MessageInfo>,
}
type ProjectStateStore = Arc<RwLock<ProjectState>>;
type ChatTasksStore = Arc<RwLock<HashMap<String, CancellationToken>>>;
type StoppedSessionsStore = Arc<RwLock<HashSet<String>>>;
const DANGEROUS_TOOLS_ENV: &str = "ATOMCODE_DAEMON_ENABLE_DANGEROUS_TOOLS";
struct SseConnectionGuard(Arc<std::sync::atomic::AtomicUsize>);
impl Drop for SseConnectionGuard {
fn drop(&mut self) {
self.0.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
}
}
#[derive(Clone)]
pub struct AppState {
pub sessions: SessionStore,
pub project: ProjectStateStore,
pub chat_tasks: ChatTasksStore,
pub stopped_sessions: StoppedSessionsStore,
pub mcp_registry: Arc<RwLock<Arc<McpRegistry>>>,
pub mcp_cache: Arc<RwLock<HashMap<PathBuf, CachedMcpRegistry>>>,
pub login_sessions: LoginSessionsStore,
pub telemetry: Arc<Telemetry>,
pub repo_origin: RepoOrigin,
pub shutdown_tx: watch::Sender<bool>,
pub last_activity: Arc<std::sync::atomic::AtomicI64>,
pub active_connections: Arc<std::sync::atomic::AtomicUsize>,
pub webui_tokens: auth_token::WebuiTokenStore,
pub enforce_token: bool,
pub pending_permissions: permission_bridge::PermissionResponders,
pub bind_host: String,
pub bind_port: u16,
}
pub struct CachedMcpRegistry {
pub registry: Arc<McpRegistry>,
pub last_used: std::time::Instant,
}
const MCP_CACHE_MAX: usize = 5;
fn default_working_dir() -> PathBuf {
std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."))
}
fn resolve_initial_working_dir(
override_dir: Option<PathBuf>,
config_default: Option<PathBuf>,
cwd: PathBuf,
) -> PathBuf {
if let Some(o) = override_dir {
if o.exists() {
return o;
}
}
if let Some(d) = config_default {
if d.exists() {
return d;
}
}
cwd
}
fn init_project_state(override_dir: Option<PathBuf>) -> ProjectState {
let config_default = Config::load(&Config::default_path())
.ok()
.and_then(|c| c.default_workdir.map(PathBuf::from));
let path = resolve_initial_working_dir(override_dir, config_default, default_working_dir());
let name = path
.file_name()
.map(|n| n.to_string_lossy().to_string())
.unwrap_or_else(|| "project".to_string());
ProjectState {
working_dir: path,
previous_dir: None,
recent_dirs: vec![],
name,
}
}
#[derive(Debug, Serialize, Clone)]
pub struct ArtifactInfo {
pub id: String,
pub artifact_type: String,
pub title: Option<String>,
pub language: Option<String>,
pub content: String,
}
#[derive(Debug, Serialize)]
pub struct ToolCallInfo {
pub id: String,
pub name: String,
pub arguments: String,
pub display: String,
}
#[derive(Debug, Serialize)]
pub struct ToolResultInfo {
pub call_id: String,
pub success: bool,
pub summary: String,
pub line_count: usize,
}
#[derive(Debug, Serialize)]
pub struct MessageInfo {
pub role: String,
pub content: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_calls: Option<Vec<ToolCallInfo>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_result: Option<ToolResultInfo>,
#[serde(skip_serializing_if = "Option::is_none")]
pub artifacts: Option<Vec<ArtifactInfo>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub images: Option<Vec<ImageData>>,
}
#[derive(Debug, Serialize)]
pub struct ImageData {
pub media_type: String,
pub data: String,
}
impl From<&atomcode_core::conversation::message::Message> for MessageInfo {
fn from(msg: &atomcode_core::conversation::message::Message) -> Self {
let role = match msg.role {
atomcode_core::conversation::message::Role::System => "system",
atomcode_core::conversation::message::Role::User => "user",
atomcode_core::conversation::message::Role::Assistant => "assistant",
atomcode_core::conversation::message::Role::Tool => "tool",
};
let (content, tool_calls, tool_result, artifacts) = match &msg.content {
atomcode_core::conversation::message::MessageContent::Text(s) => {
(s.clone(), None, None, None)
}
atomcode_core::conversation::message::MessageContent::AssistantWithToolCalls {
text,
tool_calls,
..
} => {
let calls: Vec<ToolCallInfo> = tool_calls
.iter()
.map(|tc| ToolCallInfo {
id: tc.id.clone(),
name: tc.name.clone(),
arguments: tc.arguments.clone(),
display: format_tool_args(&tc.name, &tc.arguments),
})
.collect();
let artifacts = extract_artifacts_from_tool_calls(tool_calls);
(
text.clone().unwrap_or_default(),
Some(calls),
None,
artifacts,
)
}
atomcode_core::conversation::message::MessageContent::ToolResult(r) => {
let lines = r.output.lines().count();
let first_line = r.output.lines().next().unwrap_or("");
let summary = if first_line.len() > 100 {
format!("{}...", first_line.chars().take(97).collect::<String>())
} else {
first_line.to_string()
};
(
r.output.clone(),
None,
Some(ToolResultInfo {
call_id: r.call_id.clone(),
success: r.success,
summary,
line_count: lines,
}),
None,
)
}
atomcode_core::conversation::message::MessageContent::ToolResultRef(r) => {
(r.summary.clone(), None, None, None)
}
atomcode_core::conversation::message::MessageContent::MultiPart { text, images } => {
let raw = text.clone().unwrap_or_default();
let display = if images.is_empty() {
raw
} else {
match raw.find("[图片内容(由").or_else(|| raw.find("[图片识别失败]")) {
Some(i) => raw[..i].trim_end().to_string(),
None => raw,
}
};
(display, None, None, None)
}
};
let images = match &msg.content {
atomcode_core::conversation::message::MessageContent::MultiPart { images, .. }
if !images.is_empty() =>
{
Some(
images
.iter()
.map(|i| ImageData {
media_type: i.media_type.clone(),
data: i.data.clone(),
})
.collect(),
)
}
_ => None,
};
Self {
role: role.to_string(),
content,
tool_calls,
tool_result,
artifacts,
images,
}
}
}
fn extract_artifacts_from_tool_calls(
tool_calls: &[atomcode_core::tool::ToolCall],
) -> Option<Vec<ArtifactInfo>> {
let mut artifacts = Vec::new();
for tc in tool_calls {
if tc.name == "create_file" || tc.name == "edit_file" {
let args: serde_json::Value = match serde_json::from_str(&tc.arguments) {
Ok(v) => v,
Err(_) => continue,
};
let path = match args.get("file_path").and_then(|v| v.as_str()) {
Some(p) => p,
None => continue,
};
let (artifact_type, language) = if path.ends_with(".html") || path.ends_with(".htm") {
("html", "html")
} else if path.ends_with(".svg") {
("svg", "xml")
} else if path.ends_with(".md") || path.ends_with(".markdown") {
("markdown", "markdown")
} else if path.ends_with(".pptx") {
("pptx", "pptx")
} else if path.ends_with(".docx") {
("docx", "docx")
} else if path.ends_with(".xlsx") {
("xlsx", "xlsx")
} else if path.ends_with(".pdf") {
("pdf", "pdf")
} else {
continue;
};
let content = args
.get("content")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let title = PathBuf::from(path)
.file_name()
.map(|n| n.to_string_lossy().to_string());
artifacts.push(ArtifactInfo {
id: format!("file-{}", artifacts.len() + 1),
artifact_type: artifact_type.to_string(),
title,
language: Some(language.to_string()),
content,
});
} else if tc.name == "bash" {
let args: serde_json::Value = match serde_json::from_str(&tc.arguments) {
Ok(v) => v,
Err(_) => continue,
};
let command = match args.get("command").and_then(|v| v.as_str()) {
Some(c) => c,
None => continue,
};
if let Some(path) = extract_output_file_from_bash(command) {
let (artifact_type, language) = if path.ends_with(".html") || path.ends_with(".htm")
{
("html", "html")
} else if path.ends_with(".svg") {
("svg", "xml")
} else if path.ends_with(".md") || path.ends_with(".markdown") {
("markdown", "markdown")
} else if path.ends_with(".pptx") {
("pptx", "pptx")
} else if path.ends_with(".docx") {
("docx", "docx")
} else if path.ends_with(".xlsx") {
("xlsx", "xlsx")
} else if path.ends_with(".pdf") {
("pdf", "pdf")
} else {
continue;
};
let title = PathBuf::from(&path)
.file_name()
.map(|n| n.to_string_lossy().to_string());
artifacts.push(ArtifactInfo {
id: format!("file-{}", artifacts.len() + 1),
artifact_type: artifact_type.to_string(),
title,
language: Some(language.to_string()),
content: String::new(),
});
}
}
}
if artifacts.is_empty() {
None
} else {
Some(artifacts)
}
}
fn extract_output_file_from_bash(command: &str) -> Option<String> {
let artifact_extensions = [
".html",
".htm",
".svg",
".md",
".markdown",
".pptx",
".docx",
".xlsx",
".pdf",
];
let chars: Vec<char> = command.chars().collect();
let mut i = 0;
while i < chars.len() {
if chars[i] == '>' {
let append_mode = i + 1 < chars.len() && chars[i + 1] == '>';
let start = if append_mode { i + 2 } else { i + 1 };
let mut j = start;
while j < chars.len() && chars[j].is_whitespace() {
j += 1;
}
let mut path_end = j;
while path_end < chars.len()
&& !chars[path_end].is_whitespace()
&& chars[path_end] != ';'
&& chars[path_end] != '&'
{
path_end += 1;
}
if j < path_end {
let path: String = chars[j..path_end].iter().collect();
let path = path.trim_matches(|c| c == '"' || c == '\'').to_string();
if artifact_extensions.iter().any(|ext| path.ends_with(ext)) {
return Some(path);
}
}
}
i += 1;
}
let mut in_single_quote = false;
let mut in_double_quote = false;
let mut quote_start = 0usize;
let chars: Vec<char> = command.chars().collect();
for (idx, &ch) in chars.iter().enumerate() {
if ch == '\'' && !in_double_quote {
if in_single_quote {
let path: String = chars[quote_start..idx].iter().collect();
if artifact_extensions.iter().any(|ext| path.ends_with(ext)) {
return Some(path);
}
in_single_quote = false;
} else {
in_single_quote = true;
quote_start = idx + 1;
}
} else if ch == '"' && !in_single_quote {
if in_double_quote {
let path: String = chars[quote_start..idx].iter().collect();
if artifact_extensions.iter().any(|ext| path.ends_with(ext)) {
return Some(path);
}
in_double_quote = false;
} else {
in_double_quote = true;
quote_start = idx + 1;
}
}
}
None
}
fn format_tool_args(tool_name: &str, args_json: &str) -> String {
let args: serde_json::Value = match serde_json::from_str(args_json) {
Ok(v) => v,
Err(_) => return String::new(),
};
match tool_name {
"read_file" => {
let path = args.get("file_path").and_then(|v| v.as_str()).unwrap_or("");
let short = short_path(path);
let mut s = short;
if let Some(offset) = args.get("offset").and_then(|v| v.as_u64()) {
if let Some(limit) = args.get("limit").and_then(|v| v.as_u64()) {
s.push_str(&format!(" L{}-{}", offset, offset + limit));
}
}
s
}
"create_file" => {
let path = args.get("file_path").and_then(|v| v.as_str()).unwrap_or("");
let size = args
.get("content")
.and_then(|v| v.as_str())
.map(|s| s.len())
.unwrap_or(0);
format!("{} ({} bytes)", short_path(path), size)
}
"edit_file" => {
let path = args.get("file_path").and_then(|v| v.as_str()).unwrap_or("");
short_path(path)
}
"bash" => {
let cmd = args.get("command").and_then(|v| v.as_str()).unwrap_or("");
if cmd.chars().count() > 80 {
format!("`{}...`", cmd.chars().take(77).collect::<String>())
} else {
format!("`{}`", cmd)
}
}
"list_directory" => {
let path = args.get("path").and_then(|v| v.as_str()).unwrap_or(".");
short_path(path)
}
"grep" => {
let pattern = args.get("pattern").and_then(|v| v.as_str()).unwrap_or("");
let path = args.get("path").and_then(|v| v.as_str()).unwrap_or(".");
format!("\"{}\" in {}", pattern, short_path(path))
}
"glob" => {
let pattern = args.get("pattern").and_then(|v| v.as_str()).unwrap_or("");
format!("\"{}\"", pattern)
}
"web_search" => {
let query = args.get("query").and_then(|v| v.as_str()).unwrap_or("");
format!("\"{}\"", query)
}
"web_fetch" => {
let url = args.get("url").and_then(|v| v.as_str()).unwrap_or("");
url.to_string()
}
_ => {
if let Some(obj) = args.as_object() {
obj.iter()
.map(|(k, v)| {
let val = match v {
serde_json::Value::String(s) if s.chars().count() > 30 => {
format!("{}...", s.chars().take(27).collect::<String>())
}
serde_json::Value::String(s) => s.clone(),
other => other.to_string(),
};
format!("{}={}", k, val)
})
.collect::<Vec<_>>()
.join(" ")
} else {
String::new()
}
}
}
}
fn short_path(path: &str) -> String {
let parts: Vec<&str> = path.rsplitn(3, '/').collect();
match parts.len() {
0 | 1 => path.to_string(),
2 => format!("{}/{}", parts[1], parts[0]),
_ => format!(".../{}/{}", parts[1], parts[0]),
}
}
fn dangerous_tools_enabled() -> bool {
std::env::var(DANGEROUS_TOOLS_ENV).ok().as_deref() == Some("1")
}
fn cors_layer() -> CorsLayer {
CorsLayer::new()
.allow_origin(AllowOrigin::predicate(is_loopback_origin))
.allow_methods([Method::GET, Method::POST, Method::PATCH, Method::DELETE])
.allow_headers([header::CONTENT_TYPE])
}
async fn activity_tracker_middleware(
req: axum::extract::Request,
next: axum::middleware::Next,
) -> axum::response::Response {
let skip = (req.method() == Method::GET && req.uri().path() == "/health")
|| (req.method() == Method::POST && req.uri().path() == "/shutdown");
if !skip {
if let Some(activity) = req.extensions().get::<Arc<std::sync::atomic::AtomicI64>>() {
activity.store(now_unix_ms(), std::sync::atomic::Ordering::Relaxed);
}
}
let client_mode = req
.headers()
.get("x-atomcode-client")
.and_then(|v| v.to_str().ok())
.map(resolve_client_mode)
.unwrap_or(SessionMode::Ide);
let mut req = req;
req.extensions_mut().insert(client_mode);
next.run(req).await
}
fn resolve_client_mode(header: &str) -> SessionMode {
match header {
"vscode" => SessionMode::Vscode,
"webui" => SessionMode::Webui,
"atomcode-air" => SessionMode::AtomcodeAir,
_ => SessionMode::Ide,
}
}
fn is_loopback_origin(origin: &HeaderValue, _request_parts: &RequestParts) -> bool {
let Ok(origin) = origin.to_str() else {
return false;
};
let Some(authority) = origin
.strip_prefix("http://")
.or_else(|| origin.strip_prefix("https://"))
else {
return false;
};
is_loopback_authority(authority)
}
fn is_loopback_authority(authority: &str) -> bool {
if let Some(rest) = authority.strip_prefix("[::1]") {
return rest.is_empty() || rest.starts_with(':');
}
let host = authority.split(':').next().unwrap_or(authority);
matches!(host, "localhost" | "127.0.0.1" | "::1")
}
pub(crate) fn hash_path(path: &std::path::Path) -> String {
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
let normalized = path.to_string_lossy();
let mut normalized = normalized.replace('\\', "/");
// Remove trailing slash (but keep root "/" or "C:/")
if normalized.len() > 1 && normalized.ends_with('/') {
normalized.pop();
}
// On Windows, paths are case-insensitive
#[cfg(windows)]
let normalized = normalized.to_lowercase();
let mut hasher = DefaultHasher::new();
normalized.hash(&mut hasher);
format!("{:016x}", hasher.finish())
}
/// List all projects (scans sessions directory)
fn list_projects() -> std::io::Result<Vec<ProjectInfo>> {
let sessions_root = SessionManager::sessions_root_dir();
let mut projects = Vec::new();
if !sessions_root.exists() {
return Ok(projects);
}
// Scan sessions directory for actual session data
for entry in std::fs::read_dir(sessions_root)? {
let entry = entry?;
let path = entry.path();
if path.is_dir() {
let hash = path.file_name().unwrap().to_string_lossy().to_string();
// Scan sessions in this project to get working_dir and stats
let mut session_count = 0;
let mut last_updated = 0u64;
let mut created_at = u64::MAX;
let mut working_dir = PathBuf::new();
for session_file in std::fs::read_dir(&path)? {
let session_file = session_file?;
let file_path = session_file.path();
if file_path.extension().map_or(false, |ext| ext == "json") {
if let Ok(json) = std::fs::read_to_string(&file_path) {
if let Ok(session) = serde_json::from_str::<Session>(&json) {
session_count += 1;
last_updated = last_updated.max(session.updated_at);
created_at = created_at.min(session.created_at);
if working_dir.to_string_lossy().is_empty() {
working_dir = session.working_dir;
}
}
}
}
}
// Only include projects with at least one session
if session_count > 0 {
let name = working_dir
.file_name()
.map(|n| n.to_string_lossy().to_string())
.unwrap_or_else(|| "unknown".to_string());
projects.push(ProjectInfo {
hash,
name,
working_dir,
description: None,
session_count,
created_at: if created_at == u64::MAX {
0
} else {
created_at
},
last_updated,
});
}
}
}
// Sort by last updated (most recent first)
projects.sort_by(|a, b| b.last_updated.cmp(&a.last_updated));
Ok(projects)
}
/// Session metadata with project hash for cross-project listing
#[derive(Debug, Serialize)]
pub struct SessionMetaWithProject {
pub project_hash: String,
#[serde(flatten)]
pub meta: SessionMeta,
}
/// List sessions for a project
fn list_sessions(project_hash: &str) -> std::io::Result<Vec<SessionMeta>> {
let project_dir = SessionManager::sessions_root_dir().join(project_hash);
if !project_dir.exists() {
return Ok(Vec::new());
}
let mut sessions = Vec::new();
for entry in std::fs::read_dir(project_dir)? {
let entry = entry?;
let path = entry.path();
if path.extension().map_or(false, |ext| ext == "json") {
let file_size = entry.metadata().map(|m| m.len()).unwrap_or(0);
if let Ok(json) = std::fs::read_to_string(&path) {
if let Ok(session) = serde_json::from_str::<Session>(&json) {
// Skip empty sessions (no messages)
if session.messages.is_empty() {
continue;
}
let mut meta = SessionMeta::from(&session);
meta.file_size = file_size;
sessions.push(meta);
}
}
}
}
// Sort by updated_at descending
sessions.sort_by(|a, b| b.updated_at.cmp(&a.updated_at));
Ok(sessions)
}
/// List all sessions across all projects
fn list_all_sessions() -> std::io::Result<Vec<SessionMetaWithProject>> {
let sessions_root = SessionManager::sessions_root_dir();
if !sessions_root.exists() {
return Ok(Vec::new());
}
let mut all_sessions = Vec::new();
for entry in std::fs::read_dir(sessions_root)? {
let entry = entry?;
let path = entry.path();
if path.is_dir() {
let project_hash = path.file_name().unwrap().to_string_lossy().to_string();
for session_file in std::fs::read_dir(&path)? {
let session_file = session_file?;
let file_path = session_file.path();
if file_path.extension().map_or(false, |ext| ext == "json") {
let file_size = session_file.metadata().map(|m| m.len()).unwrap_or(0);
if let Ok(json) = std::fs::read_to_string(&file_path) {
if let Ok(session) = serde_json::from_str::<Session>(&json) {
// Skip empty sessions (no messages)
if session.messages.is_empty() {
continue;
}
let mut meta = SessionMeta::from(&session);
meta.file_size = file_size;
all_sessions.push(SessionMetaWithProject {
project_hash: project_hash.clone(),
meta,
});
}
}
}
}
}
}
// Sort by updated_at descending
all_sessions.sort_by(|a, b| b.meta.updated_at.cmp(&a.meta.updated_at));
// Limit to first 50 sessions
all_sessions.truncate(50);
Ok(all_sessions)
}
/// Load a specific session
fn load_session(project_hash: &str, session_id: &str) -> std::io::Result<Session> {
let path = SessionManager::sessions_root_dir()
.join(project_hash)
.join(format!("{}.json", session_id));
let json = std::fs::read_to_string(path)?;
serde_json::from_str(&json).map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))
}
// ============== HTTP Handlers ==============
/// Health check response
#[derive(Debug, Serialize)]
pub struct HealthResponse {
pub status: &'static str,
pub version: &'static str,
pub service: &'static str,
}
async fn health() -> impl IntoResponse {
Json(HealthResponse {
status: "ok",
version: env!("CARGO_PKG_VERSION"),
service: "atomcode-daemon",
})
}
async fn shutdown_handler(State(state): State<AppState>) -> impl IntoResponse {
state.shutdown_tx.send(true).ok();
Json(serde_json::json!({"success": true}))
}
async fn get_project_state(State(state): State<AppState>) -> impl IntoResponse {
let state = state.project.read().await;
Json(ProjectState {
working_dir: state.working_dir.clone(),
previous_dir: state.previous_dir.clone(),
recent_dirs: state.recent_dirs.clone(),
name: state.name.clone(),
})
}
async fn change_dir(
State(state): State<AppState>,
axum::Extension(client_mode): axum::Extension<SessionMode>,
Json(req): Json<ChangeDirRequest>,
) -> impl IntoResponse {
let state_clone = state.clone();
daemon_scope(&state, None, client_mode, || async move {
let state = state_clone;
let mut project = state.project.write().await;
let new_path = if req.path == "-" {
match &project.previous_dir {
Some(prev) => prev.clone(),
None => {
return Json(ChangeDirResponse {
success: false,
message: "No previous directory to go back to".to_string(),
current_dir: project.working_dir.clone(),
project_hash: hash_path(&project.working_dir),
});
}
}
} else {
let expanded = if req.path.starts_with('~') {
atomcode_core::tool::real_home_dir()
.map(|h| {
h.join(
req.path
.strip_prefix('~')
.unwrap_or("")
.trim_start_matches('/'),
)
})
.unwrap_or_else(|| PathBuf::from(&req.path))
} else {
PathBuf::from(&req.path)
};
let resolved = if expanded.is_absolute() {
expanded
} else {
project.working_dir.join(&expanded)
};
if !resolved.exists() {
return Json(ChangeDirResponse {
success: false,
message: format!("Directory does not exist: {}", resolved.display()),
current_dir: project.working_dir.clone(),
project_hash: hash_path(&project.working_dir),
});
}
if !resolved.is_dir() {
return Json(ChangeDirResponse {
success: false,
message: format!("Not a directory: {}", resolved.display()),
current_dir: project.working_dir.clone(),
project_hash: hash_path(&project.working_dir),
});
}
resolved
};
let old_dir = project.working_dir.clone();
project.previous_dir = Some(old_dir);
project.working_dir = new_path.clone();
project.name = new_path
.file_name()
.map(|n| n.to_string_lossy().to_string())
.unwrap_or_else(|| "project".to_string());
project.recent_dirs.retain(|d| d != &new_path);
project.recent_dirs.insert(0, new_path.clone());
project.recent_dirs.truncate(5);
if req.set_default {
let config_path = Config::default_path();
if let Ok(mut config) = Config::load(&config_path) {
config.default_workdir = Some(new_path.to_string_lossy().to_string());
let _ = config.save(&config_path);
}
}
let hash = hash_path(&new_path);
state.telemetry.track(Event::UseCommand { type_: "cd".into(), success: Some(true), error_kind: None, error_data: None });
Json(ChangeDirResponse {
success: true,
message: format!("Changed to {}", new_path.display()),
current_dir: new_path,
project_hash: hash,
})
})
.await
}
async fn get_projects() -> impl IntoResponse {
match list_projects() {
Ok(projects) => Json(projects).into_response(),
Err(e) => {
let msg = format!("Failed to list projects: {}", e);
(StatusCode::INTERNAL_SERVER_ERROR, Json(msg)).into_response()
}
}
}
async fn get_project_sessions(Path(hash): Path<String>) -> impl IntoResponse {
match list_sessions(&hash) {
Ok(sessions) => Json(sessions).into_response(),
Err(e) => {
let msg = format!("Failed to list sessions: {}", e);
(StatusCode::INTERNAL_SERVER_ERROR, Json(msg)).into_response()
}
}
}
async fn get_session_detail(Path((hash, id)): Path<(String, String)>) -> impl IntoResponse {
match load_session(&hash, &id) {
Ok(session) => {
let detail = SessionDetail {
id: session.id.to_string(),
name: session.name,
working_dir: session.working_dir,
created_at: session.created_at,
updated_at: session.updated_at,
message_count: session.messages.len(),
messages: session.messages.iter().map(MessageInfo::from).collect(),
};
Json(detail).into_response()
}
Err(e) => {
let msg = format!("Failed to load session: {}", e);
(StatusCode::NOT_FOUND, Json(msg)).into_response()
}
}
}
async fn get_all_sessions() -> impl IntoResponse {
match list_all_sessions() {
Ok(sessions) => Json(sessions).into_response(),
Err(e) => {
let msg = format!("Failed to list sessions: {}", e);
(StatusCode::INTERNAL_SERVER_ERROR, Json(msg)).into_response()
}
}
}
async fn create_session(
State(state): State<AppState>,
Json(req): Json<CreateSessionRequest>,
) -> impl IntoResponse {
let working_dir = match req.working_dir {
Some(dir) => dir,
None => {
let project = state.project.read().await;
project.working_dir.clone()
}
};
if !working_dir.exists() {
let home = atomcode_core::tool::real_home_dir().unwrap_or_else(|| PathBuf::from("."));
let atomchat_dir = home.join("atomchat");
if atomchat_dir.exists() || std::fs::create_dir_all(&atomchat_dir).is_ok() {
} else {
let msg = format!("Working directory does not exist: {:?}", working_dir);
return (StatusCode::BAD_REQUEST, Json(msg)).into_response();
}
}
let manager = SessionManager::new(&working_dir);
let mut session = Session::new(working_dir.clone());
if let Some(title) = req.title {
session.rename(title);
}
if let Err(e) = manager.save(&session) {
let msg = format!("Failed to save session: {}", e);
return (StatusCode::INTERNAL_SERVER_ERROR, Json(msg)).into_response();
}
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
let mut hasher = DefaultHasher::new();
working_dir.hash(&mut hasher);
let project_hash = format!("{:016x}", hasher.finish());
let response = CreateSessionResponse {
id: session.id.to_string(),
name: session.name.clone(),
working_dir: session.working_dir.clone(),
project_hash,
created_at: session.created_at,
};
(StatusCode::CREATED, Json(response)).into_response()
}
fn search_sessions_by_name(keyword: &str) -> std::io::Result<Vec<SessionMetaWithProject>> {
let sessions_root = SessionManager::sessions_root_dir();
if !sessions_root.exists() {
return Ok(Vec::new());
}
let keyword_lower = keyword.to_lowercase();
let mut results = Vec::new();
for entry in std::fs::read_dir(sessions_root)? {
let entry = entry?;
let path = entry.path();
if path.is_dir() {
let project_hash = path.file_name().unwrap().to_string_lossy().to_string();
for session_file in std::fs::read_dir(&path)? {
let session_file = session_file?;
let file_path = session_file.path();
if file_path.extension().map_or(false, |ext| ext == "json") {
let file_size = session_file.metadata().map(|m| m.len()).unwrap_or(0);
if let Ok(json) = std::fs::read_to_string(&file_path) {
if let Ok(session) = serde_json::from_str::<Session>(&json) {
if session.messages.is_empty() {
continue;
}
if session.name.to_lowercase().contains(&keyword_lower) {
let mut meta = SessionMeta::from(&session);
meta.file_size = file_size;
results.push(SessionMetaWithProject {
project_hash: project_hash.clone(),
meta,
});
}
}
}
}
}
}
}
results.sort_by(|a, b| b.meta.updated_at.cmp(&a.meta.updated_at));
Ok(results)
}
async fn search_sessions(Query(query): Query<SearchQuery>) -> impl IntoResponse {
if query.q.trim().is_empty() {
return (
StatusCode::BAD_REQUEST,
Json("Search keyword cannot be empty"),
)
.into_response();
}
match search_sessions_by_name(&query.q) {
Ok(sessions) => Json(sessions).into_response(),
Err(e) => {
let msg = format!("Failed to search sessions: {}", e);
(StatusCode::INTERNAL_SERVER_ERROR, Json(msg)).into_response()
}
}
}
fn delete_session_file(project_hash: &str, session_id: &str) -> std::io::Result<()> {
let path = SessionManager::sessions_root_dir()
.join(project_hash)
.join(format!("{}.json", session_id));
if !path.exists() {
return Err(std::io::Error::new(
std::io::ErrorKind::NotFound,
format!("Session not found: {}/{}", project_hash, session_id),
));
}
std::fs::remove_file(path)
}
async fn delete_session(
State(state): State<AppState>,
axum::Extension(client_mode): axum::Extension<SessionMode>,
Path((hash, id)): Path<(String, String)>,
) -> impl IntoResponse {
let session_uuid = uuid::Uuid::parse_str(&id).ok();
let state_clone = state.clone();
daemon_scope(&state, session_uuid, client_mode, || async move {
match delete_session_file(&hash, &id) {
Ok(()) => {
state_clone.telemetry.track(Event::UseCommand { type_: "delete_session".into(), success: Some(true), error_kind: None, error_data: None });
let msg = format!("Session {} deleted successfully", id);
(StatusCode::OK, Json(msg)).into_response()
}
Err(e) => {
let msg = format!("Failed to delete session: {}", e);
(StatusCode::NOT_FOUND, Json(msg)).into_response()
}
}
})
.await
}
#[derive(Debug, Deserialize)]
pub struct RenameRequest {
pub name: String,
}
fn rename_session_file(
project_hash: &str,
session_id: &str,
new_name: &str,
) -> std::io::Result<()> {
let path = SessionManager::sessions_root_dir()
.join(project_hash)
.join(format!("{}.json", session_id));
if !path.exists() {
return Err(std::io::Error::new(
std::io::ErrorKind::NotFound,
format!("Session not found: {}/{}", project_hash, session_id),
));
}
let json = std::fs::read_to_string(&path)?;
let mut session: Session = serde_json::from_str(&json)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
session.rename(new_name.to_string());
let manager = SessionManager::new(&PathBuf::from(&session.working_dir));
manager.save(&session)
}
async fn rename_session(
State(state): State<AppState>,
axum::Extension(client_mode): axum::Extension<SessionMode>,
Path((hash, id)): Path<(String, String)>,
Json(req): Json<RenameRequest>,
) -> impl IntoResponse {
let session_uuid = uuid::Uuid::parse_str(&id).ok();
let state_clone = state.clone();
daemon_scope(&state, session_uuid, client_mode, || async move {
match rename_session_file(&hash, &id, &req.name) {
Ok(()) => {
state_clone.telemetry.track(Event::UseCommand { type_: "rename".into(), success: Some(true), error_kind: None, error_data: None });
let msg = format!("Session {} renamed to '{}'", id, req.name);
(StatusCode::OK, Json(msg)).into_response()
}
Err(e) => {
let msg = format!("Failed to rename session: {}", e);
(StatusCode::NOT_FOUND, Json(msg)).into_response()
}
}
})
.await
}
#[derive(Debug, Serialize)]
pub struct ModelInfo {
pub provider: String,
pub model: String,
pub provider_type: String,
pub is_default: bool,
}
async fn get_models() -> impl IntoResponse {
let config_path = Config::default_path();
let config = match Config::load(&config_path) {
Ok(c) => c,
Err(_e) => {
return (
StatusCode::INTERNAL_SERVER_ERROR,
Json(Vec::<ModelInfo>::new()),
)
.into_response();
}
};
let models: Vec<ModelInfo> = config
.providers
.iter()
.map(|(name, p)| ModelInfo {
provider: name.clone(),
model: p.model.clone(),
provider_type: p.provider_type.clone(),
is_default: name == &config.default_provider,
})
.collect();
(StatusCode::OK, Json(models)).into_response()
}
#[derive(Debug, Deserialize)]
pub struct ChatRequest {
pub message: String,
#[serde(default)]
pub working_dir: Option<PathBuf>,
#[serde(default)]
pub provider: Option<String>,
#[serde(default)]
pub session_id: Option<String>,
#[serde(default)]
pub images: Vec<ImageInput>,
}
#[derive(Debug, Deserialize)]
pub struct ImageInput {
pub media_type: String,
pub data: String,
}
#[derive(Debug, Serialize)]
#[serde(tag = "type")]
pub enum ChatEvent {
#[serde(rename = "tool_batch")]
ToolBatchStarted {
calls: Vec<atomcode_core::turn::event::ToolBatchCall>,
},
#[serde(rename = "text")]
TextDelta { content: String },
#[serde(rename = "reasoning")]
ReasoningDelta { content: String },
#[serde(rename = "tool_start")]
ToolCallStarted { id: String, name: String, arguments: String },
#[serde(rename = "tool_output")]
ToolOutputChunk { chunk: String },
#[serde(rename = "tool_result")]
ToolCallResult {
id: String,
name: String,
output: String,
success: bool,
duration_ms: u64,
},
#[serde(rename = "tokens")]
TokenUsage {
prompt: usize,
completion: usize,
total: usize,
},
#[serde(rename = "artifact_start")]
ArtifactStart {
id: String,
artifact_type: String,
language: Option<String>,
title: Option<String>,
},
#[serde(rename = "artifact_content")]
ArtifactContent { id: String, content: String },
#[serde(rename = "artifact_end")]
ArtifactEnd { id: String },
#[serde(rename = "done")]
Done {
tokens: usize,
tool_calls: usize,
session_id: String,
},
#[serde(rename = "permission_request")]
PermissionRequest {
session_id: String,
tool_name: String,
reason: String,
call_id: String,
arguments: String,
},
#[serde(rename = "stopped")]
Stopped,
#[serde(rename = "error")]
Error { message: String },
}
struct ArtifactDetector {
artifact_counter: usize,
state: ArtifactDetectorState,
}
#[derive(Debug, Clone)]
enum ArtifactDetectorState {
Normal,
InCodeBlock { id: String, content: String },
InHtml { id: String, content: String },
InSvg { id: String, content: String },
}
impl ArtifactDetector {
fn new() -> Self {
Self {
artifact_counter: 0,
state: ArtifactDetectorState::Normal,
}
}
fn next_id(&mut self) -> String {
self.artifact_counter += 1;
format!("artifact_{}", self.artifact_counter)
}
fn artifact_type_for_language(language: &str) -> (String, Option<String>) {
let lang_lower = language.to_lowercase();
let artifact_type = match lang_lower.as_str() {
"mermaid" => "mermaid",
"html" | "htm" => "html",
"svg" | "xmlsvg" => "svg",
"markdown" | "md" => "markdown",
_ => "code",
};
let title = if artifact_type == "code" && !language.is_empty() {
Some(language.to_string())
} else {
None
};
(artifact_type.to_string(), title)
}
fn process(&mut self, text: &str) -> Vec<ChatEvent> {
let mut events = Vec::new();
match &mut self.state {
ArtifactDetectorState::Normal => {
if text.starts_with("```") {
let rest = &text[3..];
let end_of_line = rest.find('\n').unwrap_or(rest.len());
let language = rest[..end_of_line].trim().to_string();
let (artifact_type, title) = Self::artifact_type_for_language(&language);
let id = self.next_id();
events.push(ChatEvent::ArtifactStart {
id: id.clone(),
artifact_type,
language: Some(language.clone()),
title,
});
self.state = ArtifactDetectorState::InCodeBlock {
id,
content: String::new(),
};
}
else if self.is_svg_start(text) {
let id = self.next_id();
events.push(ChatEvent::ArtifactStart {
id: id.clone(),
artifact_type: "svg".to_string(),
language: None,
title: None,
});
events.push(ChatEvent::ArtifactContent {
id: id.clone(),
content: text.to_string(),
});
self.state = ArtifactDetectorState::InSvg {
id,
content: text.to_string(),
};
}
else if self.is_html_start(text) {
let id = self.next_id();
events.push(ChatEvent::ArtifactStart {
id: id.clone(),
artifact_type: "html".to_string(),
language: None,
title: None,
});
events.push(ChatEvent::ArtifactContent {
id: id.clone(),
content: text.to_string(),
});
self.state = ArtifactDetectorState::InHtml {
id,
content: text.to_string(),
};
} else {
events.push(ChatEvent::TextDelta {
content: text.to_string(),
});
}
}
ArtifactDetectorState::InCodeBlock { id, content } => {
if text.trim() == "```" {
if !content.is_empty() {
events.push(ChatEvent::ArtifactContent {
id: id.clone(),
content: content.clone(),
});
}
events.push(ChatEvent::ArtifactEnd { id: id.clone() });
self.state = ArtifactDetectorState::Normal;
} else {
content.push_str(text);
events.push(ChatEvent::ArtifactContent {
id: id.clone(),
content: text.to_string(),
});
}
}
ArtifactDetectorState::InHtml { id, content } => {
let trimmed = text.trim();
if trimmed.ends_with("</html>")
|| trimmed.ends_with("</HTML>")
|| trimmed.ends_with("</body>")
|| trimmed.ends_with("</BODY>")
{
content.push_str(text);
events.push(ChatEvent::ArtifactContent {
id: id.clone(),
content: text.to_string(),
});
events.push(ChatEvent::ArtifactEnd { id: id.clone() });
self.state = ArtifactDetectorState::Normal;
} else {
content.push_str(text);
events.push(ChatEvent::ArtifactContent {
id: id.clone(),
content: text.to_string(),
});
}
}
ArtifactDetectorState::InSvg { id, content } => {
let trimmed = text.trim();
if trimmed.ends_with("</svg>") || trimmed.ends_with("</SVG>") {
content.push_str(text);
events.push(ChatEvent::ArtifactContent {
id: id.clone(),
content: text.to_string(),
});
events.push(ChatEvent::ArtifactEnd { id: id.clone() });
self.state = ArtifactDetectorState::Normal;
} else {
content.push_str(text);
events.push(ChatEvent::ArtifactContent {
id: id.clone(),
content: text.to_string(),
});
}
}
}
events
}
fn is_html_start(&self, text: &str) -> bool {
let trimmed = text.trim();
trimmed.starts_with("<!DOCTYPE html")
|| trimmed.starts_with("<!DOCTYPE HTML")
|| trimmed.starts_with("<html")
|| trimmed.starts_with("<HTML")
}
fn is_svg_start(&self, text: &str) -> bool {
let trimmed = text.trim();
trimmed.starts_with("<svg") || trimmed.starts_with("<SVG")
}
fn finish(&mut self) -> Option<ChatEvent> {
match &self.state {
ArtifactDetectorState::InCodeBlock { id, .. } => {
let id = id.clone();
self.state = ArtifactDetectorState::Normal;
Some(ChatEvent::ArtifactEnd { id })
}
ArtifactDetectorState::InHtml { id, .. } => {
let id = id.clone();
self.state = ArtifactDetectorState::Normal;
Some(ChatEvent::ArtifactEnd { id })
}
ArtifactDetectorState::InSvg { id, .. } => {
let id = id.clone();
self.state = ArtifactDetectorState::Normal;
Some(ChatEvent::ArtifactEnd { id })
}
ArtifactDetectorState::Normal => None,
}
}
}
type SessionStore = Arc<RwLock<std::collections::HashMap<String, Conversation>>>;
async fn chat_stream(
State(state): State<AppState>,
axum::Extension(client_mode): axum::Extension<SessionMode>,
Json(mut req): Json<ChatRequest>,
) -> impl IntoResponse {
let session_uuid = req.session_id.as_deref().and_then(|s| uuid::Uuid::parse_str(s).ok());
if req.working_dir.is_none() {
let project = state.project.read().await;
req.working_dir = Some(project.working_dir.clone());
}
let (tx, rx) = mpsc::unbounded_channel::<ChatEvent>();
let cancel_token = CancellationToken::new();
let session_id = req.session_id.clone();
if let Some(ref sid) = session_id {
state
.chat_tasks
.write()
.await
.insert(sid.clone(), cancel_token.clone());
}
let chat_tasks = state.chat_tasks.clone();
let stopped_sessions = state.stopped_sessions.clone();
let mcp_cache = state.mcp_cache.clone();
let telemetry = state.telemetry.clone();
let pending_permissions = state.pending_permissions.clone();
let webui_mode = state.enforce_token;
let chat_repo_origin = detect_repo_origin(
req.working_dir.as_deref().unwrap_or_else(|| std::path::Path::new("."))
);
let ctx_for_task = CurrentContext {
mode: Some(client_mode),
repo_origin: Some(chat_repo_origin),
session_id: session_uuid,
..CurrentContext::current()
};
tokio::spawn(async move {
CurrentContext::scope(ctx_for_task, || async move {
if let Err(e) = process_chat_request(
req,
tx.clone(),
cancel_token,
stopped_sessions.clone(),
mcp_cache,
telemetry,
pending_permissions,
webui_mode,
)
.await
{
let _ = tx.send(ChatEvent::Error {
message: e.to_string(),
});
}
if let Some(sid) = session_id {
chat_tasks.write().await.remove(&sid);
}
}).await;
});
let active_conns = state.active_connections.clone();
active_conns.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let stream = UnboundedReceiverStream::new(rx).map(|event| {
let json = serde_json::to_string(&event).unwrap_or_default();
Ok::<_, std::convert::Infallible>(axum::response::sse::Event::default().data(json))
});
let conn_guard = SseConnectionGuard(active_conns);
let guarded_stream = stream.chain(futures::stream::once(async move {
drop(conn_guard);
Ok(axum::response::sse::Event::default().comment("bye"))
}));
Sse::new(guarded_stream).keep_alive(
axum::response::sse::KeepAlive::new()
.interval(Duration::from_secs(15))
.text("ping"),
)
}
async fn process_chat_request(
req: ChatRequest,
event_tx: mpsc::UnboundedSender<ChatEvent>,
cancel_token: CancellationToken,
stopped_sessions: StoppedSessionsStore,
mcp_cache: Arc<RwLock<HashMap<PathBuf, CachedMcpRegistry>>>,
telemetry: Arc<Telemetry>,
pending_permissions: permission_bridge::PermissionResponders,
webui_mode: bool,
) -> anyhow::Result<()> {
use atomcode_core::tool::{
bash::BashTool, edit::EditFileTool, glob::GlobTool, grep::GrepTool, list_dir::ListDirTool,
read::ReadFileTool, search_replace::SearchReplaceTool, todo::TodoTool,
web_fetch::WebFetchTool, web_search::WebSearchTool, write::WriteFileTool,
};
let config_path = Config::default_path();
let config = Config::load(&config_path)?;
let provider_name = req
.provider
.unwrap_or_else(|| config.default_provider.clone());
let provider_config = config
.providers
.get(&provider_name)
.ok_or_else(|| anyhow::anyhow!("Provider '{}' not found", provider_name))?;
let provider = provider::create_provider(provider_config)?;
let working_dir = req
.working_dir
.unwrap_or_else(|| std::env::current_dir().unwrap_or_default());
let session_manager = SessionManager::new(&working_dir);
let mut session = if let Some(ref session_id_str) = req.session_id {
let session_id = SessionId::from_string(session_id_str.clone());
match session_manager.load(&session_id) {
Ok(session) => session,
Err(_) => {
Session::new(working_dir.clone())
}
}
} else {
Session::new(working_dir.clone())
};
let perm_session_key = session.id.to_string();
let conversation = Arc::new(tokio::sync::Mutex::new({
let mut conv = Conversation::new();
conv.messages = session.messages.clone();
conv
}));
{
use atomcode_core::conversation::message::{ImagePart, Message, MessageContent, Role};
use atomcode_core::vision_preprocessor::{maybe_preprocess, PreprocessOutcome};
let images: Vec<ImagePart> = req
.images
.iter()
.map(|i| ImagePart {
media_type: i.media_type.clone(),
data: i.data.clone(),
})
.collect();
let mut conv = conversation.lock().await;
if images.is_empty() {
conv.add_user_message(&req.message);
} else {
let text = match maybe_preprocess(&config, &*provider, &req.message, &images).await {
PreprocessOutcome::Skipped => req.message.clone(),
PreprocessOutcome::Replaced { text, vl_key } => {
if req.message.trim().is_empty() {
format!("[图片内容(由 {vl_key} 识别)]\n{text}")
} else {
format!("{}\n\n[图片内容(由 {vl_key} 识别)]\n{text}", req.message)
}
}
PreprocessOutcome::Failed { .. } => {
if req.message.trim().is_empty() {
"[图片识别失败]".to_string()
} else {
format!("{}\n\n[图片识别失败]", req.message)
}
}
};
let idx = conv.messages.len();
conv.messages.push(Message {
role: Role::User,
content: MessageContent::MultiPart {
text: if text.is_empty() { None } else { Some(text) },
images,
},
synthetic: false,
});
conv.turn_tracker.on_user_message(idx);
}
}
let mut tool_context =
ToolContext::with_telemetry(working_dir.clone(), req.session_id.as_deref().unwrap_or("default"), telemetry);
let mut tool_registry = ToolRegistry::new();
let disabled_tools: std::collections::HashSet<String> = std::env::var("ATOMCODE_DISABLE_TOOLS")
.ok()
.map(|v| {
v.split(',')
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
.collect()
})
.unwrap_or_default();
let enabled = |name: &str| !disabled_tools.contains(name);
if enabled("read_file") {
tool_registry.register_sync(Box::new(ReadFileTool));
}
if enabled("write_file") {
tool_registry.register_sync(Box::new(WriteFileTool));
}
if enabled("edit_file") {
tool_registry.register_sync(Box::new(EditFileTool));
}
if enabled("bash") {
tool_registry.register_sync(Box::new(BashTool));
}
if enabled("grep") {
tool_registry.register_sync(Box::new(GrepTool));
}
if enabled("glob") {
tool_registry.register_sync(Box::new(GlobTool));
}
if enabled("list_directory") {
tool_registry.register_sync(Box::new(ListDirTool));
}
if enabled("web_search") {
tool_registry.register_sync(Box::new(WebSearchTool));
}
if enabled("web_fetch") {
tool_registry.register_sync(Box::new(WebFetchTool));
}
if enabled("search_replace") {
tool_registry.register_sync(Box::new(SearchReplaceTool));
}
if enabled("todo") {
tool_registry.register_sync(Box::new(TodoTool::new()));
}
let mut skill_registry = atomcode_core::skill::SkillRegistry::new();
skill_registry.reload(&working_dir);
let has_skills = !skill_registry.is_empty();
let skill_registry = Arc::new(std::sync::RwLock::new(skill_registry));
if has_skills && enabled("use_skill") {
tool_registry.register_sync(Box::new(atomcode_core::tool::use_skill::UseSkillTool {
registry: skill_registry.clone(),
}));
}
let mcp_registry: Arc<McpRegistry> = {
let cache = mcp_cache.read().await;
if let Some(cached) = cache.get(&working_dir) {
cached.registry.clone()
} else {
drop(cache);
let new_registry = Arc::new(McpRegistry::from_config_background(&working_dir));
new_registry.wait_for_initial_connections(Duration::from_secs(5)).await;
let mut cache = mcp_cache.write().await;
if cache.len() >= MCP_CACHE_MAX {
if let Some(oldest_key) = cache
.iter()
.min_by_key(|(_, v)| v.last_used)
.map(|(k, _)| k.clone())
{
cache.remove(&oldest_key);
}
}
cache.insert(working_dir.clone(), CachedMcpRegistry {
registry: new_registry.clone(),
last_used: std::time::Instant::now(),
});
new_registry
}
};
{
let mut cache = mcp_cache.write().await;
if let Some(entry) = cache.get_mut(&working_dir) {
entry.last_used = std::time::Instant::now();
}
}
let mcp_tools = mcp_registry.list_all_tools().await;
if !mcp_tools.is_empty() {
register_mcp_tools(&mut tool_registry, mcp_registry.clone(), mcp_tools);
}
let lsp_manager = build_lsp_manager(&config.lsp, &working_dir);
if lsp_manager.is_some() && enabled("diagnostics") {
tool_registry.register_sync(Box::new(DiagnosticsTool));
}
tool_context.lsp = lsp_manager;
let shared_tools = Arc::new(tool_registry);
let (permission, perm_req_rx): (Box<dyn PermissionDecider>, Option<_>) = if webui_mode {
let (perm_req_tx, perm_req_rx) =
tokio::sync::mpsc::unbounded_channel::<ApprovalRequest>();
let (perm_resp_tx, perm_resp_rx) =
tokio::sync::mpsc::unbounded_channel::<atomcode_core::tool::PermissionDecision>();
let perm_store = std::sync::Arc::new(std::sync::RwLock::new(
atomcode_core::tool::PermissionStore::new(),
));
pending_permissions.register(perm_session_key.clone(), perm_resp_tx);
(
Box::new(InteractivePermissionDecider::new(perm_req_tx, perm_resp_rx, perm_store)),
Some(perm_req_rx),
)
} else {
(
Box::new(AutoPermissionDecider::new(AutoPermissionMode::BypassAll)),
None,
)
};
let daemon_ctx = match config.providers.get(&config.default_provider) {
Some(pc) => atomcode_core::ctx::for_provider(pc),
None => {
atomcode_core::ctx::for_provider(&atomcode_core::config::provider::ProviderConfig {
provider_type: String::new(),
api_key: None,
model: String::new(),
base_url: None,
system_prompt: None,
user_agent: None,
context_window: 128_000,
max_tokens: None,
thinking_type: None,
thinking_keep: None,
reasoning_history: None,
thinking_enabled: None,
thinking_budget: None,
skip_tls_verify: false,
ephemeral: true,
})
}
};
let mut hook_engine = atomcode_core::hook::HookEngine::new();
hook_engine.load_all(&working_dir);
let mut turn_runner = TurnRunner {
provider: provider.into(),
tools: shared_tools,
context: tool_context,
config: config.clone(),
ctx: daemon_ctx,
permission,
recently_edited_files: Vec::new(),
hook_engine: std::sync::Arc::new(hook_engine),
loop_guard: Default::default(),
current_turn_number: 0,
};
let system_prompt = build_api_system_prompt(&working_dir, &config, provider_config, &skill_registry);
let (turn_tx, mut turn_rx) = mpsc::unbounded_channel::<TurnEvent>();
let session_id_str = req.session_id.clone().unwrap_or_default();
if stopped_sessions
.write()
.await
.take(&session_id_str)
.is_some()
{
{
let conv = conversation.lock().await;
session.messages = conv.messages.clone();
session.auto_name_from_messages();
session.touch();
if let Err(e) = session_manager.save(&session) {
eprintln!("Warning: Failed to save session after early stop: {}", e);
}
}
let _ = event_tx.send(ChatEvent::Stopped);
let _ = event_tx.send(ChatEvent::Done {
tokens: 0,
tool_calls: 0,
session_id: session.id.to_string(),
});
if webui_mode {
pending_permissions.unregister(&perm_session_key);
}
return Ok(());
}
let conversation_clone = conversation.clone();
let tel_ctx = CurrentContext::current();
tokio::spawn(async move {
let _keep_perm_req_rx = perm_req_rx;
CurrentContext::scope(tel_ctx, || async move {
let mut conv = conversation_clone.lock().await;
loop {
let result = turn_runner
.run(&mut conv, &system_prompt, &turn_tx, cancel_token.clone())
.await;
match result {
TurnResult::Responded { .. } => {
break;
}
TurnResult::UsedTools { .. } => {
continue;
}
TurnResult::Failed(e) => {
let _ = turn_tx.send(TurnEvent::Error(e));
break;
}
TurnResult::Cancelled => {
break;
}
}
}
}).await;
});
let mut total_tokens = 0usize;
let mut tool_call_count = 0usize;
let mut artifact_detector = ArtifactDetector::new();
while let Some(event) = turn_rx.recv().await {
match event {
TurnEvent::TextDelta(text) => {
for chat_event in artifact_detector.process(&text) {
let _ = event_tx.send(chat_event);
}
}
TurnEvent::ReasoningDelta(text) => {
let _ = event_tx.send(ChatEvent::ReasoningDelta { content: text });
}
TurnEvent::ToolCallStarted {
id,
name,
arguments,
} => {
tool_call_count += 1;
let _ = event_tx.send(ChatEvent::ToolCallStarted {
id: id.clone(),
name: name.clone(),
arguments: arguments.clone(),
});
if name == "create_file" || name == "edit_file" {
if let Ok(args) = serde_json::from_str::<serde_json::Value>(&arguments) {
if let Some(path) = args.get("file_path").and_then(|v| v.as_str()) {
let artifact_type = if path.ends_with(".html") || path.ends_with(".htm")
{
"html"
} else if path.ends_with(".svg") {
"svg"
} else {
""
};
if !artifact_type.is_empty() {
if let Some(content) = args.get("content").and_then(|v| v.as_str())
{
let id = format!("file-{}", uuid::Uuid::new_v4());
let title = std::path::PathBuf::from(path)
.file_name()
.map(|n| n.to_string_lossy().to_string());
let _ = event_tx.send(ChatEvent::ArtifactStart {
id: id.clone(),
artifact_type: artifact_type.to_string(),
language: Some("html".to_string()),
title,
});
let _ = event_tx.send(ChatEvent::ArtifactContent {
id: id.clone(),
content: content.to_string(),
});
let _ = event_tx.send(ChatEvent::ArtifactEnd { id });
}
}
}
}
}
}
TurnEvent::ToolOutputChunk { call_id: _, chunk } => {
let _ = event_tx.send(ChatEvent::ToolOutputChunk { chunk });
}
TurnEvent::ToolCallResult {
call_id,
name,
output,
success,
duration,
} => {
let _ = event_tx.send(ChatEvent::ToolCallResult {
id: call_id,
name,
output,
success,
duration_ms: duration.as_millis() as u64,
});
}
TurnEvent::TokenUsage {
prompt_tokens,
completion_tokens,
total_tokens: tt,
cached_tokens: _,
} => {
total_tokens = tt;
let _ = event_tx.send(ChatEvent::TokenUsage {
prompt: prompt_tokens,
completion: completion_tokens,
total: tt,
});
}
TurnEvent::Error(e) => {
let _ = event_tx.send(ChatEvent::Error { message: e });
}
TurnEvent::Warning(w) => {
let _ = event_tx.send(ChatEvent::Error {
message: format!("[warning] {}", w),
});
}
TurnEvent::ContextStats { .. } => {
}
TurnEvent::ToolCallStreaming { .. } => {
}
TurnEvent::ToolBatchStarted { calls, .. } => {
let _ = event_tx.send(ChatEvent::ToolBatchStarted { calls });
}
TurnEvent::ToolBatchCompleted { .. } => {
}
TurnEvent::WorkingDirChanged(_) => {
}
TurnEvent::ApprovalRequested {
tool_name,
reason,
call,
..
} => {
let _ = event_tx.send(ChatEvent::PermissionRequest {
session_id: perm_session_key.clone(),
tool_name,
reason,
call_id: call.id,
arguments: call.arguments,
});
}
}
}
if let Some(event) = artifact_detector.finish() {
let _ = event_tx.send(event);
}
let session_id_str = req.session_id.clone().unwrap_or_default();
let was_stopped = stopped_sessions.read().await.contains(&session_id_str);
{
let mut conv = conversation.lock().await;
if was_stopped {
conv.cancel_current_turn();
}
session.messages = conv.messages.clone();
}
session.auto_name_from_messages();
session.touch();
if let Err(e) = session_manager.save(&session) {
eprintln!("Warning: Failed to save session: {}", e);
}
if was_stopped {
stopped_sessions.write().await.remove(&session_id_str);
}
let _ = event_tx.send(ChatEvent::Done {
tokens: total_tokens,
tool_calls: tool_call_count,
session_id: session.id.to_string(),
});
if webui_mode {
pending_permissions.unregister(&perm_session_key);
}
Ok(())
}
pub(crate) fn build_api_system_prompt(
working_dir: &PathBuf,
_config: &Config,
provider_config: &atomcode_core::config::provider::ProviderConfig,
skill_registry: &Arc<std::sync::RwLock<atomcode_core::skill::SkillRegistry>>,
) -> String {
let rules = if let Some(custom) = provider_config.system_prompt.as_deref() {
custom.to_string()
} else {
atomcode_core::config::prompt_sections::build_rules().to_string()
};
let shell = if cfg!(target_os = "windows") {
std::env::var("COMSPEC").unwrap_or_else(|_| "cmd.exe".into())
} else {
std::env::var("SHELL").unwrap_or_else(|_| "bash".into())
};
let env_info = format!("Platform: {} | Shell: {}", std::env::consts::OS, shell);
let model_display = &provider_config.model;
let mut prompt = format!(
"You are AtomCode. When asked who you are, say you are AtomCode \
(an AI coding agent by AtomGit) running the {} model. \
Never claim to be another product.\n\
Working directory: {wd}\n\
All file paths in tool calls must be absolute, resolved under {wd}. \
Verify file existence before editing.\n{env_info}\n",
model_display,
wd = working_dir.display(),
env_info = env_info,
);
prompt.push_str(&format!(
"\n=== GIT COMMITS ===\n\
When you create a git commit on the user's behalf, end the commit \
message with this trailer (preceded by a blank line):\n\
\n\
Co-Authored-By: AtomCode ({}) <noreply@atomgit.com>\n\
\n\
Use a HEREDOC for `git commit -m` so the trailer's blank line is \
preserved verbatim. Skip this trailer for `git commit --amend` \
and `git revert` (those operate on existing commits whose \
attribution shouldn't change).\n",
model_display
));
let instructions = atomcode_core::config::instructions::LayeredInstructions::load(working_dir);
let merged_instructions = instructions.merged();
if !merged_instructions.is_empty() {
prompt.push_str(&format!("\n{}\n", merged_instructions));
}
{
use atomcode_core::config::memory::MemoryStore;
let project_name = working_dir
.file_name()
.map(|n| n.to_string_lossy().to_string())
.unwrap_or_else(|| "project".to_string());
let global = MemoryStore::global();
let project = MemoryStore::project(working_dir);
let memory_block = MemoryStore::merged_for_prompt(&global, &project, &project_name);
if !memory_block.is_empty() {
prompt.push_str(&format!("\n{}\n", memory_block));
}
}
if let Ok(registry) = skill_registry.read() {
let skills: Vec<String> = registry
.invocable_by_llm()
.map(|s| {
let hint = s
.argument_hint
.as_ref()
.map(|h| format!(" {}", h))
.unwrap_or_default();
format!("- /{}{}: {}", s.name, hint, s.description)
})
.collect();
if !skills.is_empty() {
prompt.push_str("\n=== AVAILABLE SKILLS ===\n");
prompt.push_str(
"Use the `use_skill` tool to invoke a skill when relevant to the task.\n",
);
prompt.push_str(&skills.join("\n"));
prompt.push('\n');
}
}
let env_snapshot = atomcode_core::ctx::EnvSnapshot::capture(working_dir);
prompt.push_str(&env_snapshot.as_prompt_section());
prompt.push_str(&format!(
"\n=== RULES (follow these strictly) ===\n{rules}\n"
));
let platform = atomcode_core::config::platform_rules();
if !platform.is_empty() {
prompt.push_str(platform);
prompt.push('\n');
}
prompt
}
#[derive(Debug, Deserialize)]
struct StopChatRequest {
session_id: String,
}
#[derive(Debug, Serialize)]
struct StopChatResponse {
success: bool,
message: String,
}
async fn stop_chat(
State(state): State<AppState>,
axum::Extension(client_mode): axum::Extension<SessionMode>,
Json(req): Json<StopChatRequest>,
) -> impl IntoResponse {
let session_uuid = uuid::Uuid::parse_str(&req.session_id).ok();
let state_clone = state.clone();
daemon_scope(&state, session_uuid, client_mode, || async move {
state_clone
.stopped_sessions
.write()
.await
.insert(req.session_id.clone());
if let Some(cancel_token) = state_clone.chat_tasks.read().await.get(&req.session_id) {
cancel_token.cancel();
state_clone.telemetry.track(Event::UseCommand { type_: "stop".into(), success: Some(true), error_kind: None, error_data: None });
(
axum::http::StatusCode::OK,
Json(StopChatResponse {
success: true,
message: format!("Chat session {} stopped", req.session_id),
}),
)
} else {
state_clone.telemetry.track(Event::UseCommand { type_: "stop".into(), success: Some(true), error_kind: None, error_data: None });
(
axum::http::StatusCode::OK,
Json(StopChatResponse {
success: true,
message: format!(
"Chat session {} marked as stopped (was not running)",
req.session_id
),
}),
)
}
})
.await
}
async fn active_chat_sessions(
State(state): State<AppState>,
) -> impl IntoResponse {
let sessions: Vec<String> = state.chat_tasks.read().await.keys().cloned().collect();
Json(sessions)
}
#[derive(Debug, serde::Deserialize)]
pub struct PermissionDecisionRequest {
pub session_id: String,
pub decision: String,
}
async fn chat_permission(
State(state): State<AppState>,
Json(req): Json<PermissionDecisionRequest>,
) -> impl IntoResponse {
use atomcode_core::tool::PermissionDecision;
let decision = match req.decision.as_str() {
"allow" => PermissionDecision::Allow,
"always_allow" => PermissionDecision::Allow,
_ => PermissionDecision::Deny,
};
if state.pending_permissions.deliver(&req.session_id, decision) {
Json(serde_json::json!({ "success": true }))
} else {
Json(serde_json::json!({ "success": false, "error": "no pending permission for session" }))
}
}
#[derive(Serialize)]
struct McpServerStatus {
name: String,
status: String,
#[serde(skip_serializing_if = "Option::is_none")]
tool_count: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
error: Option<String>,
}
#[derive(Serialize)]
struct McpStatusResponse {
servers: Vec<McpServerStatus>,
}
async fn mcp_status(State(state): State<AppState>) -> Json<McpStatusResponse> {
let registry = state.mcp_registry.read().await.clone();
let statuses = registry.server_statuses().await;
let mut servers = Vec::new();
for (name, status) in statuses {
let (status_str, error) = match &status {
atomcode_core::mcp::ServerStatus::Connecting => ("connecting".to_string(), None),
atomcode_core::mcp::ServerStatus::Connected => ("connected".to_string(), None),
atomcode_core::mcp::ServerStatus::Failed(e) => ("error".to_string(), Some(e.clone())),
atomcode_core::mcp::ServerStatus::Disconnected => ("disconnected".to_string(), None),
};
let tool_count = if matches!(status, atomcode_core::mcp::ServerStatus::Connected) {
let tools = registry.list_all_tools().await;
Some(tools.iter().filter(|t| t.server_name == name).count())
} else {
None
};
servers.push(McpServerStatus {
name,
status: status_str,
tool_count,
error,
});
}
Json(McpStatusResponse { servers })
}
async fn mcp_reload(State(state): State<AppState>) -> Json<serde_json::Value> {
let project = state.project.read().await;
let project_dir = project.working_dir.clone();
drop(project);
let new_registry = McpRegistry::from_config_background(&project_dir);
*state.mcp_registry.write().await = Arc::new(new_registry);
Json(serde_json::json!({"status": "reloading"}))
}
async fn shutdown_signal(mut shutdown_rx: watch::Receiver<bool>) {
let ctrl_c = async {
tokio::signal::ctrl_c().await.ok();
};
#[cfg(unix)]
let terminate = async {
use tokio::signal::unix::{signal, SignalKind};
if let Ok(mut s) = signal(SignalKind::terminate()) {
s.recv().await;
}
};
#[cfg(not(unix))]
let terminate = std::future::pending::<()>();
let http_shutdown = async {
while !*shutdown_rx.borrow_and_update() {
if shutdown_rx.changed().await.is_err() {
break;
}
}
};
tokio::select! {
_ = ctrl_c => { tracing::info!("Received Ctrl-C, starting graceful shutdown"); }
_ = terminate => { tracing::info!("Received SIGTERM, starting graceful shutdown"); }
_ = http_shutdown => { tracing::info!("Received /shutdown request, starting graceful shutdown"); }
}
}
fn install_panic_hook(telemetry: Arc<Telemetry>) {
let default_hook = std::panic::take_hook();
std::panic::set_hook(Box::new(move |info| {
let home = atomcode_telemetry::identity::real_home_dir();
let cwd = std::env::current_dir().ok();
let loc = info
.location()
.map(|l| format!("{}:{}", l.file(), l.line()))
.unwrap_or_else(|| "unknown".into());
let msg = info
.payload()
.downcast_ref::<&str>()
.map(|s| s.to_string())
.or_else(|| info.payload().downcast_ref::<String>().cloned())
.unwrap_or_default();
let bt = std::backtrace::Backtrace::force_capture().to_string();
let scrubbed_loc =
atomcode_telemetry::scrub::scrub_path(&loc, home.as_deref(), cwd.as_deref());
let scrubbed_msg = atomcode_telemetry::scrub::truncate_head(
&atomcode_telemetry::scrub::scrub_path(&msg, home.as_deref(), cwd.as_deref()),
atomcode_telemetry::scrub::HEAD_MAX,
);
let frames =
atomcode_telemetry::scrub::backtrace_top_k(&bt, 5, home.as_deref(), cwd.as_deref());
telemetry.track(Event::Panic {
location: scrubbed_loc,
message_head: scrubbed_msg,
thread: std::thread::current().name().unwrap_or("unknown").into(),
backtrace_top_5: frames,
error_kind: Some("panic".to_string()),
error_data: Some(serde_json::json!({
"session_duration_secs": telemetry.uptime().as_secs() as u32,
"turns_completed": null,
"last_tool_name": null,
"last_event": null,
}).to_string()),
});
default_hook(info);
}));
}
fn now_unix_ms() -> i64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_millis() as i64)
.unwrap_or(0)
}
fn spawn_idle_timeout_task(
idle_timeout_secs: u64,
last_activity: Arc<std::sync::atomic::AtomicI64>,
active_connections: Arc<std::sync::atomic::AtomicUsize>,
shutdown_tx: watch::Sender<bool>,
) {
if idle_timeout_secs == 0 {
return;
}
let timeout_ms = (idle_timeout_secs * 1000) as i64;
tokio::spawn(async move {
let mut interval = tokio::time::interval(Duration::from_secs(60));
interval.tick().await;
loop {
interval.tick().await;
let conns = active_connections.load(std::sync::atomic::Ordering::Relaxed);
if conns > 0 {
continue;
}
let last = last_activity.load(std::sync::atomic::Ordering::Relaxed);
let elapsed = now_unix_ms() - last;
if elapsed >= timeout_ms {
tracing::info!(
elapsed_mins = elapsed / 60_000,
timeout_mins = idle_timeout_secs / 60,
"Daemon idle timeout reached, shutting down"
);
shutdown_tx.send(true).ok();
break;
}
}
});
}
struct WebuiHandle {
tokens: auth_token::WebuiTokenStore,
port: u16,
host: String,
abort: tokio::task::AbortHandle,
}
static WEBUI: std::sync::Mutex<Option<WebuiHandle>> = std::sync::Mutex::new(None);
async fn bind_scanning(
host: &str,
start_port: u16,
max_tries: u16,
) -> anyhow::Result<(tokio::net::TcpListener, u16)> {
let mut last_err: Option<std::io::Error> = None;
for offset in 0..max_tries {
let Some(port) = start_port.checked_add(offset) else {
break;
};
let addr = format!("{host}:{port}");
match tokio::net::TcpListener::bind(&addr).await {
Ok(listener) => {
let actual = listener.local_addr()?.port();
return Ok((listener, actual));
}
Err(e) if e.kind() == std::io::ErrorKind::AddrInUse => {
last_err = Some(e);
continue;
}
Err(e) => return Err(e.into()),
}
}
Err(anyhow::anyhow!(
"no free port in [{}, {}){}",
start_port,
start_port.saturating_add(max_tries),
last_err
.map(|e| format!(": {e}"))
.unwrap_or_default()
))
}
fn primary_lan_ipv4() -> Option<String> {
let sock = std::net::UdpSocket::bind("0.0.0.0:0").ok()?;
sock.connect("8.8.8.8:80").ok()?;
match sock.local_addr().ok()?.ip() {
std::net::IpAddr::V4(v4) if !v4.is_loopback() && !v4.is_unspecified() => {
Some(v4.to_string())
}
_ => None,
}
}
pub const WEBUI_DEFAULT_PORT: u16 = 13457;
pub async fn ensure_server_and_open(host: &str, port: u16, sync: bool) -> String {
let reuse = {
let guard = WEBUI.lock().unwrap();
match guard.as_ref() {
Some(handle) if !handle.abort.is_finished() => {
Some((handle.tokens.clone(), handle.port, handle.host.clone()))
}
_ => None,
}
};
let (tokens, actual_port, bound_host) = if let Some((tokens, p, h)) = reuse {
(tokens, p, h)
} else {
let (listener, actual_port) = match bind_scanning(host, port, 100).await {
Ok(v) => v,
Err(e) => {
return format!("webui 启动失败:{host}:{port} 起的端口绑定失败({e})");
}
};
let tokens = auth_token::WebuiTokenStore::new();
let opts = ServerOpts {
host: host.to_string(),
port: actual_port,
cli_override: CliOverride::default(),
idle_timeout_secs: 0,
startup_mode: SessionMode::Webui,
webui_tokens: Some(tokens.clone()),
quiet: true,
working_dir_override: std::env::current_dir().ok(),
prebound_listener: Some(listener),
};
let task = tokio::spawn(async move {
if let Err(e) = run_server(opts).await {
eprintln!("webui server error: {e}");
}
});
{
let mut guard = WEBUI.lock().unwrap();
*guard = Some(WebuiHandle {
tokens: tokens.clone(),
port: actual_port,
host: host.to_string(),
abort: task.abort_handle(),
});
}
(tokens, actual_port, host.to_string())
};
let token = tokens.mint();
let open_host: &str = if is_loopback_authority(&bound_host)
|| bound_host == "0.0.0.0"
|| bound_host == "::"
{
"127.0.0.1"
} else {
bound_host.as_str()
};
let sync_suffix = if sync { "&sync=1" } else { "" };
let local_url = format!("http://{}:{}/?token={}{}", open_host, actual_port, token, sync_suffix);
let opened = atomcode_core::auth::oauth::open_browser(&local_url).is_ok();
let mut msg = if opened {
format!("已在浏览器打开 webui:{local_url}")
} else {
format!("请手动在浏览器打开:{local_url}")
};
if bound_host.as_str() != host {
msg.push_str(&format!(
"\n(webui 已在运行,绑定 {bound_host};如需改绑 {host},请先 /webui stop 再重试)"
));
}
if !is_loopback_authority(&bound_host) {
if bound_host == "0.0.0.0" || bound_host == "::" {
if let Some(ip) = primary_lan_ipv4() {
msg.push_str(&format!("\n局域网访问:http://{ip}:{actual_port}/?token={token}"));
}
msg.push_str(
"\n⚠️ 上面是局域网 IP,仅同一网络内的设备可访问;公网访问请用隧道(如 cloudflared / Tailscale)。无 TLS,凡能访问者凭 token 即可进入。",
);
} else {
msg.push_str(&format!("\n访问地址:http://{bound_host}:{actual_port}/?token={token}"));
msg.push_str(
"\n⚠️ 已绑定非回环地址:凡能访问该地址者凭此 token 即可进入,请仅在可信网络使用(无 TLS)。",
);
}
}
msg
}
pub fn stop_server() -> String {
let mut guard = WEBUI.lock().unwrap();
if let Some(handle) = guard.take() {
handle.abort.abort();
"已停止 webui server".to_string()
} else {
"webui server 未在运行".to_string()
}
}
#[derive(serde::Serialize)]
struct PgyInfo {
installed: bool,
ipv4: Option<String>,
}
#[derive(serde::Serialize)]
struct TunnelStatus {
bind_host: String,
port: u16,
reachable: bool,
pgy: PgyInfo,
remote_url: Option<String>,
qr_svg: Option<String>,
}
fn pgy_ipv4_candidates(ifconfig_output: &str) -> Vec<String> {
let mut out = Vec::new();
let mut in_p2p = false;
for line in ifconfig_output.lines() {
let is_header = line.chars().next().map(|c| !c.is_whitespace()).unwrap_or(false);
if is_header {
in_p2p = line.contains("POINTOPOINT");
continue;
}
if !in_p2p {
continue;
}
if let Some(rest) = line.trim_start().strip_prefix("inet ") {
if let Some(addr) = rest.split_whitespace().next() {
if let Ok(ip) = addr.parse::<std::net::Ipv4Addr>() {
if ip.is_private() {
out.push(ip.to_string());
}
}
}
}
}
out
}
fn extract_ip_eq(line: &str) -> Option<String> {
let idx = line.find("ip=")?;
let rest = &line[idx + 3..];
let end = rest
.find(|c: char| !(c.is_ascii_digit() || c == '.'))
.unwrap_or(rest.len());
let cand = &rest[..end];
cand.parse::<std::net::Ipv4Addr>().ok().map(|_| cand.to_string())
}
fn pgy_ipv4_from_log() -> Option<String> {
let home = std::env::var_os("HOME")?;
let dir = std::path::Path::new(&home).join("Library/Logs/PgyVisitor");
let mut latest: Option<String> = None;
for entry in std::fs::read_dir(&dir).ok()?.flatten() {
let path = entry.path();
if path.extension().and_then(|e| e.to_str()) != Some("log") {
continue;
}
let Ok(content) = std::fs::read_to_string(&path) else {
continue;
};
for line in content.lines() {
if let Some(ip) = extract_ip_eq(line) {
latest = Some(ip);
}
}
}
latest
}
fn pgy_pick_ipv4(candidates: Vec<String>, log_ip: Option<String>) -> Option<String> {
match candidates.len() {
1 => candidates.into_iter().next(),
0 => None,
_ => log_ip.filter(|ip| candidates.contains(ip)),
}
}
fn pgy_installed() -> bool {
if std::path::Path::new("/Applications/PgyVisitor_download.app").exists() {
return true;
}
if let Ok(entries) = std::fs::read_dir("/Library/LaunchDaemons") {
for e in entries.flatten() {
if let Some(name) = e.file_name().to_str() {
let n = name.to_ascii_lowercase();
if n.starts_with("com.oray.") && n.contains("pgy") && n.ends_with(".plist") {
return true;
}
}
}
}
false
}
fn pgy_probe() -> PgyInfo {
let installed = pgy_installed();
let candidates = std::process::Command::new("ifconfig")
.output()
.ok()
.filter(|o| o.status.success())
.map(|o| pgy_ipv4_candidates(&String::from_utf8_lossy(&o.stdout)))
.unwrap_or_default();
let ipv4 = pgy_pick_ipv4(candidates, pgy_ipv4_from_log());
PgyInfo { installed, ipv4 }
}
async fn get_tunnel_status(
State(state): State<AppState>,
headers: axum::http::HeaderMap,
) -> impl IntoResponse {
let pgy = tokio::task::spawn_blocking(pgy_probe)
.await
.unwrap_or(PgyInfo { installed: false, ipv4: None });
let reachable = !is_loopback_authority(&state.bind_host);
let pgy_reachable = matches!(state.bind_host.as_str(), "0.0.0.0" | "::")
|| pgy.ipv4.as_deref() == Some(state.bind_host.as_str());
let token = auth_token::token_from_header(
headers
.get(axum::http::header::AUTHORIZATION)
.and_then(|h| h.to_str().ok()),
);
let (remote_url, qr_svg) = match (&pgy.ipv4, &token) {
(Some(ip), Some(tok)) if pgy_reachable => {
let url = format!("http://{}:{}/?token={}&sync=1", ip, state.bind_port, tok);
let qr = qrcode::QrCode::new(url.as_bytes()).ok().map(|code| {
code.render::<qrcode::render::svg::Color>()
.min_dimensions(200, 200)
.quiet_zone(true)
.build()
});
(Some(url), qr)
}
_ => (None, None),
};
Json(TunnelStatus {
bind_host: state.bind_host.clone(),
port: state.bind_port,
reachable,
pgy,
remote_url,
qr_svg,
})
}
#[derive(serde::Serialize)]
pub struct SkillInfo {
pub name: String,
pub description: String,
}
async fn get_skills(State(state): State<AppState>) -> impl IntoResponse {
let working_dir = { state.project.read().await.working_dir.clone() };
let mut registry = atomcode_core::skill::SkillRegistry::new();
registry.reload(&working_dir);
let skills: Vec<SkillInfo> = registry
.user_invocable()
.map(|s| SkillInfo {
name: s.name.clone(),
description: s.description.clone(),
})
.collect();
Json(skills)
}
pub fn normalize_dir_arg(arg: &str) -> PathBuf {
if let Some(rest) = arg.strip_prefix('~') {
if let Some(home) = atomcode_core::tool::real_home_dir() {
return home.join(rest.trim_start_matches('/'));
}
}
PathBuf::from(arg)
}
pub fn list_subdirs(dir: &std::path::Path) -> anyhow::Result<Vec<String>> {
let mut out = Vec::new();
for entry in std::fs::read_dir(dir)? {
let entry = entry?;
if entry.file_type()?.is_dir() {
if let Some(name) = entry.file_name().to_str() {
if !name.starts_with('.') {
out.push(name.to_string());
}
}
}
}
out.sort();
Ok(out)
}
pub fn list_files(dir: &std::path::Path) -> anyhow::Result<Vec<String>> {
let mut out = Vec::new();
for entry in std::fs::read_dir(dir)? {
let entry = entry?;
if entry.file_type()?.is_dir() {
continue;
}
if let Some(name) = entry.file_name().to_str() {
if !name.starts_with('.') {
out.push(name.to_string());
}
}
}
out.sort();
Ok(out)
}
#[derive(serde::Deserialize)]
pub struct FsListQuery {
pub path: String,
}
async fn fs_list(
State(_state): State<AppState>,
Query(q): Query<FsListQuery>,
) -> impl IntoResponse {
let expanded = normalize_dir_arg(&q.path);
let dir = expanded.canonicalize().unwrap_or(expanded);
match list_subdirs(&dir) {
Ok(dirs) => Json(serde_json::json!({
"path": dir.to_string_lossy(),
"dirs": dirs,
"files": list_files(&dir).unwrap_or_default(),
}))
.into_response(),
Err(e) => json_error(StatusCode::BAD_REQUEST, format!("{e}")).into_response(),
}
}
#[derive(serde::Deserialize)]
pub struct FsMkdirRequest { pub path: String }
async fn fs_mkdir(
State(_state): State<AppState>,
Json(req): Json<FsMkdirRequest>,
) -> impl IntoResponse {
let dir = normalize_dir_arg(&req.path);
match std::fs::create_dir_all(&dir) {
Ok(()) => {
let canon = dir.canonicalize().unwrap_or(dir);
Json(serde_json::json!({ "path": canon.to_string_lossy() })).into_response()
}
Err(e) => json_error(StatusCode::BAD_REQUEST, format!("{e}")).into_response(),
}
}
pub struct ServerOpts {
pub host: String,
pub port: u16,
pub cli_override: CliOverride,
pub idle_timeout_secs: u64,
pub startup_mode: SessionMode,
pub webui_tokens: Option<auth_token::WebuiTokenStore>,
pub working_dir_override: Option<PathBuf>,
pub quiet: bool,
pub prebound_listener: Option<tokio::net::TcpListener>,
}
pub async fn run_server(opts: ServerOpts) -> anyhow::Result<()> {
use axum::routing::patch;
let ServerOpts {
host,
port,
cli_override,
idle_timeout_secs,
startup_mode,
webui_tokens,
quiet,
working_dir_override,
prebound_listener,
} = opts;
let cfg_telemetry = match Config::load(&Config::default_path()) {
Ok(c) => c.telemetry,
Err(e) => {
tracing::warn!(?e, "Failed to load config, using defaults");
atomcode_telemetry::TelemetryConfig::default()
}
};
let resolved = resolve(&cfg_telemetry, &cli_override, Config::config_dir(), &ProcessEnv);
if !quiet {
match &resolved.state {
TelemetryState::Enabled => println!("Telemetry: enabled"),
TelemetryState::Disabled(reason) => {
println!("Telemetry: disabled (reason: {})", reason)
}
}
}
let atomcode_dir = resolved.atomcode_dir.clone();
let telemetry = Telemetry::init(resolved, env!("CARGO_PKG_VERSION").into());
install_panic_hook(telemetry.clone());
telemetry.maybe_emit_install_completed(&atomcode_dir).await;
let project_state = init_project_state(working_dir_override);
let repo_origin = detect_repo_origin(&project_state.working_dir);
telemetry.set_account_id(auth::get_stored_auth().map(|a| a.user.id));
let mcp_registry = McpRegistry::from_config_background(&project_state.working_dir);
let (shutdown_tx, shutdown_rx) = watch::channel(false);
let last_activity = Arc::new(std::sync::atomic::AtomicI64::new(now_unix_ms()));
let active_connections = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let state = AppState {
sessions: Arc::new(RwLock::new(std::collections::HashMap::new())),
project: Arc::new(RwLock::new(project_state)),
chat_tasks: Arc::new(RwLock::new(HashMap::new())),
stopped_sessions: Arc::new(RwLock::new(HashSet::new())),
mcp_registry: Arc::new(RwLock::new(Arc::new(mcp_registry))),
mcp_cache: Arc::new(RwLock::new(HashMap::new())),
login_sessions: Arc::new(RwLock::new(HashMap::new())),
telemetry: telemetry.clone(),
repo_origin: repo_origin.clone(),
shutdown_tx: shutdown_tx.clone(),
last_activity: last_activity.clone(),
active_connections: active_connections.clone(),
enforce_token: webui_tokens.is_some(),
webui_tokens: webui_tokens.unwrap_or_default(),
pending_permissions: permission_bridge::PermissionResponders::new(),
bind_host: host.clone(),
bind_port: port,
};
let public = Router::new()
.route("/health", get(health))
.route("/", axum::routing::get(webui::serve_webui))
.fallback(webui::serve_webui);
let protected = Router::new()
.route("/shutdown", post(shutdown_handler))
.route("/sessions", get(get_all_sessions).post(create_session))
.route("/sessions/search", get(search_sessions))
.route("/project", get(get_project_state))
.route("/cd", post(change_dir))
.route("/projects", get(get_projects))
.route("/projects/:hash/sessions", get(get_project_sessions))
.route(
"/projects/:hash/sessions/:id",
get(get_session_detail).delete(delete_session),
)
.route("/projects/:hash/sessions/:id/rename", patch(rename_session))
.route("/models", get(get_models))
.route("/chat", post(chat_stream))
.route("/chat/stop", post(stop_chat))
.route("/chat/active", get(active_chat_sessions))
.route("/chat/permission", post(chat_permission))
.route("/tunnel/status", get(get_tunnel_status))
.route("/live", get(live_api::live_stream))
.route("/live/message", post(live_api::live_message))
.route("/live/permission", post(live_api::live_permission))
.route("/live/provider", post(live_api::live_provider))
.route("/skills", get(get_skills))
.route("/fs/list", get(fs_list))
.route("/fs/mkdir", post(fs_mkdir))
.route("/mcp/status", get(mcp_status))
.route("/mcp/reload", post(mcp_reload))
.route("/config", get(api_config::get_config))
.route("/config/reload", post(api_config::reload_config))
.route(
"/providers",
get(api_provider::get_providers).post(api_provider::create_provider),
)
.route(
"/providers/:name",
patch(api_provider::patch_provider).delete(api_provider::delete_provider),
)
.route(
"/providers/:name/default",
post(api_provider::set_default_provider),
)
.route(
"/providers/:name/thinking",
patch(api_provider::patch_thinking),
)
.route("/auth/status", get(api_auth::auth_status))
.route("/auth/login/start", post(api_auth::auth_login_start))
.route(
"/auth/login/:login_id/poll",
post(api_auth::auth_login_poll),
)
.route("/auth/login/:login_id", delete(api_auth::auth_login_cancel))
.route("/auth/logout", post(api_auth::auth_logout))
.route("/codingplan/setup", post(api_codingplan::codingplan_setup))
.route_layer(axum::middleware::from_fn_with_state(
state.clone(),
auth_token::require_webui_token,
));
let app = public
.merge(protected)
.with_state(state)
.layer(axum::middleware::from_fn(activity_tracker_middleware))
.layer(axum::Extension(last_activity.clone()))
.layer(cors_layer());
spawn_idle_timeout_task(
idle_timeout_secs,
last_activity,
active_connections,
shutdown_tx,
);
if !quiet {
if idle_timeout_secs > 0 {
println!("Idle timeout: {} minutes", idle_timeout_secs / 60);
} else {
println!("Idle timeout: disabled");
}
}
let addr = format!("{host}:{port}");
if host != "127.0.0.1" && host != "localhost" && host != "::1" {
eprintln!(
"Warning: binding to non-loopback address '{}'. \
The daemon exposes sensitive endpoints (chat, file-edit, tool-execution). \
Ensure the network is trusted or use a reverse proxy with authentication.",
host
);
}
if dangerous_tools_enabled() {
eprintln!(
"Warning: {}=1 enables bash and write-capable daemon tools.",
DANGEROUS_TOOLS_ENV
);
}
if !quiet {
println!("AtomCode API server listening on http://{}", addr);
println!("\nAPI endpoints:");
println!(" GET /health - Health check");
println!(" GET /project - Get current working directory");
println!(
" POST /cd - Change working directory (like /cd command)"
);
println!(" GET /projects - List historical projects");
println!(" GET /projects/:hash/sessions - List sessions in a project");
println!(" GET /projects/:hash/sessions/:id - Get session detail");
println!(" DELETE /projects/:hash/sessions/:id - Delete a session");
println!(" PATCH /projects/:hash/sessions/:id/rename - Rename a session");
println!(" GET /sessions - List all sessions (cross-project)");
println!(" GET /sessions/search?q=<keyword> - Search sessions by name");
println!(" GET /models - List available models");
println!(" POST /chat - Stream chat response (SSE)");
println!(" GET /config - Get sanitized config");
println!(" POST /config/reload - Reload config from disk");
println!(" GET /providers - List providers");
println!(" POST /providers - Create/replace provider");
println!(" PATCH /providers/:name - Partially update provider");
println!(" DELETE /providers/:name - Delete provider");
println!(" POST /providers/:name/default - Set default provider");
println!(" PATCH /providers/:name/thinking - Update thinking settings");
println!(" GET /skills - List user-invocable skills");
println!(" GET /auth/status - Auth status");
println!(" POST /auth/login/start - Start OAuth login");
println!(" POST /auth/login/:login_id/poll - Poll login session");
println!(" DELETE /auth/login/:login_id - Cancel login session");
println!(" POST /auth/logout - Logout");
println!(" POST /codingplan/setup - Run CodingPlan setup");
println!("\nChange directory body:");
println!(" {{\"path\": \"/path/to/project\"}} or {{\"path\": \"-\"}} to go back");
println!("\nChat request body:");
println!(" {{\"message\": \"your question\", \"provider\": \"optional\"}}");
}
let listener = match prebound_listener {
Some(l) => l,
None => match tokio::net::TcpListener::bind(&addr).await {
Ok(l) => l,
Err(e) => {
eprintln!("Fatal: failed to bind to {}: {}", addr, e);
CurrentContext::scope(
CurrentContext {
mode: Some(startup_mode),
repo_origin: Some(repo_origin.clone()),
session_id: None,
..CurrentContext::default()
},
|| async {
telemetry.track(Event::OpenAtomcode { dangerously_skip_permissions: false });
},
)
.await;
telemetry.shutdown(Duration::from_millis(500)).await;
std::process::exit(1);
}
},
};
CurrentContext::scope(
CurrentContext {
mode: Some(startup_mode),
repo_origin: Some(repo_origin.clone()),
session_id: None,
..CurrentContext::default()
},
|| async {
telemetry.track(Event::OpenAtomcode { dangerously_skip_permissions: false });
},
)
.await;
axum::serve(listener, app)
.with_graceful_shutdown(shutdown_signal(shutdown_rx))
.await
.unwrap_or_else(|e| tracing::error!(?e, "axum::serve error"));
telemetry.shutdown(Duration::from_millis(500)).await;
Ok(())
}
#[cfg(test)]
mod fs_list_tests {
use super::*;
#[test]
fn expands_tilde() {
if let Some(home) = atomcode_core::tool::real_home_dir() {
assert_eq!(normalize_dir_arg("~"), home);
assert_eq!(normalize_dir_arg("~/x"), home.join("x"));
}
}
#[test]
fn lists_subdirs_of_temp() {
let base = std::env::temp_dir()
.join(format!("atomcode_fslist_test_{}", std::process::id()));
let _ = std::fs::create_dir_all(base.join("childdir"));
let _ = std::fs::write(base.join("afile.txt"), b"x");
let dirs = list_subdirs(&base).unwrap();
assert!(dirs.contains(&"childdir".to_string()));
assert!(!dirs.iter().any(|d| d == "afile.txt"));
let _ = std::fs::remove_dir_all(&base);
}
#[test]
fn errors_on_missing_dir() {
assert!(list_subdirs(std::path::Path::new("/no/such/dir/xyz123")).is_err());
}
}
#[cfg(test)]
mod tests {
use super::*;
fn origin_is_allowed(origin: &str) -> bool {
let origin = HeaderValue::from_str(origin).unwrap();
let request = axum::http::Request::builder().body(()).unwrap();
let (parts, _) = request.into_parts();
is_loopback_origin(&origin, &parts)
}
#[test]
fn cors_allows_loopback_origins() {
assert!(origin_is_allowed("http://localhost:3000"));
assert!(origin_is_allowed("http://127.0.0.1:3000"));
assert!(origin_is_allowed("http://[::1]:3000"));
assert!(origin_is_allowed("https://localhost"));
}
#[test]
fn initial_workdir_override_wins_over_config_default() {
let here = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
let resolved = resolve_initial_working_dir(
Some(here.clone()),
Some(PathBuf::from("/tmp")),
PathBuf::from("/nonexistent_atomcode_cwd"),
);
assert_eq!(resolved, here);
}
#[test]
fn initial_workdir_falls_back_to_config_then_cwd() {
let here = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
assert_eq!(
resolve_initial_working_dir(None, Some(here.clone()), PathBuf::from("/x")),
here
);
assert_eq!(
resolve_initial_working_dir(
None,
Some(PathBuf::from("/nonexistent_atomcode_default")),
here.clone()
),
here
);
}
#[test]
fn initial_workdir_ignores_nonexistent_override() {
let here = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
let resolved = resolve_initial_working_dir(
Some(PathBuf::from("/nonexistent_atomcode_override")),
Some(here.clone()),
PathBuf::from("/x"),
);
assert_eq!(resolved, here);
}
#[test]
fn list_files_returns_files_skips_dirs_and_hidden() {
let tmp = std::env::temp_dir().join(format!(
"atomcode_list_files_{}",
std::process::id()
));
let _ = std::fs::remove_dir_all(&tmp);
std::fs::create_dir_all(&tmp).unwrap();
std::fs::write(tmp.join("b.txt"), b"x").unwrap();
std::fs::write(tmp.join("a.txt"), b"x").unwrap();
std::fs::write(tmp.join(".hidden"), b"x").unwrap();
std::fs::create_dir_all(tmp.join("subdir")).unwrap();
let files = list_files(&tmp).unwrap();
assert_eq!(files, vec!["a.txt".to_string(), "b.txt".to_string()]);
assert!(list_subdirs(&tmp).unwrap().contains(&"subdir".to_string()));
let _ = std::fs::remove_dir_all(&tmp);
}
#[tokio::test]
async fn bind_scanning_returns_a_free_port() {
let (listener, port) = bind_scanning("127.0.0.1", 0, 1).await.unwrap();
assert_ne!(port, 0);
assert_eq!(listener.local_addr().unwrap().port(), port);
}
#[tokio::test]
async fn bind_scanning_skips_occupied_port() {
let occupied = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let busy = occupied.local_addr().unwrap().port();
let (listener, port) = bind_scanning("127.0.0.1", busy, 50).await.unwrap();
assert_ne!(port, busy);
assert!(port > busy);
drop(listener);
drop(occupied);
}
#[test]
fn cors_rejects_remote_and_opaque_origins() {
assert!(!origin_is_allowed("http://192.168.1.10:3000"));
assert!(!origin_is_allowed("http://localhost.evil.example"));
assert!(!origin_is_allowed("null"));
assert!(!origin_is_allowed("file://local/index.html"));
}
const IFCONFIG_SAMPLE: &str = "\
en0: flags=8863<UP,BROADCAST,SMART,RUNNING,SIMPLEX,MULTICAST> mtu 1500
\tinet 172.20.23.187 netmask 0xfffffc00 broadcast 172.20.23.255
utun6: flags=8051<UP,POINTOPOINT,RUNNING,MULTICAST> mtu 1280
\tinet 100.117.29.83 --> 100.117.29.83 netmask 0xffffffff
utun7: flags=8051<UP,POINTOPOINT,RUNNING,MULTICAST> mtu 1300
\tinet 172.16.2.14 --> 172.16.2.14 netmask 0xfffffc00
";
#[test]
fn pgy_candidates_picks_only_p2p_rfc1918() {
assert_eq!(pgy_ipv4_candidates(IFCONFIG_SAMPLE), vec!["172.16.2.14"]);
}
#[test]
fn pgy_candidates_excludes_rfc1918_on_non_p2p() {
let s = "en5: flags=8863<UP,BROADCAST,RUNNING,MULTICAST> mtu 1500\n\
\tinet 172.16.5.9 netmask 0xffff0000 broadcast 172.16.255.255\n";
assert!(pgy_ipv4_candidates(s).is_empty());
}
#[test]
fn pgy_candidates_is_segment_agnostic() {
let s = "utun9: flags=8051<UP,POINTOPOINT,RUNNING,MULTICAST> mtu 1300\n\
\tinet 10.8.0.3 --> 10.8.0.3 netmask 0xffffff00\n\
utun10: flags=8051<UP,POINTOPOINT,RUNNING,MULTICAST> mtu 1300\n\
\tinet 192.168.50.2 --> 192.168.50.2 netmask 0xffffff00\n";
assert_eq!(pgy_ipv4_candidates(s), vec!["10.8.0.3", "192.168.50.2"]);
}
#[test]
fn pgy_candidates_empty_when_none() {
let s = "lo0: flags=8049<UP,LOOPBACK,RUNNING,MULTICAST> mtu 16384\n\
\tinet 127.0.0.1 netmask 0xff000000\n";
assert!(pgy_ipv4_candidates(s).is_empty());
}
#[test]
fn pgy_pick_single_candidate() {
assert_eq!(
pgy_pick_ipv4(vec!["172.16.2.14".into()], None),
Some("172.16.2.14".into())
);
}
#[test]
fn pgy_pick_zero_candidates_is_none() {
assert_eq!(pgy_pick_ipv4(vec![], Some("172.16.2.14".into())), None);
}
#[test]
fn pgy_pick_multi_disambiguates_via_log() {
let cands = vec!["172.16.2.14".to_string(), "10.99.0.5".to_string()];
assert_eq!(
pgy_pick_ipv4(cands.clone(), Some("10.99.0.5".into())),
Some("10.99.0.5".into())
);
assert_eq!(pgy_pick_ipv4(cands.clone(), Some("172.31.9.9".into())), None);
assert_eq!(pgy_pick_ipv4(cands, None), None);
}
#[test]
fn extract_ip_eq_parses_log_line() {
assert_eq!(
extract_ip_eq("2026-06-01 worker connected ip=172.16.2.14 gw=172.16.0.159"),
Some("172.16.2.14".into())
);
assert_eq!(extract_ip_eq("no ip here"), None);
}
}