[Enh] Set the session cwd to use server root_dir instead of chat directory

jupyter_ai_acp_client/base_acp_persona.py

@@ -42,7 +42,7 @@ class BaseAcpPersona(BasePersona):
     Developers should always use `self.get_client()`.
     """
 
-    _client_session_future: Task[NewSessionResponse | LoadSessionResponse]
+    _client_session_future: Task[NewSessionResponse | LoadSessionResponse] | None
     """
     The future that yields the ACP client session info. Each instance of an ACP
     persona has a unique session ID, i.e. each chat reserves a unique session.
@@ -93,9 +93,11 @@ def __init__(self, *args, executable: list[str], **kwargs):
                 self._init_client()
             )
 
-        self._client_session_future = self.event_loop.create_task(
-            self._init_client_session()
-        )
+        # Session creation is deferred to the first process_message call
+        # (via _ensure_client_session) rather than started here eagerly.
+        # This gives the frontend time to sync YChat metadata — including
+        # the file browser cwd — over Y.js before the session reads it.
+        self._client_session_future = None
         self._acp_slash_commands = []
 
     async def before_agent_subprocess(self) -> None:
@@ -216,17 +218,30 @@ async def get_client(self) -> JaiAcpClient:
         """
         return await self.__class__._client_future
 
+    async def _ensure_client_session(self):
+        """Lazily create the ACP session on first access.
+
+        Session creation is deferred from __init__ so that YChat metadata
+        (e.g. the file browser cwd set by the frontend) has time to sync
+        over Y.js before the session's working directory is resolved.
+        """
+        if self._client_session_future is None:
+            self._client_session_future = self.event_loop.create_task(
+                self._init_client_session()
+            )
+        return await self._client_session_future
+
     async def get_session_response(self) -> NewSessionResponse | LoadSessionResponse:
         """
         Safely returns the ACP session response for this chat.
         """
-        return await self._client_session_future
+        return await self._ensure_client_session()
 
     async def get_session_id(self) -> str:
         """
         Safely returns the ACP session ID assigned to this chat.
         """
-        await self._client_session_future
+        await self._ensure_client_session()
         # session ID should always be stored in chat metadata after client
         # session was created or loaded.
         session_ids = self._get_existing_sessions()

jupyter_ai_acp_client/default_acp_client.py

@@ -1,5 +1,6 @@
 import asyncio
 import logging
+import os
 from pathlib import Path
 from typing import Any, Awaitable
 from time import time
@@ -163,6 +164,13 @@ async def _get_mcp_servers(self, persona: BasePersona) -> list[AcpMcpServerStdio
 
         return mcp_servers
 
+    def _resolve_cwd(self, persona: BasePersona) -> str:
+        root_dir = persona.parent.root_dir
+        chat_cwd = persona.ychat.get_metadata().get("cwd")
+        if root_dir and chat_cwd and chat_cwd != "/":
+            return os.path.join(root_dir, chat_cwd)
+        return root_dir or persona.get_chat_dir()
+
     async def create_session(self, persona: BasePersona) -> NewSessionResponse:
         """
         Create an ACP agent session through this client scoped to a
@@ -171,7 +179,8 @@ async def create_session(self, persona: BasePersona) -> NewSessionResponse:
         """
         conn = await self.get_connection()
         mcp_servers = await self._get_mcp_servers(persona)
-        session = await conn.new_session(cwd=persona.get_chat_dir(), mcp_servers=mcp_servers)
+        cwd = self._resolve_cwd(persona)
+        session = await conn.new_session(cwd=cwd, mcp_servers=mcp_servers)
         self._personas_by_session[session.session_id] = persona
         return session
 
@@ -211,7 +220,8 @@ async def _load_session_rpc(self, persona: BasePersona, session_id: str) -> Load
         """
         conn = await self.get_connection()
         mcp_servers = await self._get_mcp_servers(persona)
-        response = await conn.load_session(cwd=persona.get_chat_dir(), session_id=session_id, mcp_servers=mcp_servers)
+        cwd = self._resolve_cwd(persona)
+        response = await conn.load_session(cwd=cwd, session_id=session_id, mcp_servers=mcp_servers)
         self._personas_by_session[session_id] = persona
         return response
 

jupyter_ai_acp_client/tests/test_default_acp_client.py

@@ -239,3 +239,127 @@ async def _failing_rpc(*args, **kwargs):
             await client.load_session(persona, "stale-session-id")
 
         assert "stale-session-id" not in client._loading_sessions
