Fix remote session context working directory

Author: 1jehuang

src/agent/provider.rs

@@ -120,6 +120,7 @@ impl Agent {
             return;
         }
         self.session.working_dir = Some(dir.to_string());
+        self.session.refresh_initial_session_context_message();
         self.log_env_snapshot("working_dir");
     }
 

src/session.rs

@@ -808,6 +808,44 @@ impl Session {
         true
     }
 
+    /// Refresh the initial immutable session-context message if the session has
+    /// not started a real conversation yet. This covers remote/client-server
+    /// startup where the server creates an Agent before the subscribing client
+    /// sends the terminal working directory that tools will use.
+    pub fn refresh_initial_session_context_message(&mut self) -> bool {
+        if self.messages.iter().any(is_visible_conversation_message) {
+            return false;
+        }
+
+        let Some(message) = self.messages.iter_mut().find(|message| {
+            message.content.iter().any(|block| match block {
+                ContentBlock::Text { text, .. } => text.starts_with(SESSION_CONTEXT_PREFIX),
+                _ => false,
+            })
+        }) else {
+            return false;
+        };
+
+        let context =
+            crate::prompt::build_session_context(self.working_dir.as_deref().map(Path::new));
+        let wrapped = format!("<system-reminder>\n{}\n</system-reminder>", context.trim());
+        for block in &mut message.content {
+            if let ContentBlock::Text { text, .. } = block
+                && text.starts_with(SESSION_CONTEXT_PREFIX)
+            {
+                if *text == wrapped {
+                    return false;
+                }
+                *text = wrapped;
+                self.mark_memory_profile_dirty();
+                self.mark_messages_full_dirty();
+                return true;
+            }
+        }
+
+        false
+    }
+
     /// Get the display name for this session (short memorable name if available)
     pub fn display_name(&self) -> &str {
         self.short_name

src/session_tests/cases.rs

@@ -167,6 +167,102 @@ fn initial_session_context_uses_current_cwd_when_inserted() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn initial_session_context_can_refresh_before_real_conversation() -> Result<()> {
+    let _env_lock = lock_env();
+    let original_cwd = std::env::current_dir().map_err(|e| anyhow!(e))?;
+    let first_dir = tempfile::Builder::new()
+        .prefix("jcode-session-context-stale-")
+        .tempdir()
+        .map_err(|e| anyhow!(e))?;
+    let second_dir = tempfile::Builder::new()
+        .prefix("jcode-session-context-real-")
+        .tempdir()
+        .map_err(|e| anyhow!(e))?;
+
+    std::env::set_current_dir(first_dir.path()).map_err(|e| anyhow!(e))?;
+    let result: std::result::Result<(), anyhow::Error> = (|| {
+        let mut session = Session::create_with_id(
+            "session_context_remote_cwd_refresh_test".to_string(),
+            None,
+            Some("Remote cwd refresh".to_string()),
+        );
+        assert!(session.ensure_initial_session_context_message());
+        assert!(session.messages[0].content_preview().contains(&format!(
+            "Working directory: {}",
+            first_dir.path().display()
+        )));
+
+        session.working_dir = Some(second_dir.path().display().to_string());
+        assert!(session.refresh_initial_session_context_message());
+        let refreshed = session.messages[0].content_preview();
+        assert!(
+            refreshed.contains(&format!(
+                "Working directory: {}",
+                second_dir.path().display()
+            )),
+            "session context should refresh to subscribed cwd, got: {refreshed}"
+        );
+        assert!(!refreshed.contains(&format!(
+            "Working directory: {}",
+            first_dir.path().display()
+        )));
+        Ok(())
+    })();
+    std::env::set_current_dir(original_cwd).map_err(|e| anyhow!(e))?;
+    result?;
+
+    Ok(())
+}
+
+#[test]
+fn initial_session_context_does_not_refresh_after_real_conversation() -> Result<()> {
+    let _env_lock = lock_env();
+    let original_cwd = std::env::current_dir().map_err(|e| anyhow!(e))?;
+    let first_dir = tempfile::Builder::new()
+        .prefix("jcode-session-context-original-")
+        .tempdir()
+        .map_err(|e| anyhow!(e))?;
+    let second_dir = tempfile::Builder::new()
+        .prefix("jcode-session-context-late-")
+        .tempdir()
+        .map_err(|e| anyhow!(e))?;
+
+    std::env::set_current_dir(first_dir.path()).map_err(|e| anyhow!(e))?;
+    let result: std::result::Result<(), anyhow::Error> = (|| {
+        let mut session = Session::create_with_id(
+            "session_context_late_cwd_refresh_test".to_string(),
+            None,
+            Some("Late cwd refresh".to_string()),
+        );
+        assert!(session.ensure_initial_session_context_message());
+        session.add_message(
+            Role::User,
+            vec![ContentBlock::Text {
+                text: "hello".to_string(),
+                cache_control: None,
+            }],
+        );
+
+        session.working_dir = Some(second_dir.path().display().to_string());
+        assert!(!session.refresh_initial_session_context_message());
+        let original = session.messages[0].content_preview();
+        assert!(original.contains(&format!(
+            "Working directory: {}",
+            first_dir.path().display()
+        )));
+        assert!(!original.contains(&format!(
+            "Working directory: {}",
+            second_dir.path().display()
+        )));
+        Ok(())
+    })();
+    std::env::set_current_dir(original_cwd).map_err(|e| anyhow!(e))?;
+    result?;
+
+    Ok(())
+}
+
 #[test]
 fn existing_non_empty_session_does_not_get_retroactive_session_context() {
     let mut session = Session::create_with_id(