octomind 0.22.0

// Copyright 2025 Muvon Un Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! OctomindAgent — implements the ACP Agent trait over Octomind's session infrastructure.

use std::cell::RefCell;
use std::collections::HashMap;
use std::path::PathBuf;
use std::rc::Rc;

use agent_client_protocol::{
	AgentCapabilities, AgentSideConnection, AuthenticateRequest, AuthenticateResponse,
	AvailableCommand, AvailableCommandInput, AvailableCommandsUpdate, BlobResourceContents,
	CancelNotification, Client, ContentBlock, ContentChunk, EmbeddedResourceResource, ExtRequest,
	ExtResponse, Implementation, InitializeRequest, InitializeResponse, LoadSessionRequest,
	LoadSessionResponse, McpCapabilities, McpServer, NewSessionRequest, NewSessionResponse,
	PromptCapabilities, PromptRequest, PromptResponse, ProtocolVersion, SessionNotification,
	SessionUpdate, StopReason, ToolCall, ToolCallStatus, ToolCallUpdate, ToolCallUpdateFields,
	UnstructuredCommandInput,
};

use crate::config::mcp::McpServerConfig;
use crate::config::Config;
use crate::session::cancellation::SessionCancellation;
use crate::session::chat::session::{
	execute_api_call_and_process_response, prepare_for_api_call, process_layers_if_enabled,
	setup_and_initialize_session, setup_system_prompt_and_cache, ChatSession, GenericSessionArgs,
};
use crate::session::output::{OutputMode, WebSocketSink};
use crate::websocket::ServerMessage;
use crate::{log_debug, log_error, log_info};

/// ACP agent implementation wrapping Octomind's session infrastructure.
///
/// Single-threaded (runs inside a `tokio::task::LocalSet`), so `Rc<RefCell<...>>` is safe.
pub struct OctomindAgent {
	/// Mutable: client-injected MCP servers are merged in on first new_session/load_session.
	config: RefCell<Config>,
	role: String,
	/// Active sessions keyed by ACP session_id, paired with their working directory.
	sessions: Rc<RefCell<HashMap<String, (ChatSession, PathBuf)>>>,
	/// Active cancellation handles keyed by ACP session_id
	cancellations: Rc<RefCell<HashMap<String, SessionCancellation>>>,
	/// Connection back to the client — used to send session/update notifications
	conn: Rc<RefCell<Option<Rc<AgentSideConnection>>>>,
}

impl OctomindAgent {
	pub fn new(config: Config, role: String) -> Self {
		Self {
			config: RefCell::new(config),
			role,
			sessions: Rc::new(RefCell::new(HashMap::new())),
			cancellations: Rc::new(RefCell::new(HashMap::new())),
			conn: Rc::new(RefCell::new(None)),
		}
	}

	/// Inject the connection after it's created (chicken-and-egg: agent needs conn, conn needs agent).
	pub fn set_connection(&self, conn: Rc<AgentSideConnection>) {
		*self.conn.borrow_mut() = Some(conn);
	}
}

/// Convert ACP MCP server list into McpServerConfig entries and inject them into a config snapshot.
///
/// Returns a modified clone of `base_config` with the injected servers merged in.
/// `self.config` is never mutated — injected servers are scoped to the session only.
fn build_config_with_injected_servers(
	base_config: &Config,
	role: &str,
	servers: &[McpServer],
) -> Config {
	let mut config = base_config.clone();
	for server in servers {
		let server_config = match server {
			McpServer::Stdio(s) => {
				let args: Vec<String> = s.args.iter().map(|a| a.to_string()).collect();
				McpServerConfig::stdin(
					&s.name,
					s.command.to_string_lossy().as_ref(),
					args,
					30,
					vec![],
				)
			}
			McpServer::Http(s) => McpServerConfig::http(&s.name, &s.url, 30, vec![], None, None),
			McpServer::Sse(s) => {
				// SSE is not a supported transport in our MCP stack — skip
				log_info!("ACP: skipping SSE MCP server '{}' (not supported)", s.name);
				continue;
			}
			_ => {
				log_info!("ACP: skipping unknown MCP server transport (not supported)");
				continue;
			}
		};
		let name = server_config.name().to_string();
		if !config.mcp.servers.iter().any(|s| s.name() == name) {
			config.mcp.servers.push(server_config);
		}
		if let Some(role_entry) = config.role_map.get_mut(role) {
			if !role_entry.mcp.server_refs.contains(&name) {
				role_entry.mcp.server_refs.push(name);
			}
		}
	}
	config
}

