refactor: improve session socket path management

crates/coco-tui/src/combo_run_server.rs

@@ -1,8 +1,8 @@
 use std::path::{Path, PathBuf};
 
 use code_combo::{
-    ClientMessage, ComboRunMessage, ComboRunPayload, ComboRunResult, ServerConnection,
-    ServerMessage, SessionEnv, SessionSocketServer,
+    ClientMessage, ComboRunMessage, ComboRunPayload, ComboRunResult, SESSION_SOCKET_ENV,
+    ServerConnection, ServerMessage, SessionEnv, SessionSocketServer,
 };
 use snafu::prelude::*;
 use tokio::{net::UnixStream, sync::mpsc, task::JoinHandle};
@@ -40,10 +40,7 @@ impl ComboRunSessionServer {
         let server = SessionSocketServer::bind(&socket_path)
             .await
             .whatever_context("failed to bind session socket")?;
-        // Safety: set once during startup before spawning any tasks.
-        unsafe {
-            std::env::set_var("COCO_SESSION_SOCK", &socket_path);
-        }
+        code_combo::global::set_session_socket_path(socket_path.clone());
 
         let shutdown = CancellationToken::new();
         let shutdown_task = shutdown.clone();
@@ -83,12 +80,13 @@ impl ComboRunSessionServer {
     pub async fn shutdown(self) {
         self.shutdown.cancel();
         let _ = self.join_handle.await;
+        code_combo::global::clear_session_socket_path();
         let _ = std::fs::remove_file(&self.socket_path);
     }
 }
 
 async fn resolve_socket_path() -> Result<(PathBuf, Option<SessionEnv>)> {
-    if let Ok(path) = std::env::var("COCO_SESSION_SOCK") {
+    if let Ok(path) = std::env::var(SESSION_SOCKET_ENV) {
         let socket_path = PathBuf::from(path);
         let parent_exists = socket_path
             .parent()
@@ -99,7 +97,8 @@ async fn resolve_socket_path() -> Result<(PathBuf, Option<SessionEnv>)> {
         }
         warn!(
             socket_path = %socket_path.display(),
-            "COCO_SESSION_SOCK parent dir missing; creating new session env"
+            env = SESSION_SOCKET_ENV,
+            "session socket env parent dir missing; creating new session env"
         );
     }
     let env = SessionEnv::builder()

crates/coco-tui/src/components/chat.rs

@@ -8,8 +8,9 @@ use code_combo::{
     Agent, Block as ChatBlock, ChatResponse, ChatStreamUpdate, ComboRunEvent, ComboRunResult,
     ComboStreamKind as ComboRunStreamKind, Config, Content as ChatContent, Message as ChatMessage,
     NotificationRule, NotificationWhen, Output, PromptPayload, RetryAttempt, RetryNotifier,
-    RetryUpdate, RuntimeOverrides, SessionSocketClient, Starter, StopReason, TextEdit, ToolUse,
-    UsageStats, discover_starters, load_runtime_overrides, save_runtime_overrides,
+    RetryUpdate, RuntimeOverrides, SESSION_SOCKET_ENV, SessionSocketClient, Starter, StopReason,
+    TextEdit, ToolUse, UsageStats, discover_starters, load_runtime_overrides,
+    save_runtime_overrides,
 };
 
 /// Holds information about a collected tool result from concurrent execution.
@@ -2161,7 +2162,7 @@ impl Chat<'static> {
         }
         if let Some(session_sock) = bash_input
             .env
-            .get("COCO_SESSION_SOCK")
+            .get(SESSION_SOCKET_ENV)
             .filter(|session_sock| !session_sock.is_empty())
         {
             let combo_id = self.pending_combo_id_for_session_sock(session_sock)?;
@@ -3445,7 +3446,7 @@ fn session_sock_from_tool_use(tool_use: &ToolUse) -> Option<String> {
     let input = serde_json::from_value::<BashInput>(tool_use.input.clone()).ok()?;
     input
         .env
-        .get("COCO_SESSION_SOCK")
+        .get(SESSION_SOCKET_ENV)
         .filter(|value| !value.is_empty())
         .cloned()
 }

crates/coco-tui/src/main.rs

@@ -5,7 +5,7 @@ use std::{
 
 use clap::{CommandFactory, FromArgMatches, Parser, Subcommand};
 use code_combo::{
-    RuntimeOverrides,
+    RuntimeOverrides, SESSION_SOCKET_ENV,
     cli::{ClientCommand, handle_client_command, init_client_logging},
     default_config_dir, load_config_with_overrides, load_runtime_overrides, workspace_config_path,
 };
@@ -270,7 +270,7 @@ async fn main() -> Result<()> {
 }
 
 fn has_session_socket() -> bool {
-    let Some(path) = std::env::var_os("COCO_SESSION_SOCK") else {
+    let Some(path) = std::env::var_os(SESSION_SOCKET_ENV) else {
         return false;
     };
     let path = PathBuf::from(path);

src/agent/executor.rs

@@ -561,10 +561,16 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn bash_execution_receives_session_socket_env() {
+    async fn bash_execution_ignores_session_socket_env_injection() {
         let mut executor = Executor::default();
-        executor.set_bash_env("COCO_SESSION_SOCK", "/tmp/coco-executor.sock");
-        let input_value = bash_input_value(r#"printf "%s" "$COCO_SESSION_SOCK""#);
+        executor.set_bash_env(crate::SESSION_SOCKET_ENV, "/tmp/coco-executor.sock");
+        let input_value = bash_input_value(
+            format!(
+                r#"if [ -n "${}" ]; then printf "set\n"; else printf "empty\n"; fi"#,
+                crate::SESSION_SOCKET_ENV
+            )
+            .as_str(),
+        );
         executor.update_pcl(
             BASH_TOOL_NAME,
             PermissionControl::Once("inject-test".to_string()),
@@ -579,7 +585,7 @@ mod tests {
         };
         let output: BashOutput =
             serde_json::from_value(value).expect("deserialize bash output json");
-        assert_eq!(output.stdout, "/tmp/coco-executor.sock\n");
+        assert_eq!(output.stdout, "empty\n");
     }
 
     #[test]

src/cmd/combo.rs

@@ -5,7 +5,8 @@ use tokio::{process::Command, time::Instant};
 use tracing::{debug, info, warn};
 
 use crate::{
-    ComboRunPayload, ComboRunResult, RunComboOutput, SessionEnv, SessionSocketClient, error::Result,
+    ComboRunPayload, ComboRunResult, RunComboOutput, SESSION_SOCKET_ENV, SessionEnv,
+    SessionSocketClient, error::Result,
 };
 
 pub async fn handle_combo_run(
@@ -22,7 +23,7 @@ pub async fn handle_combo_run(
 
     let client = SessionSocketClient::from_env()
         .await
-        .whatever_context("failed to read COCO_SESSION_SOCK")?;
+        .whatever_context(format!("failed to read {SESSION_SOCKET_ENV}"))?;
 
     match client {
         Some(client) => {

src/combo.rs

@@ -30,7 +30,7 @@ pub use session::{
     ServerMessage, SessionClientError, SessionServerError, SessionSocketClient,
     SessionSocketServer, ThinkingConfig,
 };
-pub use session_env::{SessionEnv, SessionEnvBuilder, SessionEnvError};
+pub use session_env::{SESSION_SOCKET_ENV, SessionEnv, SessionEnvBuilder, SessionEnvError};
 pub use starter::PromptResponseSender;
 pub use starter::{Starter, StarterCommand, StarterError, StarterEvent, StarterExecution};
 pub use types::{

src/combo/runner.rs

@@ -31,9 +31,9 @@ pub type ExecuteResult = Result<Output, Final>;
 
 use crate::{
     Agent, Block, ChatStreamUpdate, Config, Content, Message, OutputChunk, PromptResponseSender,
-    PromptSchema, SessionEnv, Starter, StarterCommand, StarterError, StarterEvent, ToolUse,
-    bash_unsafe_ranges, bash_unsafe_reason, discover_starters, exec::StreamKind,
-    parse_primary_command, workspace_dir,
+    PromptSchema, SESSION_SOCKET_ENV, SessionEnv, Starter, StarterCommand, StarterError,
+    StarterEvent, ToolUse, bash_unsafe_ranges, bash_unsafe_reason, discover_starters,
+    exec::StreamKind, parse_primary_command, workspace_dir,
 };
 
 // Import types from combo module
@@ -45,6 +45,28 @@ pub const RUN_COMBO_TOOL_NAME: &str = "run_combo";
 
 type ComboEventCallback = Arc<std::sync::Mutex<Box<dyn FnMut(&ComboEvent) + Send>>>;
 
+struct SessionSocketPathGuard {
+    previous: Option<PathBuf>,
+}
+
+impl SessionSocketPathGuard {
+    fn install(path: PathBuf) -> Self {
+        let previous = crate::global::session_socket_path();
+        crate::global::set_session_socket_path(path);
+        Self { previous }
+    }
+}
+
+impl Drop for SessionSocketPathGuard {
+    fn drop(&mut self) {
+        if let Some(path) = self.previous.take() {
+            crate::global::set_session_socket_path(path);
+        } else {
+            crate::global::clear_session_socket_path();
+        }
+    }
+}
+
 /// Execute combo logic with event callback support.
 pub async fn run_combo<F>(
     context: Arc<Mutex<RunComboContext>>,
@@ -216,6 +238,7 @@ async fn execute_combo(
         .build()
         .map_err(|e| format!("Failed to create session env: {}", e))?;
     let session_socket_path = session_env.socket_path().to_path_buf();
+    let _session_socket_guard = SessionSocketPathGuard::install(session_socket_path.clone());
     let mut exec_envs = match prepare_mcp_envs().await {
         Ok(envs) => envs,
         Err(err) => {
@@ -428,7 +451,6 @@ async fn execute_combo(
                                 &schemas,
                                 &combo_name,
                                 cancel_token.clone(),
-                                session_socket_path.clone(),
                                 &on_event_for_emit,
                             )
                             .await;
@@ -1274,7 +1296,7 @@ async fn handle_interactive_combo_reply(
                 if matches!(command_kind, OffloadCommandKind::Coco) {
                     let mut bash_input_for_emit = bash_input.clone();
                     bash_input_for_emit.env.insert(
-                        "COCO_SESSION_SOCK".to_string(),
+                        SESSION_SOCKET_ENV.to_string(),
                         session_socket_path.to_string_lossy().to_string(),
                     );
                     emitted_tool_use.input =
@@ -1321,13 +1343,9 @@ async fn handle_interactive_combo_reply(
                                     message: err.to_string(),
                                 }
                             })?;
-                        let extra_envs = vec![(
-                            "COCO_SESSION_SOCK".to_string(),
-                            session_socket_path.to_string_lossy().to_string(),
-                        )];
                         let output = run_bash_chunked(
                             Input::Starter(bash_input_value),
-                            &extra_envs,
+                            &[],
                             cancel_token.clone(),
                             |_| {},
                         )
@@ -1615,7 +1633,6 @@ async fn handle_offload_combo_reply_with_retry(
     schemas: &[PromptSchema],
     combo_name: &str,
     cancel_token: CancellationToken,
-    session_socket_path: PathBuf,
     on_event: &ComboEventCallback,
 ) -> Result<(), ComboReplyError> {
     let max_retries = agent.combo_reply_retries();
@@ -1634,7 +1651,6 @@ async fn handle_offload_combo_reply_with_retry(
             schemas,
             combo_name,
             cancel_token.clone(),
-            session_socket_path.clone(),
             on_event,
             &directive,
         )
@@ -1656,7 +1672,6 @@ async fn handle_offload_combo_reply(
     schemas: &[PromptSchema],
     combo_name: &str,
     cancel_token: CancellationToken,
-    session_socket_path: PathBuf,
     on_event: &ComboEventCallback,
     directive: &str,
 ) -> Result<(), ComboReplyError> {
@@ -1787,13 +1802,9 @@ async fn handle_offload_combo_reply(
             message: err.to_string(),
         })?;
 
-    let extra_envs = vec![(
-        "COCO_SESSION_SOCK".to_string(),
-        session_socket_path.to_string_lossy().to_string(),
-    )];
     let output = run_bash_chunked(
         Input::Starter(bash_input_value),
-        &extra_envs,
+        &[],
         cancel_token.clone(),
         |_| {},
     )
@@ -1995,15 +2006,15 @@ mod tests {
         assert!(is_coco_reply_command(
             "/usr/local/bin/coco reply --message='hello world'"
         ));
-        assert!(is_coco_reply_command(
-            "COCO_SESSION_SOCK='/tmp/coco.sock' coco reply --message='hello world'"
-        ));
+        assert!(is_coco_reply_command(&format!(
+            "{SESSION_SOCKET_ENV}='/tmp/coco.sock' coco reply --message='hello world'"
+        )));
         assert!(is_coco_reply_command(
             "/run/current-system/sw/bin/zsh -lc \"coco reply --message='hello world'\""
         ));
-        assert!(is_coco_reply_command(
-            "bash -lc \"COCO_SESSION_SOCK=/tmp/coco.sock coco reply --message='hello world'\""
-        ));
+        assert!(is_coco_reply_command(&format!(
+            "bash -lc \"{SESSION_SOCKET_ENV}=/tmp/coco.sock coco reply --message='hello world'\""
+        )));
         assert!(!is_coco_reply_command("coco ask hello"));
         assert!(!is_coco_reply_command("echo coco reply --message=hello"));
         assert!(!is_coco_reply_command(

src/combo/session.rs

@@ -10,8 +10,8 @@ use tokio::{
 use tracing::{debug, warn};
 
 use crate::{
-    MCP_SOCKET_ENV, McpRequest, McpResponse, Message, OutputChunk, Starter, ToolUse,
-    exec::StreamKind, tools::Final,
+    MCP_SOCKET_ENV, McpRequest, McpResponse, Message, OutputChunk, SESSION_SOCKET_ENV, Starter,
+    ToolUse, exec::StreamKind, tools::Final,
 };
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
@@ -352,7 +352,7 @@ impl SessionSocketClient {
     }
 
     pub async fn from_env() -> ClientResult<Option<Self>> {
-        Self::from_env_key("COCO_SESSION_SOCK").await
+        Self::from_env_key(SESSION_SOCKET_ENV).await
     }
 
     pub async fn from_mcp_env() -> ClientResult<Option<Self>> {
@@ -367,9 +367,9 @@ impl SessionSocketClient {
         use snafu::prelude::*;
         let Some(client) = Self::from_env()
             .await
-            .whatever_context("failed to connect to COCO_SESSION_SOCK")?
+            .whatever_context(format!("failed to connect to {SESSION_SOCKET_ENV}"))?
         else {
-            whatever!("env COCO_SESSION_SOCK is not set");
+            whatever!("env {SESSION_SOCKET_ENV} is not set");
         };
         Ok(client)
     }

src/combo/session_env.rs

@@ -7,6 +7,8 @@ use std::{
 use snafu::prelude::*;
 use tempfile::{Builder, TempDir};
 
+pub const SESSION_SOCKET_ENV: &str = "COCO_SESSION_SOCK";
+
 #[derive(Debug, Snafu)]
 pub enum SessionEnvError {
     #[snafu(display("failed to locate coco executable"))]
@@ -57,7 +59,7 @@ impl Default for SessionEnvBuilder {
             binary_path: None,
             command_name: "coco".to_string(),
             socket_name: "coco.sock".to_string(),
-            socket_env_name: "COCO_SESSION_SOCK".to_string(),
+            socket_env_name: SESSION_SOCKET_ENV.to_string(),
         }
     }
 }
@@ -208,4 +210,9 @@ mod tests {
         assert_eq!(socket_path, env.socket_path().as_os_str());
         Ok(())
     }
+
+    #[test]
+    fn session_socket_env_differs_from_mcp_socket_env() {
+        assert_ne!(SESSION_SOCKET_ENV, crate::MCP_SOCKET_ENV);
+    }
 }

src/global.rs

@@ -1,10 +1,14 @@
-use std::sync::{Arc, OnceLock};
+use std::{
+    path::PathBuf,
+    sync::{Arc, OnceLock, RwLock},
+};
 
 use tokio::sync::Mutex;
 
 use crate::Config;
 
 static CONFIG: OnceLock<Arc<Mutex<Option<Config>>>> = OnceLock::new();
+static SESSION_SOCKET_PATH: OnceLock<RwLock<Option<PathBuf>>> = OnceLock::new();
 
 pub async fn set_config(config: Config) {
     let cell = CONFIG.get_or_init(|| Arc::new(Mutex::new(None)));
@@ -17,3 +21,25 @@ pub async fn config() -> Option<Config> {
     let guard = cell.lock().await;
     guard.as_ref().cloned()
 }
+
+pub fn set_session_socket_path(path: impl Into<PathBuf>) {
+    let lock = SESSION_SOCKET_PATH.get_or_init(|| RwLock::new(None));
+    let mut guard = lock
+        .write()
+        .unwrap_or_else(|poisoned| poisoned.into_inner());
+    *guard = Some(path.into());
+}
+
+pub fn clear_session_socket_path() {
+    let lock = SESSION_SOCKET_PATH.get_or_init(|| RwLock::new(None));
+    let mut guard = lock
+        .write()
+        .unwrap_or_else(|poisoned| poisoned.into_inner());
+    *guard = None;
+}
+
+pub fn session_socket_path() -> Option<PathBuf> {
+    let lock = SESSION_SOCKET_PATH.get_or_init(|| RwLock::new(None));
+    let guard = lock.read().unwrap_or_else(|poisoned| poisoned.into_inner());
+    guard.clone()
+}