diff --git a/src/harbor/agents/installed/bitfun_cli.py b/src/harbor/agents/installed/bitfun_cli.py index 5acc949942e..c5417663c9f 100644 --- a/src/harbor/agents/installed/bitfun_cli.py +++ b/src/harbor/agents/installed/bitfun_cli.py @@ -53,6 +53,7 @@ _REMOTE_CP_BACK_MANIFEST_PATH = "/logs/agent/bitfun/cp-back-manifest.json" _WINDOWS_PROMPT_FILE_NAME = "bitfun-prompt.txt" _WINDOWS_RUN_SCRIPT_NAME = "bitfun-run.bat" +_WINDOWS_CP_BACK_SCRIPT_NAME = "bitfun-cp-back.bat" _WINDOWS_BITFUN_USER_ROOT = "C:/bitfun-user" _WINDOWS_BITFUN_HOME = "C:/bitfun-home" _REDACTED_CONFIG_VALUE = "[REDACTED]" @@ -302,6 +303,10 @@ def _prompt_path(cls, environment: BaseEnvironment) -> str: def _run_script_path(cls, environment: BaseEnvironment) -> str: return str(cls._env_paths(environment).agent_dir / _WINDOWS_RUN_SCRIPT_NAME) + @classmethod + def _cp_back_script_path(cls, environment: BaseEnvironment) -> str: + return str(cls._env_paths(environment).agent_dir / _WINDOWS_CP_BACK_SCRIPT_NAME) + @classmethod def _patch_logs_dir_in_env(cls, environment: BaseEnvironment) -> str: return str(cls._env_paths(environment).agent_dir / PATCH_ARTIFACTS_SUBDIR) @@ -2265,6 +2270,74 @@ async def _upload_windows_run_script(self, environment: BaseEnvironment) -> None finally: script_path.unlink(missing_ok=True) + def _windows_cp_back_script(self, environment: BaseEnvironment) -> str: + env_paths = self._env_paths(environment) + agent_dir = self._windows_cmd_path(str(env_paths.agent_dir)) + user_root = self._windows_user_root_for(environment) + home_root = self._windows_home_for(environment) + patch_path = self._output_patch_path_for(environment) + + patch_part = "" + if patch_path: + patch_q = self._windows_cmd_path(patch_path) + patch_part = ( + f'set "PATCH_PATH={patch_q}"\r\n' + 'set "PATCH_META_PATH=%PATCH_PATH%.meta.json"\r\n' + 'if exist "%PATCH_PATH%" (\r\n' + ' > "%PATCH_META_PATH%" echo {"present":true,"created_empty_placeholder":false}\r\n' + ") else (\r\n" + ' type nul > "%PATCH_PATH%"\r\n' + ' > "%PATCH_META_PATH%" echo {"present":false,"created_empty_placeholder":true}\r\n' + ")\r\n" + ) + + return ( + "@echo off\r\n" + "setlocal EnableExtensions\r\n" + f'set "AGENT_DIR={agent_dir}"\r\n' + 'set "BITFUN_DIR=%AGENT_DIR%\\bitfun"\r\n' + 'set "SESSIONS_DIR=%BITFUN_DIR%\\sessions"\r\n' + 'set "REQUEST_TRACES_DIR=%BITFUN_DIR%\\request-traces"\r\n' + f'if not defined BITFUN_USER_ROOT set "BITFUN_USER_ROOT={user_root}"\r\n' + f'if not defined BITFUN_HOME set "BITFUN_HOME={home_root}"\r\n' + 'if not exist "%BITFUN_DIR%" mkdir "%BITFUN_DIR%"\r\n' + 'if not exist "%SESSIONS_DIR%" mkdir "%SESSIONS_DIR%"\r\n' + 'if not exist "%REQUEST_TRACES_DIR%" mkdir "%REQUEST_TRACES_DIR%"\r\n' + "for /L %%I in (1,1,6) do (\r\n" + ' for /D %%P in ("%BITFUN_HOME%\\projects\\*") do (\r\n' + ' if exist "%%P\\sessions" xcopy /E /I /Y "%%P\\sessions" "%SESSIONS_DIR%" >nul 2>nul\r\n' + ' if exist "%%P\\request-traces" xcopy /E /I /Y "%%P\\request-traces" "%REQUEST_TRACES_DIR%" >nul 2>nul\r\n' + " )\r\n" + ' dir /B /AD "%SESSIONS_DIR%\\*" >nul 2>nul && goto after_project_copy\r\n' + " ping -n 2 127.0.0.1 >nul\r\n" + ")\r\n" + ":after_project_copy\r\n" + 'if exist "%BITFUN_USER_ROOT%\\data\\token_usage" xcopy /E /I /Y "%BITFUN_USER_ROOT%\\data\\token_usage" "%BITFUN_DIR%\\token_usage" >nul 2>nul\r\n' + 'if exist "%BITFUN_USER_ROOT%\\cli-logs" xcopy /E /I /Y "%BITFUN_USER_ROOT%\\cli-logs" "%BITFUN_DIR%\\cli-logs" >nul 2>nul\r\n' + 'if exist "%BITFUN_USER_ROOT%\\logs\\bitfun-cli.log" copy /Y "%BITFUN_USER_ROOT%\\logs\\bitfun-cli.log" "%BITFUN_DIR%\\cli.log" >nul 2>nul\r\n' + 'if exist "%BITFUN_USER_ROOT%\\logs\\ai-request-audit.jsonl" copy /Y "%BITFUN_USER_ROOT%\\logs\\ai-request-audit.jsonl" "%BITFUN_DIR%\\ai-request-audit.jsonl" >nul 2>nul\r\n' + '> "%BITFUN_DIR%\\cp-back-manifest.json" echo {"windows_cp_back":true}\r\n' + f"{patch_part}" + "exit /b 0\r\n" + ) + + async def _upload_windows_cp_back_script( + self, environment: BaseEnvironment + ) -> None: + script_path = self._new_app_config_capture_temp_path(".bitfun-cp-back.bat") + try: + script_path.write_text( + self._windows_cp_back_script(environment), + encoding="utf-8", + newline="", + ) + await environment.upload_file( + script_path, + self._cp_back_script_path(environment), + ) + finally: + script_path.unlink(missing_ok=True) + async def _register_windows_config(self, environment: BaseEnvironment) -> None: if self._bitfun_config is None: return @@ -2473,83 +2546,7 @@ def _log_cp_back_gaps(self) -> None: def _cp_back_command(self, environment: BaseEnvironment | None = None) -> str: if environment is not None and self._task_os(environment) == TaskOS.WINDOWS: task_os = self._task_os(environment) - env_paths = self._env_paths(environment) - bitfun_dir = quote_shell_arg(str(env_paths.agent_dir / "bitfun"), task_os) - bitfun_dir_probe = quote_shell_arg( - str(env_paths.agent_dir / "bitfun/"), task_os - ) - sessions_dir = quote_shell_arg( - str(env_paths.agent_dir / "bitfun/sessions"), task_os - ) - sessions_dir_probe = quote_shell_arg( - str(env_paths.agent_dir / "bitfun/sessions/"), task_os - ) - request_traces_dir = quote_shell_arg( - str(env_paths.agent_dir / "bitfun/request-traces"), task_os - ) - request_traces_dir_probe = quote_shell_arg( - str(env_paths.agent_dir / "bitfun/request-traces/"), task_os - ) - token_usage_dir = quote_shell_arg( - str(env_paths.agent_dir / "bitfun/token_usage"), task_os - ) - cli_logs_dir = quote_shell_arg( - str(env_paths.agent_dir / "bitfun/cli-logs"), task_os - ) - cli_log_path = quote_shell_arg( - str(env_paths.agent_dir / "bitfun/cli.log"), task_os - ) - audit_log_path = quote_shell_arg( - str(env_paths.agent_dir / "bitfun/ai-request-audit.jsonl"), - task_os, - ) - manifest_path = quote_shell_arg( - self._remote_cp_back_manifest_path(environment), task_os - ) - user_root = self._windows_user_root_for(environment) - home_root = self._windows_home_for(environment) - commands = [ - f"if not exist {bitfun_dir_probe} mkdir {bitfun_dir}", - f"if not exist {sessions_dir_probe} mkdir {sessions_dir}", - ( - f'for /d %P in ("{home_root}\\projects\\*") do ' - f'if exist "%P\\sessions" xcopy /E /I /Y "%P\\sessions" {sessions_dir} >nul 2>nul' - ), - ( - f'if exist "{user_root}\\data\\token_usage" ' - f'xcopy /E /I /Y "{user_root}\\data\\token_usage" {token_usage_dir} >nul 2>nul' - ), - ( - f'if exist "{user_root}\\cli-logs" ' - f'xcopy /E /I /Y "{user_root}\\cli-logs" {cli_logs_dir} >nul 2>nul' - ), - ( - f'if exist "{user_root}\\logs\\bitfun-cli.log" ' - f'copy /Y "{user_root}\\logs\\bitfun-cli.log" {cli_log_path} >nul 2>nul' - ), - ( - f'if exist "{user_root}\\logs\\ai-request-audit.jsonl" ' - f'copy /Y "{user_root}\\logs\\ai-request-audit.jsonl" {audit_log_path} >nul 2>nul' - ), - f"if not exist {request_traces_dir_probe} mkdir {request_traces_dir}", - ( - f'for /d %P in ("{home_root}\\projects\\*") do ' - f'if exist "%P\\request-traces" xcopy /E /I /Y "%P\\request-traces" {request_traces_dir} >nul 2>nul' - ), - f'echo {{"windows_cp_back":true}} > {manifest_path}', - ] - patch_path = self._output_patch_path_for(environment) - if patch_path: - patch_q = quote_shell_arg(patch_path, task_os) - meta_q = quote_shell_arg(f"{patch_path}.meta.json", task_os) - commands.append( - f"if exist {patch_q} " - f'(echo {{"present":true,"created_empty_placeholder":false}} > {meta_q}) ' - f"else (type nul > {patch_q} & " - f'echo {{"present":false,"created_empty_placeholder":true}} > {meta_q})' - ) - commands.append("exit /b 0") - return " & ".join(commands) + return quote_shell_arg(self._cp_back_script_path(environment), task_os) command = _CP_BACK_COMMAND if self._output_patch_path: @@ -2652,6 +2649,8 @@ async def run( f"Failed to capture BitFun final repo state: {exc}" ) try: + if self._task_os(environment) == TaskOS.WINDOWS: + await self._upload_windows_cp_back_script(environment) await self.exec_as_agent( environment, command=self._cp_back_command(environment), diff --git a/tests/unit/agents/installed/test_bitfun_cli.py b/tests/unit/agents/installed/test_bitfun_cli.py index 7dbe2d413a8..cac6f1d9e86 100644 --- a/tests/unit/agents/installed/test_bitfun_cli.py +++ b/tests/unit/agents/installed/test_bitfun_cli.py @@ -1273,11 +1273,11 @@ async def test_windows_run_uploads_prompt_and_uses_windows_paths(self, temp_dir) commands = _exec_commands(mock_env) run_cmd = _first_command_containing(commands, "bitfun-run.bat") - cp_cmd = _first_command_containing(commands, "windows_cp_back") + cp_cmd = _first_command_containing(commands, "bitfun-cp-back.bat") probe_cmd = _first_command_containing(commands, "source=") assert "set -o pipefail" not in run_cmd assert run_cmd == "C:\\logs\\agent\\bitfun-run.bat" - assert "C:\\logs\\agent\\bitfun" in cp_cmd + assert cp_cmd == "C:\\logs\\agent\\bitfun-cp-back.bat" assert "C:\\bitfun-user\\config\\app.json" in probe_cmd assert "%BITFUN_USER_ROOT%" not in probe_cmd uploaded_targets = [ @@ -1285,6 +1285,12 @@ async def test_windows_run_uploads_prompt_and_uses_windows_paths(self, temp_dir) ] assert "C:/logs/agent/bitfun-prompt.txt" in uploaded_targets assert "C:/logs/agent/bitfun-run.bat" in uploaded_targets + assert "C:/logs/agent/bitfun-cp-back.bat" in uploaded_targets + cp_back_script = agent._windows_cp_back_script(mock_env) + assert 'set "AGENT_DIR=C:\\logs\\agent"' in cp_back_script + assert "C:\\bitfun-user" in cp_back_script + assert "C:\\bitfun-home" in cp_back_script + assert "windows_cp_back" in cp_back_script @pytest.mark.asyncio async def test_run_attempts_final_app_config_capture_after_cp_back(self, temp_dir):