/// Build the list of available slash commands to advertise to ACP clients.
///
/// Command names are sent WITHOUT the leading `/` — the client prepends it when displaying.
fn build_available_commands() -> Vec<AvailableCommand> {
	let unstructured =
		|hint: &str| AvailableCommandInput::Unstructured(UnstructuredCommandInput::new(hint));

	vec![
		AvailableCommand::new("help", "Show available commands"),
		AvailableCommand::new("role", "View or change current role")
			.input(unstructured("<role_name>")),
		AvailableCommand::new("model", "View or change current AI model")
			.input(unstructured("<provider:model>")),
		AvailableCommand::new(
			"done",
			"Finalize task with memorization, summarization, and auto-commit",
		),
		AvailableCommand::new("save", "Save the current session"),
		AvailableCommand::new("info", "Display token and cost breakdown for this session"),
		AvailableCommand::new("clear", "Clear the screen"),
		AvailableCommand::new("copy", "Copy last response to clipboard"),
		AvailableCommand::new("context", "Display session context")
			.input(unstructured("[all|assistant|user|tool|large]")),
		AvailableCommand::new("truncate", "Smart context truncation to reduce token usage"),
		AvailableCommand::new(
			"summarize",
			"Summarize entire conversation to reduce token usage",
		),
		AvailableCommand::new("cache", "Manage cache checkpoints")
			.input(unstructured("[stats|clear|threshold]")),
		AvailableCommand::new("list", "List all available sessions").input(unstructured("[page]")),
		AvailableCommand::new("session", "Switch to or create a session")
			.input(unstructured("[session_name]")),
		AvailableCommand::new("run", "Execute a command layer")
			.input(unstructured("<command_name>")),
		AvailableCommand::new("workflow", "Execute a workflow")
			.input(unstructured("<workflow_name> [input]")),
		AvailableCommand::new("mcp", "MCP server management")
			.input(unstructured("[info|list|full|health|dump|validate]")),
		AvailableCommand::new("plan", "Display current plan stored in MCP plan tool"),
		AvailableCommand::new("prompt", "Manage prompt templates")
			.input(unstructured("[template_name]")),
		AvailableCommand::new("image", "Attach image to next message")
			.input(unstructured("<path>")),
		AvailableCommand::new("video", "Attach video to next message")
			.input(unstructured("<path>")),
		AvailableCommand::new("loglevel", "Set logging level")
			.input(unstructured("[none|info|debug]")),
		AvailableCommand::new("report", "Generate detailed usage report for this session"),
		AvailableCommand::new("exit", "Exit the session"),
	]
}

/// Send the available commands list to the ACP client for the given session.
async fn send_available_commands(conn: Option<std::rc::Rc<AgentSideConnection>>, session_id: &str) {
	if let Some(conn) = conn {
		let update = SessionUpdate::AvailableCommandsUpdate(AvailableCommandsUpdate::new(
			build_available_commands(),
		));
		let notif = SessionNotification::new(std::sync::Arc::<str>::from(session_id), update);
		if let Err(e) = conn.session_notification(notif).await {
			log_error!("ACP: failed to send available_commands_update: {}", e);
		}
	}
}

