|
| 1 | +"""Tool for dynamically updating the active workspace directory.""" |
| 2 | + |
| 3 | +import logging |
| 4 | +import os |
| 5 | +from google.adk.tools import FunctionTool, ToolContext |
| 6 | +from google.adk.tools.mcp_tool.mcp_session_manager import MCPSessionManager |
| 7 | + |
| 8 | +import hitl_agent.config as config |
| 9 | +from hitl_agent.tools.filesystem_tools import filesystem_tool_r, filesystem_tool_rw |
| 10 | + |
| 11 | + |
| 12 | +async def set_working_directory_fn(path: str, tool_context: ToolContext) -> str: |
| 13 | + """Set the active workspace/working directory path for the session. |
| 14 | +
|
| 15 | + All filesystem tools (list_directory, read_file, write_file) will use |
| 16 | + the new path. |
| 17 | +
|
| 18 | + Args: |
| 19 | + path: The absolute path of the directory. |
| 20 | +
|
| 21 | + Returns: |
| 22 | + A string containing the operation result description: |
| 23 | + - If successful: "Successfully switched working directory to: <abs_path>" |
| 24 | + - If failed: "Error: Path '<path>' does not exist." or "Error: Path '<path>' is not a directory." |
| 25 | + """ |
| 26 | + if not os.path.exists(path): |
| 27 | + return f"Error: Path '{path}' does not exist." |
| 28 | + if not os.path.isdir(path): |
| 29 | + return f"Error: Path '{path}' is not a directory." |
| 30 | + |
| 31 | + abs_path = os.path.abspath(path) |
| 32 | + |
| 33 | + # Update in-memory config variables |
| 34 | + config.WORKDIR = abs_path |
| 35 | + os.environ["WORKDIR"] = abs_path |
| 36 | + |
| 37 | + # Update .env file |
| 38 | + try: |
| 39 | + env_path = ".env" |
| 40 | + if os.path.exists(env_path): |
| 41 | + with open(env_path, "r") as f: |
| 42 | + lines = f.readlines() |
| 43 | + with open(env_path, "w") as f: |
| 44 | + updated = False |
| 45 | + for line in lines: |
| 46 | + if line.strip().startswith("WORKDIR="): |
| 47 | + f.write(f'WORKDIR="{abs_path}"\n') |
| 48 | + updated = True |
| 49 | + else: |
| 50 | + f.write(line) |
| 51 | + if not updated: |
| 52 | + f.write(f'WORKDIR="{abs_path}"\n') |
| 53 | + except Exception as e: |
| 54 | + logging.error(f"Failed to update .env: {e}") |
| 55 | + |
| 56 | + # Update active filesystem tools |
| 57 | + # Update read-only filesystem tool |
| 58 | + params_r = filesystem_tool_r._connection_params |
| 59 | + server_params_r = getattr(params_r, "server_params", None) or params_r |
| 60 | + if hasattr(server_params_r, "args") and len(server_params_r.args) >= 2: |
| 61 | + server_params_r.args[-1] = abs_path |
| 62 | + await filesystem_tool_r.close() |
| 63 | + filesystem_tool_r._mcp_session_manager = MCPSessionManager( |
| 64 | + connection_params=filesystem_tool_r._connection_params, |
| 65 | + errlog=filesystem_tool_r._errlog, |
| 66 | + sampling_callback=filesystem_tool_r._sampling_callback, |
| 67 | + sampling_capabilities=filesystem_tool_r._sampling_capabilities, |
| 68 | + ) |
| 69 | + |
| 70 | + # Update read-write filesystem tool |
| 71 | + params_rw = filesystem_tool_rw._connection_params |
| 72 | + server_params_rw = getattr(params_rw, "server_params", None) or params_rw |
| 73 | + if hasattr(server_params_rw, "args") and len(server_params_rw.args) >= 2: |
| 74 | + server_params_rw.args[-1] = abs_path |
| 75 | + await filesystem_tool_rw.close() |
| 76 | + filesystem_tool_rw._mcp_session_manager = MCPSessionManager( |
| 77 | + connection_params=filesystem_tool_rw._connection_params, |
| 78 | + errlog=filesystem_tool_rw._errlog, |
| 79 | + sampling_callback=filesystem_tool_rw._sampling_callback, |
| 80 | + sampling_capabilities=filesystem_tool_rw._sampling_capabilities, |
| 81 | + ) |
| 82 | + |
| 83 | + return f"Successfully switched working directory to: {abs_path}" |
| 84 | + |
| 85 | + |
| 86 | +set_working_directory = FunctionTool(set_working_directory_fn) |
0 commit comments