+
+
+class TestSessionCwd:
+    """Tests for session cwd resolution in create_session and _load_session_rpc.
+
+    The session cwd determines the working directory for the ACP agent
+    subprocess (e.g. where Claude Code starts). It's resolved by
+    ``JaiAcpClient._resolve_cwd()`` using this priority:
+
+    1. ``chat_cwd`` from YChat metadata joined with ``root_dir`` — the file
+       browser path at the time the chat was opened. Set by the frontend in
+       jupyterlab-chat-extension (see jupyterlab/jupyter-chat#420).
+    2. ``root_dir`` alone — the Jupyter server's root directory, used when
+       the chat has no cwd metadata (e.g. older chats created before this
+       feature).
+    3. ``get_chat_dir()`` — the parent directory of the chat file, used as
+       a last resort when root_dir is also unavailable.
+
+    Each test mocks ``persona.ychat.get_metadata()`` to simulate different
+    chat metadata states, and ``persona.parent.root_dir`` to simulate the
+    server's root directory. The assertion checks what ``cwd`` kwarg was
+    passed to the ACP ``new_session`` or ``load_session`` RPC call.
+    """
+
+    async def test_create_session_uses_chat_metadata_cwd(self):
+        """When the chat has cwd metadata, join it with root_dir.
+
+        Simulates a user who opened the chat while browsing projects/my-repo
+        in the file browser. The session should start in that subdirectory.
+        """
+        client, conn, _ = _make_client_and_persona()
+        client._get_mcp_servers = AsyncMock(return_value=[])
+
+        persona = MagicMock()
+        persona.parent.root_dir = "/home/user/notebooks"
+        persona.ychat.get_metadata.return_value = {"cwd": "projects/my-repo"}
+        persona.get_chat_dir.return_value = "/home/user/notebooks/.jupyter/chats"
+
+        conn.new_session = AsyncMock(return_value=MagicMock(session_id="s1"))
+
+        await client.create_session(persona)
+
+        conn.new_session.assert_called_once()
+        assert conn.new_session.call_args.kwargs["cwd"] == "/home/user/notebooks/projects/my-repo"
+
+    async def test_create_session_uses_root_dir_when_no_chat_cwd(self):
+        """When the chat has no cwd metadata, fall back to root_dir.
+
+        This covers older chats created before the frontend started writing
+        cwd metadata, or chats created without a file browser (e.g. via API).
+        """
+        client, conn, _ = _make_client_and_persona()
+        client._get_mcp_servers = AsyncMock(return_value=[])
+
+        persona = MagicMock()
+        persona.parent.root_dir = "/home/user/notebooks"
+        persona.ychat.get_metadata.return_value = {}
+        persona.get_chat_dir.return_value = "/home/user/notebooks/.jupyter/chats"
+
+        conn.new_session = AsyncMock(return_value=MagicMock(session_id="s1"))
+
+        await client.create_session(persona)
+
+        conn.new_session.assert_called_once()
+        assert conn.new_session.call_args.kwargs["cwd"] == "/home/user/notebooks"
+
+    async def test_create_session_uses_root_dir_when_chat_cwd_is_root(self):
+        """When the chat cwd is "/", treat it as the file browser root.
+
+        The JupyterLab file browser reports "/" when at the top level.
+        os.path.join(root_dir, "/") would incorrectly return "/" (the
+        filesystem root), so we skip the join and use root_dir directly.
+        """
+        client, conn, _ = _make_client_and_persona()
+        client._get_mcp_servers = AsyncMock(return_value=[])
+
+        persona = MagicMock()
+        persona.parent.root_dir = "/home/user/notebooks"
+        persona.ychat.get_metadata.return_value = {"cwd": "/"}
+        persona.get_chat_dir.return_value = "/home/user/notebooks/.jupyter/chats"
+
+        conn.new_session = AsyncMock(return_value=MagicMock(session_id="s1"))
+
+        await client.create_session(persona)
+
+        conn.new_session.assert_called_once()
+        assert conn.new_session.call_args.kwargs["cwd"] == "/home/user/notebooks"
+
+    async def test_create_session_falls_back_to_chat_dir(self):
+        """When root_dir is unavailable, fall back to get_chat_dir().
+
+        This is the last resort — the directory containing the chat file.
+        """
+        client, conn, _ = _make_client_and_persona()
+        client._get_mcp_servers = AsyncMock(return_value=[])
+
+        persona = MagicMock()
+        persona.parent.root_dir = None
+        persona.ychat.get_metadata.return_value = {}
+        persona.get_chat_dir.return_value = "/home/user/.jupyter/chats"
+
+        conn.new_session = AsyncMock(return_value=MagicMock(session_id="s1"))
+
+        await client.create_session(persona)
+
+        conn.new_session.assert_called_once()
+        assert conn.new_session.call_args.kwargs["cwd"] == "/home/user/.jupyter/chats"
+
+    async def test_load_session_uses_chat_metadata_cwd(self):
+        """load_session resolves cwd the same way as create_session."""
+        client, conn, _ = _make_client_and_persona()
+        client._get_mcp_servers = AsyncMock(return_value=[])
+
+        persona = MagicMock()
+        persona.parent.root_dir = "/home/user/notebooks"
+        persona.ychat.get_metadata.return_value = {"cwd": "projects/my-repo"}
+        persona.get_chat_dir.return_value = "/home/user/notebooks/.jupyter/chats"
+
+        conn.load_session = AsyncMock(return_value=MagicMock(session_id="s1"))
+
+        await client._load_session_rpc(persona, "s1")
+
+        conn.load_session.assert_called_once()
+        assert conn.load_session.call_args.kwargs["cwd"] == "/home/user/notebooks/projects/my-repo"