#[async_trait::async_trait(?Send)]
impl agent_client_protocol::Agent for OctomindAgent {
	async fn initialize(
		&self,
		args: InitializeRequest,
	) -> agent_client_protocol::Result<InitializeResponse> {
		log_debug!("ACP: initialize from {:?}", args.client_info);

		// Advertise extension capabilities in _meta per ACP spec
		let mut meta = agent_client_protocol::Meta::new();
		meta.insert(
			"octomind.dev".to_string(),
			serde_json::json!({
				"commands": true
			}),
		);

		let response = InitializeResponse::new(ProtocolVersion::LATEST)
			.agent_capabilities(
				AgentCapabilities::default()
					.load_session(true)
					// Advertise HTTP MCP transport support so clients offer us HTTP servers.
					// SSE is not supported — we skip those servers silently in inject_acp_mcp_servers.
					.mcp_capabilities(McpCapabilities::new().http(true))
					.prompt_capabilities(
						PromptCapabilities::default()
							.image(true)
							.embedded_context(true),
					)
					.meta(meta),
			)
			.agent_info(Implementation::new("octomind", env!("CARGO_PKG_VERSION")));
		Ok(response)
	}

	async fn authenticate(
		&self,
		_args: AuthenticateRequest,
	) -> agent_client_protocol::Result<AuthenticateResponse> {
		Ok(AuthenticateResponse::default())
	}

	async fn new_session(
		&self,
		args: NewSessionRequest,
	) -> agent_client_protocol::Result<NewSessionResponse> {
		// Set per-session working directory via thread-local (safe: single-threaded LocalSet)
		crate::mcp::set_session_working_directory(args.cwd.clone());
		let session_cwd = args.cwd.clone();

		// Build a per-session config snapshot with injected servers merged in.
		// self.config is never mutated — injected servers are scoped to this session only.
		let config_snapshot = build_config_with_injected_servers(
			&self.config.borrow(),
			&self.role,
			&args.mcp_servers,
		);

		// Start any newly injected servers and register their tools in the tool map.
		// initialize_mcp_for_role is idempotent: already-running servers and already-registered
		// tools are skipped via config-hash and is_server_already_running checks.
		crate::mcp::initialize_mcp_for_role(&self.role, &config_snapshot)
			.await
			.map_err(|e| agent_client_protocol::Error::internal_error().data(e.to_string()))?;

		let session_args = GenericSessionArgs {
			role: self.role.clone(),
			mode: "websocket".into(),
			..Default::default()
		};
		let (mut chat_session, config_for_role, session_role, _) =
			setup_and_initialize_session(&session_args, &config_snapshot)
				.await
				.map_err(|e| agent_client_protocol::Error::internal_error().data(e.to_string()))?;

		setup_system_prompt_and_cache(&mut chat_session, &config_for_role, &session_role, false)
			.await
			.map_err(|e| agent_client_protocol::Error::internal_error().data(e.to_string()))?;

		let session_id = chat_session.session.info.name.clone();
		log_debug!("ACP: new_session created: {}", session_id);

		self.sessions
			.borrow_mut()
			.insert(session_id.clone(), (chat_session, session_cwd));
		self.cancellations
			.borrow_mut()
			.insert(session_id.clone(), SessionCancellation::new());

		let conn = self.conn.borrow().clone();
		send_available_commands(conn, &session_id).await;

		Ok(NewSessionResponse::new(session_id))
	}

	async fn prompt(&self, args: PromptRequest) -> agent_client_protocol::Result<PromptResponse> {
		let session_id = args.session_id.to_string();

		// Extract text, images, and videos from prompt content blocks
		let mut text_parts = Vec::new();
		let mut images = Vec::new();
		let mut videos = Vec::new();
		for block in &args.prompt {
			match block {
				ContentBlock::Text(t) => text_parts.push(t.text.as_str()),
				ContentBlock::Image(img) => {
					images.push(crate::session::image::ImageAttachment {
						data: crate::session::image::ImageData::Base64(img.data.clone()),
						media_type: img.mime_type.clone(),
						source_type: crate::session::image::SourceType::Url, // ACP images are inline data, closest match
						dimensions: None,
						size_bytes: None,
					});
				}
				ContentBlock::Resource(res) => {
					// Extract video from embedded blob resources (ACP has no native video block)
					if let EmbeddedResourceResource::BlobResourceContents(BlobResourceContents {
						blob,
						mime_type: Some(mime),
						..
					}) = &res.resource
					{
						if mime.starts_with("video/") {
							videos.push(crate::session::video::VideoAttachment {
								data: crate::session::video::VideoData::Base64(blob.clone()),
								media_type: mime.clone(),
								source_type: crate::session::video::SourceType::Url,
								dimensions: None,
								size_bytes: None,
								duration_secs: None,
							});
						}
					}
				}
				_ => {} // Skip audio, resource links, etc.
			}
		}
		let input: String = text_parts.join("\n");

		if input.trim().is_empty() && images.is_empty() && videos.is_empty() {
			return Ok(PromptResponse::new(StopReason::EndTurn));
		}

		// Slash commands are sent as regular prompts per the ACP spec.
		// Intercept them here before the AI pipeline, execute via process_command,
		// and stream the result back as an AgentMessageChunk.
		if input.trim_start().starts_with('/') {
			let (mut chat_session, session_cwd) =
				match self.sessions.borrow_mut().remove(&session_id) {
					Some(s) => s,
					None => {
						return Err(agent_client_protocol::Error::invalid_params()
							.data(format!("session not found: {session_id}")));
					}
				};

			crate::mcp::set_session_working_directory(session_cwd.clone());

			let operation_rx = self
				.cancellations
				.borrow_mut()
				.entry(session_id.clone())
				.or_default()
				.new_operation();

			let mut config = self.config.borrow().clone();
			let result = crate::session::chat::session::commands::process_command(
				&mut chat_session,
				input.trim(),
				&mut config,
				&self.role,
				operation_rx,
			)
			.await;
			// Write back any config mutations (e.g. model/role changes)
			*self.config.borrow_mut() = config;

			self.sessions
				.borrow_mut()
				.insert(session_id.clone(), (chat_session, session_cwd));

			let text = match result {
				Ok(crate::session::chat::session::commands::CommandResult::HandledWithOutput(
					output,
				)) => serde_json::to_string_pretty(&output.to_json())
					.unwrap_or_else(|_| "Command executed.".to_string()),
				Ok(crate::session::chat::session::commands::CommandResult::Handled) => {
					"Command executed.".to_string()
				}
				Ok(crate::session::chat::session::commands::CommandResult::Exit) => {
					"Session exit requested.".to_string()
				}
				Ok(crate::session::chat::session::commands::CommandResult::TreatAsUserInput) => {
					let available: Vec<&str> = crate::session::chat::COMMANDS.to_vec();

					format!(
						"The {} command is not supported by Octomind.\n\nAvailable commands: {}",
						input.trim(),
						available.join(", ")
					)
				}
				Err(e) => format!("Command failed: {e}"),
			};

			let conn = self.conn.borrow().clone();
			if let Some(conn) = conn {
				let update = SessionUpdate::AgentMessageChunk(ContentChunk::new(text.into()));
				let notif = SessionNotification::new(
					std::sync::Arc::<str>::from(session_id.as_str()),
					update,
				);
				if let Err(e) = conn.session_notification(notif).await {
					log_error!("ACP: failed to send command result: {}", e);
				}
			}

			return Ok(PromptResponse::new(StopReason::EndTurn));
		}

		// Take session out of map for exclusive access
		let (mut chat_session, session_cwd) = match self.sessions.borrow_mut().remove(&session_id) {
			Some(s) => s,
			None => {
				return Err(agent_client_protocol::Error::invalid_params()
					.data(format!("session not found: {session_id}")));
			}
		};

		// Restore this session's working directory for tool calls
		crate::mcp::set_session_working_directory(session_cwd.clone());

		let config_for_role = self.config.borrow().get_merged_config_for_role(&self.role);
		let current_dir = session_cwd.clone();

		// Get or create cancellation for this session
		let mut cancellation = self
			.cancellations
			.borrow_mut()
			.remove(&session_id)
			.unwrap_or_default();
		cancellation.reset();
		let operation_rx = cancellation.new_operation();
		// Re-insert cancellation so cancel() can find it during prompt execution
		self.cancellations
			.borrow_mut()
			.insert(session_id.clone(), cancellation);

		// Process through layers (pre-processing step)
		let first_message_processed = !chat_session.session.messages.is_empty();
		let (processed_input, layers_modified_session, layer_cancelled) =
			process_layers_if_enabled(
				&input,
				&mut chat_session,
				&config_for_role,
				&self.role,
				first_message_processed,
				operation_rx.clone(),
			)
			.await
			.map_err(|e| agent_client_protocol::Error::internal_error().data(e.to_string()))?;

		if layer_cancelled {
			self.sessions
				.borrow_mut()
				.insert(session_id.clone(), (chat_session, session_cwd.clone()));
			return Ok(PromptResponse::new(StopReason::Cancelled));
		}

		// Attach ACP images/videos as pending so add_user_message picks them up
		if let Some(first_image) = images.into_iter().next() {
			chat_session.pending_image = Some(first_image);
		}
		if let Some(first_video) = videos.into_iter().next() {
			chat_session.pending_video = Some(first_video);
		}

		// Add user message if layers didn't modify session
		if !layers_modified_session {
			let final_input = crate::session::chat::session::utils::append_constraints_if_exists(
				&processed_input,
				&config_for_role.custom_constraints_file_name,
				&current_dir,
			);
			if let Err(e) = chat_session.add_user_message(&final_input) {
				self.sessions
					.borrow_mut()
					.insert(session_id.clone(), (chat_session, session_cwd));
				return Err(agent_client_protocol::Error::internal_error().data(e.to_string()));
			}
		}

		// Prepare for API call
		if let Err(e) =
			prepare_for_api_call(&mut chat_session, &config_for_role, operation_rx.clone()).await
		{
			self.sessions
				.borrow_mut()
				.insert(session_id.clone(), (chat_session, session_cwd));
			return Err(agent_client_protocol::Error::internal_error().data(e.to_string()));
		}

		// Channel-based sink: session pipeline emits ServerMessages, we forward them as ACP notifications
		let (ws_tx, mut ws_rx) = tokio::sync::mpsc::unbounded_channel::<ServerMessage>();
		let ws_sink = WebSocketSink::new(ws_tx.clone());

		// Forward MCP server notifications through the same channel.
		// Safe: prompt() holds exclusive access to the session (removed from map above),
		// so no two prompts for the same session can race on this global sender.
		crate::mcp::process::set_notification_sender(Some(session_id.clone()), ws_tx);

		// Spawn a local task to stream notifications to the client in real-time
		// while the API call runs concurrently. The channel closes when ws_sink drops.
		// Use Arc<str> so each SessionNotification::new() call clones the Arc pointer
		// rather than allocating a new String per notification.
		let session_id_for_task: std::sync::Arc<str> = session_id.as_str().into();
		let conn_for_task = self.conn.borrow().as_ref().cloned();
		let forward_task = tokio::task::spawn_local(async move {
			while let Some(msg) = ws_rx.recv().await {
				let update = match msg {
					ServerMessage::Assistant(p) => Some(SessionUpdate::AgentMessageChunk(
						ContentChunk::new(p.content.into()),
					)),
					ServerMessage::Thinking(p) => Some(SessionUpdate::AgentThoughtChunk(
						ContentChunk::new(p.content.into()),
					)),
					ServerMessage::ToolUse(p) => {
						let tool_call = ToolCall::new(p.tool_id.clone(), p.tool.clone())
							.status(ToolCallStatus::InProgress)
							.raw_input(p.params.clone());
						Some(SessionUpdate::ToolCall(tool_call))
					}
					ServerMessage::ToolResult(p) => {
						let status = if p.success {
							ToolCallStatus::Completed
						} else {
							ToolCallStatus::Failed
						};
						let update = ToolCallUpdate::new(
							p.tool_id.clone(),
							ToolCallUpdateFields::new().status(status).raw_output(
								serde_json::from_str::<serde_json::Value>(&p.content)
									.unwrap_or(serde_json::Value::String(p.content)),
							),
						);
						Some(SessionUpdate::ToolCallUpdate(update))
					}
					_ => None,
				};
				if let (Some(update), Some(conn)) = (update, conn_for_task.as_ref()) {
					let notif = SessionNotification::new(session_id_for_task.clone(), update);
					if let Err(e) = conn.session_notification(notif).await {
						log_error!("ACP: failed to send session notification: {}", e);
					}
				}
			}
		});

		// Execute the AI call
		let api_result = execute_api_call_and_process_response(
			&mut chat_session,
			&config_for_role,
			&self.role,
			operation_rx.clone(),
			// Reuse WebSocket output mode — ACP and WebSocket both use the same
			// channel-based ServerMessage sink; the transport layer differs, not the pipeline.
			OutputMode::WebSocket,
			ws_sink,
		)
		.await;

		// Clear the global notification sender so the channel can close.
		// Without this, forward_task.await hangs forever because NOTIFICATION_SENDER
		// holds a clone of ws_tx, preventing the channel from closing.
		crate::mcp::process::clear_notification_sender(Some(session_id.clone()));

		// Wait for the forwarding task to drain any remaining messages
		let _ = forward_task.await;

		// Put session back
		self.sessions
			.borrow_mut()
			.insert(session_id.clone(), (chat_session, session_cwd));
		// Note: cancellation was already inserted at the start of prompt() so cancel() can find it

		match api_result {
			Ok(_) => {
				if *operation_rx.borrow() {
					Ok(PromptResponse::new(StopReason::Cancelled))
				} else {
					Ok(PromptResponse::new(StopReason::EndTurn))
				}
			}
			Err(e) => {
				log_error!("ACP: prompt API call failed: {}", e);
				Err(agent_client_protocol::Error::internal_error().data(e.to_string()))
			}
		}
	}

	async fn cancel(&self, args: CancelNotification) -> agent_client_protocol::Result<()> {
		let session_id = args.session_id.to_string();
		log_debug!("ACP: cancel requested for session: {}", session_id);
		// Safe to borrow while prompt() may be awaiting: we run inside a LocalSet
		// (single-threaded), so cancel() only executes when prompt() yields at an await
		// point — the RefCell is never doubly borrowed on the same call stack.
		if let Some(cancellation) = self.cancellations.borrow().get(&session_id) {
			cancellation.shutdown();
		}
		Ok(())
	}

	async fn load_session(
		&self,
		args: LoadSessionRequest,
	) -> agent_client_protocol::Result<LoadSessionResponse> {
		let session_id = args.session_id.to_string();
		log_debug!("ACP: load_session requested: {}", session_id);

		// Set per-session working directory via thread-local
		crate::mcp::set_session_working_directory(args.cwd.clone());
		let session_cwd = args.cwd.clone();

		// Build a per-session config snapshot with injected servers merged in.
		// self.config is never mutated — injected servers are scoped to this session only.
		let config_snapshot = build_config_with_injected_servers(
			&self.config.borrow(),
			&self.role,
			&args.mcp_servers,
		);
		crate::mcp::initialize_mcp_for_role(&self.role, &config_snapshot)
			.await
			.map_err(|e| agent_client_protocol::Error::internal_error().data(e.to_string()))?;

		// Resume the existing session from disk by its ID
		let session_args = GenericSessionArgs {
			resume: Some(session_id.clone()),
			role: self.role.clone(),
			mode: "websocket".into(),
			..Default::default()
		};
		let (mut chat_session, config_for_role, session_role, _) =
			setup_and_initialize_session(&session_args, &config_snapshot)
				.await
				.map_err(|e| agent_client_protocol::Error::internal_error().data(e.to_string()))?;

		setup_system_prompt_and_cache(&mut chat_session, &config_for_role, &session_role, false)
			.await
			.map_err(|e| agent_client_protocol::Error::internal_error().data(e.to_string()))?;

		self.sessions
			.borrow_mut()
			.insert(session_id.clone(), (chat_session, session_cwd));
		self.cancellations
			.borrow_mut()
			.insert(session_id.clone(), SessionCancellation::new());

		let conn = self.conn.borrow().clone();
		send_available_commands(conn, &session_id).await;

		Ok(LoadSessionResponse::new())
	}

	async fn ext_method(&self, args: ExtRequest) -> agent_client_protocol::Result<ExtResponse> {
		super::commands::handle_ext_method(
			args,
			&self.sessions,
			&self.config,
			&self.role,
			&self.cancellations,
		)
		.await
	}
}